Porting a 128-expert MoE (Gemma-4 26B-A4B) to AWS Inferentia2 — where every rank weighted the wrong experts

Of the five Gemma-4 models I ported to AWS Inferentia2, the 26B-A4B was the one that was supposed to
be impossible: a Mixture of Experts. It compiled on the first try — then produced nothing. The
device output was empty. The CPU reference was perfect. My unit tests all passed. The bug was in a
place I'd have sworn couldn't have one.

This is the MoE entry in the series. The dense models (E2B/E4B/12B/31B) are hard because of
architecture; the 26B is hard because it's sparse — 128 experts, top-8 routing — and none of the
AWS vendor stack has ever traced an MoE for Gemma-4.

"A4B" is a lie your memory budget will not forgive

The name says 26B-A4B: ~4 billion active parameters per token. That sounds like a 4B model. It
is not. All 128 experts (~49 GB) must be resident in HBM — top-8 routing only reduces compute, not
footprint. So you budget for a 26B, not a 4B, and it needs the 24xlarge's 192 GB, same class as the
2×-larger dense 31B. (This bites again later when we try to shrink it onto a smaller box.)

The architecture: a dual-path FFN nobody warned me about

I inspected the transformers-5.13 module on a meta device before writing a line of port code, and it's
weirder than "MLP replaced by MoE." Every one of the 30 layers runs a shared dense MLP in parallel
with the 128-expert MoE, combined and passed through four feed-forward layernorms:

residual (post-attention)
├─ dense: pre_feedforward_layernorm  → mlp (2112)                → post_feedforward_layernorm_1  ─┐
└─ moe:   pre_feedforward_layernorm_2 → router → 128 experts (704) → post_feedforward_layernorm_2 ─┤
                                            hidden = dense + moe ────────────────────────────────────┘
   → post_feedforward_layernorm → + residual → × layer_scalar

• Router (Gemma4TextRouter): its own RMSNorm + a scale + a per_expert_scale, then softmax(128, fp32) → top-8 → renormalize → × per_expert_scale.
• Experts (Gemma4TextExperts): a fused gate_up_proj [128,1408,2816] + down_proj [128,2816,704], and its forward is a sparse gather/scatter loop (torch.where, index_add_).
• Attention is the same mixed sliding/global layout as the 31B (which I'd already solved).

That sparse expert loop is the crux: it's data-dependent and will not trace to a static Neuron
graph.

The traceable trick: compute all experts, mask to top-8

You can't put a per-token, per-expert gather/scatter in a static HLO graph. So don't. Compute all 128
experts on every token, weight each by its router weight — which is zero for the non-top-8 — and
sum. Because a non-selected expert contributes 0 × expert(x) = 0, this is mathematically identical
to HF's sparse top-8, but it's a fixed-shape sequence of matmuls the compiler loves. Wasteful in FLOPs
(you compute 128 experts to use 8), correct in every bit, and traceable.

And the beautiful part: the whole thing collapses onto two standard parallel linears that
ModelBuilder already knows how to shard — no custom 3D-parameter sharding:

• gate_up_proj [128,1408,2816] → [180224, 2816] as one ColumnParallelLinear (rank r gets experts 16*r*…16*r*+15)
• down_proj [128,2816,704] → [2816, 90112] as one RowParallelLinear (input-sharded → all-reduce)

I keep HF's dual-path layer forward and its router unchanged (router replicated; its top_k_weights
already carry the renorm + per-expert scale) and swap only self.experts → my DenseExperts. Before
spending a cent on hardware, I unit-tested the math and the TP=8 sharding against HF on CPU:
MAXDIFF ≈ 2e-6, cosine 1.0. The math was right.

The bug: SEQ_MATCH-perfect math, empty device output

First device run: it compiled (the first-of-its-kind MoE trace on Neuron — 30 MoE layers,
MB_TRACED). And the output was ''. Empty. First token = end-of-turn, immediately.

• CPU reference: "The capital of France is Paris." ✅
• My all-experts-dense math: unit-tested identical to HF ✅
• My TP=8 sharding logic: unit-tested identical to HF ✅
• Device: nothing ❌

When your math is proven, your sharding logic is proven, and the device is still wrong, the bug is in
the gap between "how I think NxD's primitives behave" and "how they actually behave under the trace."

