TorchTitan update: Broader model support, improved large-scale training performance

1. Less need to maintain separate model code

TorchTitan is a platform for generative AI pretraining and post-training, which adapts pretrained models for specific purposes. The update supports recent models including Qwen3.5 and DeepSeek-V4. Its configuration system avoids maintaining separate model code when changing model components, parallelization methods, or kernels that perform GPU computations for specific hardware.

Parallelization distributes training work across multiple GPUs. Teams can change how work is divided or swap computational components while using the same model. This reduces the burden of rebuilding the entire training infrastructure whenever a new model is added.

2. TitanRL unifies model implementations for training and inference

TitanRL uses the same models for pretraining, post-training, and inference inside vLLM. This reduces duplicate work updating both training and inference implementations when model architectures change. Response generation for training runs asynchronously, without waiting for completion, and training and generation can each scale through multiple parallelization methods.

The team reported stable training over hundreds of steps, each a training-data processing iteration. Tests covered models that activate only some expert modules per input and dense models that use all parameters for every input. Tasks included context, the input range a model can reference at once, and code execution. One published example was a software problem-solving experiment that trained and evaluated the Qwen3.5-35B-A3B base model on 40 GB300 GPUs with a 64K context.

The team reported bitwise agreement between trainer and generator log probabilities in supported on-policy configurations using batch-invariant mode. Log probabilities are logarithms of probabilities; on-policy training uses data generated by the current policy. The team used this to study numerical discrepancies in reinforcement learning for mathematics and terminal agents.

3. Large-scale throughput comes with execution conditions

The company reported MLPerf training benchmark results from its internal GB300 cluster using DeepSeek V3 671B with an MTP1 module that assists next-token prediction. It applied quantization, representing numbers at lower precision, using the 8-bit floating-point format MXFP8. MLPerf is a public benchmark for machine learning performance. The 256-GPU configuration achieved 5,324.21 tokens per second per GPU. It used 2 pipeline stages, which split computation into stages, and data parallelism of 128, which distributes data across GPUs. Expert parallelism, which divides expert-module computation, was 64. Local batch size, the data processed at once on one device, and the number of microbatches, smaller portions of a batch, were each 120.

For Muse Spark 1.2 OSS, the company reported 1,500.61 tokens per second per GPU with learned routing at 96K context on 256 GPUs. Routing that forces an even load across experts achieved 1,839.52 tokens per second per GPU. At 192K context on 512 GPUs, the figures were 1,335.86 and 1,588.22 tokens per second, respectively. Results with different contexts, GPU counts, and parallel configurations cannot be ranked as performance under identical conditions.

4. Sharding boundaries are key to reducing communication

FlexShard provides finer control over the boundaries used to divide model states across GPUs. The company tested DeepSeek V3 on 256 H100 GPUs with Owned and BlockShard layouts for the distributed Muon optimizer, which updates training states across GPUs. This removed All Gather communication before updates. Compared with an equivalent implementation using only FSDP, the existing fully sharded data parallel training method, optimizer time fell by up to 79%. The optimizer is the algorithm that updates model parameters.

Benefits and costs depend on how data is divided. Previously, misaligned sharding boundaries and quantization blocks required sending data in a larger format or exchanging extra correction values for compressed transfers. The new method aligns boundaries with quantization blocks, removing that extra communication. The company reported up to a 20% reduction in total training step time for Llama3 405B on 256 H100 GPUs.

5. Shorter waits for checkpoint saving

The new checkpoint library shares metadata, information describing the data, before the first save and skips collective communication among processes on subsequent saves. This reduces the coordination overhead each time training states are saved. The company reported shorter trainer waits and save and load times than the previous method on 512 H100 GPUs, as shown below.

Table: Time taken by the previous and new checkpoint methods

Conditions: 512 H100 GPUs, DSv3 671B, checkpoint interval 50 (company-reported)

Units: Trainer wait time in ms; save and load times in s

P99 is the value at the 99% point when measurements are sorted in ascending order. Less training downtime for saving can allow more frequent checkpoints. To get started, clone the TorchTitan repository and follow its documentation.

The update's value lies in reducing duplicate model implementations and distributed training waits, refining the development path from new-model experiments to large-scale reinforcement learning.