Gemma-4 на Inferentia2: где ломается путь от модели до стабильного inference

Соло-разработчику, который выносит LLM из ноутбука в свой endpoint под агентов или IDE-воркфлоу, важнее не свежесть чекпоинта, а узкие места между «модель скачана» и побайтно верным decode. В полевом отчёте xbill все три размера Gemma-4 — E2B, E4B и 12B — завёл на AWS Inferentia2 только после обхода vendor stack; при этом тот же двухъядерный accelerator можно обслужить на inf2.xlarge примерно в четыре раза дешевле inf2.8xlarge.

Почему vLLM, optimum-neuron и NxD не довели Gemma-4 до рабочего вывода

Автор намеренно не пошёл по пути AWS vendor inference. Причины конкретны и воспроизводимы для любого, кто пытается поднять Gemma-4 «как в документации».

В optimum-neuron нет классов gemma4_text и gemma4_unified — vLLM через этот слой просто нечего инстанцировать. NxD ModelBuilder не выражает cross-layer KV-sharing: граф либо падает, либо молча теряет sharing между слоями.

Публичный Neuron vLLM/NxD endpoint поднимался healthy, но выдавал fluent-looking gibberish — текст выглядел грамматично и был семантически пуст. Отдельный удар по доверию: автопорт фреймворка отчитался «100% PASS», хотя golden reference строился из PLE-stripped checkpoint, и обе стороны сравнения были неверны одинаково.

Вывод для inference-практики: «зелёный healthcheck» и «правдоподобный токен-стрим» не заменяют CPU reference на той же машине.

Три разных «mixed heads» — не одна архитектурная ошибка

В заголовке и snippet фигурируют mixed attention heads. В полном тексте автор разводит это на три независимые проблемы шардирования — стратегия «просто поделить всё на tensor parallelism» ломается во всех трёх случаях.

E2B и E4B — KV-sharing между слоями. Часть слоёв помечена self_attn.is_kv_shared_layer и не считает собственные K/V, а читает KV «владельца». NxD не может выразить такую зависимость в статическом графе.

E4B — GQA под TP=2. Восемь query heads и два KV head; наивное деление num_key_value_groups даёт неверный repeat_kv и «правдоподобный, но неправильный» вывод.

12B — чередование sliding-window и global слоёв. Sliding-window слои с nkv=8 делятся на TP=2; global слои с num_global_key_value_heads=1 — нет. Для global K/V остаются реплицированными, а num_key_value_groups ужимается под sharded query heads.

Option B: прямой trace Hugging Face transformers 5.13

Рабочая замена — не дожимать vendor path, а трассировать live forward pass из transformers 5.13.

• E2B, single core: torch_neuronx.trace() на полном HF forward.
• E4B и 12B, TP=2: neuronx_distributed.trace.parallel_model_trace и ModelBuilder.

KV-sharing, Per-Layer Embeddings, GQA и mix sliding/global attention идут в граф как обычные live dependencies — без попытки упростить архитектуру до того, что понимает NxD из коробки.

По итогам порта greedy decode побайтно совпадает с CPU reference на всех трёх размерах (SEQ_MATCH True; пример: «The capital of France is Paris.»). Артефакты выложены на Hugging Face (xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2) и Docker Hub (xbill9/gemma4-optb{,-e4b,-12b}).

Лимиты neuronx-cc: SBUF, softcap и буферы, которые NxD не подхватывает

Повторяющееся узкое место — on-chip state buffer (SBUF): лимит 196 608 B на partition. Остальное — цепочка частных фиксов под каждую модель.

Проблема	Модели	Симптом	Fix
PLE table на device	E2B/E4B	262144 × hidden не влезает в 16 GB core	PLE и word embedding на host, gather на CPU
Logit soft-capping	12B	[NCC_INLA001] Allocated memory out of bound на tanh over vocab 262144	softcap на host при sampling; greedy без on-device cap
Fused/SDPA attention	12B	sliding_window=1024 переполняет SBUF	cfg._attn_implementation = "eager"
fp32 constants	E4B	~15.4 GB resident constants → NRT_RESOURCE / status=4	bf16 weights и TP=2
layer_scalar buffer	все три	NxD shard_children грузит parameters, не buffers → garbage (cos ≈ 0)	ручное копирование layer_scalar до trace (commit e785f6d)

Компиляция neuronx-cc идёт на CPU; пик host RAM при TP=2 compile — >128 GB. Автор рекомендует ≥55 GB swap на время сборки. Docker-образы 70+ GB разумнее собирать на inf2.8xlarge, не на inf2.xlarge.

Закреплённые версии toolchain из поста:

Компонент	Версия
Neuron SDK	2.23
torch-neuronx	2.8.0
neuronx-cc	2.23.6484
neuronx-distributed	0.17
transformers	5.13.0

Скорость decode и соло-serving на дешёвом inf2

Железо: Inferentia2 на inf2.xlarge (1 chip / 2 cores / 32 GB HBM; host RAM 16 GB) и inf2.8xlarge (тот же accelerator, 128 GB host RAM). Accelerator идентичен — разница в vCPU, host RAM и цене.

Модель	Build	First token	Decode
E2B	two-graph static KV	~0.06 s	~44 tok/s
E4B	tp_alias	~1.4–1.6 s	~33 tok/s
E4B	device-prefill (bf16)	~0.11–0.16 s	~36–39 tok/s
12B	device-prefill (bf16)	~0.1 s	~15 tok/s

E2B (~5B реальный footprint при маркетинговых ~2B) помещается на одном 16 GB core в bf16. E4B (~8B) и 12B требуют TP=2. Все три размера автор поднимает на одном inf2.xlarge через slim servers — embedding на CPU, KV cache device-resident.

Trade-off prefill: host-seed (tp_alias) тянет first token до ~1.4–1.6 s, но decode на E4B slim доходит до ~60 tok/s; device-prefill даёт first token ~0.1–0.16 s при decode ~36–39 tok/s.

Ограничения порта честны: batch size 1, single-stream; нет continuous batching и paged attention; context зашит в neff (увеличение — recompile); multimodal пути 12B не подключены к device graph.

Для соло-workflow «свой LLM под агентов» практический вывод такой: сначала CPU reference и проверка tokenizer/buffers (layer_scalar, mismatch SDK), потом trace — а не отладка «сломанного accelerator» по fluent gibberish.

Дешёвый inf2.xlarge с swapfile на время neff load (~14.5 GB пик на 16 GB host) — рабочая база, если не нужен тяжёлый compile на той же машине.

Источники

• xbill — «Porting Gemma-4 (2B / 4B / 12B) to AWS Inferentia2» (Dev.to): Dev.to — доступ 2026-07-13