---
title: "JAX"
description: "6 posts about JAX, summarised, each linking to the original."
---

# JAX
> 6 posts about JAX, summarised, each linking to the original.

## Articles

### [Reproducing OLMo 3 7B Pre-training in MaxText: case study of large scale training on TPUs](https://yomu.fyi/post/reproducing-olmo-3-7b-pre-training-in-maxtext-case-study-of-large-scal.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Gagik Amirkhanyan, Ran Ran, Aireen Mei, Matt Davidow, TPU Inference Software Engineering Team, Google Cloud, AI2 Team
- Published: Sep 24, 2026
- Type: Problem & solution

The post describes the reproduction of AI2’s OLMo 3 7B pre-training using MaxText on Google Cloud TPUs, covering both stages, matching the reference loss curve and held-out metrics. PyTorch weights were converted to JAX with logit parity, KL divergence 1.5e-3 and 98.75% top-1 agreement, and held-out evaluations caught a data loader bug that made incorrect gains appear. The run demonstrated reliability via exact checkpoint-resume with zero loss delta, resilience to a 75% capacity loss via resize with 1% throughput preservation, and seamless TPU generation change. Performance optimizations achieved 44.5% MFU on Ironwood, and head reshaping from 32 heads by 128 to 16 by 256 yielded 12.4% faster throughput at identical parameters and FLOPs. The post emphasizes that faithful reproduction requires verification on generalization metrics and provides instructions for readers to reproduce the run with open sourced code.


### [Accelerating Dropless MoE Training in JAX with NVIDIA Transformer Engine](https://yomu.fyi/post/accelerating-dropless-moe-training-in-jax-with-nvidia-transformer-engi.md)
- Company: [NVIDIA](https://yomu.fyi/company/nvidia-developer-blog.md)
- Author: Tanya Lenz
- Published: Sep 14, 2026
- Type: Explainer

Production-scale mixture-of-experts (MoE) training must handle uneven, data-dependent expert loads, creating ragged tensors and expensive dispatch, combine, and all-to-all communication. Dropless MoE processes every selected token, preserving model quality but requiring kernels that support variable token counts without padding or dropping. The post presents JAX and NVIDIA Transformer Engine optimizations centered on group-aware MXFP8 quantization, grouped GEMM, and fused expert-parallel (EP) dispatch and combine using NCCL EP. Grouped GEMM computes each expert’s actual workload in one kernel path, while EP reduces movement overhead and can deduplicate tokens crossing networks. For DeepSeek-V3 training on NVIDIA GB200, the reported throughput increased from 103 to 1,068 TFLOPS/GPU, a 10.4x improvement; the post also reports roughly 10x throughput improvement and 97% scaling efficiency at 1,024 GPUs.


### [HeyGen x Google Cloud: Bringing Avatar IV to TPUs](https://yomu.fyi/post/heygen-x-google-cloud-bringing-avatar-iv-to-tpus.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Published: Aug 13, 2026
- Type: Benchmark

HeyGen and Google Cloud ported Avatar IV, an 18B-plus-parameter diffusion stack for streamed talking-head video, to an eight-chip Trillium (v6e) host. The pipeline uses two diffusion transformers and a VAE decoder, with FSDP weight sharding and Ulysses sequence parallelism across the mesh, and runs through torchax, JAX, XLA, and shape-specific Pallas kernels. Optimization targeted exposed all-to-all collectives, frame-misaligned sparse-attention blocks, serial softmax work, layout copies, scheduler choices, and custom-kernel cost estimates; changes were measured per video chunk against output-quality gates. The resulting pipeline was 1.86× faster than the first working TPU version, comparable in performance to the 8×H100 production setup and up to 25% more cost efficient per generated-video minute, while shipped changes preserved byte-identical output or passed defined similarity review.


### [Run Ray on TPU, Part 2: Ray AI libraries](https://yomu.fyi/post/run-ray-on-tpu-part-2-ray-ai-libraries.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Ivan Nardini, Spencer Peterson
- Published: Jul 24, 2026
- Type: Tutorial

Running Ray on TPUs requires multi-host workloads to remain on one intact TPU slice whose hosts share a high-speed ICI link. This guide presents Ray's AI libraries on top of GKE and Ray Core: Ray Serve for serving, Ray Data's iter\_jax\_batches() for input pipelines and batch inference, and Ray Train's JaxTrainer for distributed JAX training. Serve uses accelerator\_config.topology to defer slice placement to the replica, while iter\_jax\_batches() returns device-sharded JAX arrays and handles ragged final batches through drop, pad, or raise choices. JaxTrainer takes a training function and slice topology, launches one worker per host, and provides checkpointing, fault-tolerant restarts, and multi-slice scale-out; official TPU images and dashboard metrics round out the workflow.


### [Scaling Agentic RL: High-Throughput Agentic Training with Tunix](https://yomu.fyi/post/scaling-agentic-rl-high-throughput-agentic-training-with-tunix.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Haoyu Gao, Lance Wang, Shadi Noghabi, Tianshu Bao, Weiren Yu
- Published: Jul 21, 2026
- Type: Explainer

Tunix is presented as a post-training library for scaling agentic reinforcement learning, where multi-turn tool use and environment latency can leave TPUs idle. Its design combines an asyncio-based, high-concurrency trajectory collector with asynchronous vLLM-TPU and SGLang-Jax integration, then streams variable-length trajectories through a producer-consumer queue rather than waiting for a complete synchronous batch. Composable agent and environment abstractions expose APIs for custom tools, benchmarks, and external systems while preserving multi-turn token handling and lifecycle management. Lightweight RL-specific metrics correlate rollout, training, weight synchronization, model-call, and environment stages with TPU timelines, complementing targeted XProf traces. The stated outcome is a more observable pipeline intended to reduce stalls, keep accelerators utilized, and support agentic RL natively in the JAX/TPU ecosystem.


### [We terminated a TPU mid-training and it recovered in seconds:  Introduction to elastic training with MaxText](https://yomu.fyi/post/we-terminated-a-tpu-mid-training-and-it-recovered-in-seconds-introduct.md)
- Company: [Google](https://yomu.fyi/company/google.md)
- Author: Luke Baumann, Abhinav Singh, Ivan Nardini
- Published: Jul 6, 2026
- Type: Tutorial

Distributed training on multiple TPU machines normally fails as a whole when one worker disappears: collective communication times out, processes exit, and the job restarts from a checkpoint. The article demonstrates elastic training with JAX, MaxText, PathwayS, Orbax, and Cloud TPUs on GKE, using a single CPU-side controller that remains alive while TPU workers execute compiled programs. A Pathways elastic\_retry wrapper catches the failure, cleans up partial state, restores the last viable checkpoint, and re-enters the training function without relaunching the workload; the demo returned to training in under two minutes, with no JobSet restarts, after losing 88 steps. It also contrasts pause-and-resume with replica resize, explains direct checkpoint persistence through ENABLE\_PATHWAYS\_PERSISTENCE=1, and describes planned snapshot-based recovery from host memory.