My Inferentia port matched its reference token-for-token — and still output garbage

Short summary

Porting the 30B dense Gemma-4 to Inferentia2 revealed that scale, not architecture, was the challenge. Manual parallel rank tracing caused OOM (8× fp32 compile graphs) or deadlock (serialized ranks hung on weight-collection rendezvous). Switching to NxD's ModelBuilder compiled one rank then loaded sharded weights, solving both. A subtle bug where layer_scalar buffers weren't moved by the checkpoint loader caused silent output corruption despite token-for-token matching — fixed by manually copying buffers from safetensors.

• •Manual parallel tracing of 8 ranks OOMs at 300+ GB; serialized tracing deadlocks on weight-collection rendezvous — use ModelBuilder instead
• •Global attention layers with 4 KV heads < TP=8 must be replicated, not sharded; sliding layers with 16 KV heads shard normally
• •layer_scalar buffers silently default to 1.0 when checkpoint loader skips them — read from safetensors and copy manually

Generated with AI, which can make mistakes.