分布式训练总览

分布式训练沿数据、张量、流水线三个正交维度切分计算与状态；选型的核心问题只有一个——模型与优化器状态撞上哪一级显存墙，就需要在哪一级引入切分。

并行的三个维度

维度	切分对象	通信模式	通信量级	适用互联
数据并行 DP	批次数据	梯度 AllReduce	每步约等于参数量	单机/跨机皆可
张量并行 TP	层内权重矩阵	每层前向/反向各 2 次 AllReduce	极高，逐层触发	仅 NVLink 机内
流水线并行 PP	层间 stage	点对点 send/recv	低，仅边界激活	适合跨机慢互联

三个维度正交，可叠加组合；先数据并行、再模型并行是成本从低到高的引入顺序。

数据并行谱系：从 DataParallel 到 FSDP

演进主线：PyTorch 1.0 引入 torch.distributed 与 DDP → 1.10 torchrun 取代 launch 脚本 → 1.11 FSDP → 2.x 以 DTensor 为基座提供原生张量并行 API。

nn.DataParallel：已弃用

nn.DataParallel（DP）是单进程多线程方案：受 GIL 与线程同步拖累，每步还要 scatter 输入、replicate 模型、gather 输出，单卡 0 号 GPU 成为带宽与显存热点，且完全不支持模型并行。官方文档明确不推荐——任何场景都应改用 DDP，迁移成本仅几行代码。

DDP：全量副本 + 梯度同步

DDP 为每张 GPU 启动独立进程，各持完整模型副本，仅在 backward 时按 bucket 重叠执行梯度 AllReduce，数学上严格等价于对全批数据做单次梯度下降：

$θ_{t + 1} = θ_{t} - η \cdot \frac{1}{N} \sum_{i = 1}^{N} \nabla_{θ} L (x_{i}, y_{i}; θ_{t})$

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
 
dist.init_process_group("nccl")  # rank/world_size 由 torchrun 注入环境变量
model = DDP(model.to(torch.cuda.current_device()))
# DataLoader 须配 DistributedSampler 切分数据，见 [[Dataset 与 DataLoader]]

进程组与集合通信 API 见 torch.distributed 通信原语；多机训练用 torchrun 拉起，不要手写 mp.spawn。

显存墙与 ZeRO / FSDP

DDP 的隐性代价：每卡冗余存放全量参数、梯度、优化器状态。以混合精度 + Adam 训练 7.5B 模型为例——fp16 参数 15 GB + fp16 梯度 15 GB + 优化器状态（fp32 副本 + 一二阶动量）90 GB，合计约 120 GB，单张 80 GB 卡必然 OOM，此即显存墙。

ZeRO（DeepSpeed）按阶段消除冗余，FSDP 是同一分片思想（源自 ZeRO-3/FairScale）的 PyTorch 原生实现：

切分内容	ZeRO 术语	FSDP 对应
仅优化器状态	ZeRO-1	—
+ 梯度	ZeRO-2	`SHARD_GRAD_OP`
+ 参数	ZeRO-3	`FULL_SHARD`，静态显存约 ÷ N

代价：前向/反向需 all-gather 临时聚合完整参数、梯度再 reduce-scatter 回分片，通信量约为 DDP 的 1.5 倍，须靠预取与通信/计算重叠掩盖。权重分片后保存与加载不能再用普通 state_dict，参见大模型检查点。

模型并行与流水线并行

张量并行（TP） 将单层权重矩阵按行/列切到多卡，典型如 Megatron 按注意力 head 切分 QKV 投影。每层前后向各引入 2 次 AllReduce，通信量随层数线性增长，因此几乎只部署在 NVLink 互联的单机内部。

流水线并行（PP） 将模型按层切成连续 stage，micro-batch 在 stage 间流水执行。空闲气泡占比为 $\frac{p - 1}{m + p - 1}$ （ $p$ 为 stage 数、 $m$ 为 micro-batch 数），靠增大 $m$ 与 1F1B 调度压低。只传边界激活，通信量小，适合跨机。

3D 混合并行：机内 TP + 机间 PP + 外层再套数据并行/FSDP，是 Megatron-LM 与 DeepSpeed 训练百亿至万亿参数模型的标准拓扑；MoE 模型还叠加专家并行（EP）。

通信后端

后端	适用场景
NCCL	NVIDIA GPU 单机/多机，默认首选
Gloo	CPU 训练、调试与回退
MPI	HPC 集群（需源码编译）

按模型规模选型

以单卡 80 GB 为粗略标尺（先配合 torch.cuda 与混合精度的 AMP、激活重计算、梯度累积把单卡余量榨干，再谈分布式）：

模型规模	推荐方案
参数 + 优化器状态单卡放得下（约 ≤2B）	DDP
单卡放不下、单层放得下（约 3–30B）	FSDP `FULL_SHARD` 或 ZeRO-3
单层也放不下 / MoE（>30B）	TP + PP + DP/FSDP（Megatron / DeepSpeed）
快速实验、多后端切换	Accelerate 一套代码切换 DDP/FSDP/DeepSpeed

决策顺序：先压单卡，再 DDP 水平扩展，跨过显存墙才引入分片；TP/PP 的工程与通信调试成本远高于 FSDP，非必要不启用。

知识花园

探索