distributed-llm-pretraining-torchtitan

Pretrains large language models at scale using PyTorch torchtitan with 4D parallelism.

13.0k|930|Updated Nov 3, 2025

Search Tags:#pytorch #distributed training #tensor parallelism #torchtitan #float8 #fsdp2 #llm pretraining

Ready to run

npx skills add https://github.com/Orchestra-Research/AI-research-SKILLs --skill distributed-llm-pretraining-torchtitan-orchestra-research

💬 Prompt for ChatGPT / Claude Chat▼

Please help me install this Agent Skill.
Skill: distributed-llm-pretraining-torchtitan
Source: https://github.com/Orchestra-Research/AI-research-SKILLs/tree/main/01-model-architecture/torchtitan
Command: npx skills add https://github.com/Orchestra-Research/AI-research-SKILLs --skill distributed-llm-pretraining-torchtitan-orchestra-research

SYSTEM DOCUMENTATION & REQUIREMENTS

💡 This Skill requires torch, torchtitan, torchao, and includes references (resource) components.

What problem does it solve? Pretraining LLMs from scratch across hundreds of GPUs requires coordinating data, tensor, pipeline, and context parallelism plus checkpointing and mixed precision, which is complex to configure correctly. This Skill provides tested workflows for distributed pretraining with torchtitan, PyTorch's native training platform. ## Core Features & Use Cases - 4D Parallelism Configuration: Compose FSDP2, tensor parallelism, pipeline parallelism, and context parallelism for models from 8B to 405B+ parameters on 8 to 512+ GPUs. - Float8 and torch.compile Acceleration: Enable Float8 training with tensorwise or rowwise scaling on H100 GPUs for up to 48% throughput gains over FSDP baselines. - Distributed Checkpointing: Use PyTorch DCP for async checkpointing, seed checkpoints for pipeline parallelism, and bidirectional HuggingFace checkpoint conversion. - Use Case: A research team wants to pretrain a Llama 3.1 70B model on 256 H100 GPUs. Use this Skill to configure FSDP plus tensor parallelism, set up a SLURM multi-node job, enable activation checkpointing, and resume automatically from distributed checkpoints. ## Quick Start Ask the AI to configure and launch a torchtitan pretraining run for Llama 3.1 8B on 8 GPUs with Float8 and torch.compile enabled.

Frequently Asked Questions about distributed-llm-pretraining-torchtitan

High-intent search queries and answers about installing and using this skill.

FAQPage Schema

How do I pretrain Llama 3.1 8B on multiple GPUs with torchtitan?▼

Download the tokenizer with the download_hf_assets script, create a TOML config specifying the llama3 model and 8B flavor, then launch with torchrun using 8 processes per node. Training logs to TensorBoard and checkpoints save at configured intervals.

TorchTitan vs Megatron-LM for LLM pretraining, which should I use?▼

TorchTitan is PyTorch-native with composable 4D parallelism and no third-party dependencies, making it suitable for research and HuggingFace-interoperable checkpoints. Megatron-LM targets maximum performance for NVIDIA-only production deployments.

Does torchtitan support Float8 training on H100 GPUs?▼

Yes, torchtitan supports Float8 training on H100 or newer GPUs via torchao, with tensorwise and rowwise dynamic scaling recipes. Combined with torch.compile, benchmarks show up to 48% throughput improvement over FSDP baselines on Llama 8B.

How do I convert torchtitan checkpoints to HuggingFace format?▼

Use the convert_to_hf.py script with the input directory, output directory, model name, and flavor, or enable last_save_in_hf in the checkpoint config to save directly in HuggingFace format during training. Reverse conversion uses convert_from_hf.py.

Why does torchtitan run out of memory on large models?▼

Out of memory occurs when activation memory exceeds GPU capacity. Enable full activation checkpointing instead of selective, reduce local batch size, or use gradient accumulation with a larger global batch size to lower peak memory.

When should I not use torchtitan for training?▼

Avoid torchtitan for fine-tuning tasks, where Axolotl, TRL, or torchtune are better suited, and for inference workloads where DeepSpeed or serving frameworks apply. It is designed specifically for large-scale pretraining from scratch.