Chaos and Order

💡 왼쪽 원문을 읽으면서 오른쪽에 따라 써보세요. Tab 키로 힌트를 받을 수 있습니다.

引言

训练大语言模型(LLM)是现代 AI 工程中最复杂的挑战之一。在单张 GPU 上训练 Llama 3 70B 根本不可能。必须高效利用从数十到数千张 GPU，而这需要多种并行化策略的配合。

由 Meta PyTorch 团队开发的 torchtitan，正是为这种复杂的大规模 LLM 训练而生的参考实现(reference implementation)。它运用最新的 PyTorch 特性，以简洁、可扩展的方式实现，旨在让 AI 研究者和工程师能够学习并应用分布式训练的最佳实践。

本指南涵盖 torchtitan 的方方面面 — 从基础理论到实际安装、进阶并行化策略，再到性能优化，一次性梳理完毕。

1. Torch-Titan 简介

什么是 torchtitan？

torchtitan 是 Meta PyTorch 团队公开的、面向大规模 LLM 训练的生产级参考实现。在 GitHub 上可以通过 pytorch/torchtitan 找到它。

现有的许多 LLM 训练代码库(Megatron-LM、DeepSpeed、NeMo 等)功能丰富,但复杂度也很高。torchtitan 采取了不同的理念:

明确性(Clarity):只使用 PyTorch 原生 API,抽象层级降到最低
模块化(Modularity):每种并行化技术都可以独立开关
现代性(Modernity):积极采用 PyTorch 2.x 的最新特性
可复现性(Reproducibility):结构上便于复现训练实验

支持的模型

目前 torchtitan 默认支持的模型:

Llama 2 (7B、13B、34B、70B)
Llama 3 (8B、70B、405B)
Llama 3.1、3.2 系列

虽然默认支持的是 Llama 系列,但只要是基于 Transformer 的模型,都可以按照同样的结构轻松添加自定义模型。

与现有框架的差异

Megatron-LM (NVIDIA):

GPU 集群上最成熟的方案
高度优化,但依赖 NVIDIA 生态
代码库复杂,定制困难

DeepSpeed (Microsoft):

以 ZeRO 优化器闻名
同时支持训练与推理
依赖 C++ 自定义内核,有时难以与 PyTorch 完全集成

torchtitan (Meta/PyTorch):

纯 PyTorch 原生 API
运用 PyTorch 2.x 的 torch.compile、FSDP2、torch.distributed.tensor
结构清晰,也适合作为教学素材
紧跟 PyTorch 版本的更新节奏

仓库结构

torchtitan/
├── torchtitan/
│   ├── models/
│   │   ├── llama/          # Llama 模型定义
│   │   └── __init__.py
│   ├── parallelisms/
│   │   ├── parallelize_llama.py  # 并行化应用逻辑
│   │   ├── pipeline_llama.py     # 流水线并行
│   │   └── __init__.py
│   ├── optimizer.py         # 优化器配置
│   ├── checkpoint.py        # 检查点
│   ├── profiling.py         # 性能分析
│   └── utils.py
├── train.py                 # 训练入口
├── train_configs/           # TOML 配置文件
│   ├── llama3_8b.toml
│   ├── llama3_70b.toml
│   └── llama3_405b.toml
└── estimation.py            # 显存/FLOPs 估算工具

2. 分布式训练范式复习

在多张 GPU 上训练大模型的方法大致分为四种。torchtitan 支持全部四种,甚至可以将它们同时组合(4D 并行)。

数据并行 (Data Parallelism, DP)

最基础的并行方式。将模型复制到每张 GPU,把数据批次拆开,各 GPU 独立处理。

GPU 0: [批次 0~15]  →  梯度 0  ┐
GPU 1: [批次 16~31] →  梯度 1  ├→ All-Reduce → 同步
GPU 2: [批次 32~47] →  梯度 2  ┘

PyTorch 的 DistributedDataParallel(DDP)是标准实现。但这要求模型能装进单张 GPU 的显存。仅按 FP16 计算,70B 参数的模型就需要 140GB — 单张 GPU 无法容纳。

