TorchTitan adds RL training with TitanRL and support for new frontier models

New models and a composable configuration system

TorchTitan is a PyTorch-native platform for large-scale pre- and post-training of generative AI models. The team says it added PyTorch-native support for Muse Spark Open, Muse Glimmer, Qwen3-VL, Qwen3.5, Kimi K2.7, Kimi K3 and DeepSeek-V4. According to PyTorch, these are full training-stack integrations that work across TorchTitan's 5D parallelism stack.

The work also introduced a composable configuration system built around typed Python recipes, model-config transforms and targeted component overrides. Users can change parallelization strategies, model components or hardware-specific kernels without maintaining separate model forks. The same interfaces integrate kernels from Attention Gym, including Gated DeltaNet, Kimi Delta Attention and sparse-attention kernels, as well as opt-in Helion kernels such as fused RoPE. All of this remains compatible with torch.compile.

TitanRL: one model for training and inference

TitanRL is the platform's new reinforcement learning stack. The PyTorch team points out that many RL stacks keep separate training and inference implementations of a model, which means architecture changes must be kept aligned and makes it hard to tell whether a problem comes from the reward, the algorithm or the two copies disagreeing. TitanRL keeps the same model from pretraining and post-training to inside vLLM, so the work stays in one codebase.

For agentic workloads, rollouts run asynchronously, and training and generation each scale with multiple forms of parallelism. The team says it tested TitanRL on several workloads and saw stable training on MoE and dense models for hundreds of steps, on tasks involving long context, code execution, compaction, harness and self-play. The shown runs used the Qwen3.5-35B-A3B base model: one on SWE-Rebench easy with 64K context on 40 GB300 GPUs, and one trained on Tmax 15k and evaluated on terminal bench 2.1 with 128K context on 20 GB300 GPUs.

A batch-invariant mode is meant to help with debugging. According to the post, in supported on-policy configurations the trainer and generator produce bitwise identical log probabilities, so the gap between them is exactly zero at every step.

Performance work for MoE models

The team describes performance work around three areas: faster kernels, more overlap of communication and compute, and finer control of memory. The listed features are:

• local_compile: applies torch.compile to selected regions instead of requiring full-graph or full-block tracing.
• torch_remat: an activation checkpointing API with fine-grained control over recomputed regions.
• Dist-MoE: fused communication-compute kernels for NVIDIA SM100+ GPUs.
• Revamped pipeline parallel package: CUDA-graphable, deterministic execution with controls to maximize overlap when composing FSDP and pipeline parallelism.
• Overrides: integration of model-specific kernels such as fused MLA, fused SwiGLU and fused Complex RoPE.

For DeepSeek V3 671B with an MTP1 module, evaluated on the MLPerf training benchmark on an internal GB300 cluster with MXFP8 quantization, the PP2-256 run on 256 GPUs reached 5,324.21 tokens per second per GPU and 1,553.89 TFLOPs per GPU. The PP2-512 run on 512 GPUs reached a global 2,694,493 tokens per second. For Muse Spark Open, a 256-GPU run with learned routing and 96K context reached 1,500.61 tokens per second per GPU and 25.84% MFU, while the forced balanced variant reached 31.68%. On 512 GPUs with 192K context, MFU was 31.27% with learned routing and 37.17% with forced balanced routing.

Ecosystem integrations and FlexShard

According to the post, TorchTitan has been integrated into veRL as a training backend and added to Miles. It is the default training backend for Meshy and inspires the core implementation of Areal's training engine Archon.

TorchTitan also now integrates with FlexShard, a new data parallelism library for sharding model and optimizer state. Unlike FSDP2, which shards row-wise or column-wise, FlexShard offers a bucket-centric API with customizable placement policies, including Owned, where one rank owns an entire parameter, and BlockShard, which allows 2D tile-aligned sharding.

This lets TorchTitan run the distributed Muon optimizer without All Gather collectives before the update. On DeepSeekV3 at 256 H100s, the team reports optimizer time improved by up to 79% compared with an equivalent FSDP-only implementation. FlexShard also supports blockwise-quantized parameter AllGathers, where FP8-quantized parameters are gathered, reducing communication volume by about 50%.

What it means for developers

The PyTorch team says the aim is for TorchTitan to be easy to use while letting users start small and scale up. For teams working on newly released architectures, the update offers a single native PyTorch route from pretraining to RL, with kernels and parallelism options swappable through configuration. The performance figures come from the team's own internal cluster and benchmarks, so results on other setups may differ.