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

Short summary

Porting the 128-expert MoE Gemma-4 26B-A4B to Inferentia2 required converting sparse expert gather/scatter into dense computation — computing all 128 experts and masking to top-8, which collapses onto standard parallel linears. The model compiled but produced empty output because a runtime collective used a Python int for rank identification, baking rank 0's expert weights into every rank's graph. The fix: register an SPMDRank int32 parameter checkpoint-loaded per rank, matching NxD's own MoE module approach.

• •Sparse expert gather/scatter can't trace to static HLO — compute all 128 experts, mask to top-8, collapse onto ColumnParallel/RowParallel linears
• •MoE footprint is total params (26B ~49GB), not active params (4B) — budget for the full model in HBM
• •Runtime collectives using Python int for rank ID get baked at compile time — use SPMDRank parameter checkpoint-loaded per rank instead

Generated with AI, which can make mistakes.