张量并行 (Tensor Parallelism, TP)

张量并行把模型的单个层分散到多张 GPU 上,沿列(column)或行(row)方向切分矩阵运算。

将 Linear 层(d_model=4096, d_ff=16384)分散到 4 张 GPU:

GPU 0: 列 0~4095   (4096 x 4096)
GPU 1: 列 4096~8191
GPU 2: 列 8192~12287
GPU 3: 列 12288~16383

每张 GPU 只处理整层的 1/4。需要 All-Reduce 或 All-Gather 来合并结果。NVLink 带宽越高,效率越好。

流水线并行 (Pipeline Parallelism, PP)

按顺序把模型的层分给多张 GPU。第一张 GPU 处理前 N 层,再把输出传给下一张 GPU。

GPU 0: 第 0~9 层   → 激活值 →
GPU 1: 第 10~19 层 → 激活值 →
GPU 2: 第 20~29 层 → 激活值 →
GPU 3: 第 30~39 层 → 损失

朴素实现一次只处理一个微批次(microbatch),导致 GPU 利用率低(流水线气泡,pipeline bubble)。GPipe、1F1B、Interleaved 等调度方式对此加以改善。

序列并行 (Sequence Parallelism, SP)

把注意力层中的序列维度分散到多张 GPU。用于解决长上下文(128K+ token)训练时注意力矩阵显存以 O(n²) 爆炸增长的问题。

把序列长度 4096 分散到 4 张 GPU:
GPU 0: token 0~1023
GPU 1: token 1024~2047
GPU 2: token 2048~3071
GPU 3: token 3072~4095

再用 Ring Attention 等算法计算出完整的注意力结果。

4D 并行:把一切组合起来

torchtitan 的核心优势在于支持把这四种并行方式同时组合的 4D 并行。

4D 并行示例(128 张 GPU):
- DP = 2(数据并行,2 个副本)
- TP = 8(张量并行,8 张 GPU 处理一层)
- PP = 4(流水线并行,4 段流水线)
- SP 随 TP 一并启用

总 GPU 数 = DP × TP × PP = 2 × 8 × 4 = 64
(若算入 SP,还可以有更多组合方式)

每种并行方式的通信模式各不相同,因此需要根据硬件拓扑找到最优组合。一般来说:

TP 应用于同一台服务器内、由 NVLink 互联的 GPU 之间
PP 应用于服务器之间
DP 应用于最外层维度

3. FSDP2 (Fully Sharded Data Parallel v2)

ZeRO 与 FSDP 的关系

Microsoft DeepSpeed 的 ZeRO(Zero Redundancy Optimizer)是一种把参数、梯度、优化器状态分散到多张 GPU 上的创新方法论。PyTorch 的 FSDP(Fully Sharded Data Parallel)则是这一思想的 PyTorch 原生实现。

ZeRO 分为三个阶段:

ZeRO-1:仅分片优化器状态
ZeRO-2:优化器状态 + 梯度分片
ZeRO-3:优化器状态 + 梯度 + 参数分片

FSDP 相当于 ZeRO-3。

FSDP1 与 FSDP2 的差异

在 PyTorch 2.0 之前,使用的是 torch.distributed.fsdp.FullyShardedDataParallel(FSDP1)。在 PyTorch 2.4+ 中,推荐使用 torch.distributed._composable.fsdp 里的 fully_shard,也就是 FSDP2。

主要差异:

特性	FSDP1	FSDP2
API 风格	包装器(wrapper)方式	可组合(composable)API
TP 集成	有限	原生集成
显存效率	较好	更好
torch.compile	部分支持	完全支持
代码可读性	复杂	清晰

# FSDP1 方式(旧式)
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(model, ...)

# FSDP2 方式(torchtitan 使用)
from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy

mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.float32
)

# 对每个 Transformer 层应用 FSDP2
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)

# 也应用到整个模型
fully_shard(model, mp_policy=mp_policy)

FSDP 的工作原理

FSDP 在前向传播、反向传播、权重更新这三个阶段的行为各不相同。

