# Porting Gemma-4 (2B / 4B / 12B) to AWS Inferentia2 > Source: > Published: 2026-07-13 13:36:03+00:00 *A field report on running Google's Gemma-4 family on AWS Inferentia2 ( inf2), covering the three architectural obstacles that break the vendor stack — **mixed attention heads* `optimum-neuron` / NxD`neuronx-cc` ) limits —| Models | `google/gemma-4-E2B-it` , `google/gemma-4-E4B-it` , `google/gemma-4-12B-it` | | Hardware | AWS Inferentia2 — `inf2.xlarge` (1 chip / 2 cores / 32 GB HBM) and `inf2.8xlarge` (same accelerator, 128 GB host RAM) | | Software | Neuron SDK 2.23 · `torch-neuronx` 2.8.0 · `neuronx-cc` 2.23.6484 · `neuronx-distributed` 0.17 · `transformers` 5.13.0 | | Result | E2B ~44 tok/s (1 core), E4B ~33–39 tok/s (TP=2), 12B ~15 tok/s (TP=2) — greedy decode token-for-token identical to the CPU reference on all three | | Artifacts | HF: `xbill9/gemma-4-{E2B,E4B,12B}-it-inferentia2` · Docker Hub: `xbill9/gemma4-optb{,-e4b,-12b}` | Gemma-4 is not a vanilla decoder. Across the family it combines several features that each map cleanly onto TPU/XLA (where the model was designed) but individually break the AWS inference path: `tanh` cap at 30) over a The AWS vendor stack (`optimum-neuron` + `neuronx-distributed` + the Neuron vLLM backend) has **no Gemma-4 model class at all**, and its graph builder cannot express KV-sharing. The public Neuron E2B | E4B | 12B | | |---|---|---|---| | HF class | `Gemma4ForConditionalGeneration` (text) | same | `Gemma4UnifiedForConditionalGeneration` | | model_type | `gemma4_text` | `gemma4_text` | `gemma4_unified` (encoder-free multimodal) | | Params (effective) | ~5B (2B) | ~8B (4B) | 12B | | Per-Layer Embeddings | yes | yes | no (`hidden_size_per_layer_input=0` ) | | Cross-layer KV-sharing | yes (15 non-shared layers own KV) | yes | no (`num_kv_shared_layers=0` , every layer owns KV) | | Query / KV heads | GQA | 8 q / 2 kv | 16 q / 8 kv, global layers nkv=1 | | Attention | sliding + global, `sliding_window=512` | sliding + global, `sliding_window=512` | sliding + global, `sliding_window=1024` , `attention_k_eq_v=true` | | head_dim | — | 256 | 256 (global 512) | | Logit softcap | 30 | 30 | 30 | | Vocab | 262 144 | 262 144 | 262 144 (tied embeddings) | | Fits one 16 GB core? | yes (bf16) | no → TP=2 | no → TP=2 | Three model sizes, three *different* reasons the vendor path fails, and — as it turns out — three distinct meanings of "mixed heads." "Mixed heads" was the single most expensive class of bug in this port. It shows up in three completely different forms as you move up the size ladder. On E2B/E4B a layer's attention is flagged `self_attn.is_kv_shared_layer` . Shared layers do **not** run their own `k_proj` /`v_proj` ; they attend against a KV tensor produced by an earlier "owner" layer. On TPU this is a free graph view. In AWS's `neuronx-distributed` (NxD) `ModelBuilder` , which wants a static, per-layer weight→buffer mapping, it **cannot be represented** — there is no place to say "this layer's K/V is that layer's K/V, computed once, live." The fix (**"Option B"**, see §3) is to stop fighting the framework and trace the Hugging Face forward pass directly. KV-sharing then falls out as an ordinary live data dependency in the traced graph. The port only has to enumerate which layers actually own a cache: ``` # discover the layers that write KV; shared layers never touch the cache NONSHARED = [] for i, lyr in enumerate(lang.layers[:cfg.num_hidden_layers]): a = lyr.self_attn if not a.is_kv_shared_layer: # <- the crux NONSHARED.append(i) LINFO[i] = (a.k_proj.out_features // a.head_dim, a.head_dim) # (n_kv, head_dim) ``` Only the `NONSHARED` layers (15 of them on E2B) allocate a KV buffer; the shared layers read through. This is why the KV cache is far smaller than a naïve "one buffer per layer" implementation would allocate — and why it fits a single core in bf16. E4B does not fit one 16 GB core (§4.3), so it must be sharded across both NeuronCores (TP=2). GQA now bites: E4B has **8 query heads and 2 KV heads**. Column-sharding the query projection across 2 ranks is trivial (4 q-heads/rank). The KV projection has only **2** heads, and the subtle failure is that `repeat_kv` inside attention broadcasts each KV head to a *group* of query heads — so you cannot independently halve the KV heads and the query heads without scrambling **which query attends which KV**. The rule that makes it correct: shard KV to `nkv // TP` heads per rank **only when it divides**, and crucially **keep num_key_value_groups unchanged** so `repeat_kv` still maps each rank's 4 query ``` python def _kv_rank_width(nkv): return nkv // TP if nkv % TP == 0 else nkv # 2 // 2 = 1 KV head / rank on E4B # ... per layer: nkv = a.k_proj.out_features // a.head_dim if nkv % TP == 0: # divisible: shard k/v, keep groups a.k_proj = col(a.k_proj) a.v_proj = col(a.v_proj) # else: see 2.3 ``` Naïvely also dividing `num_key_value_groups` produces the classic "4-vs-8 `repeat_kv` mismatch": plausible-looking but wrong output. The KV cache is then kept **device-resident and sliced per rank** (head *r* → rank *r*), because NxD's `forward()` only returns rank 0's KV (§5.2). The 12B is where "mixed heads" becomes structural. `gemma4_unified` interleaves two attention types with **different KV-head counts**: `nkv = 8` → divisible by TP=2 → shard to 4 heads/rank (the §2.2 path).`num_global_key_value_heads = 1` → You cannot column-shard a single KV head across two ranks. The working answer is to **leave the global layers' K/V replicated** (a plain `nn.Linear` , which NxD loads identically on every rank) and`repeat_kv` matches the ``` a = lyr.self_attn hd = a.head_dim nq = a.q_proj.out_features // hd # capture BEFORE replacing q_proj (see gotcha) a.q_proj = col(a.q_proj); a.o_proj = row(a.o_proj) nkv = a.k_proj.out_features // hd if nkv % TP == 0: # sliding layers: nkv=8 -> 4/rank, keep groups a.k_proj = col(a.k_proj); a.v_proj = col(a.v_proj) else: # global layers: nkv=1, indivisible # leave k/v REPLICATED; shrink groups to the per-rank sharded q-head count a.num_key_value_groups = (nq // TP) // nkv ``` Two gotchas made this a multi-hour fight (commit `d70dc94` ): `nq = q_proj.out_features // head_dim` `q_proj` with a `ColumnParallelLinear` , whose `.out_features` reports the `AttributeError` /wrong groups.`else` branch had literally never been exercised by E4B (all E4B layers are `nkv=2` , divisible), so it was dead, untested code until the 12B walked into it.**Takeaway:** "mixed heads" is not one problem. It is (a) mixed KV *sharing* across layers, (b) mixed query/KV *counts* within a layer (GQA), and (c) mixed attention *types* with different KV counts per type. Each needs a different sharding rule, and the naïve "just divide everything by TP" is wrong in all three. `optimum-neuron` / NxD couldn't do it The port deliberately does **not** use the AWS vendor inference path. Here is why each vendor layer failed, and what replaced it. ** optimum-neuron has no Gemma-4.** There is simply no `gemma4_text` / `gemma4_unified` model class`optimum-neuron` . The Neuron vLLM backend dispatches through `optimum-neuron` , so vLLM has nothing**NxD ModelBuilder cannot represent KV-sharing.** NxD wants a static graph where each layer maps **The vendor endpoint served gibberish.** The public Neuron vLLM/NxD deployment we began with came up "healthy" and produced fluent-looking but semantically empty text — the worst failure mode, because it *looks* like it works. **An automated port falsely reported success.** A framework auto-port pass claimed **"100% PASS"**. On inspection its golden reference had been built from a **PLE-stripped** checkpoint, so both sides of the comparison were wrong in the same way. A green check on a broken oracle is worse than a red one. **The replacement — "Option B."** Instead of teaching NxD about Gemma-4, `torch_neuronx.trace()` (for single-core E2B) and `neuronx_distributed.trace.parallel_model_trace` / `ModelBuilder` (for TP E4B/12B) are pointed **straight at the Hugging Face transformers 5.13 text forward pass**. KV-sharing, There is still a TP-launch wrinkle: `parallel_model_trace` spawns rank workers, and `transformers` 5.13's FX utilities plus the spawn need a `transformers.utils.fx` shim and a `_TP_CHILD` environment sentinel to guard against infinite re-spawn. `neuronx-cc` ) wall Once the graph is traceable, `neuronx-cc` imposes its own hard limits. The recurring one is **SBUF**, the on-chip state buffer (SRAM), capped at **196 608 B per partition**. Several Gemma-4 ops blow straight through it, each with a distinct fix. The Per-Layer Embedding table is `262144 × hidden` . Materialised on device it trips the compiler and by itself over-commits the 16 GB core. **Fix:** keep the PLE (and word) embedding **on the host**, gather on CPU, and feed the result in **as activations**. The device graph never sees the table. Gemma-4 applies `softcap * tanh(logits / softcap)` over the full 262 144-token vocabulary. In fp32 that `tanh` is a custom-call whose working set is `128 × 524288` bytes — **524 288 > 196 608 B/partition**: ``` [NCC_INLA001] Allocated memory out of bound (128x524288) 524288 vs 196608 B/partition ``` **Fix (commit 328ca88):** `argmax(softcap(x)) == argmax(x)` — greedy decode is completely unaffected. The device returns `temperature > 0` ). Correct and free.`tanh` cap does fit and is kept in-graph.)`sliding_window=1024` (12B) Dropping softcap was necessary but not sufficient — the real overflow on the 12B was the **fused attention custom-call**, because the 12B's `sliding_window = 1024` is `a0a2a75` ):`cfg._attn_implementation = "eager"` on both the config and its `text_config` ). Eager tiles attention`sliding_window=512` the fused pathA single Option-B neff for E4B materialises **~15.4 GB of fp32 model constants**; a second neff (prefill + decode) on the same core tips past the 16 GB NeuronCore budget → `NRT_RESOURCE / status=4 Allocation Failure` . Two things are needed: `MB_WDTYPE=bf16` ): NxD's `shard_children` casts the checkpoint to the layer dtype, halving the neff. E2B is the exception: in bf16 it fits one core, which is why it stays single-core at ~44 tok/s. `neuronx-cc` is happiest with plain arithmetic. Several things are written out by hand rather than left to library helpers or dynamic control flow: `0.5*x*(1+tanh(0.7978845608*(x + 0.044715*x^3)))` .`DynamicCache` at trace time`buf*(1-oh) + k*oh` — pure arithmetic, trace-safe, no scatter op or data-dependent indexing.`layer_scalar` is a Each Gemma-4 layer does `hidden_states *= self.layer_scalar` . `layer_scalar` is a registered **buffer** (default `1.0` , real value ≈ 0.06). NxD's weight-sharding (`shard_children` / `get_sharded_checkpoint` ) loads **parameters only, never buffers** — so without intervention every layer over-scales ~16×, compounding across all layers into `cos ≈ 0` garbage. **Fix (commit e785f6d):** read the per-layer `layer_scalar` from the checkpoint and copy it into the module by`neuronx-cc` runs on **CPU** — no NeuronCore needed to compile, so you can build neffs on any box. But a TP=2 compile runs `neuronx-cc` for **both ranks concurrently** and peaks **past 128 GB host RAM**; add ≥55 GB of swap before compiling (the resulting neff runs fine without swap). Correctness and SBUF are only half the story; throughput comes from keeping the KV cache **on the device** and never round-tripping it through the host. Three designs evolved: Two neffs share one static KV buffer: **prefill** (padded prompt ≤ `KV_BUCKET` , returns the 15 non-shared layers' K/V) and **decode** (single-token forward against a fixed `KV_MAX` buffer, one-hot masked write). KV tensors are graph inputs/outputs. ~44 tok/s, ~0.06 s prefill after a one-time ~100 s neff load. Under TP a neff spans **both** cores, so you cannot park prefill on core 0 and decode on core 1. Two sub-designs handle prefill: `tp_alias` `nn.Parameter` s aliased as graph I/O`input_output_aliases` ), updated in place each step, so the cache never leaves the cores. Each rank's KV is seeded with head `ModelBuilder` ) ``` # KV parameters aliased as graph I/O so the cache is never round-tripped through the host aliases = {} for j in range(NK): aliases[w.kbuf[j]] = 1 + j for j in range(NK): aliases[w.vbuf[j]] = 1 + NK + j ``` Per-token cost is essentially **flat across context length** because the cache is device-resident. The traced executor is saved with `torch.jit.save(model, path)` and reloaded with `torch.jit.load` + `nxd_model.initialize_with_saved_weights(torch.tensor(start_rank))` . (The executor's own `.save` is TorchScript's, not NxD's — a subtle trap.) All three ports match the CPU float reference **token-for-token** on greedy decode (`SEQ_MATCH True` ), e.g. *"The capital of France is **Paris**."* | Model | Build | Hardware | TP | Context | First token | Decode | Host RAM | |---|---|---|---|---|---|---|---| E2B | two-graph static KV | `inf2.8xlarge` / `inf2.xlarge` (slim) | 1 core | 512 / 128 | ~0.06 s | ~44 tok/s | ~6 GB (slim) | E4B | `tp_alias` | `inf2.8xlarge` | 2 | 512 / 128 · 2048 / 512 | ~1.4–1.6 s | ~33 tok/s | — | E4B | device-prefill (bf16) | `inf2.8xlarge` / `inf2.xlarge` (slim) | 2 | 512 / 128 | ~0.11–0.16 s | ~36–39 tok/s | ~8 GB (slim) | 12B | device-prefill (bf16) | `inf2.8xlarge` / `inf2.xlarge` (slim) | 2 | 256 / 64 | ~0.1 s | ~15 tok/s | ~8 GB (slim) | The 12B's lower decode rate is inherent (dense 12B, all weights active per token), not a port defect; per-token cost matches the 8xlarge full server. Notably, **all three run on a single inf2.xlarge** `tokenizer.json` masquerades as a device bug.`hf download` once left `tokenizer.json` absent; `GemmaTokenizer` loaded with no vocab and mapped `` . The model then faithfully emitted garbage (unused high-id tokens), which looked exactly like a device/reload/precision failure and consumed hours. The decisive diagnostic was a `tok("hello world").input_ids` before suspecting the compiler.`inf2.xlarge` (16 GB) needs swap.`inf2.xlarge` .`inf2.8xlarge` (700 GB root); the push only uploads the tiny server layer since the base is already on Hub.The three obstacles above were the entry point. The port also produced a set of findings that generalize beyond Gemma-4 and were, in aggregate, worth more than any single fix. Every time this port produced garbage, the NeuronCore was **innocent**. The causes, in order of how often they bit: a broken/missing **tokenizer** (§7), a mis-restored **weight reload** (§8.2), a wrong **buffer** (`layer_scalar` , §4.6), and a mismatched **driver/SDK** version (a precompiled neff can *mis-execute into garbage rather than error* on the wrong runtime). All four are **cheaper to rule out than a device bug**, and all four The fastest oracle turned out to be a **CPU-reference run on the same box**: load the model in bf16 with `from_pretrained` (≈5 s off the page cache once the weights are downloaded) and run one forward. If the CPU reference produces the *same* garbage as the device, the accelerator is exonerated and the bug is upstream (tokenizer, input, weights). This single technique collapsed a multi-hour "device bug" into a two-minute tokenizer fix. **Reach for the CPU oracle before profiling the neff.** The most dangerous illusion in the whole project was a green `SEQ_MATCH` . Correctness was validated **in-process** on the freshly traced model (`ModelBuilder.trace(initialize_model_weights=True)` ) — but the server loads a **saved** model in a **fresh process** and calls `initialize_with_saved_weights()` . Those are different code paths, and the in-process pass never exercised the one the server actually uses. (`torch.jit.load` *without* the init call fails outright — "This model is not initialized" — so the init step is load-bearing, not cosmetic.) Combined with the auto-port's **false "100% PASS"** against a PLE-stripped golden (§3), the lesson is blunt: **a passing test against the wrong oracle, or on the wrong execution path, is worse than a failing one.** Validate (a) the exact artifact you will ship, (b) in the exact way you will load it, E2B is marketed as "2B" and E4B as "4B" — the MatFormer/PLE *effective* parameter counts. But the **device footprint is the full parameter count** (~5B and ~8B), and that is what has to fit 16 GB. This is exactly why E4B, the "4B" model, does **not** fit one core while E2B, the "2B" model, does. **Plan capacity from real parameters, never the effective headline.** The second half of the trap: **bf16 halves the on-disk neff but not the on-device constant count.** The intuition "just use lower precision to fit" is wrong here — E4B's ~15.4 GB of resident fp32 constants stay ~15.4 GB of *slots* regardless of dtype tricks at the boundaries, and a second neff still tips past 16 GB. **Tensor parallelism, not precision, is the lever that actually fits the model on the core** (§4.4). The same cross-layer KV-sharing that NxD can't represent (§2.1, §3) is what makes E2B fit one core in the first place. Only the **15 non-shared layers allocate a KV buffer**; the rest read through. The device-resident cache is therefore far smaller than a naïve one-buffer-per-layer implementation, which is a large part of why the memory budget closes. An architectural feature that breaks the vendor abstraction can still be a **net win** once you stop fighting it. Two of the compiler wins came from *math*, not from the compiler: `argmax` — greedy decode is identical whether or not it runs. Moving it host-side (only needed for sampling) freed the SBUF it was overflowing (§4.2). Generalize: `layer_scalar` (§4.6). When a model multiplies by a learned scalar that lives in a buffer, that scalar will be wrong after any param-only load — a bug with no error message, only degraded output.`inf2.xlarge` and `inf2.8xlarge` carry the **identical 2-NeuronCore / 32 GB-HBM accelerator**; they differ only in **host** vCPU (4 vs 32) and RAM (16 vs 128 GB). Since Gemma-4's transformer runs entirely on the cores, the only thing standing between the ~4× cheaper box and full performance is **host memory** — solved by "slim" servers that keep just the embedding table on the CPU and drop the transformer layers (they live in the neff). Result: **all three models serve on a single inf2.xlarge**, and because throughput barely drops, the cheap box is Two prefill strategies coexist and trade off cleanly, so the "right" one depends on the workload: host-seed (`tp_alias` ) | device-prefill (`ModelBuilder` ) | | |---|---|---| | First token | ~1.4–1.6 s (CPU prefill) | ~0.1–0.16 s | | Decode | up to ~60 tok/s (E4B slim) | ~36–39 tok/s (E4B) | | Best for | long generations | chat / first-token-bound | Because the KV cache is **device-resident** in both, per-token latency is essentially **flat across context length** — there is no host round-trip that grows with the sequence. The knob to turn is **Hugging Face** (recipe + compiled neffs + Dockerfiles + model cards): `xbill9/gemma-4-E2B-it-inferentia2` `xbill9/gemma-4-E4B-it-inferentia2` `xbill9/gemma-4-12B-it-inferentia2` **Docker Hub** (prebuilt, run with `--device /dev/neuron0 --ipc=host` ): `xbill9/gemma4-optb` — E2B: `latest` /`512-128` , `slim` , `tp2-slim` , `tp2-2048` `xbill9/gemma4-optb-e4b` — E4B: `latest` /`512-128` , `tp2-devprefill-512` , `slim-devprefill` `xbill9/gemma4-optb-12b` — 12B: `tp2-devprefill-256` , `slim-devprefill` **Key source** (branch `gemma4-inf2-nxd-kvshare` ): `optb_kv.py` (E2B two-graph), `tp_alias_trace.py` (E4B TP+alias), `tp_mb.py` (E4B/12B device-prefill `ModelBuilder` ), `optb_server_*.py` (full / slim / device-prefill HTTP servers with OpenAI routes). Fix history: `e785f6d` (layer_scalar), `d70dc94` (12B global-layer sharding), `328ca88` (drop on-device softcap), `a0a2a75` (force eager attention), `78967ac` (bf16 weights). `gemma4_unified` audio (640 samples/40 ms) and image (48×48 patch) projection paths exist but are not yet wired to the device graph.Natural next steps: static batching, a 12B 512-context recompile, a `tp_alias` -style 12B decode path (the E4B data suggests ~2× throughput headroom), and wiring the 12B multimodal projections. *Base models © Google, Apache-2.0. The compiled neffs embed the base weights in bf16 and are redistributed as Apache-2.0 derivatives with attribution. The Option-B recipe, TP sharding rules, device-resident KV design, and servers are Apache-2.0. Not affiliated with or endorsed by Google or AWS; provided for research/testing.*