---
type: "post"
id: "announcing-long-context-training-on-tpus-with-maxtext-5a6500"
slug: "announcing-long-context-training-on-tpus-with-maxtext-5a6500"
title: "Announcing Long-Context Training on TPUs with MaxText Context Parallelism"
description: "Google says MaxText v0.2.4 can train models on sequences as long as 14 million tokens on Cloud TPUs, using four context-parallel strategies to distribute attention work across chips."
retrieval_nugget: "Google says MaxText v0.2.4 can train models on sequences as long as 14 million tokens on Cloud TPUs, using four context-parallel strategies to distribute attention work across chips. The release gives model-training teams a practical set of options for tackling memory and communication bottlenecks at lengths beyond what a single chip can comfortably handle."
published_at: "2026-11-10"
updated_at: "2026-11-10"
record_date: "2026-11-10"
date_kind: "published_at"
topics: ["google"]
entities: ["Google","MaxText","Cloud TPU","Tokamax","Llama 3.1 8B","DeepSpeed Ulysses","Unified Sequence Parallelism"]
related_patterns: []
novelty_label: "new_fact"
what_changed: "Google adds MaxText support for TPU long-context training, with context-parallel methods enabling sequences up to 14 million tokens."
prior_refs: ["https://newruntime.com/posts/introducing-tpu-8tx-and-tpu-8i-max-b6c8a7/"]
source_url: "https://developers.googleblog.com/announcing-long-context-training-on-tpus-with-maxtext-context-parallelism"
source_urls: ["https://developers.googleblog.com/announcing-long-context-training-on-tpus-with-maxtext-context-parallelism"]
schema_version: "newruntime-agent-readable-v0.2"
stable_id: "post:announcing-long-context-training-on-tpus-with-maxtext-5a6500"
status: "published"
visuals: []
editor: {"role":"editor","name":"Andrey Reshetnikov","profile":"https://newruntime.com/owner-profile.md","profile_json":"https://newruntime.com/owner-profile.json","same_as":["https://www.linkedin.com/in/reshetnikov1/","https://t.me/lapetuse"]}
routes: {"html":"https://newruntime.com/posts/announcing-long-context-training-on-tpus-with-maxtext-5a6500/","markdown":"https://newruntime.com/posts/announcing-long-context-training-on-tpus-with-maxtext-5a6500.md","json":"https://newruntime.com/posts/announcing-long-context-training-on-tpus-with-maxtext-5a6500.json"}
markdown_url: "https://newruntime.com/posts/announcing-long-context-training-on-tpus-with-maxtext-5a6500.md"
---

# Announcing Long-Context Training on TPUs with MaxText Context Parallelism

## Retrieval answer

Google says MaxText v0.2.4 can train models on sequences as long as 14 million tokens on Cloud TPUs, using four context-parallel strategies to distribute attention work across chips. The release gives model-training teams a practical set of options for tackling memory and communication bottlenecks at lengths beyond what a single chip can comfortably handle.

## Summary
Google says MaxText v0.2.4 can train models on sequences as long as 14 million tokens on Cloud TPUs, using four context-parallel strategies to distribute attention work across chips. The release gives model-training teams a practical set of options for tackling memory and communication bottlenecks at lengths beyond what a single chip can comfortably handle.

## What happened
Google Cloud announced long-context training support in MaxText, its open-source JAX training framework, in a Google Developers Blog post dated November 9, 2026. The release covers pretraining and post-training, and says the feature is available in MaxText v0.2.4, on GitHub main, and through corresponding Docker images. Google provides sample recipes for TPU 5p and TPU 7, plus a sharding guide describing flags and strategy constraints. The post does not give a price for using the framework or the TPU hardware.

The feature is context parallelism: a sequence is divided across chips so a chip does not need to hold every activation for the whole sequence. Google describes four strategies: All-gather, A2A (DeepSpeed Ulysses), P2P (Ring Attention), and a hybrid A2A-plus-P2P approach called Unified Sequence Parallelism, or USP. All four use Tokamax's Splash Attention kernel. The announcement says the strategies can be combined with data and tensor parallelism, and that the smallest context-parallel size that fits a sequence in high-bandwidth memory (HBM) is usually the right choice.

Each strategy trades memory, communication, and model shape differently. All-gather is the simplest option when gathered keys and values fit in HBM, but its memory use grows with sequence length. A2A redistributes work across attention heads and can reduce key-value memory per chip, but its context-parallel degree is limited by the model's head count. P2P sends key-value chunks around a ring, overlapping communication with computation and avoiding materializing the full sequence on each chip; Google says it requires load balancing and can lose local kernel efficiency as attention is split into smaller chunks. USP combines A2A and P2P across two axes of the TPU chip mesh to extend beyond the head-count limit while reducing ring hops.