前向传播(Forward pass):

执行层之前:通过 All-Gather 在所有 GPU 上重建被分片的参数
执行层:用完整参数计算
执行层之后:丢弃参数(节省显存)

反向传播(Backward pass):

层反向传播之前:通过 All-Gather 重建参数
计算梯度
通过 Reduce-Scatter 把梯度分片、分散到各 GPU
丢弃参数

权重更新:

每张 GPU 只更新自己负责的分片所对应的参数、梯度和优化器状态

在 4 张 GPU 上训练 70B 模型的示例:

显存节省:权重 140GB / 4 = 每张 GPU 35GB(+ 分散的优化器状态)
代价:产生通信开销(All-Gather、Reduce-Scatter)

CPU Offload

当 GPU 显存不足时,可以把优化器状态和梯度卸载到 CPU 内存(RAM)。

from torch.distributed._composable.fsdp import fully_shard, CPUOffloadPolicy

# CPU 卸载策略
cpu_offload = CPUOffloadPolicy(offload_params=True)

for layer in model.layers:
    fully_shard(
        layer,
        offload_policy=cpu_offload
    )

CPU 卸载能大幅减少显存占用,但由于 PCIe 传输开销,训练速度会变慢(2-5 倍)。只在显存绝对不够用时,才把它当作最后的手段。

混合精度与 FSDP2

FSDP2 可以按层应用不同的精度。

from torch.distributed._composable.fsdp import MixedPrecisionPolicy

# 基本混合精度配置
mp_policy = MixedPrecisionPolicy(
    param_dtype=torch.bfloat16,   # 参数:BF16(通信效率)
    reduce_dtype=torch.float32    # 梯度规约:FP32(稳定性)
)

# 特定层保持 FP32(例如嵌入层)
fully_shard(model.embed_tokens)  # 不传 mp_policy → FP32
for layer in model.layers:
    fully_shard(layer, mp_policy=mp_policy)
fully_shard(model, mp_policy=mp_policy)

4. Tensor Parallelism (张量并行)

Transformer 中的张量并行

在 Megatron-LM 中首次被系统化的张量并行,应用于 Transformer 的两个核心模块。

MLP (前馈网络,Feed-Forward Network):

FFN(x) = GELU(x * W1) * W2

列并行(Column Parallel,切分 W1):
GPU 0: x * W1[0:d/4]  → GELU → y[0:d/4]
GPU 1: x * W1[d/4:d/2] → GELU → y[d/4:d/2]
...

行并行(Row Parallel,切分 W2 + All-Reduce):
GPU 0: y[0:d/4] * W2[0:d/4, :] → 部分和 0
GPU 1: y[d/4:d/2] * W2[d/4:d/2, :] → 部分和 1
...
All-Reduce → 最终输出

Self-Attention: 将 Query、Key、Value 矩阵按注意力头(head)切分。

# torchtitan 中的 TP 应用(基于 DTensor)
from torch.distributed.tensor.parallel import (
    parallelize_module,
    ColwiseParallel,
    RowwiseParallel,
    SequenceParallel,
    PrepareModuleInput,
)

# 定义并行化方案
plan = {
    "attention": PrepareModuleInput(
        input_layouts=(Shard(1),),
        desired_input_layouts=(Replicate(),),
    ),
    "attention.wq": ColwiseParallel(),
    "attention.wk": ColwiseParallel(),
    "attention.wv": ColwiseParallel(),
    "attention.wo": RowwiseParallel(output_layouts=Shard(1)),
    "feed_forward.w1": ColwiseParallel(),
    "feed_forward.w2": RowwiseParallel(output_layouts=Shard(1)),
    "feed_forward.w3": ColwiseParallel(),
}

parallelize_module(model_layer, tp_mesh, plan)

DTensor:分布式张量的基础

torchtitan 基于 PyTorch 的 torch.distributed.tensor(DTensor)实现张量并行。DTensor 是一种新的抽象 — 在逻辑上是单个张量,但实际上分布在多个设备上。

