基于《大规模语言模型：从理论到实践（第2版）》第4章 分布式模型训练
爆款小标题：数据并行、张量并行、流水并行怎么选？原书第4章+DeepSpeed 落地
为什么这一节重要
当模型参数量达到 7B、70B 甚至更大时，单卡显存无法容纳完整模型与优化器状态，必须依赖分布式训练把计算与显存分摊到多卡或多机。分布式训练有多种并行方式：数据并行、张量并行、流水并行，以及 ZeRO 等显存优化技术。选错策略会导致显存溢出、训练极慢或通信成为瓶颈。本节基于原书第 4 章与 DeepSpeed 实践，把各类并行的适用场景、ZeRO 的分片逻辑与使用顺序讲清，并给出「先跑通再调优」的落地建议。
学习目标
学完本节，你将能够：
- 区分三种并行：准确说明数据并行（DP）、张量并行（TP）、流水并行（PP）各自如何切分计算与数据、适用什么规模的模型、以及通信特点。
- 理解 ZeRO：说明 ZeRO-1/2/3 分别对优化器状态、梯度、参数做了哪些分片与通信，以及 ZeRO-3 相比 ZeRO-2 的额外收益与通信开销。
- 能做初步选型：根据模型规模（如 7B/70B）与单卡显存、卡数，初步选择「DP+ZeRO」或「DP+ZeRO+TP/PP」的组合，并能在 DeepSpeed 配置上跑通一个示例。
一、数据并行（Data Parallelism, DP）（原书第 4 章）
做法：每张卡上各保留一份完整的模型副本，每卡处理不同的数据（batch 的不同分片）；前向与反向在各卡本地完成，梯度在卡间做 AllReduce（或等价操作）得到全局梯度，再各自用相同梯度更新参数，从而保证各卡参数一致。
优点：实现简单、与单卡训练逻辑最接近；扩展性好，卡数增加时只需保证数据分片与梯度同步正确。缺点：单卡必须能放下完整模型与一份优化器状态等，因此模型不能超过单卡显存。
适用：模型能完整放入单卡（或通过 ZeRO 等优化后能放入）时，优先用数据并行扩规模。原书第 4 章将其作为基础并行方式介绍。
通信：每步反向结束后一次 AllReduce（或 ReduceScatter + AllGather），通信量与模型参数量成正比；卡数多时通信可能成为瓶颈，需结合梯度累积与通信重叠等手段。
二、张量并行（Tensor Parallelism, TP）（原书第 4 章）
做法：把单层的权重与计算切分到多卡。例如把一层 (Y = XW) 的 (W) 按行或列切分，每卡持有一部分 (W)，计算时每卡算局部结果，再通过 AllGather 或 ReduceScatter 拼成完整输出。这样单层不再依赖单卡显存，但单层计算过程中需要多卡通信。
优点：能训练「单层就很大」的模型（如 70B 的 FFN 层）。缺点：通信频繁（每层都可能通信）、对带宽要求高；实现复杂，通常依赖 Megatron 等框架。
适用：当单卡放不下单层（例如单层参数量超过单卡显存）时，需要用张量并行或流水并行。原书第 4 章对 TP 的切分方式与通信有进一步说明。
工程注意：TP 通常与 DP 组合使用（例如 4 卡 TP × 2 组 = 8 卡），组内做张量并行、组间做数据并行；集群拓扑（如 NVLink、PCIe）会影响通信效率，需结合带宽规划。
三、流水并行（Pipeline Parallelism, PP）（原书第 4 章）
做法：把模型按层切分到多卡，形成「流水线」：第 1 卡算第 1 批的前几层，算完传给第 2 卡算接下来几层，依次类推；当第 1 卡把第 1 批传走后，可以开始算第 2 批的前几层，从而实现多批同时在 pipeline 上流动。
优点：单卡不需要存完整模型，只需存若干层；适合层数多、单层不算特别大的模型。缺点：存在「流水线气泡」——部分卡在等前一批数据时空闲，利用率会下降；需要调 micro batch 等参数平衡效率与显存。
适用：层数多、但单层能放进单卡时，可用 PP 把层摊到多卡。常与 DP、TP 组合成 3D 并行（原书第 4 章）。
四、ZeRO：优化器、梯度与参数的分片（原书 DeepSpeed 实践）
问题：数据并行时，每卡都存完整模型 + 完整优化器状态（如 Adam 的动量与方差）+ 完整梯度，显存占用是单卡训练的「卡数倍」，反而限制了可用的 batch size 或模型规模。
ZeRO-1：只对优化器状态分片。每卡只存 1/N 的优化器状态（N 为 DP 卡数），更新时按需 Gather 再分片，从而把优化器状态占用的显存降到约 1/N。
ZeRO-2：在 ZeRO-1 基础上再对梯度分片。每卡只存 1/N 的梯度，AllReduce 时用 ReduceScatter 等操作直接得到分片结果，梯度显存也降到约 1/N。
ZeRO-3：在 ZeRO-2 基础上再对参数分片。每卡只存 1/N 的模型参数，前向与反向时按需 Gather 完整参数（或分层 Gather），计算完再释放。这样单卡显存不再需要存完整模型，可训练远大于单卡显存的模型；代价是前向与反向中多了参数 AllGather 等通信，通信量比 ZeRO-2 大。
选型建议：模型能放进单卡时，ZeRO-2 即可显著降低显存并加速（因梯度与优化器状态分片）；模型放不进单卡时，需要 ZeRO-3 或 ZeRO-3 + TP/PP。原书配套的 DeepSpeed 实践部分给出了 ZeRO-2/3 的配置示例，建议先跑通再按显存与速度调 batch size 与并行度。
五、工程实战要点
1. 从小规模组合开始
不要一上来就 3D 并行（DP+TP+PP）。建议顺序：先数据并行 + ZeRO-2（或 ZeRO-3 若模型大）；若单卡仍放不下（例如单层就很大），再考虑加 TP；若层数多、单层能放下，可考虑 PP。先在小规模（如 2 卡、4 卡）跑通配置与数据管线，再放大。3D 并行的配置复杂（如 TP 与 PP 的组划分、通信拓扑），调试成本高，只有在「必须」时才上。
2. 梯度累积与有效 batch
当单卡 batch size 受显存限制较小时，可通过梯度累积扩大有效 batch：多步前向—反向，累积梯度后再做一次 AllReduce 与优化器更新。有效 batch = 每卡 batch × 梯度累积步数 × DP 卡数。预训练时常用较大有效 batch（如数百万 token/step），需合理设置梯度累积步数以平衡显存与吞吐。
3. 使用 DeepSpeed 时先跑通再调参
用官方或原书给出的 ZeRO 配置文件（如 ds_config.json）先跑通一个 epoch；再根据显存占用与速度调整 train_batch_size、gradient_accumulation_steps、以及 ZeRO 的 stage。遇到 OOM 时优先考虑增大梯度累积、减小 batch、或开 ZeRO-3。
DeepSpeed ZeRO-2 配置示例：
{
"zero_optimization": {
"stage": 2,
"offload_optimizer": {"device": "none"},
"allgather_partitions": true
},
"train_batch_size": 32,
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 2
}
stage: 2表示 ZeRO-2（优化器状态 + 梯度分片）；stage: 3则加上参数分片。- 有效 batch = 4 × 2 × 卡数；若 4 卡则 4×2×4=32。OOM 时减小
train_micro_batch_size_per_gpu或增大gradient_accumulation_steps。
4. 通信带宽与拓扑
张量并行与 ZeRO-3 的 AllGather 都依赖卡间带宽。多机时要注意网络带宽；单机多卡时 NVLink 比 PCIe 快很多。集群规划时把「通信量」与「带宽」一起考虑，避免通信成为瓶颈。
六、常见误区与避坑指南
误区一：一上来就 3D 并行
3D 并行配置复杂、调试成本高，且未必需要。避坑：先评估「单卡能放多少」「模型多大」，从 DP+ZeRO 开始，再按需加 TP/PP。
误区二：忽视通信带宽
TP 与梯度同步都依赖带宽；带宽不足时加卡数可能反而变慢。避坑：做扩展实验时监控步时与通信占比；多机时优先保证网络带宽与拓扑。
误区三：ZeRO-3 一定比 ZeRO-2 慢
ZeRO-3 多了参数通信，在模型不大、单卡能放下时，ZeRO-2 往往更快。避坑：在「能放下」的前提下优先 ZeRO-2；只有显存不够时才上 ZeRO-3。
七、小结与衔接
本节结合原书第 4 章与 DeepSpeed，梳理了数据并行、张量并行、流水并行的适用场景与通信特点，以及 ZeRO-1/2/3 对优化器状态、梯度、参数的分片与对应显存与通信权衡；并给出了「先 DP+ZeRO，再按需 TP/PP」的选型顺序与工程要点。下一节将进入指令微调：如何构建指令数据、如何用 LoRA 做参数高效微调，以及长上下文扩展的常见做法（原书第 5 章）。
课后思考题
- 在 8 张 A100 上训练 70B 模型，单卡显存不足以放下完整模型时，你会优先考虑哪种并行组合？简要说明理由。
- ZeRO-3 相比 ZeRO-2 多做了哪类分片？带来的好处与额外通信开销是什么？
