Портирование Gemma-4 на AWS Inferentia2: полное руководство

Зачем портировать Gemma-4 на Inferentia2?

Семейство Gemma-4 включает модели от 2B до 31B параметров, включая разреженные MoE. Официальный стек AWS (optimum-neuron, Neuron vLLM) не поддерживает эти архитектуры — ни общие KV-кэши, ни MoE. В этой статье я покажу, как вручную портировать каждую модель, используя единый подход: трассировку реального forward, ручной тензорный параллелизм и ModelBuilder. В конце вы сможете запустить любую модель Gemma-4 на Inferentia2 с точностью до токена.

Обзор семейства Gemma-4

Все пять моделей имеют общую основу: softcap 30, словарь на 262144 токена, связанные эмбеддинги, GQA и перемежающееся скользящее + глобальное внимание. Различия — в плотности, наличии Per-Layer Embeddings (PLE) и KV-шеринга.

Все пять моделей декодируют токен-в-токен идентично CPU fp32-эталону и отвечают: «Столица Франции — Париж». Это жёсткий критерий корректности.

Общие проблемы при портировании

1. Отсутствие поддержки в vendor-стеке

optimum-neuron и Neuron vLLM не имеют класса модели Gemma-4. Их графовые построители не могут выразить KV-шеринг или MoE. Все порты — ручные, на основе transformers 5.13.0.

2. Одинаковый скелет

Softcap 30, 262144-токенный словарь, связанные эмбеддинги, GQA, перемежающееся скользящее + глобальное внимание. Схема обслуживания: host-side эмбеддинги, устройство выполняет prefill + decode с on-device KV-кэшем, жадный decode сравнивается с host fp32-эталоном.

3. Окружение

Используйте aws_neuronx_venv_pytorch_2_8_nxd_inference venv + pip install transformers==5.13.0. Добавьте shim для transformers.utils.fx, чтобы neuronx_distributed импортировался. Установите export PATH=$VENV/bin:/opt/aws/neuron/bin — это фиксит ошибку импорта libneuronpjrt-path.

4. Softcap — только на хосте

30·tanh(x/30) монотонен, поэтому не меняет жадный argmax. Удалите его из графа устройства и применяйте на хосте при семплировании.

Особенности каждой модели

E2B / E4B — «эффективные» модели

Самые сложные архитектурно: Per-Layer Embeddings (отдельный вклад эмбеддингов на каждый слой), cross-layer KV-шеринг (многие слои заимствуют K/V у соседей) и MatFormer (вложенность). KV-шеринг не выражается в vendor-построителе, но отлично трассируется как обычные графовые зависимости через torch_neuronx.trace. E2B помещается на одно ядро 16 ГБ; E4B требует TP=2.

12B — мультимодальная без энкодера

Поставляется как Gemma4UnifiedForConditionalGeneration, но без загрузки аудио/визуального энкодера — используйте model.language_model. PLE и KV-шеринг отсутствуют. Проблема: sliding_window=1024 переполняет SBUF Neuron, поэтому принудительно используйте eager-реализацию внимания.

31B — плотный гигант

Архитектурно прост: нет PLE, нет KV-шеринга. Сложность в масштабе: нужно 8 ядер, ручной тензорно-параллельный трейсер для E4B падает с OOM. Здесь обязателен NxD ModelBuilder: компилируете один ранг, шардите веса по рангам. Компиляция ~39 минут, пик памяти 182 ГБ. Глобальные слои (nkv=4) реплицируются, а не шардируются.

26B-A4B — MoE

Единственная разреженная модель: общий плотный MLP + 128 экспертов MoE, четыре layernorm, роутер с собственным масштабом. Экспертный forward — data-dependent gather/scatter, который не трассируется. Решение: вычисляйте все 128 экспертов, маскируйте top-8 (статическая форма). «A4B» означает ~4B активных параметров, но все 128 экспертов (~49 ГБ) reside — экономия на compute, не на памяти.

Эволюция рецепта сборки

1. torch_neuronx.trace (E2B): трассировка реального forward на одном ядре. Доказал, что KV-шеринг и PLE компилируются. Не масштабируется на несколько ядер.
2. tp_alias (E4B): ручной тензорный параллелизм с алиасированными KV-буферами и CPU-seed prefill. Шардирует GQA до 1 KV-головы на ранг. Работает на двух ядрах, но падает на 31B.
3. NxD ModelBuilder (12B, 31B, 26B): компиляция одного ранга, загрузка весов порангово через checkpoint_loader. Этот рецепт выдерживает и плотного гиганта, и MoE. MoE расширяет ModelBuilder: все эксперты — два auto-shardable параллельных линейных слоя + SPMD-rank scatter.

Баги, которые повторяются в каждом порте

Главный урок: равенство устройства и эталона не гарантирует корректность. Примеры:

• 31B: ручной chat prompt токенизировал как буквальные символы (правильные токены — <|turn|>/). Устройство и CPU выдали одинаковый мусор — SEQ_MATCH True на бессмыслице.
• 26B MoE: устройство выдавало пустой вывод, CPU — правильный. Оказалось, tensor-parallel scatter «запёк» срез ранга 0 в трассировку ModelBuilder, и каждый ранг взвешивал не те эксперты.
• E-family: оригинальный эндпоинт выдавал уверенную бессмыслицу из-за неправильной обвязки (PLE, масштаб эмбеддингов, KV-соединения), а не из-за акселератора.

Вывод: проверяйте не только равенство, но и осмысленность вывода («Париж»).

Практические выводы

1. Трассируйте реальный forward модели. То, что ломает vendor-построители (KV-шеринг, MoE-роутинг), часто отлично трассируется как обычные графовые зависимости.
2. Эскалируйте рецепт по мере роста модели: trace → ручной TP → ModelBuilder. Переходите на single-rank-compile до того, как ручной подход упадёт с OOM.
3. «Активные параметры» и «эффективный размер» — числа для compute/маркетинга, не для памяти. MatFormer и MoE не уменьшают HBM, которую нужно зарезервировать.
4. Равенство с эталоном необходимо, но недостаточно. Всегда проверяйте, что вывод осмыслен. Самые дорогие баги — те, что проходят SEQ_MATCH.

Начните с портирования E2B на одном ядре, затем переходите к более крупным моделям. Используйте готовые образы Docker: xbill9/gemma4-optb и аналоги. Полный код и deep-dive — в сопутствующих статьях серии.