from torch.distributed.tensor import DTensor, Shard, Replicate
import torch.distributed as dist

# 创建 1D mesh(用于 TP)
tp_mesh = dist.init_device_mesh("cuda", (8,), mesh_dim_names=("tp",))

# 创建 2D mesh(DP + TP)
mesh_2d = dist.init_device_mesh(
    "cuda", (2, 8), mesh_dim_names=("dp", "tp")
)

# 把张量按特定维度分片
# Shard(0): 行方向分片
# Shard(1): 列方向分片
# Replicate(): 复制(不分片)

序列并行 (Sequence Parallelism)

作为张量并行的自然延伸,把 LayerNorm 和 Dropout 运算沿序列维度并行化。

标准 TP 流程:
Input(复制) → LayerNorm(复制) → Attention(TP) → [All-Reduce] → Output(复制)

SP + TP 流程:
Input(序列分片) → LayerNorm(序列分片) →
[All-Gather] → Attention(TP) → [Reduce-Scatter] →
Output(序列分片)

SP 与 TP 通信联动,用 All-Gather + Reduce-Scatter 取代 All-Reduce。通信总量相同,但在显存层面更高效 — LayerNorm 的激活值被分散到 N 张 GPU 上,从而节省显存。

5. Pipeline Parallelism (流水线并行)

流水线气泡问题

朴素的流水线并行会带来严重的低效。

朴素流水线(4 张 GPU,4 个微批次):

时间 →
GPU 0: [M0 F][M0 B]      [      气泡      ]
GPU 1:        [M0 F][M0 B][      气泡      ]
GPU 2:               [M0 F][M0 B][  气泡  ]
GPU 3:                      [M0 F][M0 B]

F = 前向(Forward),B = 反向(Backward)。GPU 0 处理完 M0 之后要长时间等待。气泡比例 = (PP - 1) / (micro_batches + PP - 1)。

GPipe 调度

GPipe 通过向流水线中注入多个微批次来减少气泡。

GPipe(4 张 GPU,4 个微批次):

时间 →
GPU 0: [M0F][M1F][M2F][M3F]               [M3B][M2B][M1B][M0B]
GPU 1:      [M0F][M1F][M2F][M3F]      [M3B][M2B][M1B][M0B]
GPU 2:           [M0F][M1F][M2F][M3F][M3B][M2B][M1B][M0B]
GPU 3:                [M0F][M1F][M2F][M3F][M3B][M2B][M1B][M0B]

所有前向传播完成之后才进行反向传播。缺点是需要大量的激活值显存。

1F1B (One Forward One Backward) 调度

1F1B 是一种显存效率更好的调度方式。每张 GPU 交替对微批次执行前向和反向传播。

# torchtitan 的流水线并行配置(TOML)
# [experimental]
# pipeline_parallel_degree = 4
# pipeline_parallel_schedule = "1f1b"

# Python 代码中的设置
from torchtitan.parallelisms.pipeline_llama import pipeline_llama

# 创建流水线阶段(stage)
stages, model_parts = pipeline_llama(
    model,
    pp_mesh,
    parallel_dims,
    job_config,
    device,
    model_config
)

Interleaved 调度

Interleaved(交错式)调度让每张 GPU 负责多个流水线阶段。能进一步减少气泡,但实现更复杂。

Interleaved 1F1B(4 张 GPU,每张 GPU 2 个阶段):

GPU 0 负责的层:[第 0~4 层] 和 [第 20~24 层]
GPU 1:[第 5~9 层] 和 [第 25~29 层]
...

6. torchtitan 的安装与基本使用

系统要求

Python 3.10+
PyTorch 2.5+(推荐使用最新 nightly)
CUDA 12.1+
GPU:推荐 H100 或 A100(最低 40GB 显存)

安装

# 克隆仓库
git clone https://github.com/pytorch/torchtitan
cd torchtitan

# 安装 PyTorch nightly(包含最新特性)
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu124

# 安装依赖
pip install -r requirements.txt

# 安装 torchtitan 包
pip install -e .

