⚡

Torchtitan分布式LLM预训练Skill distributed-llm-pretraining-torchtitan

Torchtitan是一个用于大规模语言模型（LLM）预训练的PyTorch原生平台，支持4D并行（FSDP2、TP、PP、CP）、Float8训练和分布式检查点，适用于从8到512+ GPU的Llama、DeepSeek等模型的预训练，提供高效的训练速度和可扩展性。关键词：分布式训练、LLM预训练、PyTorch、Torchtitan、4D并行、Float8、深度学习、人工智能。

深度学习 0 次安装 0 次浏览 更新于 3/21/2026

名称：分布式-llm-预训练-torchtitan 描述：提供使用 torchtitan 的 PyTorch 原生分布式 LLM 预训练，支持 4D 并行（FSDP2、TP、PP、CP）。适用于从 8 到 512+ GPU 的大规模 Llama 3.1、DeepSeek V3 或自定义模型的预训练，支持 Float8、torch.compile 和分布式检查点。 版本：1.0.0 作者：Orchestra Research 许可证：MIT 标签：[模型架构, 分布式训练, TorchTitan, FSDP2, 张量并行, 管道并行, 上下文并行, Float8, Llama, 预训练] 依赖项：[torch>=2.6.0, torchtitan>=0.2.0, torchao>=0.5.0]

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

快速入门

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

安装:

# 从 PyPI（稳定版）
pip install torchtitan

# 从源码（最新功能，需要 PyTorch nightly）
git clone https://github.com/pytorch/torchtitan
cd torchtitan
pip install -r requirements.txt

下载分词器:

# 从 https://huggingface.co/settings/tokens 获取 HF token
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

复制此清单：

单节点预训练：
- [ ] 步骤 1：下载分词器
- [ ] 步骤 2：配置训练
- [ ] 步骤 3：启动训练
- [ ] 步骤 4：监控和检查点

步骤 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 训练"

[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  # 使用所有 GPU 进行 FSDP

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

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

步骤 3：启动训练

# 单节点上的 8 GPU
CONFIG_FILE="./llama3_8b_custom.toml" ./run_train.sh

# 或使用 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 进行多节点训练

多节点训练：
- [ ] 步骤 1：为规模配置并行
- [ ] 步骤 2：设置 SLURM 脚本
- [ ] 步骤 3：提交作业
- [ ] 步骤 4：从检查点恢复

步骤 1：为规模配置并行

对于在 256 GPU（32 节点）上的 70B 模型：

[parallelism]
data_parallel_shard_degree = 32  # 跨 32 个 rank 的 FSDP
tensor_parallel_degree = 8        # 节点内的 TP
pipeline_parallel_degree = 1      # 70B 不使用 PP
context_parallel_degree = 1       # 增加以处理长序列

步骤 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 训练：
- [ ] 步骤 1：安装 torchao
- [ ] 步骤 2：配置 Float8
- [ ] 步骤 3：使用编译启动

步骤 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"]  # 排除输出层

[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 并行（FSDP + TP + PP + CP）：
- [ ] 步骤 1：创建种子检查点
- [ ] 步骤 2：配置 4D 并行
- [ ] 步骤 3：在 512 GPU 上启动

步骤 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
pipeline_parallel_degree = 8     # 跨节点的 PP
context_parallel_degree = 1      # 用于长序列的 CP

[training]
local_batch_size = 32
seq_len = 8192

步骤 3：在 512 GPU 上启动

# 64 节点 x 8 GPU = 512 GPU
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）
  • 在 H100 上使用 Float8 支持进行训练
  • 希望与 torchtune/HuggingFace 互操作的检查点

使用替代方案代替:

  • Megatron-LM: 针对 NVIDIA 专属部署的最大性能
  • DeepSpeed: 更广泛的 ZeRO 优化生态系统，支持推理
  • Axolotl/TRL: 用于微调而非预训练
  • LitGPT: 教育用途，小规模训练

常见问题

问题：大型模型内存不足

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

[activation_checkpoint]
mode = "full"  # 代替 "selective"

[training]
local_batch_size = 1

或使用梯度累积：

[training]
local_batch_size = 1
global_batch_size = 32  # 累积梯度

问题：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 的重分片能力：

# 将分片检查点转换为单文件
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+编译+FP8 8,532 +48%
Llama 70B 256 FSDP+TP+异步TP 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 协议。

资源