Frontier Pretraining on TPU: GPT-OSS with MaxText
Patrick Toulme’s deep dive into running a GPT-OSS MoE pretraining job on TPU using MaxText, tracing the full compilation pipeline from Python to fused TPU kernels.
Key Claims
- MaxText + JAX/XLA does automatically what teams spend months rebuilding by hand in PyTorch: FSDP sharding, kernel fusion, collective overlap, MoE routing
- 3.5B parameter MoE (16 experts, top-2 routing) trained for under $10 on 4 TPU v6e chips
- Single
python3 -m maxtext.trainers.pre_train.traincommand with config flags produces: 887 fused kernels, 104 async all-gathers, 12 ragged all-to-all collectives, 8 Pallas attention kernels, 24 MegaBlox grouped matmuls - 22.2 TFLOP/s per device at steady state (182ms per step)
- The parallelism specification is a configuration parameter (
ici_fsdp_parallelism=2 ici_expert_parallelism=2), not code
Technical Highlights
- XLA compiles entire training step (forward + backward + optimizer) into a single binary — enables cross-boundary fusion and compute-communication overlap
- RMSNorm: 10 separate HLO instructions fused into one kernel (saves 320MB memory traffic without fusion × 32 instances)
- Adam optimizer: 15 instructions per parameter fused into monster kernels
- Step 0: 30s (compilation). Steps 1-199: 180ms each (execution)
Takeaways
- MaxText is dramatically underused — “I rarely hear MaxText mentioned as even an option”
- Teams rebuilding training infra from scratch in PyTorch often don’t achieve as high MFU
- The XLA compilation approach (whole-program optimization) is fundamentally different from PyTorch’s approach
- This is the open-source equivalent of frontier pretraining infrastructure
ici_fsdp_parallelism=2→ici_fsdp_parallelism=128— same code, no changes needed
Connections
- Scaling & Compute — infrastructure for frontier pretraining
- MoE Architectures — MoE training specifics
- Kernel Engineering — XLA automated vs manual kernel engineering
- RL Infrastructure — similar distributed systems challenges