# 下载训练数据的分词器
python torchtitan/datasets/download_tokenizer.py \
    --repo_id meta-llama/Meta-Llama-3-8B \
    --tokenizer_path "original" \
    --hf_token YOUR_HF_TOKEN

配置文件 (TOML)

torchtitan 使用 TOML 格式的配置文件。

# train_configs/llama3_8b.toml

[job]
dump_folder = "./outputs"
description = "Llama 3 8B training"

[profiling]
enable_profiling = true
save_traces_folder = "profile_trace"
profile_freq = 100

[metrics]
log_freq = 10
enable_tensorboard = true

[model]
name = "llama3"
flavor = "8B"
norm_type = "rmsnorm"
tokenizer_path = "./torchtitan/datasets/tokenizer/original/tokenizer.model"

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

[training]
batch_size = 8
seq_len = 2048
warmup_steps = 200
max_norm = 1.0
steps = 1000
data_parallel_replicate_degree = 1
data_parallel_shard_degree = -1  # 自动使用剩余的 GPU 数
tensor_parallel_degree = 1
enable_loss_parallel = false

[experimental]
pipeline_parallel_degree = 1

[checkpoint]
enable_checkpoint = true
folder = "checkpoint"
interval_type = "steps"
interval = 500
model_weights_only = false
export_dtype = "float32"
async_mode = "disabled"

[activation_checkpoint]
mode = "selective"  # full, selective, none
selective_ac_option = "op"  # 选择性检查点选项

[float8]
enable_float8_linear = false

运行 Llama 3 8B 训练

# 单节点、8 张 GPU 训练 Llama 3 8B
torchrun --nproc_per_node=8 \
    train.py \
    --job.config_file train_configs/llama3_8b.toml

# TP=4、DP=2 的配置训练
torchrun --nproc_per_node=8 \
    train.py \
    --job.config_file train_configs/llama3_8b.toml \
    --training.tensor_parallel_degree 4 \
    --training.data_parallel_shard_degree 2

# 多节点训练(2 个节点,各 8 张 GPU)
torchrun \
    --nproc_per_node=8 \
    --nnodes=2 \
    --rdzv_id=101 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:29400 \
    train.py \
    --job.config_file train_configs/llama3_70b.toml

显存/FLOPs 估算工具

在开始训练之前,可以预先估算显存和计算量。

# 估算显存与 FLOP
python estimation.py \
    --job.config_file train_configs/llama3_8b.toml

# 输出示例:
# Estimated model size: 15.01 GB
# Estimated optimizer state size: 30.02 GB
# Total estimated GPU memory: 65.23 GB
# Estimated FLOP per step: 1.61e+14

7. Flash Attention 集成

什么是 Flash Attention?

Flash Attention 是斯坦福大学 Tri Dao 等人于 2022 年发表的注意力算法。它把标准注意力的显存复杂度从 O(N²) 降到 O(N),实际速度也提升了 2-4 倍。

标准注意力的问题:

标准注意力运算:
S = Q * K^T      # (seq_len, seq_len) 矩阵 — 显存爆炸!
P = softmax(S)
O = P * V

seq_len=4096: S 矩阵 = 4096 × 4096 × 4 bytes = 64MB (FP32)
seq_len=32768: S 矩阵 = 32768 × 32768 × 4 bytes = 4GB(单层就有这么大!)

Flash Attention 的核心思想:

把 Q、K、V 按块(block)处理(分块,tiling)
在 SRAM(共享内存)而不是 HBM 中计算
不把注意力矩阵存进 HBM
数值上精确的结果(不是近似值)

Flash Attention 2 与 3

Flash Attention 2 (2023):

改进了注意力计算的并行化
掩码处理更高效
相较 A100,注意力速度快 2-4 倍

Flash Attention 3 (2024):

针对 H100 Hopper 架构做了特化
运用 WGMMA(Warpgroup Matrix Multiply Accumulate)
在 H100 上相较 FA2 再提升 1.5-2 倍
支持 FP8

在 torchtitan 中使用 Flash Attention

# torchtitan/models/llama/model.py 中的注意力实现
import torch.nn.functional as F

