TorchTitan bekommt mit TitanRL Reinforcement Learning und unterstützt neue Frontier-Modelle

Neue Modelle und ein modulares Konfigurationssystem

TorchTitan ist eine PyTorch-native Plattform für das Pre- und Post-Training generativer KI-Modelle im großen Maßstab. Das Team gibt an, PyTorch-native Unterstützung für Muse Spark Open, Muse Glimmer, Qwen3-VL, Qwen3.5, Kimi K2.7, Kimi K3 und DeepSeek-V4 ergänzt zu haben. Laut PyTorch handelt es sich dabei um vollständige Integrationen in den Trainings-Stack, die über den gesamten 5D-Parallelism-Stack von TorchTitan funktionieren.

Außerdem wurde ein modulares Konfigurationssystem eingeführt, das auf typisierten Python-Rezepten, Transformationen der Modellkonfiguration und gezielten Überschreibungen einzelner Komponenten beruht. Nutzer können Parallelisierungsstrategien, Modellkomponenten oder hardwarespezifische Kernel ändern, ohne separate Modell-Forks pflegen zu müssen. Über dieselben Schnittstellen lassen sich Kernel aus Attention Gym einbinden, darunter Gated DeltaNet, Kimi Delta Attention und Sparse-Attention-Kernel, ebenso optionale Helion-Kernel wie fused RoPE. All das bleibt mit torch.compile kompatibel.

TitanRL: ein Modell für Training und Inferenz

TitanRL ist der neue Reinforcement-Learning-Stack der Plattform. Das PyTorch-Team weist darauf hin, dass viele RL-Stacks getrennte Implementierungen eines Modells für Training und Inferenz vorhalten. Das bedeutet, dass Architekturänderungen synchron gehalten werden müssen, und es erschwert die Klärung, ob ein Problem an der Belohnung, am Algorithmus oder an Abweichungen zwischen den beiden Kopien liegt. TitanRL behält dasselbe Modell vom Pretraining und Post-Training bis hinein in vLLM bei, sodass die Arbeit in einer einzigen Codebasis bleibt.

Bei agentischen Workloads laufen Rollouts asynchron, und Training und Generierung skalieren jeweils mit mehreren Formen der Parallelisierung. Das Team gibt an, TitanRL an mehreren Workloads getestet zu haben und über Hunderte von Schritten stabiles Training auf MoE- und Dense-Modellen beobachtet zu haben, und zwar bei Aufgaben mit langem Kontext, Codeausführung, Compaction, Harness und Self-Play. Die gezeigten Läufe nutzten das Basismodell Qwen3.5-35B-A3B: einen auf SWE-Rebench easy mit 64K Kontext auf 40 GB300-GPUs und einen, der auf Tmax 15k trainiert und auf terminal bench 2.1 mit 128K Kontext auf 20 GB300-GPUs ausgewertet wurde.

Ein batch-invarianter Modus soll beim Debugging helfen. Laut dem Beitrag liefern Trainer und Generator in unterstützten On-Policy-Konfigurationen bitweise identische Log-Wahrscheinlichkeiten, sodass die Lücke zwischen beiden in jedem Schritt exakt null beträgt.

Performance-Arbeit für MoE-Modelle

Das Team beschreibt Performance-Arbeiten in drei Bereichen: schnellere Kernel, stärkere Überlappung von Kommunikation und Berechnung sowie feinere Kontrolle über den Speicher. Zu den aufgeführten Funktionen gehören:

• local_compile: wendet torch.compile auf ausgewählte Bereiche an, statt das Tracing des gesamten Graphen oder ganzer Blöcke zu verlangen.
• torch_remat: eine API für Activation Checkpointing mit feingranularer Kontrolle über neu berechnete Bereiche.
• Dist-MoE: fusionierte Kommunikations-Compute-Kernel für NVIDIA-GPUs ab SM100.
• Überarbeitetes Pipeline-Parallel-Paket: CUDA-Graph-fähige, deterministische Ausführung mit Steuerungsmöglichkeiten, um die Überlappung bei der Kombination von FSDP und Pipeline-Parallelismus zu maximieren.
• Overrides: Integration modellspezifischer Kernel wie fused MLA, fused SwiGLU und fused Complex RoPE.

Für DeepSeek V3 671B mit einem MTP1-Modul, ausgewertet mit dem MLPerf-Training-Benchmark auf einem internen GB300-Cluster mit MXFP8-Quantisierung, erreichte der PP2-256-Lauf auf 256 GPUs 5.324,21 Token pro Sekunde und GPU sowie 1.553,89 TFLOPs pro GPU. Der PP2-512-Lauf auf 512 GPUs kam auf global 2.694.493 Token pro Sekunde. Bei Muse Spark Open erreichte ein Lauf auf 256 GPUs mit gelerntem Routing und 96K Kontext 1.500,61 Token pro Sekunde und GPU bei 25,84 % MFU, die Variante mit erzwungener Balancierung kam auf 31,68 %. Auf 512 GPUs mit 192K Kontext lag die MFU bei 31,27 % mit gelerntem Routing und bei 37,17 % mit erzwungen ausbalanciertem Routing.

Integrationen im Ökosystem und FlexShard

Laut dem Beitrag wurde TorchTitan als Trainings-Backend in veRL integriert und zu Miles hinzugefügt. Es ist das Standard-Trainings-Backend für Meshy und diente als Vorbild für die Kernimplementierung von Archon, der Trainings-Engine von Areal.

Außerdem ist TorchTitan nun mit FlexShard verzahnt, einer neuen Data-Parallelism-Bibliothek zum Sharding von Modell- und Optimizer-Zustand. Anders als FSDP2, das zeilen- oder spaltenweise aufteilt, bietet FlexShard eine Bucket-zentrierte API mit anpassbaren Platzierungsrichtlinien, darunter Owned, bei der ein Rank einen ganzen Parameter besitzt, und BlockShard, das ein an 2D-Kacheln ausgerichtetes Sharding erlaubt.

Dadurch kann TorchTitan den verteilten Muon-Optimizer ohne All-Gather-Kollektive vor dem Update ausführen. Auf DeepSeekV3 mit 256 H100-GPUs verbesserte sich die Optimizer-Zeit nach Angaben des Teams um bis zu 79 % gegenüber einer äquivalenten Implementierung, die nur auf FSDP setzt. FlexShard unterstützt zudem blockweise quantisierte Parameter-AllGathers, bei denen FP8-quantisierte Parameter gesammelt werden, was das Kommunikationsvolumen um etwa 50 % senkt.

Was das für Entwickler bedeutet

Das PyTorch-Team erklärt, TorchTitan solle einfach zu benutzen sein und es Nutzern ermöglichen, klein anzufangen und dann hochzuskalieren. Für Teams, die mit neu veröffentlichten Architekturen arbeiten, bietet das Update einen einzigen nativen PyTorch-Weg vom Pretraining bis zum RL, bei dem sich Kernel und Parallelisierungsoptionen per Konfiguration austauschen lassen. Die Leistungszahlen stammen vom internen Cluster und aus den Benchmarks des Teams, daher können die Ergebnisse auf anderen Setups abweichen.