PyTorch 2.x 深度学习专题【左扬精讲】—— 大规模分布式训练核心要素：从并行策略到通信原语的理论详解
PyTorch 2.x 深度学习专题【左扬精讲】—— 大规模分布式训练核心要点：从并行策略到通信原语的理论知识
分布式训练 数据并行 模型并行 流水线并行 张量并行 通信原语 All-Reduce NCCL
学习重点提示
本文涵盖分布式训练的核心知识，内容深度依次递进。建议读者根据自身背景选择性深入：
- 必须掌握（面试/工程基准线）：
- 并行策略的分类体系：理解"切什么"和"怎么切"是掌握并行策略的核心。切分的维度决定了通信模式、资源利用率和适用场景，三大并行策略（数据并行、流水线并行、张量并行）本质上是不同的切分角度，并非互相替代而是互补组合
- 数据并行的梯度聚合机制：梯度聚合是数据并行正确性的保障，All-Reduce 保证了所有 GPU 在每步更新后持有相同的模型参数。Ring-All-Reduce 通过去中心化设计避免了通信瓶颈，是工业界的事实标准
- 集合通信与点对点通信的区别：集合通信（Collective）涉及组内所有节点，如 All-Reduce、All-Gather；点对点通信（P2P）仅涉及两个节点，如 Send/Recv。前者用于梯度同步等全局操作，后者用于流水线并行的 stage 间激活传递
- 需要理解（深度优化/系统设计）：
- 流水线并行的气泡问题：气泡产生的根本原因是上下游 stage 之间的计算依赖。气泡开销与 stage 数量 S 成正比，与 micro-batch 数量 K 成反比（气泡比例 ≈ (S-1)/(K+S-1)）。1F1B 调度通过让反向传播紧跟前向传播，降低了在制品数量和显存峰值
- 张量并行的分片策略：Column Parallel 将权重按列切分，输出需要 All-Gather 拼接；Row Parallel 将权重按行切分，输出需要 All-Reduce 累加。LLM 的 Attention QKV 投影天然适合 Column Parallel，因为切分后的 Q/K/V 可以独立计算注意力分数
- 通信开销与计算 overlap：通信开销直接影响 GPU 利用率。通过 CUDA Stream 将通信和计算分配到不同流水线，在等待通信完成的同时执行不依赖该通信结果的计算，可以有效隐藏通信延迟。这是现代训练系统的核心性能优化手段
- 了解即可（进阶专题，可后续深入）：
- 通信调度算法：包括梯度 bucketing、延迟调度、自适应通信策略等，这些是 NCCL 库和框架层面的实现细节
- 混合并行的工程实现细节：3D Parallelism（TP+PP+DP）的配置调优、多机网络拓扑感知、故障恢复机制等，通常在大模型训练框架（Megatron-LM、DeepSpeed）中封装好
目录导航
一、为什么需要分布式训练？
What — 分布式训练是什么？
分布式训练是指利用多个计算设备（GPU/TPU/多台机器）并行执行深度学习模型的训练过程，以突破单卡显存的限制或加速训练收敛。随着 Transformer 架构的崛起，模型参数规模从 ResNet 的约 6000 万飞速增长到 GPT-4 的约 1.8 万亿（业界广泛估计，OpenAI 未官方公布），训练一个如此庞大的模型需要消耗数十到数百 PetaFLOPS（每秒千万亿次浮点运算）的算力，单卡根本无法承载。
Why — 为什么需要分布式训练？
分布式训练解决两个核心问题：
- 显存墙：一个 1750 亿参数的 GPT-3 模型，按 FP16 半精度存储需要约 350GB 显存，而当时最先进的 A100 GPU 单卡仅有 80GB 显存
- 算力墙：训练 1 万亿参数规模的模型需要约 10^25 次浮点运算，假设用单卡 A100（SXM 版本，FP16/BF16 Tensor Core 算力 312 TFLOPS）需要数十年
没有分布式训练会发生什么？
- 无法训练参数量超过单卡显存上限的模型
- 训练周期从几周延长到几年甚至几十年
- 大模型研究的大门对学术界和中小企业彻底关闭
核心结论：分布式训练是训练大模型的必要条件，它通过将计算和存储分散到多个设备上来突破单卡瓶颈。
二、并行策略全景图
What — 并行策略是什么？
并行策略是指将深度学习训练任务拆分到多个设备上的方法。根据"拆什么"和"怎么拆"，业界主流分为三大类：数据并行（Data Parallelism）、流水线并行（Pipeline Parallelism）、张量并行（Tensor Parallelism）。
图 1 并行策略分类
┌──────────────────────────────────────────────────────────────────────────────┐
│ 分布式训练并行策略 │
│ （从单卡 → 多卡 → 多机扩展） │
├──────────────────────────────────────────────────────────────────────────────┤
│ │
│ 单卡瓶颈 │
│ ┌─────────────────────────────────────┐ │
│ │ GPU 0 │ │
│ │ ┌───────────────────────────────┐ │ │
│ │ │ Embedding │ │ │
│ │ ├───────────────────────────────┤ │ │
│ │ │ Transformer Block × N │ │ 显存不足 / 计算太慢 │
│ │ │ Layer 0, Layer 1, ... │ │ │
│ │ ├───────────────────────────────┤ │ │
│ │ │ Output Layer (LM Head) │ │ │
│ │ └───────────────────────────────┘ │ │
│ └─────────────────────────────────────┘ │
│ │ │
│ ▼ │
├──────────────────────────────────────────────────────────────────────────────┤
│ │
│ ┌──────────┐ ┌──────────────┐ ┌──────────────────┐ │
│ │ 数据并行 │ │ 流水线并行 │ │ 张量并行 │ │
│ │ (DP) │ │ (PP) │ │ (TP) │ │
│ └────┬─────┘ └──────┬───────┘ └────────┬─────────┘ │
│ │ │ │ │
│ 多卡扩展 多卡扩展 多卡扩展 │
│ ┌────┴────┐ ┌────┴────┐ ┌────┴────┐ │
│ │ GPU0 GPU1 │ │ GPU0 GPU1 │ │ GPU0 GPU1 │ │
│ │ Layer0-11│ │ Layer0-5 │ │ W[:,0:H/2]│ │
│ │ Layer0-11│ │ Layer6-11│ │ W[:,H/2:H]│ │
│ └──────────┘ └──────────┘ └──────────┘ │
│ │ │ │ │
│ ┌────┴────┐ ┌────┴────┐ ┌────┴────┐ │
│ │ 多机扩展 │ │ 多机扩展 │ │ 通信密集 │ │
│ │ Rack 0 │ │ Rack 0 │ │ 需要 NVLink│ │
│ │ Rack 1 │ │ Rack 1 │ │ 单机内TP │ │
│ └──────────┘ └──────────┘ └──────────┘ │
│ │
│ 适用：参数量 < 单卡显存 │ 适用：层数多显存不足 │ 适用：单层参数 > 单卡显存 │
│ 但数据量大 但数据量中等 且需要 All-Gather │
└──────────────────────────────────────────────────────────────────────────────┘图 1：分布式训练并行策略全景图，从单卡到多卡、多机扩展的层次关系
| 并行策略 | 拆分维度 | 通信模式 | 适用场景 | 主要挑战 |
|---|---|---|---|---|
| 数据并行 | 数据批次（Batch） | All-Reduce | 参数量 < 显存上限，数据量大 | 梯度同步开销 |
| 流水线并行 | 模型层（Layer） | Send/Recv（P2P） | 参数量 >> 单卡显存 | 流水线气泡（Bubble） |
| 张量并行 | 单层内部（Tensor） | All-Gather / All-Reduce | 单层 >> 单卡显存 | 通信量巨大 |
三种并行策略本质上是"切分方向"的不同选择：
- 数据并行：横向切数据，每个设备有完整模型副本，适合数据密集场景
- 流水线并行：纵向切模型，按层分组到不同设备，适合超深网络
- 张量并行：在单层内部进行矩阵分块，适合超宽层（如大语言模型的 Attention 和 FFN）
实际训练系统通常将三者组合使用，即 3D Parallelism（3D 并行）。
三、通信原语详解
What — 通信原语是什么？
通信原语是分布式系统中定义数据如何在多个节点之间传输和变换的基础操作。根据参与者的数量和角色，分为两大类：集合通信（Collective Communication）和点对点通信（Point-to-Point Communication）。
图 2 通信原语分类
┌──────────────────────────────────────────────────────────────────────────────┐
│ 分布式通信原语 │
│ │
│ ┌─────────────────────────────┐ ┌─────────────────────────────────┐ │
│ │ 集合通信 Collective │ │ 点对点通信 P2P │ │
│ │ (所有节点参与) │ │ (两节点间) │ │
│ └──────────────┬──────────────┘ └───────────────┬─────────────────┘ │
│ │ │ │
│ ┌─────────────┼─────────────┐ ┌─────────────┼─────────────┐ │
│ │ │ │ │ │ │ │
│ ▼ ▼ ▼ │ ▼ │ │
│ Broadcast All-Reduce All-Gather │ Send ──────▶ Recv │
│ 一发全收 求和归约 收集拼接 │ ◀────── ──────▶ │
│ ▼ ▼ ▼ │ Send Recv │
│ GPU0 GPU1 GPU2 GPU3 ... │ │
│ ═══ ═══ ═══ ═══ │ GPU0 GPU1 │
│ └────── All-Reduce ──────┘ │
└──────────────────────────────────────────────────────────────────────────────┘图 2：通信原语分类图，展示集合通信（Collective）与点对点通信（P2P）的区别
3.1 Broadcast — 一对多广播
What — Broadcast 是什么？
Broadcast（广播）是指将一个节点上的数据完整复制到组内所有其他节点的操作。典型场景是模型参数初始化和全局学习率广播。
设有 N 个节点，源节点 s 持有数据 x，广播后所有节点 i 都持有 x：
forall i in [0, N): receive(x) from s
实际实现中，数据通过树形结构分发，时间复杂度为 O(log N)。
3.2 Scatter — 一对多分发
What — Scatter 是什么？
Scatter（分发）与 Broadcast 不同，Scatter 将一份数据切分成 N 个分片，每个节点收到不同的分片。典型场景是张量并行中权重矩阵按列分片后的分发。
| 原语 | 源节点 | 目标节点 | 数据变化 |
|---|---|---|---|
| Broadcast | 1 个 | N 个 | 相同数据复制 N 份 |
| Scatter | 1 个 | N 个 | 数据切分，每个节点拿 1/N |
3.3 Gather — 多对一收集
What — Gather 是什么？
Gather（收集）将多个节点上的数据汇聚到一个节点。与 Reduce 不同，Gather 不做任何计算，只是简单拼接。
3.4 Reduce — 多对一归约
What — Reduce 是什么？
Reduce（归约）将多个节点上的数据通过某种二元运算（如求和、求最大值）合并到单个节点。这是数据并行中梯度聚合的核心操作。
设有 N 个节点，每个节点 i 持有梯度 g_i，归约操作将结果存到根节点 r：
g_r = g_0 ⊕ g_1 ⊕ ... ⊕ g_{N-1}
其中 ⊕ 是二元运算符。在深度学习中通常是加法运算，用于计算所有节点的梯度之和。
3.5 All-Gather — 多对多收集
What — All-Gather 是什么？
All-Gather（全局收集）是 Gather 的扩展版本：将每个节点的数据收集起来，然后分发给所有节点。换言之，每个节点最终都持有所有节点数据的完整副本。
Why — 为什么需要 All-Gather？
张量并行中，单个线性层的输入 X 需要广播到所有分片，而输出 Y 的各分片需要重新拼接。没有 All-Gather，层之间的数据流就会断裂。
3.6 All-Reduce — 多对多归约
What — All-Reduce 是什么？
All-Reduce（全局归约）是分布式训练中最重要的通信原语。它结合了 Reduce 和 Broadcast：先对所有节点的数据做归约操作（通常是求和），然后将结果广播给所有节点。最终，所有节点持有相同的归约结果。
设有 N 个节点，每个节点 i 持有梯度 g_i，执行 All-Reduce 后，所有节点都得到相同的聚合结果：
forall i: result_i = g_0 ⊕ g_1 ⊕ ... ⊕ g_{N-1}
在 Ring-All-Reduce 算法中，通信被组织成两步：
- Reduce-Scatter：在 N 个节点间做 N-1 次数据交换，每次让每个节点累加一部分数据
- All-Gather：将累加后的完整结果广播给所有节点
3.7 Scatter-Gather — 分发-收集组合
What — Scatter-Gather 是什么？
Scatter-Gather（分发-收集组合）先执行 Scatter 再执行 Gather。它实际上是 Reduce-Scatter 和 All-Gather 的组合，在某些集合通信库（如 NCCL）中作为独立原语暴露给用户。
3.8 P2P（Point-to-Point）— 点对点通信
What — P2P 通信是什么？
P2P（点对点通信）是指两个节点之间的直接数据交换，仅涉及一个发送者（Sender）和一个接收者（Receiver）。Send 和 Recv 是最基本的 P2P 操作。
Why — P2P 在流水线并行中的作用
流水线并行中，上游 stage 完成后需要将中间激活（Activation）传递给下游 stage。这种传递是严格的顺序依赖关系，只能用 P2P 通信。每个 micro-batch 的前向传播会触发一次 Send，反向传播也会触发一次 Recv。
通信原语小结：
- 一对多：Broadcast（复制）、Scatter（切分分发）
- 多对一：Gather（收集）、Reduce（归约）
- 多对多：All-Gather（全局收集）、All-Reduce（全局归约）、Scatter-Gather（分发收集组合）
- 点对点：Send/Recv（P2P）
四、数据并行（Data Parallelism）
What — 数据并行是什么？
数据并行（Data Parallelism）是分布式训练中最广泛使用的并行策略。其核心思想是：每个设备保存一份完整的模型副本，不同设备处理不同的数据分片（Data Shard），最后通过通信同步梯度，确保所有设备上的模型参数保持一致。
图 3 数据并行梯度聚合流程（4 GPU 示例）
┌──────────────────────────────────────────────────────────────────────────────┐
│ Step 1: 本地前向+反向，各自计算梯度 │
│ │
│ GPU 0 GPU 1 GPU 2 GPU 3 │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ ✕ │ │ ✕ │ │ ✕ │ │ ✕ │ │
│ │ Grad0│ │ Grad1│ │ Grad2│ │ Grad3│ │
│ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │ │ │ │
│ └───────────────┼───────────────┼───────────────┘ │
│ ▼ │
│ Step 2: Ring-All-Reduce 同步梯度 │
│ │
│ GPU0 ──▶ Grad0+Grad1 ──▶ GPU1 ──▶ Grad1+Grad2 ──▶ GPU2 │
│ ▲ │ │
│ │ Ring 环式传递 │ │
│ │ ▼ │
│ Grad3+Grad0 ◀── GPU3 ◀── Grad2+Grad3 ◀── GPU2 ◀─────────────── │
│ │
│ 最终：所有 GPU 持有 Sum(Grad0, Grad1, Grad2, Grad3) / 4 = 平均梯度 │
│ │
│ Step 3: 更新模型参数（各 GPU 独立进行，参数一致） │
│ │
│ GPU 0 GPU 1 GPU 2 GPU 3 │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ ✓ │ │ ✓ │ │ ✓ │ │ ✓ │ │
│ │ 全模型 │ │ 全模型 │ │ 全模型 │ │ 全模型 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ │
└──────────────────────────────────────────────────────────────────────────────┘图 3：数据并行的梯度聚合流程，展示从本地梯度计算到 All-Reduce 同步的关键步骤
4.1 数据并行的主流实现
| 实现方式 | 通信库 | 通信时机 | 特点 |
|---|---|---|---|
| DDP（DistributedDataParallel） | NCCL/Gloo | 每轮反向传播后 | 异步梯度同步，Ring-All-Reduce，带 bucketing 优化 |
| ZeRO（Zero Redundancy Optimizer） | NCCL | 按需通信 | 分片优化器状态/梯度/参数，显存极度节省 |
| DP（DataParallel） | NCCL/Gloo | 每轮前反向后 | PyTorch 早期实现，单进程多卡，通信效率低 |
DDP 是 PyTorch 分布式训练的事实标准，其核心初始化和训练流程如下：
import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 1. 初始化分布式环境，必须在任何 DDP API 调用之前
dist.init_process_group(backend="nccl") # NCCL 是 NVIDIA 的 GPU 通信库
# 2. 获取本地秩（local_rank），即当前进程所在的 GPU 编号
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank) # 将当前 GPU 设置为该进程的默认设备
# 3. 构建模型并移动到 GPU
model = nn.Linear(512, 512).cuda(local_rank)
# 4. 用 DDP 包装模型，关键：模型必须在此之前已经移动到 GPU
model = DDP(model, device_ids=[local_rank])
# 5. 训练循环：与单卡训练完全一致，无需手动同步梯度
for batch in dataloader:
inputs = batch.cuda(local_rank)
targets = batch["labels"].cuda(local_rank)
outputs = model(inputs) # 前向传播
loss = nn.functional.cross_entropy(outputs, targets)
loss.backward() # 反向传播：DDP 会在内部自动触发 All-Reduce
optimizer.step()
optimizer.zero_grad()设计意图：DDP 将梯度同步封装在 loss.backward() 内部，对用户代码透明。相比 DP，DDP 不需要每步等待一个进程完成前反向再交给下一个，真正实现了多卡并行计算。
4.2 梯度聚合机制
What — 梯度聚合是什么？
梯度聚合（Gradient Aggregation）是指在数据并行中，多个设备分别计算出本地梯度后，如何合并成全局梯度并同步到所有设备的技术。核心通信原语是 All-Reduce。
设 N 个 GPU 并行处理，每个 GPU i 计算得到本地梯度 g_i。对于随机梯度下降（SGD），全局参数更新公式为：
theta_{t+1} = theta_t - eta * (1/N) * sum_{i=0}^{N-1} g_i
其中 eta 是学习率。All-Reduce 操作的本质就是计算 sum_{i=0}^{N-1} g_i 并将结果广播给所有 GPU。
以 Ring-All-Reduce 为例，N 个 GPU 排成一个环，分两步执行：
- Reduce-Scatter：每个 GPU 将自己的梯度分 N-1 次发送给环中的下一个 GPU，每次累加收到的数据。完成后，每个 GPU 持有 1/N 的全局累加和。
- All-Gather：每个 GPU 将自己持有的 1/N 结果分 N-1 次发送给环中的下一个 GPU，每次接收并拼接。最终所有 GPU 持有完整的全局累加和。
Why — 为什么需要 Ring-All-Reduce？
朴素 All-Reduce 需要一个中心节点作为聚合枢纽，通信量和时间复杂度均为 O(N)。Ring-All-Reduce 将聚合分散到环中每个节点，总通信量仍为 O(N)，但无中心瓶颈，带宽利用率更高。
没有 Ring-All-Reduce 会发生什么？
- 中心节点的网卡带宽成为瓶颈，所有 GPU 等待最慢的链路
- 大规模多卡训练时，梯度同步时间占总训练时间的 20%-30%
- 无法充分利用 NVLink/NIC 的多通道带宽
4.3 ZeRO：超越传统数据并行
What — ZeRO 是什么？
ZeRO（Zero Redundancy Optimizer）是微软提出的显存优化技术，本质上是在数据并行基础上对模型状态（优化器状态、梯度、参数）进行分片存储，避免每个 GPU 都保存完整副本。
| ZeRO Stage | 分片内容 | 显存节省比例 | 通信量增加 |
|---|---|---|---|
| ZeRO-1 | 优化器状态（Adam moments） | ~4x | 基本不变 |
| ZeRO-2 | 优化器状态 + 梯度 | ~8x | 小幅增加 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | ~Nx（N=设备数） | 显著增加 |
注意：ZeRO-3 的参数分片意味着每次前向传播需要按需从其他 GPU 获取缺失的参数分片，会显著增加通信量。它适合参数量极大但多机互联带宽充足的场景（如 400G InfiniBand）。
数据并行小结：
- DDP 是 PyTorch 分布式训练的标准实现，梯度同步对用户透明
- All-Reduce 是梯度聚合的核心通信原语，Ring-All-Reduce 是工业界主流实现
- ZeRO 通过分片存储突破显存墙，是训练超大模型的关键技术
五、模型并行：流水线并行（Pipeline Parallelism）
What — 流水线并行是什么？
流水线并行（Pipeline Parallelism）是将模型的各层按顺序分配到不同设备上的并行策略。例如，一个 24 层Transformer模型，可以将前 12 层放在 GPU 0，后 12 层放在 GPU 1。前向传播时，GPU 0 算完前 12 层后将激活传给 GPU 1；反向传播时则反向进行。
图 4 流水线并行示意图（4 stage，8 micro-batch）
┌──────────────────────────────────────────────────────────────────────────────┐
│ GPipe 调度（1F1B） │
│ │
│ Stage 0 Stage 1 Stage 2 Stage 3 时间 │
│ (Layer 0-5) (Layer 6-11)(Layer12-17)(Layer18-23) │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ F0 │ ───▶ │ F0 │ ───▶ │ F0 │ ───▶ │ F0 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │ │
│ Warm-up (气泡产生阶段) ▼ │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ F1 │ ───▶ │ F1 │ ───▶ │ F1 │ ───▶ │ F1 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │ │ │
│ └────────────┴────────────┘ │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ B0 │ ◀─── │ F2 │ ───▶ │ F2 │ ───▶ │ F2 │ ───▶ │ F2 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │ │
│ 1F1B 阶段 ▼ │
│ (前向与反向交错) │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ B1 │ ◀─── │ B0 │ ◀─── │ F3 │ ───▶ │ F3 │ ───▶ │ F3 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │ │ │
│ └────────────┴────────────┴──────────────────────────────────▶ │
│ │
│ ... │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ B7 │ ◀─── │ B6 │ ◀─── │ B5 │ ◀─── │ B4 │ ◀─── │ F7 │ │
│ └─────┘ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐ │
│ │ B7 │ │ B7 │ ◀─── │ B6 │ ◀─── │ B5 │ ◀─── ... │
│ └─────┘ └─────┘ └─────┘ └─────┘ │
│ │ │
│ Cool-down (气泡消散) │
│ │
│ ───▶ = 前向传播 F ◀─── = 反向传播 B ░░ = 流水线气泡 (Bubble) │
│ │
│ 气泡比例 ≈ (S-1)/(K+S-1) (S=4 stage, K=8 micro-batch → ~27%) │
└──────────────────────────────────────────────────────────────────────────────┘图 4：流水线并行示意图，展示 micro-batch 在多 stage 间的流水线式流动
5.1 流水线气泡（Pipeline Bubble）
What — 流水线气泡是什么？
流水线气泡（Pipeline Bubble）是指在流水线并行中，由于下游设备必须等待上游完成才能开始计算，导致部分设备处于空闲状态的资源浪费现象。气泡的大小直接影响 GPU 利用率。
设模型分为 S 个 stage（流水线阶段），每个 micro-batch 的前向时间为 F，反向时间为 B。使用 GPipe（经典流水线调度）时，流水线启动阶段（Warm-up）有 S-1 个气泡，流水线填满后每轮只有 1 个气泡。
气泡占总时间的比例为（以 K 个 micro-batch 为例）：
bubble_ratio = (S-1) / (K + S - 1)
当 K（micro-batch 数量）远大于 S 时，气泡比例趋近于 0。因此，增加 micro-batch 数量是减少气泡开销的关键策略。
5.2 1F1B 调度策略
What — 1F1B 是什么？
1F1B（One Forward One Backward，一前向一反向）调度是 GPipe 提出的优化策略：在流水线启动后，稳定状态中每个 stage 每轮只同时运行 1 个前向和 1 个反向 micro-batch。相比传统的"先跑完所有前向再跑反向"，1F1B 显著降低了在制品（In-Flight Micro-batches）数量，从而降低显存占用。
Why — 为什么 1F1B 能节省显存？
朴素流水线调度需要同时保存所有 micro-batch 的激活（因为反向传播需要等所有前向完成）。1F1B 策略在流水线稳定后，反向传播立即跟在同一 micro-batch 的前向传播之后执行，前向时保存的激活可以在反向后立即释放，显存峰值约为朴素调度的 1/S（S 为 stage 数）。
5.3 Interleaved 调度
What — Interleaved 调度是什么？
Interleaved（交叉）调度是 1F1B 的扩展，每个 GPU 处理多个不连续的模型块（chunk），而非连续的层分组。例如，将模型分为 4 个 chunk（每 chunk 包含连续的若干层），朴素分配是 GPU 0 处理 chunk 0，GPU 1 处理 chunk 1；Interleaved 分配是 GPU 0 处理 chunk 0 和 chunk 2，GPU 1 处理 chunk 1 和 chunk 3，每个 GPU 以交叉方式处理多个 chunk。
注意：Interleaved 调度会增加通信次数（每个 GPU 需要与更多其他 GPU 通信），但能进一步减少气泡比例。它是 PipeDream 和 Megatron-LM 等框架采用的高级调度策略。
流水线并行小结：
- 流水线并行按层拆分模型，适合参数量远超单卡显存的超深网络
- Pipeline Bubble 是核心挑战，增加 micro-batch 数量可缓解
- 1F1B 调度在保持计算效率的同时显著降低显存峰值
六、模型并行：张量并行（Tensor Parallelism）
What — 张量并行是什么？
张量并行（Tensor Parallelism）是在单层（如单个矩阵乘法或自注意力层）内部进行分片，将计算分散到多个设备上。与流水线并行"按层切分"不同，张量并行"在层内部切分"，通信发生在每层的前向和反向传播中。
图 5 张量并行列并行（Column Parallel）
┌──────────────────────────────────────────────────────────────────────────────┐
│ Y = X · A (标准矩阵乘法) │
│ │
│ X [B, H] A [H, H] Y [B, H] │
│ ┌──────┐ ┌────────┐ ┌──────┐ │
│ │ 输入 │ × │ 权重的 │ ───▶ │ 输出 │ │
│ │ 激活 │ │ H×H │ │ │ │
│ └──────┘ └────────┘ └──────┘ │
│ │
│ │
│ 按列切分权重矩阵 │
│ │
│ X [B, H] │
│ ┌──────┐ │
│ │ 输入 │ │
│ └──────┘ │
│ │ ┌──────────────────────────────────┐ │
│ ├──────────────────▶│ A_0 [H, H/N] │ A_1 [H, H/N] │ │
│ │ │ GPU 0 │ GPU 1 │ │
│ │ └───────┬──────────┴───────┬────────┘ │
│ │ │ │ │
│ │ ▼ ▼ │
│ │ Y_0 [B, H/N] Y_1 [B, H/N] │
│ │ GPU 0 GPU 1 │
│ │ │ │ │
│ │ └─────────┬─────────┘ │
│ │ ▼ │
│ │ All-Gather 拼接 │
│ │ │ │
│ │ ▼ │
│ │ Y [B, H] = concat(Y_0, Y_1) │
│ │
│ ───────────────────────────────────────────────────────────────────────── │
│ 通信原语：前向传播后需要 All-Gather 拼接各 GPU 的局部输出 │
│ 输出维度：从 [B, H/N] → [B, H] │
└──────────────────────────────────────────────────────────────────────────────┘图 5：张量并行列并行示意图，展示单层矩阵乘法 Y = XA 按列分片的计算流程
6.1 Column Parallel Linear（列并行）
What — 列并行是什么？
列并行（Column Parallel）将线性层的权重矩阵 A 按列切分为 [H, H/N] 的 N 个分片 A_i，每个 GPU 计算 Y_i = X * A_i，输出 Y 的维度从 [B, H] 变为 [B, H/N]。最后需要通过 All-Gather 将各分片拼接成完整的 [B, H] 输出。
6.2 Row Parallel Linear（行并行）
What — 行并行是什么？
行并行（Row Parallel）将权重矩阵 A 按行切分为 [H/N, H] 的 N 个分片，每个 GPU 计算 Y_i = X_i * A_i（其中 X 也需要相应拆分）。由于部分和需要累加，行并行的通信原语是 All-Reduce。
| 分片方式 | 权重分片 | 输出通信 | 典型应用 |
|---|---|---|---|
| Column Parallel | 按列切分 A[:, :] | All-Gather | 自注意力 QKV 投影、FFN 第一层 |
| Row Parallel | 按行切分 A[:, :] | All-Reduce | 自注意力输出投影、FFN 第二层 |
6.3 Megatron-LM 的张量并行策略
What — Megatron-LM 是什么？
Megatron-LM 是 NVIDIA 提出的超大规模语言模型训练框架，其核心贡献是系统性地将 Transformer 架构中的所有层分解为列并行和行并行的组合。自注意力层的 QKV 投影使用 Column Parallel，输出投影使用 Row Parallel；FFN 的第一层（up-projection）使用 Column Parallel，第二层（down-projection）使用 Row Parallel。
以单层 Transformer 为例，完整的张量并行布局如下：
GPU 0 GPU 1 GPU N-1
| | |
v v v
Q_proj_col Q_proj_col Q_proj_col # Column Parallel，All-Gather 输出
K_proj_col K_proj_col K_proj_col # Column Parallel，All-Gather 输出
V_proj_col V_proj_col V_proj_col # Column Parallel，All-Gather 输出
| | |
v v v
Q @ K^T Q @ K^T Q @ K^T # 各自计算，不需要通信
| | |
v v v
softmax softmax softmax # 各自计算，不需要通信
| | |
v v v
attn @ V attn @ V attn @ V # 各自计算，不需要通信
| | |
v v v
out_proj_row out_proj_row out_proj_row # Row Parallel，All-Reduce 输出
| | |
v v v
MLP_col MLP_col MLP_col # Column Parallel，All-Gather 输出
| | |
v v v
MLP_row MLP_row MLP_row # Row Parallel，All-Reduce 输出设计精髓：Megatron 的核心洞察是 Transformer 架构中注意力机制的 QKV 投影天然适合列切分，因为切分后的 Q/K/V 可以独立计算注意力分数，无需通信。只有在投影回原空间时（All-Gather）或层间残差连接时（All-Reduce）才需要通信。
注意：张量并行假设所有 GPU 在物理上紧密互联（如 NVLink），因为每层前反向都需要多次 All-Gather/All-Reduce。NVLink 提供 900GB/s 的单向带宽（HGX A100 8-GPU 配置），远高于 PCIe Gen4 的 32GB/s 和 InfiniBand HDR 的 50GB/s 单向带宽。在多机场景下，如果 GPU 间通信必须经过网卡，张量并行的效率会急剧下降。
张量并行小结：
- 张量并行在单层内部进行矩阵分片，适合超宽层（如大语言模型的 Attention 和 FFN）
- Column Parallel 需要 All-Gather 拼接输出，Row Parallel 需要 All-Reduce 累加结果
- 张量并行通信密集，必须部署在高速互联硬件（NVLink）上
七、混合并行与通信调度
What — 混合并行是什么？
混合并行（Hybrid Parallelism）是指在实际大模型训练系统中，同时使用数据并行、流水线并行、张量并行三种策略的组合。例如，NVIDIA 的 PTD-P（Pipeline+Tensor+Data）框架将多机内使用张量并行，多机间使用流水线并行和数据并行。
7.1 通信与计算 Overlap
What — 通信计算 Overlap 是什么？
通信计算 Overlap（Overlap）是指在等待通信完成的同时，让 GPU 执行其他不依赖该通信结果的计算任务，从而隐藏通信延迟。在现代大模型训练系统中，Overlapping 是提升 GPU 利用率的关键技术。
PyTorch DDP 原生支持梯度同步与反向传播的 overlap，通过设置 bucket_cap_mb 参数和 CUDA Stream 调度实现：
import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 1. 初始化分布式环境
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
# 2. 构建模型并用 DDP 包装
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank], broadcast_buffers=False)
# 3. 训练循环：DDP 原生支持 overlap，梯度同步与反向传播并行执行
for batch in dataloader:
inputs = batch["input"].cuda(local_rank)
targets = batch["target"].cuda(local_rank)
outputs = model(inputs)
loss = loss_fn(outputs, targets)
loss.backward() # DDP 在内部自动 overlap 梯度同步与反向传播设计意图：DDP 将梯度同步封装在 loss.backward() 内部，通过延迟 All-Reduce（将多个梯度张量打包成 bucket）和 CUDA Stream 机制，自动实现梯度同步与反向传播的 overlap，对用户代码透明。
7.2 NCCL 通信优化
What — NCCL 是什么？
NCCL（NVIDIA Collective Communications Library）是 NVIDIA 提供的 GPU 专用集合通信库，支持 All-Reduce、Broadcast、All-Gather 等高效集合通信原语。相比通用的 Gloo 库，NCCL 针对 NVIDIA GPU 的 NVLink、PCIe、NIC 等硬件拓扑做了深度优化。
| 通信库 | 适用场景 | 硬件支持 | 性能 |
|---|---|---|---|
| NCCL | GPU 间通信 | NVLink, PCIe, InfiniBand | 最高 |
| Gloo | 通用 CPU/GPU | 任意网络 | 中等 |
| MPI | 高性能计算 | InfiniBand, Omni-Path | 高 |
混合并行与通信调度小结：
- 实际大模型训练系统使用数据并行+流水线并行+张量并行的 3D 混合并行
- 通信计算 Overlap 是提升 GPU 利用率的关键技术
- NCCL 是 GPU 分布式训练的事实标准通信库
八、总结与 Roadmap 预告
本篇核心要点回顾：
- 分布式训练是训练大模型的必要条件，解决显存墙和算力墙两个核心问题
- 并行策略三剑客：数据并行（横向切数据）、流水线并行（纵向切模型）、张量并行（层内切张量）
- 通信原语体系：一对多（Broadcast/Scatter）、多对一（Gather/Reduce）、多对多（All-Gather/All-Reduce）、点对点（Send/Recv）
- 数据并行核心：DDP 封装梯度同步，All-Reduce 聚合梯度，ZeRO 分片优化显存
- 流水线并行核心：按层拆分，Pipeline Bubble 是主要挑战，1F1B 调度降低显存
- 张量并行核心：单层内矩阵分片，Column Parallel + Row Parallel 组合，All-Gather/All-Reduce 拼接结果
Roadmap 预告：下篇文章我们将深入 PyTorch 分布式训练的核心 API，讲解 torch.distributed 的初始化流程、进程组创建、集合通信调用，并手把手实现一个完整的多机多卡训练脚本。敬请期待！
FAQ（常见问题速查）
以下是 20 组关于分布式训练核心要素的常见问题与解答，帮助读者快速巩固知识点。
Q1. 数据并行和模型并行的核心区别是什么？
一句话结论：数据并行每个设备有完整模型副本，模型并行每个设备只有部分模型。数据并行通过切分数据批次让多卡并行处理，通信聚合梯度；模型并行通过切分模型结构让多卡协同计算单层，通信传递激活/梯度。
Q2. 为什么 All-Reduce 是数据并行中最重要的通信原语？
一句话结论：因为它完成了梯度同步的核心操作——将所有设备的本地梯度归约为全局梯度。All-Reduce 保证所有设备在每步更新后持有相同的模型参数，这是数据并行正确性的基础。
Q3. Ring-All-Reduce 相比中心化 All-Reduce 的优势是什么？
一句话结论：去中心化设计消除了通信瓶颈。中心化 All-Reduce 需要所有 GPU 与单个根节点通信，带宽压力集中在根节点；Ring-All-Reduce 将通信分散到环中每个节点，负载均衡。
Q4. Pipeline Bubble 产生的原因是什么？
一句话结论：上下游 stage 之间的计算依赖导致设备空闲等待。下游 GPU 必须等待上游完成当前 micro-batch 的前向传播才能开始计算，这种顺序依赖产生了流水线气泡。
Q5. 1F1B 调度相比朴素流水线调度的优势是什么？
一句话结论：在保持高 GPU 利用率的同时显著降低显存峰值。朴素调度需要同时保存所有 micro-batch 的激活；1F1B 让反向传播紧跟同一 micro-batch 的前向传播，激活可以在反向后立即释放。
Q6. 张量并行为什么必须使用 All-Gather？
一句话结论：因为列并行后各 GPU 只持有输出的部分列，需要拼接成完整输出。Column Parallel 将权重矩阵按列切分，每个 GPU 计算 Y_i = X * A_i，输出维度变为 H/N，需要 All-Gather 将 [B, H/N] 的 N 个分片拼接为 [B, H]。
Q7. 为什么张量并行适合部署在 NVLink 互联的 GPU 上？
一句话结论：张量并行每层都需要多次 All-Gather/All-Reduce，对带宽要求极高。NVLink 提供 900GB/s 的单向带宽（HGX A100 8-GPU 配置），远高于 PCIe Gen4 的 32GB/s 和 InfiniBand HDR 的 50GB/s 单向带宽，能够支撑张量并行的密集通信。
Q8. ZeRO-3 和张量并行都能切分模型，区别是什么？
一句话结论：ZeRO-3 是参数分片存储（按需获取），张量并行是参数分片计算（每 GPU 计算部分）。ZeRO-3 的每个 GPU 在任何时刻只需要持有完整的参数，但可以按需从其他 GPU 获取缺失的分片；张量并行的每个 GPU 在计算时需要持有完整输入，但只计算参数的一部分，两者的通信模式和使用场景完全不同。
Q9. DDP 为什么比 DP（PyTorch DataParallel）更高效？
一句话结论：DDP 实现真正的多卡并行计算，DP 实际上是串行前反向。DP 使用单进程多线程，每个 GPU 必须等待前一个 GPU 完成前反向才能处理下一个；DDP 使用多进程，每个 GPU 独立计算，梯度同步异步进行。
Q10. Broadcast 和 Scatter 的区别是什么？
一句话结论：Broadcast 复制相同数据，Scatter 分发不同数据。Broadcast 将源节点的完整数据复制到所有目标节点；Scatter 将源节点的数据切分为 N 个分片，每个目标节点收到不同的分片。
Q11. 为什么张量并行中自注意力层不需要通信计算 Q@K^T？
一句话结论：因为每个 GPU 持有完整的 Q、K、V，可以独立计算注意力分数。Column Parallel 的 QKV 投影之后，通过 All-Gather 拼接了完整的 Q/K/V，所以各自计算 Q@K^T 是正确的，无需额外通信。
Q12. 什么是通信计算 Overlap？为什么它能提升 GPU 利用率？
一句话结论：在等待通信结果的同时执行不依赖该通信的计算，隐藏通信延迟。GPU 可以利用 CUDA Stream 同时调度通信流和计算流，当通信阻塞时切换到不依赖通信结果的计算任务。
Q13. 为什么增加 micro-batch 数量可以减少 Pipeline Bubble？
一句话结论：气泡比例 = (S-1) / (K+S-1)，K 越大气泡比例越小。流水线启动阶段有 S-1 个气泡是固定的，气泡占总时间比例随 K 增大而减小。GPipe 论文建议 K >= 4S 以保证气泡比例低于 20%。
Q14. NCCL 和 Gloo 的适用场景有什么区别？
一句话结论：NCCL 专用于 GPU 间高速通信，Gloo 适用于通用 CPU/GPU 场景。NCCL 针对 NVIDIA 硬件拓扑（NVLink、PCIe）深度优化，GPU 间通信性能远高于 Gloo；Gloo 是 PyTorch 内置的通用库，支持任意网络拓扑。
Q15. 为什么 FFN 的第一层适合列并行，第二层适合行并行？
一句话结论：因为 MLP 的形状是 up-projection（窄到宽）和 down-projection（宽到窄），组合使用列+行并行可以实现完整的两层映射。Column Parallel 将 [B, H] 投影到 [B, 4H]，Row Parallel 将 [B, 4H] 投影回 [B, H]，All-Gather/All-Reduce 嵌入其中保证正确性。
Q16. 什么是 Reduce-Scatter？它和 All-Reduce 的关系是什么？
一句话结论：Reduce-Scatter 是 All-Reduce 的前半步，先归约再分发，每个节点得到归约结果的不同部分。Ring-All-Reduce 由 Reduce-Scatter（前半步）和 All-Gather（后半步）组成。
Q17. 为什么 Interleaved 调度会增加通信次数？
一句话结论：因为每个 GPU 需要与更多其他 GPU 交换中间激活。朴素分配中 GPU 0 只与 GPU 1 通信（两层相邻）；Interleaved 分配中 GPU 0 可能需要与 GPU 1、GPU 2、GPU 3 都通信（每层被分配到不同 GPU）。
Q18. 数据并行中"同步 SGD"和"异步 SGD"的区别是什么？
一句话结论：同步 SGD 等所有 GPU 算完梯度再更新，异步 SGD 让 GPU 各自独立更新参数。同步 SGD 是 DDP 的默认行为，保证所有设备参数严格一致；异步 SGD 允许各设备参数存在短暂不一致，可能加速但会引入梯度陈旧问题。
Q19. 为什么张量并行不适合层归一化（LayerNorm）？
一句话结论：LayerNorm 需要统计输入的全局均值和方差，切分后无法正确计算。LayerNorm 的归一化涉及全局统计量（mean、variance），如果对输入进行张量并行分片，每个 GPU 只看到部分数据，计算出的统计量是错误的。
Q20. 3D 并行（数据+流水线+张量）的典型配置是什么？
一句话结论：机内 8 卡使用张量并行（TP=8），机间使用流水线并行（PP），数据并行（DP）跨所有机器。以 64 卡 A100（8 台机器，每台 8 卡）为例：TP=8（机内）、PP=4（4 stage/机）、DP=2（2 组流水线），总参数量 8*4*2=64 的并行度。
预告：下一篇我们将深入 PyTorch 分布式训练核心 API，从 torch.distributed 的初始化到 NCCL 通信原语调用，手把手实现多机多卡训练脚本。敬请期待！
浙公网安备 33010602011771号