distributed-llm-pretraining-torchtitan

TorchTitan - PyTorch 原生分布式 LLM 预训练

快速开始

TorchTitan 是 PyTorch 官方的大规模 LLM 预训练平台，支持可组合的 4D 并行（FSDP2、TP、PP、CP），在 H100 GPU 上的速度比基线提升 65% 以上。

安装：

# From PyPI (stable)
pip install torchtitan

# From source (latest features, requires PyTorch nightly)
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

下载分词器：

# Get HF token from https://huggingface.co/settings/tokens
python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=...

在 8 块 GPU 上开始训练：

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh

常见工作流

工作流 1：在单节点上预训练 Llama 3.1 8B

复制此清单：

Single Node Pretraining:
- [ ] Step 1: Download tokenizer
- [ ] Step 2: Configure training
- [ ] Step 3: Launch training
- [ ] Step 4: Monitor and checkpoint

步骤 1：下载分词器

python scripts/download_hf_assets.py \
  --repo_id meta-llama/Llama-3.1-8B \
  --assets tokenizer \
  --hf_token=YOUR_HF_TOKEN

步骤 2：配置训练

编辑或创建 TOML 配置文件：

# llama3_8b_custom.toml
[job]
dump_folder = "./outputs"
description = "Llama 3.1 8B training"

[model]
name = "llama3"
flavor = "8B"
hf_assets_path = "./assets/hf/Llama-3.1-8B"

[optimizer]
name = "AdamW"
lr = 3e-4

[lr_scheduler]
warmup_steps = 200

[training]
local_batch_size = 2
seq_len = 8192
max_norm = 1.0
steps = 1000
dataset = "c4"

[parallelism]
data_parallel_shard_degree = -1  # Use all GPUs for FSDP

[activation_checkpoint]
mode = "selective"
selective_ac_option = "op"

[checkpoint]
enable = true
folder = "checkpoint"
interval = 500

步骤 3：启动训练

# 8 GPUs on single node
CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh

# Or explicitly with torchrun
torchrun --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_8b_custom.toml

步骤 4：监控并保存检查点

TensorBoard 日志保存在 ./outputs/tb/：

tensorboard --logdir ./outputs/tb

工作流 2：使用 SLURM 进行多节点训练

Multi-Node Training:
- [ ] Step 1: Configure parallelism for scale
- [ ] Step 2: Set up SLURM script
- [ ] Step 3: Submit job
- [ ] Step 4: Resume from checkpoint

步骤 1：配置并行策略以扩展规模

在 256 块 GPU（32 个节点）上运行 70B 模型：

[parallelism]
data_parallel_shard_degree = 32  # FSDP across 32 ranks
tensor_parallel_degree = 8        # TP within node
pipeline_parallel_degree = 1      # No PP for 70B
context_parallel_degree = 1       # Increase for long sequences

步骤 2：设置 SLURM 脚本

#!/bin/bash
#SBATCH --job-name=llama70b
#SBATCH --nodes=32
#SBATCH --ntasks-per-node=8
#SBATCH --gpus-per-node=8

srun torchrun \
  --nnodes=32 \
  --nproc_per_node=8 \
  --rdzv_backend=c10d \
  --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
  -m torchtitan.train \
  --job.config_file ./llama3_70b.toml

步骤 3：提交作业

sbatch multinode_trainer.slurm

步骤 4：从检查点恢复

如果配置的文件夹中存在检查点，训练会自动恢复。

工作流 3：为 H100 启用 Float8 训练

Float8 可在 H100 GPU 上带来 30%–50% 的速度提升。

Float8 Training:
- [ ] Step 1: Install torchao
- [ ] Step 2: Configure Float8
- [ ] Step 3: Launch with compile

步骤 1：安装 torchao

USE_CPP=0 pip install git+https://github.com/pytorch/ao.git

步骤 2：配置 Float8

添加到 TOML 配置中：

[model]
converters = ["quantize.linear.float8"]

[quantize.linear.float8]
enable_fsdp_float8_all_gather = true
precompute_float8_dynamic_scale_for_fsdp = true
filter_fqns = ["output"]  # Exclude output layer

[compile]
enable = true
components = ["model", "loss"]

步骤 3：启用编译并启动

CONFIG_FILE="./llama3_8b.toml" ./run_train.sh \
  --model.converters="quantize.linear.float8" \
  --quantize.linear.float8.enable_fsdp_float8_all_gather \
  --compile.enable

工作流 4：为 405B 模型启用 4D 并行

4D Parallelism (FSDP + TP + PP + CP):
- [ ] Step 1: Create seed checkpoint
- [ ] Step 2: Configure 4D parallelism
- [ ] Step 3: Launch on 512 GPUs

步骤 1：创建种子检查点

确保 PP 阶段之间的初始化一致，必须执行此步骤：

NGPU=1 CONFIG_FILE=./llama3_405b.toml ./run_train.sh \
  --checkpoint.enable \
  --checkpoint.create_seed_checkpoint \
  --parallelism.data_parallel_shard_degree 1 \
  --parallelism.tensor_parallel_degree 1 \
  --parallelism.pipeline_parallel_degree 1

