单卡装不下的大模型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 优化数据并行的显存优化ZeROZero Redundancy Optimizer不是一种新的并行策略而是对数据并行的显存优化。它分三个阶段逐步消除数据并行中的冗余阶段切分对象显存节省通信量ZeRO-1优化器状态4x与 DP 相同ZeRO-2优化器状态 梯度8x与 DP 相同ZeRO-3优化器状态 梯度 参数N 倍NGPU 数增加 1.5xZeRO-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 并行配置TP8, PP4, DP8 # 总 GPU 数 8 * 4 * 8 256 GPUS_PER_NODE8 NNODES32 MASTER_ADDRnode01 MASTER_PORT6000 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 10003.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: # 非 NVLinkTP 限制在节点内 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 通信量约 128MBFP16。在 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 DP70B 以上模型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。每增加一层并行调试和运维成本就上一个台阶。