---
title: "MaxText"
description: "2 posts about MaxText, summarised, each linking to the original."
---

# MaxText
> 2 posts about MaxText, 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.


### [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.