def forward(
    self,
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> torch.Tensor:
    bs, seqlen, _ = x.shape
    xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

    # 应用 RoPE 编码
    xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

    # 使用 Flash Attention(PyTorch 内置 SDPA)
    # torch.nn.functional.scaled_dot_product_attention
    # 会自动使用 Flash Attention
    output = F.scaled_dot_product_attention(
        xq, xk, xv,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True,  # 因果掩码(语言模型)
    )

    return self.wo(output.view(bs, seqlen, -1))

从 PyTorch 2.0 开始,torch.nn.functional.scaled_dot_product_attention(SDPA)会自动使用 Flash Attention。无需额外安装任何软件包,就能享受到 Flash Attention 的优势。

若要直接使用 Flash Attention 原始软件包:

pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func

# 直接调用 Flash Attention
output = flash_attn_func(
    q, k, v,
    dropout_p=0.0,
    softmax_scale=None,  # 自动设为 1/sqrt(head_dim)
    causal=True,
)

8. Async Tensor Parallelism

为什么需要异步 TP

在标准张量并行中,All-Reduce 或 All-Gather/Reduce-Scatter 会阻塞下一层的计算。GPU 在等待通信期间处于空闲(idle)状态。

同步 TP:
[GPU 计算] → [等待通信] → [GPU 计算] → [等待通信] ...

异步 TP 让通信与计算重叠执行:

异步 TP:
[GPU 计算 A] ──────────────────→
[通信(A 的结果)]──────→
                  [GPU 计算 B] ─→
                       [通信(B)]

torchtitan 中的异步 TP 实现

# 在配置文件中启用异步 TP
# [training]
# enable_async_tensor_parallel = true

# 代码中的对应实现
from torchtitan.parallelisms import parallelize_llama

# async_tp 在内部利用 torch.distributed.tensor.parallel 的
# async_all_gather
model = parallelize_llama(
    model,
    world_mesh,
    parallel_dims,
    job_config
)

异步 TP 通过计算-通信重叠(compute-communication overlap),在较高的 TP degree(8 以上)时尤其有效。在 H100 + NVLink 环境下,有报告显示可带来额外 5-15% 的性能提升。

9. 检查点与恢复训练

Distributed Checkpoint (dcp)

在大规模分布式训练中,检查点(checkpoint)远不止「保存模型」这么简单。需要在数千张 GPU 上同时保存与加载。

PyTorch 的 torch.distributed.checkpoint(dcp)原生支持分布式检查点。

# torchtitan 的检查点实现
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.state_dict import (
    get_model_state_dict,
    get_optimizer_state_dict,
    set_model_state_dict,
    set_optimizer_state_dict,
    StateDictOptions,
)

# 保存
def save_checkpoint(model, optimizer, step, output_dir):
    state_dict = {
        "model": get_model_state_dict(model),
        "optimizer": get_optimizer_state_dict(model, optimizer),
        "extra": {"step": step},
    }

    dcp.save(
        state_dict,
        checkpoint_id=f"{output_dir}/step-{step}",
    )

# 加载
def load_checkpoint(model, optimizer, checkpoint_path):
    state_dict = {
        "model": get_model_state_dict(model),
        "optimizer": get_optimizer_state_dict(model, optimizer),
        "extra": {},
    }

    dcp.load(
        state_dict,
        checkpoint_id=checkpoint_path,
    )

    set_model_state_dict(
        model,
        model_state_dict=state_dict["model"],
        options=StateDictOptions(strict=True),
    )
    set_optimizer_state_dict(
        model,
        optimizer,
        optim_state_dict=state_dict["optimizer"],
    )

    return state_dict["extra"]["step"]

异步检查点

大模型的检查点保存可能耗时数分钟,期间 GPU 会停滞。异步检查点让训练继续进行的同时,在后台完成保存。

# 异步检查点配置
[checkpoint]
async_mode = "async"  # "disabled", "async", "async_with_pinned_mem"

# 异步检查点在内部由独立线程执行
# 训练循环不会等待检查点保存完成
# 但在保存下一个检查点时,会确认上一次保存是否已完成

from torchtitan.checkpoint import CheckpointManager

checkpoint_manager = CheckpointManager(
    dataloader=train_dataloader,
    model_parts=model_parts,
    optimizer=optimizer,
    lr_scheduler=lr_scheduler,
    states={"train_state": train_state},
    job_config=job_config,
)

# 在训练循环中
for step in range(num_steps):
    # ... 训练代码 ...

    # 异步保存检查点(非阻塞)
    checkpoint_manager.save(curr_step=step, force=False)

检查点格式转换

分布式训练保存的检查点,可以转换成单文件形式,或转换成 HuggingFace 格式。

# 把分布式检查点转换为单个 HuggingFace 模型
python scripts/convert_checkpoint.py \
    --checkpoint_path outputs/checkpoint/step-1000 \
    --output_path outputs/hf_model \
    --model_type llama3 \
    --model_flavor 8B

10. 性能分析

使用 PyTorch Profiler

要找出性能瓶颈,首先要知道时间都花在了哪里。PyTorch Profiler 正是为此而生的强大工具。

import torch
from torch.profiler import profile, record_function, ProfilerActivity

# 基本性能分析
with profile(
    activities=[
        ProfilerActivity.CPU,
        ProfilerActivity.CUDA,
    ],
    record_shapes=True,    # 记录张量 shape
    profile_memory=True,   # 记录显存使用量
    with_stack=True,       # 记录调用栈
) as prof:
    with record_function("model_inference"):
        output = model(input)

# 输出结果
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20
))

