distributed-llm-pretraining-torchtitan

Pretrain large language models on H100 GPUs by distributing the workload across PyTorch's FSDP2, Tensor Parallelism, Pipeline Parallelism, and Context Parallelism. It supports Float8 optimization and distributed checkpointing for efficient training at scale.

Yes, torchtitan supports large-scale LLM pretraining for research and development. The framework is designed to handle pretraining custom models like Llama, DeepSeek, and GPT-OSS on up to 512 GPUs.

Supported parallelism strategies include PyTorch-native FSDP2, Tensor Parallelism (TP), Pipeline Parallelism (PP), and Context Parallelism (CP). These strategies combine to efficiently distribute large-scale LLM pretraining across multiple GPUs.

Yes, PyTorch FSDP2 is optimized for H100 GPUs with native support for Float8. This optimization reduces memory bandwidth and footprint, enabling efficient large-scale LLM pretraining.

You need PyTorch 2.6.0 or later, torchtitan 0.2.0 or later, and torchao 0.5.0 or later. Once dependencies are met, you can use the provided configuration files and scripts to start pretraining models like Llama 3.1 8B on 8 GPUs.

PyTorch-native parallelism integrates FSDP2, TP, PP, and CP directly without external dependencies. This approach provides native distributed checkpointing and Float8 support tailored for H100 GPUs, streamlining large-scale LLM research.