Five Gemma-4 models, one accelerator: what porting E2B 31B to AWS Inferentia2 taught me

I ported the whole Gemma-4 family — E2B, E4B, 12B, 31B, and the 26B-A4B MoE — to run on...

title: "Five Gemma-4 models, one accelerator: what porting E2B → 31B to AWS Inferentia2 taught me" published: true

I ported the whole Gemma-4 family — E2B, E4B, 12B, 31B, and the 26B-A4B MoE — to run on AWS Inferentia2. Each has its own write-up in this series; this is the map. What's shared, what's different, how the recipe evolved from "trace it and pray" to a single-rank-compile pipeline that carries a 30-billion-parameter dense model and a 128-expert MoE — and the one bug that appears, in some costume, in every port.

The whole family at a glance

All five decode token-for-token identically to their CPU fp32 reference, and all five answer "The capital of France is Paris." That equality bar — not "looks fluent" — is the thread that ties the series together, for a reason I'll come back to.

What every port shares

1. The vendor stack can't do any of them. optimum-neuron and the Neuron vLLM backend have no Gemma-4 model class at all, and their graph builders can't express the architecture (KV-sharing on the E-family, MoE on the 26B). The public endpoint I started from loaded something and served fluent-looking gibberish. Every one of these is a hand-rolled port of transformers-5.13's own forward.

2. The same skeleton underneath. Softcap 30, a 262,144-token vocabulary, tied embeddings, GQA, and interleaved sliding + global attention show up across the family. So does the serving shape: host-side word embeddings, a device prefill + decode with an on-device KV cache, greedy decode compared against a host fp32 reference.

3. The same environment quirks. The aws_neuronx_venv_pytorch_2_8_nxd_inference venv + pip install transformers==5.13.0, a transformers.utils.fx shim so neuronx_distributed imports under the new transformers, and export PATH=$VENV/bin:/opt/aws/neuron/bin (that PATH bit fixes a libneuronpjrt-path import error that will waste your afternoon).

4. Softcap is a host-only concern. 30·tanh(x/30) is monotonic, so it never changes the greedy argmax. Every port drops it from the device graph (no tanh over 262K logits) and re-applies it host-side only when sampling needs real probabilities.

What makes each one different

E2B / E4B — the "effective" models. These are the hard-architecture ports. Per-Layer Embeddings (a separate embedding contribution per layer), cross-layer KV-sharing (many layers don't compute their own K/V — they borrow a neighbor's), and MatFormer nesting (the "2B/4B" is a sub-network of a bigger model). The whole reason this project started: KV-sharing is exactly what the vendor graph builder cannot express — but it traces fine as live graph dependencies when you just torch_neuronx.trace the real forward. E2B fits one 16 GB core; E4B is ~2× and needs TP=2, which forces the GQA question below.

12B — the encoder-free multimodal one. Ships as Gemma4UnifiedForConditionalGeneration (multimodal class) with no audio/vision encoder loaded — you take model.language_model and move on. Drops PLE and KV-sharing, but its sliding_window=1024 attention overflows Neuron's fused-attention SBUF, so you force the eager attention implementation to make it compile.

31B — the dense giant. No PLE, no KV-sharing — architecturally simple. The hard part is scale: it doesn't fit fewer than 8 cores, and the hand-driven tensor-parallel tracer that served E4B OOMs or deadlocks at 31B. This is where NxD ModelBuilder becomes mandatory (compile one rank, shard weights per-rank). A ~39-minute compile, 182 GB host peak, and a mixed-attention layout where the 10 global layers (nkv=4) get replicated instead of sharded.

26B-A4B — the MoE. The only sparse one, and the weirdest: a shared dense MLP running in parallel with a 128-expert MoE, four feed-forward layernorms, a router with its own scale + per-expert scale. The expert forward is a data-dependent gather/scatter that won't trace — so you compute all 128 experts, mask to the top-8 (identical math, static shape). "A4B" means ~4B active params but all 128 experts (~49 GB) are resident — it saves compute, not memory, and needs the same 24xlarge as the 2×-larger 31B.

The recipe evolved — and that's the real story

You can read the family as one recipe maturing under pressure:

1. torch_neuronx.trace (E2B). Trace the real HF forward on one core. Proves KV-sharing/PLE compile at all. Doesn't scale past one core.
2. tp_alias (E4B). Hand-driven tensor parallelism: on-device aliased KV buffers + a CPU-seed prefill, sharding GQA down to 1 KV head/rank. Serves two cores. Doesn't survive contact with 31B (the all-ranks-in-process tracer OOMs; serializing it deadlocks).
3. NxD ModelBuilder (12B, 31B, 26B). Compile one rank, load weights per-rank from a checkpoint_loader. This is the recipe that finally carries everything — the dense giant and the MoE — because the per-rank compile is what makes 8-way TP tractable.

Each step was forced by the next model refusing to run on the previous step's recipe. The MoE then extends ModelBuilder one more notch (all-experts-dense experts as two auto-shardable parallel linears + an SPMD-rank scatter).

The bug that's in every port

Here's the thread from the top. Every one of these ports, at some point, produced output that was wrong while the CPU reference and my unit tests were perfect — and every time, the cause was the same shape: the device faithfully computed a broken input, and only a coherence check caught it.

• 31B: a manually-built chat prompt tokenized <start_of_turn> into literal characters (Gemma-4's turn tokens are actually <|turn>/<turn|>). Device output matched the CPU reference token for token — because both faithfully processed the same garbage. SEQ_MATCH True on nonsense.
• 26B MoE: the device output was empty while CPU was perfect. A standalone tensor-parallel scatter baked rank 0's slice into the single-rank ModelBuilder trace, so every rank weighted the wrong experts. My math and sharding logic unit-tested identical to HF — the bug lived entirely in "how the primitive behaves under the trace."
• E-family: the original "gibberish" endpoint — a real model loaded and producing confident nonsense because the harness around it (PLE path, embedding scale, KV wiring) was wrong, not the accelerator.

The lesson that survives all five: device == reference proves the accelerator is faithful. It does not prove the reference is correct. A shared upstream bug — chat template, PLE, a runtime-rank scatter — passes an equality check in dead silence. So you assert the output is sensible ("Paris"), not only that device equals host. It's why every result in this series is quoted as an actual sentence, not just a green SEQ_MATCH.

If you only take four things

1. Trace the model's real forward. The features that break vendor graph builders (KV-sharing, MoE routing) often trace fine as ordinary graph dependencies. Don't reimplement what you can trace.
2. Let the recipe escalate with the model. trace → hand-rolled TP → ModelBuilder. Reach for the single-rank-compile path before the manual one OOMs, not after.
3. "Active params" and "effective size" are compute/marketing numbers, not memory numbers. MatFormer and MoE don't shrink the HBM you must reserve — budget for the whole thing.
4. Equality with a reference is necessary, not sufficient. Always also check the output is coherent. The most expensive bugs in this whole family were the ones that passed SEQ_MATCH.

Artifacts

• Docker Hub: xbill9/gemma4-optb{,-e4b,-31b,-26b,-26b-int8}
• Hugging Face: xbill9/gemma-4-{E2B,E4B,31B,26B-A4B}-it-inferentia2 (+ -int8)
• Per-model deep-dives: the other posts in this series.

Written with AI assistance (the ports, the debugging, and this write-up were done across Claude Code sessions); every number quoted is from a real run on real hardware.