# 保存为 Chrome Trace(在 chrome://tracing 中可视化)
prof.export_chrome_trace("trace.json")

torchtitan 内置性能分析

torchtitan 通过配置文件控制性能分析。

[profiling]
enable_profiling = true
save_traces_folder = "profile_trace"
profile_freq = 100      # 每 100 步做一次性能分析
enable_memory_snapshot = true  # 保存显存快照

# 参考 torchtitan/profiling.py 内部实现
from contextlib import contextmanager
import torch.profiler as profiler

@contextmanager
def maybe_enable_profiling(config, global_step=0):
    if not config.profiling.enable_profiling:
        yield
        return

    with profiler.profile(
        activities=[
            profiler.ProfilerActivity.CPU,
            profiler.ProfilerActivity.CUDA,
        ],
        schedule=profiler.schedule(
            skip_first=10,
            wait=5,
            warmup=5,
            active=1,
            repeat=1,
        ),
        on_trace_ready=profiler.tensorboard_trace_handler(
            config.profiling.save_traces_folder
        ),
        record_shapes=True,
        profile_memory=True,
        with_stack=True,
    ) as p:
        yield p

TensorBoard 集成

# 记录指标
from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir="runs/experiment_1")

# 在训练循环中
for step, (input, target) in enumerate(dataloader):
    loss = train_step(model, optimizer, input, target)

    writer.add_scalar("Loss/train", loss.item(), step)
    writer.add_scalar("LR", optimizer.param_groups[0]["lr"], step)

    # 监控 GPU 显存
    writer.add_scalar(
        "GPU/memory_allocated_GB",
        torch.cuda.memory_allocated() / 1e9,
        step
    )
    writer.add_scalar(
        "GPU/memory_reserved_GB",
        torch.cuda.memory_reserved() / 1e9,
        step
    )

# 运行 TensorBoard
# tensorboard --logdir=runs/

显存使用量分析

# GPU 显存快照
torch.cuda.memory._record_memory_history(max_entries=100000)

# 执行一部分训练
for step in range(100):
    loss = train_step(model, optimizer, batch)

# 保存显存快照
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

# 用可视化工具分析
# 在 https://pytorch.org/memory_viz 上传 .pickle 文件

# 用显存分析器分析每层的显存使用情况
from torch.cuda._memory_viz import profile_plot

# 把分析结果保存为 HTML
with open("memory_profile.html", "w") as f:
    f.write(profile_plot(prof))

训练效率指标:MFU

