pytorch/torchtitan

TorchTitan is a PyTorch native platform designed for rapid experimentation and large-scale training of generative AI models. It serves as a minimal clean-room implementation of PyTorch native scaling techniques, providing a flexible foundation for developers to build upon. The project is currently under extensive development, with the README recommending users employ the most recent PyTorch nightly to access the latest features. The platform has achieved significant milestones, including acceptance of its research paper at ICLR 2025, a GPU MODE lecture in December 2024, and a presentation at PyTorch Conference 2024.

The core mission of TorchTitan is to accelerate innovation in generative AI by empowering researchers and developers to explore new modeling architectures and infrastructure techniques. The platform is designed around three guiding principles: ease of understanding, use, and extension for different training purposes; minimal changes to model code when applying multi-dimensional parallelism; and a bias towards a clean, minimal codebase while providing basic reusable and swappable components. TorchTitan has been showcasing PyTorch's latest distributed training features through support for pretraining Llama 3.1 LLMs of various sizes, including 8B, 70B, and 405B parameter models.

The platform offers comprehensive support for multi-dimensional composable parallelisms, including FSDP2 with per-parameter sharding, Tensor Parallel with async variants, Pipeline Parallel with zero-bubble capabilities, and Context Parallel for training long-context LLMs. Additional key features include meta device initialization, selective and full activation checkpointing, distributed checkpointing with async support, torch.compile integration, Float8 and MXFP8 quantization support, supervised fine-tuning with chat-formatted datasets, and flexible learning rate scheduling. The platform provides extensive debugging tools, structured logging, and helper scripts for tokenizer downloads, checkpoint conversion, and distributed inference.

TorchTitan supports installation through multiple methods including direct source code execution, nightly builds, and stable releases via pip or conda. The platform includes an experiments folder to accelerate contributions and innovations, with dedicated guidelines for both experimental contributions and core fixes. Performance has been reported on up to 512 GPUs, with verified loss convergence correctness across various techniques. The source code is made available under a BSD 3 license, though users may have other legal obligations governing their use of linked third-party data and models.