Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang, Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, Stratos Idreos

TL;DR

TorchTitan introduces a PyTorch-native, open-source distributed framework that unifies 4D parallelism (DP/FSDP, TP, PP, CP) with hardware-aware optimizations (Float8, AsyncTP) and production tools (DTensor-based DCP, Flight Recorder) to accelerate large-language-model pretraining. It provides a modular, composable architecture and a meta-device initialization flow, enabling scalable training across thousands of GPUs with elastic growth and long-context capabilities. Empirical results on Llama 3.1 models (8B–405B) show up to 65% throughput gains at smaller scales and substantial improvements at larger scales, plus demonstrated long-context training enabled by CP within a 4D framework. By delivering a unified, production-ready test bed and clear training recipes, TorchTitan lowers engineering overhead and accelerates exploration of distributed training strategies for cutting-edge LLMs.

Abstract

The development of large language models (LLMs) has been instrumental in advancing state-of-the-art natural language processing applications. Training LLMs with billions of parameters and trillions of tokens require sophisticated distributed systems that enable composing and comparing several state-of-the-art techniques in order to efficiently scale across thousands of accelerators. However, existing solutions are complex, scattered across multiple libraries/repositories, lack interoperability, and are cumbersome to maintain. Thus, curating and empirically comparing training recipes require non-trivial engineering effort. This paper introduces TorchTitan, an open-source, PyTorch-native distributed training system that unifies state-of-the-art techniques, streamlining integration and reducing overhead. TorchTitan enables 3D parallelism in a modular manner with elastic scaling, providing comprehensive logging, checkpointing, and debugging tools for production-ready training. It also incorporates hardware-software co-designed solutions, leveraging features like Float8 training and SymmetricMemory. As a flexible test bed, TorchTitan facilitates custom recipe curation and comparison, allowing us to develop optimized training recipes for Llama 3.1 and provide guidance on selecting techniques for maximum efficiency based on our experiences. We thoroughly assess TorchTitan on the Llama 3.1 family of LLMs, spanning 8 billion to 405 billion parameters, and showcase its exceptional performance, modular composability, and elastic scalability. By stacking training optimizations, we demonstrate accelerations of 65.08% with 1D parallelism at the 128-GPU scale (Llama 3.1 8B), an additional 12.59% with 2D parallelism at the 256-GPU scale (Llama 3.1 70B), and an additional 30% with 3D parallelism at the 512-GPU scale (Llama 3.1 405B) on NVIDIA H100 GPUs over optimized baselines.

TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training

TL;DR

TorchTitan introduces a PyTorch-native, open-source distributed framework that unifies 4D parallelism (DP/FSDP, TP, PP, CP) with hardware-aware optimizations (Float8, AsyncTP) and production tools (DTensor-based DCP, Flight Recorder) to accelerate large-language-model pretraining. It provides a modular, composable architecture and a meta-device initialization flow, enabling scalable training across thousands of GPUs with elastic growth and long-context capabilities. Empirical results on Llama 3.1 models (8B–405B) show up to 65% throughput gains at smaller scales and substantial improvements at larger scales, plus demonstrated long-context training enabled by CP within a 4D framework. By delivering a unified, production-ready test bed and clear training recipes, TorchTitan lowers engineering overhead and accelerates exploration of distributed training strategies for cutting-edge LLMs.

Abstract

The development of large language models (LLMs) has been instrumental in advancing state-of-the-art natural language processing applications. Training LLMs with billions of parameters and trillions of tokens require sophisticated distributed systems that enable composing and comparing several state-of-the-art techniques in order to efficiently scale across thousands of accelerators. However, existing solutions are complex, scattered across multiple libraries/repositories, lack interoperability, and are cumbersome to maintain. Thus, curating and empirically comparing training recipes require non-trivial engineering effort. This paper introduces TorchTitan, an open-source, PyTorch-native distributed training system that unifies state-of-the-art techniques, streamlining integration and reducing overhead. TorchTitan enables 3D parallelism in a modular manner with elastic scaling, providing comprehensive logging, checkpointing, and debugging tools for production-ready training. It also incorporates hardware-software co-designed solutions, leveraging features like Float8 training and SymmetricMemory. As a flexible test bed, TorchTitan facilitates custom recipe curation and comparison, allowing us to develop optimized training recipes for Llama 3.1 and provide guidance on selecting techniques for maximum efficiency based on our experiences. We thoroughly assess TorchTitan on the Llama 3.1 family of LLMs, spanning 8 billion to 405 billion parameters, and showcase its exceptional performance, modular composability, and elastic scalability. By stacking training optimizations, we demonstrate accelerations of 65.08% with 1D parallelism at the 128-GPU scale (Llama 3.1 8B), an additional 12.59% with 2D parallelism at the 256-GPU scale (Llama 3.1 70B), and an additional 30% with 3D parallelism at the 512-GPU scale (Llama 3.1 405B) on NVIDIA H100 GPUs over optimized baselines.
Paper Structure (45 sections, 5 figures, 9 tables)

This paper contains 45 sections, 5 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Elasticity through composability
  3. Composable N-D parallelism training
  4. Large-scale model initialization using meta device
  5. Fully Sharded Data Parallel
  6. Tensor Parallel
  7. Tensor and Sequence Parallel (TP/SP)
  8. Loss Parallel
  9. Pipeline Parallel
  10. Context Parallelism
  11. Optimizing training efficiencies
  12. Navigating compute-memory trade-offs using activation checkpointing
  13. Regional compilation to exploit torch.compile optimizations
  14. Asynchronous Tensor Parallel to maximally overlap communication
  15. Boosting throughput with mixed precision training and Float8 support
  16. ...and 30 more sections

Figures (5)

  • Figure 1: Composable and Modular TorchTitan initialization workflow.
  • Figure 2: Scaling with 4D Parallelism
  • Figure 3: Tensor Parallel in detail (2 GPUs, data moves from left to right).
  • Figure 4: FSDP2 + Tensor Parallel (TP degree 4) sharding layout, with 2 nodes of 4 GPUs.
  • Figure 5: Loss converging tests on Llama 3.1 8B. C4 dataset. Local batch size 4, global batch size 32. 3000 steps, 600 warmup steps.