Research map / Large language models / Reasoning and chain of thought
TorchTitan: One-stop PyTorch native solution for production ready LLM pretraining
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents TorchTitan, a PyTorch-native, composable 4D parallelism framework for production LLM pretraining, evaluated on Llama 3.1 8B/70B/405B from 8 to 512 H100 GPUs.
Problem
Existing distributed LLM training systems are non-composable, inflexible, inefficient in hardware use, weak on production needs (checkpointing, failure recovery, debugging), and depend on external poorly maintained dependencies.
Approach
Builds a modular PyTorch-native training system on DTensor and DeviceMesh that composes FSDP2, TP/SP, PP, and CP as n-D parallelism, combined with torch.compile, Float8, AsyncTP, selective activation checkpointing, distributed checkpointing, and Flight Recorder.
Key results
- Float8 + torch.compile on Llama 3.1 8B at 128 H100s raises FSDP throughput from 5,645 to 9,319 tok/s/GPU (+65.08%).
- AsyncTP on Llama 3.1 70B at 256 GPUs (2D) improves throughput from 897 to 1,010 tok/s/GPU (+12.59%).
- Interleaved 1F1B on Llama 3.1 405B at 512 GPUs (3D) improves throughput from 100 to 130 tok/s/GPU (+30%).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add torch.compile (ablation) | FSDP · Llama 3.1 8B, 8 GPUs | throughput (Tok/Sec): 6,258 -> 6,674 (+ 6.64%) | Helped |
| add torch.compile (ablation) | FSDP · Llama 3.1 8B, 128 GPUs | throughput (Tok/Sec): 5,645 -> 6,482 (+ 14.82%) | Helped |
| add Float8 on top of torch.compile (ablation) | FSDP + torch.compile · Llama 3.1 8B, 8 GPUs | throughput (Tok/Sec): 6,674 -> 9,409 (+ 50.35%) | Helped |
| add Float8 on top of torch.compile (ablation) | FSDP + torch.compile · Llama 3.1 8B, 128 GPUs | throughput (Tok/Sec): 6,482 -> 9,319 (+ 65.08%) | Helped |
| add AsyncTP (ablation) | 2D parallelism with torch.compile and Float8 · Llama 3.1 70B, 256 GPUs | throughput (Tok/Sec): 897 -> 1,010 (+ 12.59%) | Helped |
| use Interleaved 1F1B instead of 1F1B (ablation) | 3D parallelism with torch.compile, Float8, and AsyncTP · Llama 3.1 405B, 512 GPUs | throughput (Tok/Sec): 100 -> 130 (+ 30.00%) | Helped |
| increase CP degree to enable longer context (ablation) | FSDP + CP · Llama 3.1 8B, 8 GPUs | maximum sequence length: 32,768 -> 262,144 | Helped |
| increase CP degree from 1 to 8 (ablation) | FSDP + CP · Llama 3.1 8B, 8 GPUs | throughput (Tok/Sec): 3,890 -> 548 | Hurt |
| increase CP degree to enable longer context (ablation) | FSDP + CP · Llama 3.1 405B, 512 GPUs | maximum sequence length: 32,768 -> 262,144 | Helped |
| increase CP degree from 1 to 8 (ablation) | 4D parallelism · Llama 3.1 405B, 512 GPUs | throughput (Tok/Sec): 76 -> 16 | Hurt |
| replace FSDP1 with FSDP2 (ablation) | FSDP training · Llama 2 7B | GPU memory: 7% | Helped |
| replace FSDP1 with FSDP2 (ablation) | FSDP training · Llama 2 7B | throughput: 1.5% | Helped |
Surprising findings
- Reports that FSDP2 uses about 7% less GPU memory and gives about 1.5% higher throughput than FSDP1 on Llama 2 7B.
Limitations
- MFU is not reported because Float8 makes the MFU definition ill-defined.
- Non-standard H100 GPUs with 95 GiB memory and unknown peak TFLOPs limit absolute comparability.
- Comparisons against related systems are mostly feature tables and LOC counts rather than head-to-head runs.
Experimental setup
Datasets: C4
Metrics: throughput (tokens/sec/GPU) peak memory (GiB) loss convergence
Baselines: FSDP1 Megatron-LM DeepSpeed veScale torch.compile-free FSDP baseline 2D parallelism baseline
Builds on: FSDP DTensor DeviceMesh Megatron-LM torch.compile
Closest papers by meaning
- ShiftAddLLM: Accelerating Pretrained LLMs via Post-Training Multiplication-Less Reparameterization (NeurIPS 2024 · similarity 0.88)
- EE-LLM: Large-Scale Training and Inference of Early-Exit Large Language Models with 3D Parallelism (ICML 2024 · similarity 0.87)
- Sheared LLaMA: Accelerating Language Model Pre-training via Structured Pruning (ICLR 2024 · similarity 0.87)
- Heterogeneous Parallelism for Multimodal Large Language Model Training (NeurIPS 2026 · similarity 0.87)
- Cost-Efficient LLM Training with Lifetime-Aware Tensor Offloading via GPUDirect Storage (NeurIPS 2025 · similarity 0.87)
- StreamBP: Memory-Efficient Exact Backpropagation for Long Sequence Training of LLMs (NeurIPS 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.