分布式训练：把模型撑到千亿参数
目标：搞清楚「单卡装不下」时，模型怎么被切开、数据怎么流动、瓶颈在哪。
你大概不会亲手训练千亿模型，但这套概念是理解所有大模型工程讨论的基础，也是高薪岗位的硬门槛。
一、先算清楚显存
训练 7B 模型（混合精度 + Adam）的显存构成：
| 项 | 计算 | 大小 |
|---|---|---|
| 权重（fp16） | 7B × 2B | 14 GB |
| 梯度（fp16） | 7B × 2B | 14 GB |
| 优化器状态（fp32 的 m, v + fp32 主权重） | 7B × 12B | 84 GB |
| 激活值 | 与 batch × seq 成正比 | 数十 GB |
| 合计（不含激活） | ≈ 112 GB |
结论：7B 全参训练需要 4×A100-80G 起步。 这就是为什么大家都用 LoRA（可训参数降到 0.5%，优化器状态随之降 200 倍）。
第一课就是这条
显存不是「装不下权重」的问题，主要是优化器状态与激活。 理解这点，你就明白 ZeRO、梯度检查点、LoRA 各自在省什么。
二、四种并行
| 并行 | 切什么 | 解决什么 | 代价 |
|---|---|---|---|
| 数据并行 DP | 切数据，模型各存一份 | 加速训练 | 显存不省 |
| 张量并行 TP | 切单层内的矩阵 | 单层太大 | 通信极密集，限机内 |
| 流水线并行 PP | 切层（前几层在 GPU1，后几层在 GPU2） | 模型太深 | 有空泡（bubble） |
| 序列并行 / 上下文并行 | 切序列长度 | 超长上下文 | 实现复杂 |
| 专家并行 EP | 切 MoE 的专家 | MoE 模型 | 路由与通信复杂 |
python
# PyTorch 原生数据并行：DDP
import torch, torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group("nccl")
model = torch.nn.Linear(4096, 4096).cuda(dist.get_rank())
model = DDP(model, device_ids=[dist.get_rank()])
# 训练循环不变；DDP 自动在所有卡间同步梯度（all-reduce）
loss.backward()
optimizer.step()DDP 与 DP 的区别：DDP 每卡一个进程、梯度自动归约（快）；老式 DP 单进程多线程、有 GIL 与复制开销（慢，已弃用）。
三、ZeRO：把冗余消掉
标准 DP 中，每张卡都存完整权重+梯度+优化器状态 —— 这是巨大浪费。
| 阶段 | 切分什么 | 7B 单卡显存（8 卡） |
|---|---|---|
| ZeRO-1 | 优化器状态 | ~30 GB |
| ZeRO-2 | + 梯度 | ~18 GB |
| ZeRO-3 / FSDP | + 权重 | ~14 GB |
| 加上 LoRA | 只存 adapter | 数 GB |
代价：通信量上升。ZeRO-3 每步要 all-gather 权重，网络必须够快。
python
# 用 accelerate 极简上手
from accelerate import Accelerator
acc = Accelerator(mixed_precision="bf16", gradient_accumulation_steps=4)
model, optimizer, dl, sched = acc.prepare(model, optimizer, dl, sched)
for batch in dl:
with acc.accumulate(model):
loss = criterion(model(**batch).logits, batch["labels"])
acc.backward(loss)
optimizer.step(); sched.step(); optimizer.zero_grad()四、梯度检查点（用时间换显存）
前向时不保存中间激活，反向时重算。
| 显存 | 速度 | |
|---|---|---|
| 默认 | 高 | 快 |
| 梯度检查点 | 降 50%–70% | 慢 20%–30% |
python
model.gradient_checkpointing_enable()这是最划算的一招：省下的显存可以直接换更大 batch 或更长序列。
五、通信：真正的瓶颈
| 通信原语 | 含义 | 用在哪 |
|---|---|---|
| Broadcast | 一传多 | 参数初始化 |
| All-reduce | 全部求和并分发 | DDP 梯度同步 |
| All-gather | 全部收集 | ZeRO-3 取权重 |
| Reduce-scatter | 求和后分摊 | ZeRO 梯度 |
| All-to-all | 全交换 | MoE 专家路由 |
| 网络 | 带宽 | 能干什么 |
|---|---|---|
| PCIe | ~32 GB/s | 只能做轻量通信 |
| NVLink（机内） | 300–900 GB/s | 适合 TP |
| InfiniBand / RDMA（机间） | 100–400 Gb/s | PP / DP / ZeRO |
经验法则：TP 放在机内（NVLink），PP/DP 跨机。跨机做 TP 会慢到无法训练。
六、并行策略怎么选（决策表）
| 情况 | 方案 |
|---|---|
| 模型能放进单卡 | DDP（纯数据并行） |
| 单卡放得下推理，但训练显存不够 | DDP + ZeRO-2/3 + 梯度检查点 |
| 单层就放不下（如 70B） | TP（机内）+ PP 或 ZeRO-3 |
| 超长序列（>32K） | 序列并行 / 上下文并行 |
| MoE 模型 | EP + DP |
| 只想微调 | LoRA/QLoRA，能不用分布式就不用 |
一条最实用的建议
大多数团队根本不需要自研分布式方案。 用 DeepSpeed / FSDP / Megatron-LM / torchtitan 的现成配置，把精力放在数据与效果上。
七、规模化训练的现实问题
| 问题 | 现象 | 处理 |
|---|---|---|
| Loss spike | loss 突然跳高甚至 nan | 降 lr、梯度裁剪、跳过坏 batch、回滚 checkpoint |
| 吞吐上不去 | GPU 利用率 < 60% | 查数据加载（dataloader 瓶颈）、通信占比 |
| 训练中断 | 节点故障 | 每 N 步存 checkpoint、自动续训 |
| 复现不了 | 换机器结果不同 | 固定种子、记录环境、接受硬件级非确定性 |
| 成本失控 | 卡时超预算 | 先小规模验证 scaling，再放大 |
大模型训练的核心工程能力 = 让一个跑 2 周的任务不中断、不跑偏、可恢复。
八、学习方式建议
| 步骤 | 做什么 |
|---|---|
| 1 | 单机 8 卡或 2 卡上跑通 DDP + AMP + 梯度累积（不需大模型，用 ResNet 即可） |
| 2 | 用 FSDP/DeepSpeed 训一个 1B 模型，观察显存变化 |
| 3 | 做一次 TP=2 的实验，看到通信开销 |
| 4 | 读 Megatron-LM 的并行配置文档，理解每种并行的约束 |
| 5 | 跑一次真实的长任务（含中断恢复） |
没有卡也能学：大部分概念可以用 2 卡甚至单卡多进程模拟（torchrun --nproc_per_node=2）。