单卡装不下的大模型：LLM 模型并行训练策略与工程选型
单卡装不下的大模型：LLM 模型并行训练策略与工程选型
一、显存墙与算力瓶颈：大模型训练的第一道坎
训练一个 70B 参数的大模型，仅模型参数就需要 140GB 显存（FP16），加上梯度、优化器状态和激活值，总显存需求超过 1TB。单张 A100 80GB 显存连参数都装不下，更不用说训练了。这就是"显存墙"——模型规模增长的速度远超单卡显存增长的速度。
即使模型能装进单卡，训练效率也是问题。70B 模型在单卡 A100 上的训练速度约为 0.1 samples/s，训练 1 万亿 Token 需要数年时间。必须通过并行策略将训练任务分布到多张卡、多台机器上，才能在可接受的时间内完成训练。
模型并行不是数据并行的替代品，而是互补关系。数据并行解决"算力不够"的问题，模型并行解决"显存不够"的问题。当模型大到单卡装不下时，必须先做模型并行，再在模型并行的基础上叠加数据并行。
二、三种并行策略：从切分维度看本质
大模型训练的并行策略可以从"切分什么"的角度分为三类：数据并行、张量并行和流水线并行。理解切分维度，才能搞清楚每种策略的通信模式和适用场景。
graph TB
subgraph "数据并行 DP"
DP_M[完整模型副本] --> DP_G0[GPU 0: 微批次 0]
DP_M --> DP_G1[GPU 1: 微批次 1]
DP_M --> DP_G2[GPU 2: 微批次 2]
DP_G0 --> |AllReduce 梯度| DP_SYNC[梯度同步]
DP_G1 --> DP_SYNC
DP_G2 --> DP_SYNC
end
subgraph "张量并行 TP"
TP_M[模型层切分] --> TP_G0[GPU 0: 列切分 Q/K/V]
TP_M --> TP_G1[GPU 1: 列切分 Q/K/V]
TP_G0 --> |AllReduce 激活| TP_SYNC[前向/反向同步]
TP_G1 --> TP_SYNC
end
subgraph "流水线并行 PP"
PP_L0[Layer 0-5: GPU 0] --> |激活传递| PP_L1[Layer 6-11: GPU 1]
PP_L1 --> |激活传递| PP_L2[Layer 12-17: GPU 2]
PP_L2 --> |激活传递| PP_L3[Layer 18-23: GPU 3]
end
style DP_SYNC fill:#e1f5fe
style TP_SYNC fill:#fff3e0
style PP_L0 fill:#e8f5e9
数据并行（DP）：每张卡持有完整模型副本，但处理不同的数据批次。反向传播后通过 AllReduce 同步梯度。通信量与模型参数量成正比。优点是实现简单，缺点是每张卡必须装下完整模型。
张量并行（TP）：将单个层的参数在维度上切分到多张卡。例如，将 Transformer 的注意力头均匀分配到 4 张卡上，每张卡计算 1/4 的注意力头，然后通过 AllReduce 合并结果。通信量与隐藏维度成正比，每层前向和反向各需要两次 AllReduce。张量并行要求高带宽互联（NVLink），跨节点性能急剧下降。
流水线并行（PP）：将模型按层切分到多张卡，每张卡负责连续的几层。数据像流水线一样依次通过各卡。核心问题是流水线气泡——前面的卡在等后面的卡完成反向传播时处于空闲状态。通过微批次（Micro-batch）调度可以减少气泡比例，但无法完全消除。
2.1 ZeRO 优化：数据并行的显存优化
ZeRO（Zero Redundancy Optimizer）不是一种新的并行策略，而是对数据并行的显存优化。它分三个阶段逐步消除数据并行中的冗余：
| 阶段 | 切分对象 | 显存节省 | 通信量 |
|---|---|---|---|
| ZeRO-1 | 优化器状态 | 4x | 与 DP 相同 |
| ZeRO-2 | 优化器状态 + 梯度 | 8x | 与 DP 相同 |
| ZeRO-3 | 优化器状态 + 梯度 + 参数 | N 倍（N=GPU 数） | 增加 1.5x |
ZeRO-3 的代价是通信量增加——每次前向计算前需要从其他卡 Gather 参数，反向传播后需要 Reduce-Scatter 梯度。在 NVLink 互联的节点内，这个通信开销可以接受；跨节点时，通信延迟可能成为瓶颈。
三、3D 并行策略的工程配置
3.1 DeepSpeed 配置示例
{
"train_batch_size": 2048,
"train_micro_batch_size_per_gpu": 4,
"gradient_accumulation_steps": 8,
"zero_optimization": {
"stage": 3,
"offload_param": {
"device": "cpu",
"pin_memory": true
},
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
"overlap_comm": true,
"contiguous_gradients": true,
"sub_group_size": 1e9
},
"tensor_parallel": {
"enabled": true,
"tp_size": 8
},
"pipeline_parallel": {
"enabled": true,
"pp_size": 4,
"pp_micro_batches": 16
},
"fp16": {
"enabled": true,
"loss_scale": 0,
"initial_scale_power": 16
},
"gradient_clipping": 1.0,
"prescale_gradients": false
}
3.2 Megatron-LM 启动脚本
#!/bin/bash
# 3D 并行配置：TP=8, PP=4, DP=8
# 总 GPU 数 = 8 * 4 * 8 = 256
GPUS_PER_NODE=8
NNODES=32
MASTER_ADDR=node01
MASTER_PORT=6000
DISTRIBUTED_ARGS="
--nproc_per_node $GPUS_PER_NODE \
--nnodes $NNODES \
--master_addr $MASTER_ADDR \
--master_port $MASTER_PORT
"
torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
--tensor-model-parallel-size 8 \
--pipeline-model-parallel-size 4 \
--num-layers 80 \
--hidden-size 8192 \
--num-attention-heads 64 \
--seq-length 4096 \
--micro-batch-size 2 \
--global-batch-size 2048 \
--train-iters 500000 \
--lr 1.5e-4 \
--min-lr 1.0e-5 \
--lr-decay-iters 400000 \
--lr-warmup-iters 2000 \
--lr-decay-style cosine \
--weight-decay 0.1 \
--adam-beta1 0.9 \
--adam-beta2 0.95 \
--clip-grad 1.0 \
--fp16 \
--bf16 \
--log-interval 100 \
--save-interval 5000 \
--eval-interval 1000
3.3 并行策略选型决策树
def choose_parallel_strategy(
model_params_b: float, # 模型参数量（十亿）
gpu_memory_gb: float, # 单卡显存（GB）
num_gpus: int, # 可用 GPU 数
inter_node_bw: str, # 节点间带宽: "nvlink" / "infiniband" / "ethernet"
) -> dict:
"""根据模型规模和硬件条件选择并行策略"""
# 估算单卡所需显存（FP16 参数 + 梯度 + 优化器 + 激活值）
mem_per_gpu_needed = model_params_b * 20 # 粗估：20 bytes/param
result = {"dp": 1, "tp": 1, "pp": 1, "zero_stage": 0}
# 第一步：判断是否需要模型并行
if mem_per_gpu_needed <= gpu_memory_gb:
# 单卡能装下，只需数据并行 + ZeRO
result["dp"] = num_gpus
result["zero_stage"] = 2 if mem_per_gpu_needed > gpu_memory_gb * 0.5 else 1
return result
# 第二步：需要张量并行（优先节点内）
if inter_node_bw == "nvlink":
# NVLink 互联，TP 可以跨卡
tp_size = min(8, int(mem_per_gpu_needed / gpu_memory_gb) + 1)
result["tp"] = tp_size
else:
# 非 NVLink，TP 限制在节点内
result["tp"] = min(8, num_gpus)
# 第三步：如果 TP 还不够，叠加流水线并行
mem_after_tp = mem_per_gpu_needed / result["tp"]
if mem_after_tp > gpu_memory_gb:
pp_size = int(mem_after_tp / gpu_memory_gb) + 1
result["pp"] = min(pp_size, num_gpus // result["tp"])
# 第四步：剩余 GPU 用于数据并行
gpus_for_dp = num_gpus // (result["tp"] * result["pp"])
result["dp"] = max(1, gpus_for_dp)
# 第五步：叠加 ZeRO 优化
if mem_after_tp / result["pp"] > gpu_memory_gb * 0.7:
result["zero_stage"] = 3 # 需要参数卸载
elif mem_after_tp / result["pp"] > gpu_memory_gb * 0.4:
result["zero_stage"] = 2
else:
result["zero_stage"] = 1
return result
四、并行训练的代价：通信开销、气泡与调试地狱
3D 并行不是免费的午餐，每一层并行都引入了通信开销和工程复杂度。
张量并行的通信瓶颈。TP 在每层 Transformer 的前向和反向传播中各需要两次 AllReduce。以 70B 模型为例，隐藏维度 8192，每次 AllReduce 通信量约 128MB（FP16）。在 NVLink 互联下（带宽 600GB/s），单次 AllReduce 约 0.2ms；跨节点 InfiniBand 下（带宽 100Gb/s），约 10ms。如果 TP 跨节点，通信开销可能占总训练时间的 30% 以上。结论：TP 必须限制在节点内，跨节点用 PP 或 DP。
流水线气泡。PP 的气泡比例约为 (PP-1) / (PP + MicroBatches - 1)。4 级 PP、16 个微批次时，气泡比例约 16%。这意味着 16% 的 GPU 时间在空转。增加微批次数可以降低气泡，但会增加显存占用（需要缓存更多激活值）。
3D 并行的调试难度。TP + PP + DP 三维组合下，一个 256 卡的训练任务可能出现：某张卡 OOM、某次 AllReduce 超时、某级流水线死锁。定位问题需要理解每张卡在 3D 网格中的位置、当前执行的微批次编号、以及通信组划分。调试成本远高于单卡训练。
检查点（Checkpoint）的复杂性。3D 并行下，模型参数分散在多张卡上，保存和加载检查点需要协调所有卡的状态。ZeRO-3 模式下，每张卡只持有部分参数，恢复训练时必须确保所有卡的参数片段一致。建议使用 DeepSpeed 的统一检查点接口，而非手动保存。
适用边界：7B 以下模型，单卡或 DP + ZeRO 即可；7B-30B 模型，TP + DP；70B 以上模型，TP + PP + DP + ZeRO-3。关键约束是节点内 NVLink 带宽——TP 的效率直接取决于互联带宽。
五、总结
LLM 模型并行训练的核心思路是：张量并行解决单层显存问题（节点内 NVLink），流水线并行解决跨层显存问题（可跨节点），数据并行叠加在模型并行之上解决算力问题，ZeRO 优化进一步压缩数据并行的显存冗余。
落地路线建议：第一步，从 DP + ZeRO-2 起步，验证训练流程和模型收敛性；第二步，当模型大到单卡装不下时，引入 TP（限制在节点内 8 卡）；第三步，当 TP 仍不够时，叠加 PP，配置微批次调度减少气泡；第四步，叠加 ZeRO-3 + CPU Offload，进一步压缩显存；第五步，建立完善的检查点和恢复机制，3D 并行下的训练中断恢复成本远高于单卡。
并行策略的选择不是"越复杂越好"，而是"刚好够用"。能用 DP + ZeRO 解决的问题，不要引入 TP；能用 TP 解决的问题，不要叠加 PP。每增加一层并行，调试和运维成本就上一个台阶。
汇聚全球AI编程工具，助力开发者即刻编程。
更多推荐
12
0- 0
已为社区贡献53条内容
所有评论(0)