Here it was. To weight rank r's experts by the router, I scattered the dense weight matrix per-rank
with scatter_to_tensor_model_parallel_region. That function picks its slice using
get_tensor_model_parallel_rank() — a Python int. ModelBuilder compiles one rank and replicates
the graph across all 8 at runtime. So the trace baked rank 0's slice (Wd[:, 0:16]) into the graph,
and at runtime every rank weighted experts 0–15 while its gate_up computed its own experts
(16–31, 32–47, …). Total misalignment → garbage → the model confidently ends the turn.

The auto-sharded ColumnParallel/RowParallel don't have this problem because their weights are
pre-sliced at load time — the forward graph is rank-agnostic. A standalone collective on a runtime
activation is not.

The fix is the exact mechanism NxD's own MoE module uses (enable_spmd_rank): register an SPMDRank
module (a [1] int32 parameter checkpoint-loaded from arange(TP), so each rank gets its own number)
and scatter with the runtime-rank-aware variant:

Wl = scatter_to_process_group_spmd(Wd, 1, self.spmd_rank.get_rank())  # each rank gets ITS experts

Recompile. CPU GEN == DEV GEN == "The capital of France is Paris.", SEQ_MATCH True, 77 ms
prefill. The last and hardest of the five Gemma-4 variants was on Inferentia.

The general lesson, now taped to my monitor: any standalone tensor-parallel collective on a
runtime activation needs the SPMD-rank variant under ModelBuilder — auto-sharded parallel linears are
fine because their weights are pre-sliced; runtime tensors are not.

A packaging trap worth 20 wasted minutes

Publishing, the server wouldn't load the model: PytorchStreamReader ... failed finding central
directory. My background "mirror to S3" loop had uploaded the 64.6 GB neff while torch.jit.save was
still writing it — a truncated 21.9 GB partial — and I'd terminated the box before the final clean
upload finished. Fix: never mirror an artifact mid-write; assert zipfile.is_zipfile(...) then do one
clean upload. (The "real" neff being 3× the size of the corrupt partial was the tell.)

Can it run on a smaller box? Yes, no, and "deeper than it looks"

The experts are ~93% of the weight, so quantizing them is the obvious lever. I built an int8-expert
variant (QuantizedColumnParallel/QuantizedRowParallel, per-channel symmetric) — two NxD autograd
quirks to fix on the way (.clone() the down output; monkeypatch scale_dequantize to be out-of-place,
it does x *= scale in place on an async-comm view). The result was genuinely good:

• int8 is numerically PERFECT — SEQ_MATCH True, token-for-token identical to fp32.
• int8 is genuinely smaller — neff 41.8 GB vs 64.6 GB bf16, exactly the 22.8 GB expert saving (the compiler keeps it int8, doesn't upcast).

But it fits only the 24xlarge. On a 2-core box (inf2.8xlarge/inf2.xlarge, 32 GB HBM, TP=2) it
Allocation Failures — experts drop to 11.4 GB/rank, but lm_head + the compiled graph + Neuron's
runtime reserve push it ~3–4 GB over the 16 GB core. And inf2 has no 4-core SKU to split the
difference. Fitting a 2-core box needs fp4 experts — which, unlike int8's five-line quantizer, is a
deep NxD microscaling integration (packed uint16, from_float hooks) with real accuracy risk. That's
a separate expedition.

So the shipped set is bf16 (24xlarge) and a validated int8 (24xlarge, smaller/faster, bf16-exact
quality) — both published — with fp4-for-2-cores documented as the next frontier.

Takeaways

1. A data-dependent op won't trace? Make it fixed-shape. All-experts-dense trades FLOPs for a static graph and exact correctness. Optimize routing later.
2. Unit-test math AND sharding logic on CPU before you pay for a NeuronCore. It's what let me say with certainty "the bug is in the primitives, not my code."
3. Runtime-rank is the MoE-on-ModelBuilder gotcha. Weights get pre-sliced; runtime activations need the SPMD-rank scatter. If the device is empty while CPU + unit tests are perfect, look here.
4. "Active params" is a compute number, not a memory number. MoE saves FLOPs, not HBM.
5. int8 per-channel was free accuracy-wise — a smaller, faster serving artifact at zero quality cost. The sub-8-bit step is where the real work (and risk) begins.

Artifacts

• Docker Hub: xbill9/gemma4-optb-26b (bf16) · xbill9/gemma4-optb-26b-int8
• HF: xbill9/gemma-4-26B-A4B-it-inferentia2 (bf16) · -int8
• Recipe: tp_mb_moe.py (the DenseExperts + SPMDRank scatter) · tp_mb_moe_int8.py

Written with AI assistance (the port, the debugging, and this write-up were done in a Claude Code
session); every log line quoted is from a real run on real hardware.