分布式训练基础
1.4.1 为什么需要分布式训练?
模型规模爆炸
大模型参数规模快速增长,单卡已经无法承载完整训练过程。
2018 2019 2020 2021 2022 2023 2024
│ │ │ │ │ │ │
▼ ▼ ▼ ▼ ▼ ▼ ▼
BERT GPT-2 GPT-3 GPT-J OPT LLaMA GPT-4
1.1B 1.5B 175B 6B 175B 65B ~1.8T显存需求计算
以训练一个 175B 参数的 GPT-3 模型 为例:
| 项目 | 计算方式 | 显存需求 |
|---|---|---|
| 模型参数 | 175B × 2 字节(FP16) | 350 GB |
| 梯度 | 175B × 2 字节 | 350 GB |
| 优化器状态(Adam) | 175B × 2 × 4 字节 | 1,400 GB |
| 激活值 | 取决于 batch size 和序列长度 | ~500 GB |
| 总计 | 约 2.6 TB |
单张 A100(80GB)只能存储约:
80GB / 2.6TB ≈ 3%因此至少需要:
2.6TB / 80GB ≈ 33 张 A100结论:超大模型训练必须依赖多 GPU、多节点协同,也就是分布式训练。
1.4.2 数据并行(Data Parallelism, DP)
核心思想
每个 GPU 保存一份完整模型副本,但处理不同的数据批次。
全局 Batch: [样本1, 样本2, 样本3, 样本4, 样本5, 样本6]
│
▼ 切分
┌───────────┼───────────┐
▼ ▼ ▼
GPU 0 GPU 1 GPU 2
[样本1,2] [样本3,4] [样本5,6]
模型 A 模型 A 模型 A
│ │ │
└───────────┼───────────┘
▼
AllReduce 梯度平均
▼
所有 GPU 同步更新模型参数伪代码示意
def data_parallel_training():
# 1. 每个 GPU 加载相同的模型副本
model = load_model().to(local_gpu)
# 2. 数据分发(Scatter)
local_batch = global_batch[rank::world_size]
# 3. 前向传播
loss = model(local_batch)
# 4. 反向传播
loss.backward()
# 5. 梯度同步(AllReduce)
all_reduce(gradients, op=AVERAGE)
# 6. 参数更新
optimizer.step()优缺点
| 优点 | 缺点 |
|---|---|
| 实现简单,易于理解 | 每个 GPU 都需要存储完整模型 |
| 理想情况下加速比接近线性 | 模型大小受限于单卡显存 |
| 适合中小模型和大 batch 训练 | 通信开销随 GPU 数量增加 |
1.4.3 模型并行(Model Parallelism, MP)
核心思想
将模型切分到不同 GPU,每个 GPU 只存储部分参数。
原始模型(4 层 Transformer)
┌─────────┬─────────┬─────────┬─────────┐
│ Layer 1 │ Layer 2 │ Layer 3 │ Layer 4 │
└─────────┴─────────┴─────────┴─────────┘
│ │ │ │
▼ ▼ ▼ ▼
GPU 0 GPU 1 GPU 2 GPU 3
Layer 1 Layer 2 Layer 3 Layer 4
前向传播:GPU 0 → GPU 1 → GPU 2 → GPU 3
反向传播:GPU 3 → GPU 2 → GPU 1 → GPU 0伪代码示意
def model_parallel_forward(x, rank):
if rank == 0:
x = layer1(x)
send(x, dest=1)
return None
elif rank == 1:
x = recv(src=0)
x = layer2(x)
send(x, dest=2)
return None
elif rank == 2:
x = recv(src=1)
x = layer3(x)
send(x, dest=3)
return None
elif rank == 3:
x = recv(src=2)
x = layer4(x)
return x优缺点
| 优点 | 缺点 |
|---|---|
| 可以训练超大模型 | GPU 利用率可能较低 |
| 不受单卡显存限制 | |
| 适合模型层数多、参数量大的场景 | 通信频繁,延迟敏感,实现复杂 |
注:流水线气泡(Pipeline Bubble,也叫 bubble / stall cycle):在流水线里人为插入一个“什么都不做(不写结果、不改变状态)”的空拍/空槽,让后面已经发射的指令原地多停一拍或多拍,等前面的依赖、资源冲突或控制风险解除后再继续。
你可以把它想成在传送带(流水线)上放了一段“假包裹”——它占着一个工位的时间,但不产生有效工作,于是整体吞吐下降。
流水线要求:同一时刻不同 stage 在处理不同指令。但当出现下面情况之一时,后续指令“现在不能安全往下走”:
- 数据冒险(Data Hazard):前一条指令还没写完目标寄存器/内存,后一条就要读它(RAW)。
- 结构冲突(Structural Hazard):某个争用。
- 控制冒险(Control Hazard):分支/跳转还没解析出来,不知道下一条该取哪条指令。
解决手段通常有三类:
- 转发/旁路(forwarding/bypass):把结果“抄近路”直接送到后面读口 → 减少气泡
- 乱序执行/重命名(OoO / rename):从根本上消解很多伪依赖 → 更少气泡
- 实在等不了:就只能 stall → 插入气泡,让后续指令卡住几拍
气泡本质上就是 stall 的一种实现形式:不是取消指令,而是让它“占着流水级但不推进有效结果”。
1.4.4 流水线并行(Pipeline Parallelism, PP)
核心思想
流水线并行结合数据并行和模型并行,将数据切成多个 micro-batch,让不同 GPU 尽量持续工作,减少空闲时间。
朴素模型并行(大量空闲):
GPU 0: [Layer1]████░░░░░░░░░░░░
GPU 1: ░░░░[Layer2]████░░░░░░░░
GPU 2: ░░░░░░░░[Layer3]████░░░░
GPU 3: ░░░░░░░░░░░░[Layer4]████
流水线并行(micro-batch 减少空闲):
GPU 0: [L1:m1][L1:m2][L1:m3][L1:m4]
GPU 1: ░░░░[L2:m1][L2:m2][L2:m3][L2:m4]
GPU 2: ░░░░░░░░[L3:m1][L3:m2][L3:m3][L3:m4]
GPU 3: ░░░░░░░░░░░░[L4:m1][L4:m2][L4:m3][L4:m4]GPipe vs PipeDream
| 特性 | GPipe | PipeDream |
|---|---|---|
| 更新方式 | 同步,等待所有 micro-batch | 异步,立即更新 |
| 内存开销 | 高,需要存储多个激活值 | 低 |
| 实现复杂度 | 相对简单 | 更复杂 |
| 收敛稳定性 | 好 | 需要额外处理 |
1.4.5 张量并行(Tensor Parallelism, TP)
核心思想
张量并行将单层内部的矩阵计算拆分到多个 GPU。
以 Linear 层为例:
Y = X × W
X: [batch, in_dim]
W: [in_dim, out_dim]按列拆分(Column Parallel)
W = [W1 | W2]
GPU 0: Y1 = X × W1
GPU 1: Y2 = X × W2
最终结果:Y = [Y1 | Y2]按行拆分(Row Parallel)
W = [W1]
[W2]
GPU 0: Y1 = X1 × W1
GPU 1: Y2 = X2 × W2
最终结果:Y = Y1 + Y2,需要 AllReduceMegatron-LM 的张量并行策略
Megatron-LM(更准确地说今天的 Megatron-Core)是 NVIDIA 开源的大规模 Transformer 训练框架/库,专门解决一个核心问题:
💡 单个 GPU 的显存根本塞不下百亿/千亿参数模型,怎么办?
它通过多层次并行拆分(切权重、切层、切数据、切序列),让你能在成百上千张 GPU上高效训练 GPT、LLaMA、DeepSeek、Qwen 这类超大模型。
在 Transformer 的 MLP 层中,Megatron-LM 通常采用:
- 第一层矩阵按列并行;
- 第二层矩阵按行并行;——这样第一层的结果就不用 Reduce 正好是第二层的输入
- 在必要位置进行通信同步。
1.4.6 3D 并行组合
核心思想
将 数据并行(DP)、流水线并行(PP)、张量并行(TP) 组合使用。
例如:
假设有 32 个 GPU,训练一个超大模型:
DP = 2
PP = 4
TP = 4
总 GPU 数 = 2 × 4 × 4 = 32Config 卡:
**DP=2 × PP=4 × TP=2 = 16 GPU/replica × 2 replicas = 32 GPU**(注:上述的 32 = 2×4×4 配置是错的,这里用更标准的 2×4×2=16 per replica × 2 replicas = 32 GPU;)
通信模式
| 并行维度 | 通信范围 | 常用网络 |
|---|---|---|
| TP 组内 | 同一层内部张量切分通信 | NVLink,最快,900GB/s |
| PP 阶段间 | 相邻流水线阶段传输激活值 | NVLink / PCIe |
| DP 副本间 | 梯度同步 | InfiniBand / RoCE |
并行策略选择指南
| 场景 | 推荐配置 | 说明 |
|---|---|---|
| 小模型(<10B) | DP only | 简单高效 |
| 中等模型(10–100B) | DP + TP | 单节点内 TP,跨节点 DP |
| 大模型(100B–1T) | DP + PP + TP | 3D 并行 |
| 超大模型(>1T) | ZeRO + 3D | 结合显存优化 |
1.4.7 专家并行(Expert Parallelism)
Expert Parallelism(专家并行) 是一种针对 MoE(Mixture of Experts,混合专家) 架构模型(如 DeepSeek-V3、Mixtral)的分布式训练/推理策略。
核心思想
MoE 模型中有多个“专家”(Expert)子网络,每次计算只激活其中的一部分。Expert Parallelism 将不同的专家分配到不同的 GPU 上,每个 GPU 只负责计算自己拥有的那部分专家,并通过 All-to-All 通信将 token 路由到对应的专家所在 GPU。
配置示例
集群规模:96 × NVIDIA H100 GPUs
模型:DeepSeek-V3 (671B参数,MoE架构)
并行策略:EP + DP + TP + PD分离
- Expert Parallelism (EP): 8
- Data Parallelism (DP): 4
- Tensor Parallelism (TP): 3
- Prefill-Decode 分离:Prefill:Decode = 1:2配置中 Expert Parallelism: 8表示:
- 模型的总专家被分成 8 组,每组放在一个 GPU 或一组 GPU 上(结合 TP 使用)。
- 当 token 经过门控(Gate)决定要去哪些专家时,会通过 All-to-All 通信将 token 发送到对应的 GPU 上进行计算。
在 DeepSeek-V3 中的具体效果 :
配置 EP=8, DP=4, TP=3, PD 分离表明:
- 总共 96 张 H100,按
8×4×3 = 96正好分配。 - Prefill 和 Decode 节点比例为 1:2,即 32 张做 Prefill,64 张做 Decode。
- 在每个节点内部,通过 EP 将 8 个专家分布在 8 个 GPU 组上,配合 TP 进一步切分单个专家的参数,DP 复制多份数据并行训练/推理。
与其他并行维度的关系
| 并行维度 | 作用 |
|---|---|
| EP (Expert Parallelism) | 拆分专家,每个 GPU 负责部分专家 |
| DP (Data Parallelism) | 复制整个模型(含所有专家),处理不同数据批次 |
| TP (Tensor Parallelism) | 拆分单个 Transformer 层的参数(如注意力头的切分) |
| PP (Pipeline Parallelism) | 按层切分,不同 GPU 负责不同层(此处未列出) |
为什么需要 EP?
- MoE 模型参数量巨大(如 671B),但每次只激活少量专家(稀疏激活)。如果所有专家都放在同一块 GPU 上,显存装不下。
- EP 通过将专家分散到多卡,使得每张卡只存储一部分专家,从而突破单卡显存限制。
- 同时,EP 让计算负载更均衡:不同 token 可能激活不同专家,通过 All-to-All 通信实现动态路由。
典型通信模式
- 前向:每个 GPU 上的 token 需要发送给目标专家所在的 GPU → All-to-All 通信。
- 反向:梯度需要按同样的路由反向传播回来 → 再次 All-to-All。
Expert Parallelism 是专门为 MoE 模型设计的并行策略,通过将专家分散到不同 GPU 上来解决超大模型显存不足的问题,同时利用稀疏激活特性减少不必要的计算。
小结
分布式训练的核心目标是:把模型、数据和计算合理拆分到多张 GPU / 多个节点上,并尽可能降低通信和空闲开销。
| 并行方式 | 拆分对象 | 主要解决的问题 |
|---|---|---|
| 数据并行 DP | 数据 batch | 提升吞吐量 |
| 模型并行 MP | 模型层 | 单卡放不下完整模型 |
| 流水线并行 PP | 模型阶段 + micro-batch | 降低模型并行空闲时间 |
| 张量并行 TP | 单层矩阵 / 张量 | 单层参数或计算过大 |
| 3D 并行 | DP + PP + TP | 超大模型训练 |
实际大模型训练往往不是选择某一种并行方式,而是组合使用多种并行策略。