A developer’s guide to training with Ironwood TPUs

This post provides a technical guide for developers on optimizing AI model training using Google's seventh-generation Ironwood TPU within the JAX and MaxText ecosystems.

• •Ironwood supports pods of up to 9,216 chips via ICI, OCS, DCN, and massive HBM, with native FP8 support in MXUs that can theoretically double throughput over BF16
• •The Qwix library enables FP8 training recipes configurable through MaxText flags without compromising model quality
• •Tokamax kernels offer Splash Attention for long-context I/O efficiency, Megablox GMM for MoE ragged tensors, and kernel tuning utilities for tile size optimization
• •Fourth-generation SparseCores can offload All-Gather and Reduce-Scatter collectives via XLA flags, freeing TensorCores for primary computation
• •