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

Short summary

An ML engineer ported all five Gemma-4 models (E2B through 26B-A4B MoE) to AWS Inferentia2, achieving token-for-token identical output to CPU fp32 references. The post details how each model required a different compilation recipe — from simple torch_neuronx.trace for the smallest, to NxD ModelBuilder with 8-way tensor parallelism for the 31B dense and 26B MoE variants. Key challenges included KV-sharing, per-layer embeddings, and Neuron's attention buffer limits.

• •All five Gemma-4 models ported to Inferentia2 with exact output match to CPU references
• •Three compilation recipes needed: torch_neuronx.trace, tp_alias, and NxD ModelBuilder depending on model size
• •Hardest ports were E2B/E4B (novel architectures) and 26B-A4B MoE (128 experts, custom routing math)

Generated with AI, which can make mistakes.