MFU(Model FLOP Utilization)是实际 GPU 性能相对于理论最大性能的比率。

def compute_mfu(
    model_num_params: int,
    batch_size: int,
    seq_len: int,
    elapsed_time: float,  # 每步耗时(秒)
    num_gpus: int,
    gpu_peak_flops: float,  # GPU 理论最大 FLOPS
) -> float:
    """
    计算 LLM 训练的 MFU。
    参考:PaLM 论文(Chowdhery et al., 2022)
    """
    # 前向 FLOPs = 2 * num_params * batch_size * seq_len
    # 反向大约是前向的 2 倍 → 总计 6 * num_params * batch_size * seq_len
    flops_per_step = 6 * model_num_params * batch_size * seq_len
    achieved_flops = flops_per_step / elapsed_time
    peak_flops_total = gpu_peak_flops * num_gpus

    return achieved_flops / peak_flops_total

# 使用示例
mfu = compute_mfu(
    model_num_params=8e9,     # 80 亿参数
    batch_size=8,
    seq_len=2048,
    elapsed_time=0.5,         # 每步 0.5 秒
    num_gpus=8,
    gpu_peak_flops=989e12,    # H100 BF16: 约 989 TFLOPS
)
print(f"MFU: {mfu:.1%}")  # 例如:MFU: 45.2%

# 一般能达到的 MFU:
# - 良好实现:40-60%
# - 优化过的实现(torchtitan、Megatron):50-65%
# - 理论最大值:约 70%(通信/显存开销不可避免)

11. 实战训练配置示例

Llama 3 8B:单节点 8× H100

# train_configs/llama3_8b_h100x8.toml

[job]
dump_folder = "./outputs/llama3_8b"

[model]
name = "llama3"
flavor = "8B"

[training]
batch_size = 4
seq_len = 8192
data_parallel_replicate_degree = 1
data_parallel_shard_degree = 8  # FSDP2: 8 张 GPU
tensor_parallel_degree = 1

[activation_checkpoint]
mode = "selective"

[float8]
enable_float8_linear = true  # 启用 FP8 训练

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

torchrun --nproc_per_node=8 train.py \
    --job.config_file train_configs/llama3_8b_h100x8.toml

Llama 3 70B:4 节点 × 8 H100(共 32 张 GPU)

# train_configs/llama3_70b_32gpu.toml

[training]
batch_size = 2
seq_len = 4096
data_parallel_replicate_degree = 2  # DDP: 2 个副本
data_parallel_shard_degree = 4      # FSDP: 4 张 GPU 分片
tensor_parallel_degree = 4          # TP: 4 张 GPU

# 总 GPU 数 = 2 × 4 × 4 = 32

[experimental]
pipeline_parallel_degree = 1  # PP 关闭

Llama 3 405B:大规模集群

# train_configs/llama3_405b_large.toml

[training]
batch_size = 1
seq_len = 2048
data_parallel_replicate_degree = 4   # DDP
data_parallel_shard_degree = 8       # FSDP
tensor_parallel_degree = 8           # TP

[experimental]
pipeline_parallel_degree = 8         # PP

# 总 GPU 数 = 4 × 8 × 8 × 8 = 2048

结语:从 torchtitan 学到的东西

torchtitan 不只是一个训练工具,更是一份清晰展示现代 LLM 分布式训练最佳实践的教学资源。

通过本指南学到的核心内容:

4D 并行:组合 DP、TP、PP、SP,高效利用数千张 GPU
FSDP2:用 PyTorch 原生 API 达到 ZeRO-3 级别的显存效率
Flash Attention:把 O(N²) 的显存降为 O(N),速度提升 2-4 倍
异步检查点:训练不中断也能保存检查点
MFU 优化:以理论 GPU 性能的 40-65% 为目标

分布式训练是硬件、软件、算法交织的复杂领域。torchtitan 尽可能把这种复杂性以简明的方式呈现出来,让学习和实验都更容易上手。希望你能亲自跑一跑代码,尝试各种并行组合,培养出对分布式训练的直觉。

引言