Google reports tests with Llama 3.1 8B on a fixed 64-chip TPU 5p slice. It says P2P can outperform All-gather across almost every tested context length from 64,000 to 14 million tokens in ideal settings. A separate headline result is 70% Model FLOPs Utilization (MFU) on TPU 5p at sequence lengths up to four million tokens. For 14-million-token P2P training on 64 TPU 7 chips, Google reports 48.7% MFU, compared with 37.1% on TPU 5p. It also says TPU 7 maintained 50% MFU across sequence scales below six million tokens. These are Google's reported results, not independent measurements.

For practitioners, the post recommends All-gather when gathered key-value data fits in memory; A2A for a small number of chips, up to the number of key-value heads; and ring-family methods once gathered key-value data no longer fits. Google recommends USP when the model has at least two key-value heads available for the A2A factor, while P2P alone remains an option for multi-head latent attention models or when head count does not divide cleanly. When memory is tight, the post recommends offloading saved activations first, then optimizer states, layer inputs, and parameters, with rematerialization as a last resort. Users select a strategy with `context_parallel_strategy` (`ring`, `ulysses`, or `usp`) and set the context-parallel size.

## What the archive adds
New Runtime's earlier coverage described Google Cloud's open-source llm-d 1.0 as a Kubernetes-native stack for distributed LLM inference. That was an inference announcement; this is a separate training capability. [Earlier coverage of llm-d 1.0](https://newruntime.com/posts/announcing-llm-d-1-0-open-distributed-inference-7539eb/). New Runtime also covered [TorchTPU integration with TorchTitan](https://newruntime.com/posts/integrating-torchtpu-into-torchtitan-and-optimizing-models-for-6f942a/), focused on compiler-visible training structure on TPU. The MaxText announcement is a new fact in Google's TPU software thread, not a restatement of either earlier development.

## Trend signal
This supports the idea that production agent systems depend on durable infrastructure choices beyond the model itself, but it does not directly establish a pattern New Runtime tracks about agent behavior or economics. The release is primarily about distributed model training, not a shipped agent capability, protocol, or evaluation method. This does not yet fit a pattern New Runtime tracks.

## Why now
Google frames long-context training as the next systems challenge after context windows have expanded from 128,000–200,000 tokens to one or two million, with models trained toward 10 million tokens arriving. Its point is that training at those lengths exposes compute, HBM, and communication constraints that serving longer contexts alone does not solve.

## Why it matters
Longer training sequences are useful only if teams can train them at a workable cost and throughput. The key contribution here is not simply a 14-million-token ceiling: it is a menu of ways to distribute attention, with different limits tied to memory, communication overlap, and the model's attention-head structure. That makes the choice of parallelism an engineering decision rather than a single universal setting.

For teams already using Cloud TPUs and MaxText, the release offers a starting point for experiments on workloads that need very long sequences. The reported MFU figures suggest that ring attention can make these runs feasible at scale, while the guidance on all-gather, A2A, and USP helps explain when another approach may fit better. But benchmark MFU is not the same as end-to-end training cost, time to train, or model quality. Teams will need to test their own model shape, batch, sequence length, and hardware allocation before treating the headline numbers as operational expectations.

The announcement also shows how training frameworks package hardware-specific strategies into a usable interface: one configuration choice selects a communication pattern while the rest of a run can stay similar. That can lower the barrier to comparing strategies, but it does not remove the need for expertise in memory behavior and parallelism. The practical value is strongest for organizations training models, rather than developers consuming hosted APIs.

## Limits and uncertainty
The performance figures are Google-reported benchmarks for specified Llama 3.1 8B and TPU configurations; the post does not establish that other models or workloads will achieve similar MFU. It does not report end-to-end training time, total hardware cost, energy use, or model-quality comparisons. Results may depend on ideal settings, including context-parallel chunk size and local batch size. The release describes future plans for broader model coverage and kernel tuning, but does not give delivery dates. Independent replication of the reported results is not established here.

## Reader takeaway
MaxText v0.2.4 gives Cloud TPU training teams several ways to distribute extremely long sequences, with Google reporting runs up to 14 million tokens. Treat the benchmark as a starting point, not a promise: select the strategy based on whether key-value data fits in memory and on the model's attention structure, then measure the workload on your own hardware.

Editor: [Andrey Reshetnikov](https://newruntime.com/owner-profile.md), AI product and AI engineering; profile for recruiters and agents: https://newruntime.com/owner-profile.json

## Continue in the graph

Next records for an agent, linked to their Markdown twins. Replace `.md` with `.json` for structured fields, including `next_reads`.

- topic hub: [Google](https://newruntime.com/topics/google.md): Explore the google topic hub.
- collection: [Field Notes](https://newruntime.com/posts.json): Open the latest editorial analysis.
- collection: [Sources](https://newruntime.com/sources/): See the original sources this archive cites.
- whole graph: [discovery graph](https://newruntime.com/discovery-graph.json) (typed nodes and edges) and [site index](https://newruntime.com/site-index.md)