分佈式 LLM 預訓練 Torchtitan

使用 torchtitan 提供基於 PyTorch 原生的分佈式 LLM 預訓練，支持 4D 並行（FSDP2、TP、PP、CP）。適用於在 8 到 512+ 張 GPU 上大規模預訓練 Llama 3.1、DeepSeek V3 或自定義模型，並支持 Float8、torch.compile 和分佈式檢查點。

技能元數據


來源	可選 — 使用 `hermes skills install official/mlops/torchtitan` 安裝
路徑	`optional-skills/mlops/torchtitan`
版本	`1.0.0`
作者	Orchestra Research
許可證	MIT
依賴項	`torch>=2.6.0`, `torchtitan>=0.2.0`, `torchao>=0.5.0`
標籤	`Model Architecture`, `Distributed Training`, `TorchTitan`, `FSDP2`, `Tensor Parallel`, `Pipeline Parallel`, `Context Parallel`, `Float8`, `Llama`, `Pretraining`

參考：完整 SKILL.md

信息

以下是 Hermes 在觸發此技能時加載的完整技能定義。這是技能激活時代理看到的指令。

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

下載 tokenizer：

# 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：下載 tokenizer

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：使用 compile 啟動

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

何時使用及替代方案對比

在以下情況使用 TorchTitan：

從頭預訓練 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 的重分片（resharding）功能：

# 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 配置：請參閱 references/fsdp.md 以獲取詳細的 FSDP2 與 FSDP1 對比以及 ZeRO 等效項。

Float8 訓練：請參閱 references/float8.md 以瞭解按張量縮放與按行縮放的方案。

檢查點保存：請參閱 references/checkpoint.md 以瞭解 HuggingFace 轉換和異步檢查點保存。

添加自定義模型：請參閱 references/custom-models.md 以瞭解 TrainSpec 協議。

資源

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

技能元數據​

參考：完整 SKILL.md​

TorchTitan - PyTorch 原生分佈式 LLM 預訓練

快速開始​

常見工作流​

工作流 1：在單節點上預訓練 Llama 3.1 8B​

工作流 2：使用 SLURM 進行多節點訓練​

工作流 3：在 H100 上啟用 Float8 訓練​

工作流 4：針對 405B 模型的 4D 並行​

何時使用及替代方案對比​

常見問題​

支持的模型​

性能基準測試 (H100)​

高級主題​

資源​