torchtitan
An open-source generative AI training platform that runs on your own NVIDIA or AMD GPU servers and supports distributed training with PyTorch.
torchtitan is a training platform for researchers and developers experimenting with generative AI models on their own GPU servers or clusters. Its replaceable components and extension points support work on model architectures and training infrastructure.
- Distributed training: Combine data, tensor, pipeline and context parallelism to spread training across GPUs, with support for multi-node Slurm clusters.
- Pretraining and fine-tuning: Train with C4 or custom datasets, or use chat-formatted datasets for supervised fine-tuning; supported model families include Llama 3, Qwen3, DeepSeek V3 and V4, GPT-OSS, Kimi, Muse Glimmer and Flux.
- Memory and compute: Use activation checkpointing and BF16 optimizer states to reduce memory use, alongside torch.compile and low-precision training with MXFP8 and NVFP4.
- Checkpoints and inference: Save distributed checkpoints asynchronously, load compatible checkpoints in torchtune for further fine-tuning, convert between Hugging Face and PyTorch distributed formats, and run inference through vLLM.
- Training diagnostics: Inspect CPU and GPU performance and memory use, and track loss, GPU memory and throughput through TensorBoard or Weights & Biases.
The Python codebase is open source under the BSD-3-Clause license. It supports NVIDIA GPUs and AMD GPUs through ROCm. The project is under extensive development and depends on PyTorch nightly builds; stable releases pin specific nightly versions of PyTorch and torchao.
Logotyper