步骤 2：配置 4D 并行

[parallelism]
data_parallel_shard_degree = 8   # FSDP
tensor_parallel_degree = 8       # TP within node
pipeline_parallel_degree = 8     # PP across nodes
context_parallel_degree = 1      # CP for long sequences

[training]
local_batch_size = 32
seq_len = 8192

步骤 3：在 512 块 GPU 上启动

# 64 nodes x 8 GPUs = 512 GPUs
srun torchrun --nnodes=64 --nproc_per_node=8 \
  -m torchtitan.train \
  --job.config_file ./llama3_405b.toml

何时使用及替代方案

适用场景：

• 从头开始预训练 LLM（8B 至 405B 以上）
• 需要不依赖第三方依赖项的 PyTorch 原生解决方案
• 需要可组合的 4D 并行（FSDP2、TP、PP、CP）
• 使用支持 Float8 的 H100 进行训练
• 需要与 torchtune/HuggingFace 互通的检查点

替代方案：

• Megatron-LM：适用于仅使用 NVIDIA 的部署，性能最高
• DeepSpeed：提供更全面的 ZeRO 优化生态系统，并支持推理
• Axolotl/TRL：适用于微调，而非预训练
• LitGPT：适用于教学和小规模训练

常见问题

问题：大型模型导致内存不足

启用激活检查点并减小批次大小：

[activation_checkpoint]
mode = "full"  # Instead of "selective"

[training]
local_batch_size = 1

也可以使用梯度累积：

[training]
local_batch_size = 1
global_batch_size = 32  # Accumulates gradients

问题：TP 与异步集合通信一起使用时导致内存占用过高

设置环境变量：

export TORCH_NCCL_AVOID_RECORD_STREAMS=1

问题：Float8 训练没有提速

Float8 只对大型 GEMM 运算有益。过滤掉较小的层：

[quantize.linear.float8]
filter_fqns = ["attention.wk", "attention.wv", "output", "auto_filter_small_kn"]

问题：更改并行策略后，检查点加载失败

使用 DCP 的重分片功能：

# Convert sharded checkpoint to single file
python -m torch.distributed.checkpoint.format_utils \
  dcp_to_torch checkpoint/step-1000 checkpoint.pt

问题：流水线并行初始化

请先创建种子检查点（参见工作流 4，步骤 1）。

支持的模型

模型	规模	状态
Llama 3.1	8B、70B、405B	生产就绪
Llama 4	多种	实验性
DeepSeek V3	16B、236B、671B（MoE）	实验性
GPT-OSS	20B、120B（MoE）	实验性
Qwen 3	多种	实验性
Flux	扩散模型	实验性

性能基准（H100）

模型	GPU 数量	并行策略	TPS/GPU	技术
Llama 8B	8	FSDP	5,762	基线
Llama 8B	8	FSDP+compile+FP8	8,532	+48%
Llama 70B	256	FSDP+TP+AsyncTP	876	2D 并行
Llama 405B	512	FSDP+TP+PP	128	3D 并行

高级主题

FSDP2 配置：详细了解 FSDP2 与 FSDP1 的比较及 ZeRO 等效方案，请参阅 references/fsdp.md。

Float8 训练：有关 tensorwise 与 rowwise 缩放方案，请参阅 references/float8.md。

检查点：有关 HuggingFace 转换和异步检查点，请参阅 references/checkpoint.md。

添加自定义模型：有关 TrainSpec 协议，请参阅 references/custom-models.md。

资源

• GitHub：https://github.com/pytorch/torchtitan
• 论文：https://arxiv.org/abs/2410.06511
• ICLR 2025：https://iclr.cc/virtual/2025/poster/29620
• PyTorch 论坛：https://discuss.pytorch.org/c/distributed/torchtitan/44

Agent 操作流程

1. 检查环境。 检查 GPU 类型、内存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确定输入。 根据数据或用户确认格式、标识符和参数。任何会影响结果的值都应先询问，不要猜测。
3. 先运行小规模版本。 执行冒烟测试：使用小型模型或子集、少量步数，并确认损失下降且检查点已保存。
4. 按照上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请这样做
CUDA 显存不足	减小批次大小、启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否存在无效值，并启用梯度裁剪。
这些说明中提到的函数、标志或端点在已安装版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），进行适配并告知用户。切勿编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所作的假设。

诚信准则

• 切勿编造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请如实说明。
• 没有日志时，切勿声称训练取得了某种结果；启动大型任务前，应估算计算成本。
• 请将此处与特定版本相关的细节视为可能已过时，并根据已安装版本的官方文档进行确认。
• 对于需要付费、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征得用户同意。

相关技能

• torchforge-rl-training：提供使用 torchforge 进行 PyTorch 原生 Agent 强化学习的指导。torchforge 是 Meta 的库，旨在将基础设施与算法分离。
• pytorch-fsdp2：为训练脚本添加 PyTorch FSDP2（fully_shard），并提供正确的初始化、分片、混合精度/卸载配置和分布式检查点…
• dask：适用于超出内存容量的 pandas/NumPy 工作流的分布式计算。