从单卡到千卡集群理解 DDP 数据并行的本质搞懂 AllReduce 通信原语 看清模型并行TP / PP / ZeRO何时登场。附完整 PyTorch 实现。一、为什么需要分布式训练⏱ 速度瓶颈单卡训练 GPT-3175B需要 355 年。8 卡 DDP 能将 epoch 时间线性缩短至约 1/N。 显存瓶颈单张 A10080GB放不下一个 70B 模型FP16 约 140GB。必须把模型切分到多卡。 数据规模训练数据动辄 TB 级单卡 DataLoader 成为 I/O 瓶颈分布式数据加载是必需项。核心权衡数据并行解决速度问题模型能放进单卡 模型并行解决显存问题模型放不进单卡。 实际训练中两者常混合使用——即3D 并行DP TP PP。二、数据并行DDP核心原理Distributed Data ParallelDDP是 PyTorch 官方推荐的数据并行方案 替代了旧的DataParallelDP。每个 GPU 拥有一份完整的模型副本 但处理不同的 mini-batch 数据切片。Batch (N×B)GPU 0 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₀GPU 1 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₁GPU 2 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₂GPU 3 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₃AllReduceSum / Mean∇ Σ∇ᵢ/4统一梯度 各 GPU 用汇总梯度同步更新参数 → 所有副本保持一致 DDP 数据并行 AllReduce 梯度同步 参数一致性DDP vs DPDataParallel特性DP旧DDP推荐通信方式主卡 Gather → Scatter有瓶颈AllReduce 环状通信无单点瓶颈负载均衡主卡承担额外 scatter/gather 开销所有卡对等负载均匀多进程模型单进程多线程GIL 瓶颈每 GPU 一个独立进程无 GIL性能显著慢于 DDP接近线性加速比使用方式model nn.DataParallel(model)需初始化进程组 包装模型三、AllReduce 通信原语AllReduce 是所有分布式训练框架的核心通信操作将每个进程的本地张量归约reduce后 把结果广播broadcast回所有进程。最常用的实现是Ring AllReduce。Ring AllReduce 算法两步法Ring AllReduce 复杂度通信量 2(N-1)/N × 数据量。 当 N 很大时趋近2×数据量与 GPU 数量无关——这是它能线性扩展的关键。 NCCL 库NVIDIA Collective Communications Library高度优化了这一过程。常用集合通信原语对比原语操作典型用途Broadcast一个进程的数据复制到所有进程分发初始模型权重Reduce所有进程的数据归约到一个进程汇总 loss / 指标到 rank 0AllReduce归约 广播回所有进程DDP 梯度同步最核心AllGather所有进程的数据拼接后广播Ring AllReduce 第二阶段ReduceScatter归约 每进程只接收一个分片Ring AllReduce 第一阶段四、PyTorch DDP 完整代码实现基础DDP#!/usr/bin/env python3 DDP 最小可运行示例 — 数据并行训练 ResNet 风格模型 import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # ═══════════════ Step 1: 初始化进程组 ═══════════════ def setup(rank, world_size): 每个进程调用一次建立通信组 dist.init_process_group( backendnccl, # NVIDIA NCCL 后端GPU init_methodenv://, # 从环境变量读取地址 rankrank, world_sizeworld_size, ) torch.cuda.set_device(rank) # 绑定当前进程到指定 GPU def cleanup(): dist.destroy_process_group() # ═══════════════ Step 2: 准备数据含 DistributedSampler ═══════════════ class DummyDataset(Dataset): def __len__(self): return 10000 def __getitem__(self, idx): return torch.randn(3, 224, 224), torch.randint(0, 10, (1,)) def get_dataloader(rank, world_size, batch_size): dataset DummyDataset() sampler DistributedSampler( dataset, num_replicasworld_size, # 总共 N 个 GPU rankrank, # 当前 GPU 编号 shuffleTrue, drop_lastTrue, ) return DataLoader(dataset, batch_sizebatch_size, samplersampler) # ═══════════════ Step 3: 训练主函数 ═══════════════ def train(rank, world_size, epochs5): setup(rank, world_size) # 模型 → GPU → DDP 包装 model nn.Sequential( nn.Conv2d(3, 64, 3, padding1), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(64, 10), ).to(rank) ddp_model DDP(model, device_ids[rank]) dataloader get_dataloader(rank, world_size, batch_size64) criterion nn.CrossEntropyLoss() optimizer optim.SGD(ddp_model.parameters(), lr0.01) for epoch in range(epochs): dataloader.sampler.set_epoch(epoch) # 关键每 epoch 重新 shuffle total_loss 0.0 for batch_idx, (data, target) in enumerate(dataloader): data, target data.to(rank), target.to(rank).squeeze() optimizer.zero_grad() output ddp_model(data) loss criterion(output, target) loss.backward() # DDP 自动触发 AllReduce 同步梯度 optimizer.step() total_loss loss.item() if rank 0: # 仅 rank 0 打印日志 print(fEpoch {epoch} | Loss: {total_loss / len(dataloader):.4f}) cleanup()torchrun启动使用 torchrun 启动 DDP 训练推荐方式 不需要手动设置 rank/world_sizetorchrun 自动注入环境变量 import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def main(): # torchrun 自动设置这些环境变量 local_rank int(os.environ[LOCAL_RANK]) world_size int(os.environ[WORLD_SIZE]) rank int(os.environ[RANK]) dist.init_process_group(backendnccl) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) ddp_model DDP(model, device_ids[local_rank]) # ... 训练循环同上 ... dist.destroy_process_group() if __name__ __main__: main() Shell # 单机 4 卡启动 torchrun --nproc_per_node4 train_ddp.py # 多机多卡2 节点 × 4 GPUs # 在 node0 上 torchrun --nnodes2 --nproc_per_node4 \ --node_rank0 --master_addr192.168.1.10 --master_port29500 \ train_ddp.py # 在 node1 上 torchrun --nnodes2 --nproc_per_node4 \ --node_rank1 --master_addr192.168.1.10 --master_port29500 \ train_ddp.py混合精度DDP混合精度训练 DDPFP16 前向/反向 FP32 主权重 显存减半速度提升利用 Tensor Core import torch from torch.cuda.amp import autocast, GradScaler from torch.nn.parallel import DistributedDataParallel as DDP def train_amp_ddp(rank, world_size): # ... 初始化进程组同上 ... model MyModel().to(rank) ddp_model DDP(model, device_ids[rank]) scaler GradScaler() # 梯度缩放器防止 FP16 下溢 optimizer torch.optim.AdamW(ddp_model.parameters(), lr1e-3) for data, target in dataloader: data, target data.to(rank), target.to(rank) optimizer.zero_grad() with autocast(): # 前向自动转为 FP16 output ddp_model(data) loss criterion(output, target) scaler.scale(loss).backward() # 反向loss × scale → 避免下溢 scaler.step(optimizer) # 更新unscale 梯度 → 正常 step scaler.update() # 动态调整 scale 因子 # ⚠️ DDP AMP 时 backward 中的 AllReduce 在 FP16 精度完成 # 梯度通信量减半进一步加速断电续训DDP 断点续训 / 模型保存的最佳实践 import torch from torch.nn.parallel import DistributedDataParallel as DDP def save_checkpoint(ddp_model, optimizer, epoch, path, rank): 仅 rank 0 保存其他进程跳过 if rank ! 0: return checkpoint { model_state_dict: ddp_model.module.state_dict(), # ↑ 注意 .module 去掉 DDP 包装前缀 optimizer_state_dict: optimizer.state_dict(), epoch: epoch, } torch.save(checkpoint, path) def load_checkpoint(model, optimizer, path, rank): 所有 rank 加载同一份 checkpoint checkpoint torch.load(path, map_locationfcuda:{rank}) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) return checkpoint[epoch] # ── 使用示例 ── ddp_model DDP(model, device_ids[rank]) # 保存 save_checkpoint(ddp_model, optimizer, epoch, checkpoint.pt, rank) # 加载在所有进程上调用 start_epoch load_checkpoint(ddp_model, optimizer, checkpoint.pt, rank)关键细节DistributedSampler确保每个 GPU 看到互不重叠的数据分片。sampler.set_epoch(epoch)必须在每个 epoch 调用否则每个 epoch 的 shuffle 结果相同失去随机性。 DDP 在loss.backward()时自动进行梯度 AllReduce——无需手动调用通信操作。常见坑点⚠️ sampler.set_epoch()忘记调用会导致每个 epoch 的 shuffle 顺序完全相同等效于没有 shuffle模型泛化能力下降。⚠️ BatchNorm 同步DDP 默认不同步 BN 统计量。跨卡 BatchNorm 需额外使用SyncBatchNorm小 batch 时尤其重要。⚠️ 随机种子各进程需设置不同种子避免生成相同数据。通常seed base_seed rank。⚠️ loss.backward() 同步DDP 在backward()时进行 AllReduce。如果某进程跳过 backward如条件分支会导致死锁。五、当模型大到放不进单卡模型并行当模型参数量超过单卡显存如 Llama-70B 需要约 140GB FP16而 A100 只有 80GB 数据并行不够用必须将模型本身切分到多张 GPU —— 这就是模型并行。三种主流方案三种模型并行详细对比方案切分粒度通信模式适用场景代表框架Tensor Parallel (TP)单层内矩阵按列/行切分每步 AllReduce通信密集单机多卡NVLink 高带宽Megatron-LMPipeline Parallel (PP)按层切分到不同 GPU仅层边界通信稀疏跨节点低带宽容忍GPipe, PipeDreamZeRO / FSDP优化器状态 / 梯度 / 参数分片等同于 DP 通信量通用性价比最高DeepSpeed, PyTorch FSDP模型并行代码实现DeepSpeed ZeROPipeline Parallel — 按层切分模型到不同 GPU 核心GPU0 持有 Layer 1~8GPU1 持有 Layer 9~16 通过微批次micro-batch流水线掩盖通信延迟 关键概念 - micro-batch将一个大 batch 切分为 M 个 micro-batch - 1F1B 调度一个 forward 一个 backward 交替执行 - bubble流水线空泡GPU 空闲等待的时间比例 (P-1)/M import torch import torch.nn as nn import torch.distributed as dist from queue import Queue # ═══════════════ 原生 PP 实现理解原理 ═══════════════ class PipelineStage(nn.Module): 每个 GPU 上的流水线阶段 def __init__(self, layers, rank, prev_rank, next_rank): super().__init__() self.layers nn.Sequential(*layers) self.rank rank self.prev prev_rank self.next next_rank def forward(self, xNone): if self.prev is not None: # 从上游 GPU 接收激活值 x torch.empty(self.input_shape, deviceself.rank) dist.recv(x, srcself.prev) out self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dstself.next) return out # ═══════════════ 生产环境torch.distributed.pipelining ═══════════════ from torch.distributed.pipelining import ( pipeline, ScheduleGPipe, Schedule1F1B, PipelineStage, SplitPoint, ) def train_pp(): # 1. 定义模型 切分点 model MyBigTransformer() split_spec { model.layers.4: SplitPoint.END, # Layer 0~3 → GPU0 model.layers.8: SplitPoint.END, # Layer 4~7 → GPU1 } # 2. 构建流水线 stages pipeline( modulemodel, num_chunks8, # micro-batch 数量越多 bubble 越小 example_args(example_input,), split_specsplit_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward再全部 backwardbubble 大 # Schedule1F1B: 一次 forward 接一次 backward 交替bubble 小推荐 schedule Schedule1F1B(stages, num_chunks8) for batch in dataloader: # 将 batch 拆成 8 个 micro-batch流水线执行 loss schedule.step(batch) # ═══════════════ Bubble 公式 ═══════════════ # bubble 比例 (P - 1) / M其中 P 流水线阶段数M micro-batch 数量 # 例4 阶段 × 32 micro-batch → bubble 3/32 ≈ 9.4% # 例4 阶段 × 4 micro-batch → bubble 3/4 75%几乎不可用 # ═══════════════ DeepSpeed Pipeline Parallel ═══════════════ # DeepSpeed 也提供 PP 支持通过配置即可启用 # pipeline: { stages: auto, partition: type:transformer } GPU 拓扑示意# 4 GPU 流水线示例8 micro-batches # # GPU0: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU1: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU2: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU3: ▓▓▓▓▓▓▓ F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 ▓▓▓▓▓▓▓ # ↑ bubble ↑ bubble # # F forward, B backward, ▓ bubble (GPU 空闲) # 1F1B 调度GPU0 做完 F1 立即做 B0不需等所有 F 完成Pytorch FSDPPyTorch FSDP — 原生 ZeRO-3 等效实现 (PyTorch ≥ 2.0) 无需安装第三方库API 与 DDP 风格一致 import torch import torch.distributed as dist from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, CPUOffload, BackwardPrefetch, ) from torch.distributed.fsdp.wrap import ( transformer_auto_wrap_policy, size_based_auto_wrap_policy, ) import functools def train_fsdp(rank, world_size): dist.init_process_group(backendnccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) # ═══════════════ 1. FSDP 混合精度配置 ═══════════════ mp_policy MixedPrecision( param_dtypetorch.bfloat16, # 参数存储精度 reduce_dtypetorch.bfloat16, # 梯度 AllReduce 精度 buffer_dtypetorch.bfloat16, # buffer 精度 ) # ═══════════════ 2. Transformer 自动包装策略 ═══════════════ from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf) auto_wrap_policy functools.partial( transformer_auto_wrap_policy, transformer_layer_cls{ type(model.model.layers[0]), # 按 Transformer Block 粒度包装 }, ) # ═══════════════ 3. FSDP 包装模型 ═══════════════ model FSDP( model.to(rank), sharding_strategyShardingStrategy.FULL_SHARD, # ZeRO-3 等效 mixed_precisionmp_policy, auto_wrap_policyauto_wrap_policy, cpu_offloadCPUOffload(offload_paramsTrue), # 可选参数 offload 到 CPU backward_prefetchBackwardPrefetch.BACKWARD_PRE, # 预取下一层参数 device_idrank, ) optimizer torch.optim.AdamW(model.parameters(), lr1e-5) for batch in dataloader: optimizer.zero_grad() loss model(batch).loss loss.backward() optimizer.step() dist.destroy_process_group() # ═══════════════ FSDP 分片策略选择 ═══════════════ # ShardingStrategy.FULL_SHARD → ZeRO-3 等效参数梯度优化器全分片 # ShardingStrategy.SHARD_GRAD_OP → ZeRO-2 等效分片梯度优化器 # ShardingStrategy.NO_SHARD → DDP 等效不分片 # ShardingStrategy.HYBRID_SHARD → 节点内 FULL_SHARD 节点间 NO_SHARDTensor ParallelTensor Parallel — 单层内矩阵切分 核心思想将 Linear 层的权重矩阵按列或行切分到多 GPU 每步需 AllReduce / AllGather 通信适合 NVLink 高带宽场景 原生实现理解原理用→ 生产环境用 Megatron-LM 或 torch.distributed.tensor import torch import torch.nn as nn import torch.distributed as dist from torch.distributed import all_reduce, ReduceOp # ═══════════════ 列并行 LinearColumn Parallel ═══════════════ class ColumnParallelLinear(nn.Module): W: (in_features, out_features) → 按列切到 N 个 GPU 每 GPU 持有 W[:, i*C : (i1)*C]输出需 AllGather 拼接 def __init__(self, in_features, out_features): super().__init__() self.rank dist.get_rank() self.world_size dist.get_world_size() assert out_features % self.world_size 0 self.out_per_gpu out_features // self.world_size self.weight nn.Parameter(torch.empty(in_features, self.out_per_gpu)) nn.init.kaiming_uniform_(self.weight) def forward(self, x): # 每 GPU 计算自己那列 local_out x self.weight # (batch, out_per_gpu) # AllGather 拼接完整输出 gathered [torch.empty_like(local_out) for _ in range(self.world_size)] dist.all_gather(gathered, local_out) return torch.cat(gathered, dim-1) # ═══════════════ 行并行 LinearRow Parallel ═══════════════ class RowParallelLinear(nn.Module): W: (in_features, out_features) → 按行切到 N 个 GPU 每 GPU 持有 W[i*R:(i1)*R, :]输入需提前分片输出 AllReduce 求和 def __init__(self, in_features, out_features): super().__init__() self.rank dist.get_rank() self.world_size dist.get_world_size() assert in_features % self.world_size 0 self.in_per_gpu in_features // self.world_size self.weight nn.Parameter(torch.empty(self.in_per_gpu, out_features)) nn.init.kaiming_uniform_(self.weight) def forward(self, x): # x 已经在前一步被 AllGather 得到完整输入但这里只取自己那行 start self.rank * self.in_per_gpu local_x x[:, start : start self.in_per_gpu] local_out local_x self.weight # AllReduce 求和得到完整输出 dist.all_reduce(local_out, opReduceOp.SUM) return local_out # ═══════════════ 标准 TP MLP Block 组合 ═══════════════ class TP_MLP(nn.Module): ColumnParallel(d → 4d) → GELU → RowParallel(4d → d) → Dropout 列并行输出已 AllGather → 行并行输入直接用 → 输出 AllReduce 求和 一个完整 TP block 只有 1 次 AllReduceRowParallel 末尾 def __init__(self, d_model, d_ff): super().__init__() self.fc1 ColumnParallelLinear(d_model, d_ff) self.fc2 RowParallelLinear(d_ff, d_model) self.act nn.GELU() self.dropout nn.Dropout(0.1) def forward(self, x): return self.dropout(self.fc2(self.act(self.fc1(x)))) # ═══════════════ 生产环境torch.distributed.tensor.parallel ═══════════════ # PyTorch 2.0 内置 TP基于 DTensor一行代码切分模型 from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel # parallelize_module(model, tp_mesh, { # fc1: ColwiseParallel(), # 列并行 # fc2: RowwiseParallel(), # 行并行自动插入 AllReduce # })Pipeline ParallelPipeline Parallel — 按层切分模型到不同 GPU 核心GPU0 持有 Layer 1~8GPU1 持有 Layer 9~16 通过微批次micro-batch流水线掩盖通信延迟 关键概念 - micro-batch将一个大 batch 切分为 M 个 micro-batch - 1F1B 调度一个 forward 一个 backward 交替执行 - bubble流水线空泡GPU 空闲等待的时间比例 (P-1)/M import torch import torch.nn as nn import torch.distributed as dist from queue import Queue # ═══════════════ 原生 PP 实现理解原理 ═══════════════ class PipelineStage(nn.Module): 每个 GPU 上的流水线阶段 def __init__(self, layers, rank, prev_rank, next_rank): super().__init__() self.layers nn.Sequential(*layers) self.rank rank self.prev prev_rank self.next next_rank def forward(self, xNone): if self.prev is not None: # 从上游 GPU 接收激活值 x torch.empty(self.input_shape, deviceself.rank) dist.recv(x, srcself.prev) out self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dstself.next) return out # ═══════════════ 生产环境torch.distributed.pipelining ═══════════════ from torch.distributed.pipelining import ( pipeline, ScheduleGPipe, Schedule1F1B, PipelineStage, SplitPoint, ) def train_pp(): # 1. 定义模型 切分点 model MyBigTransformer() split_spec { model.layers.4: SplitPoint.END, # Layer 0~3 → GPU0 model.layers.8: SplitPoint.END, # Layer 4~7 → GPU1 } # 2. 构建流水线 stages pipeline( modulemodel, num_chunks8, # micro-batch 数量越多 bubble 越小 example_args(example_input,), split_specsplit_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward再全部 backwardbubble 大 # Schedule1F1B: 一次 forward 接一次 backward 交替bubble 小推荐 schedule Schedule1F1B(stages, num_chunks8) for batch in dataloader: # 将 batch 拆成 8 个 micro-batch流水线执行 loss schedule.step(batch) # ═══════════════ Bubble 公式 ═══════════════ # bubble 比例 (P - 1) / M其中 P 流水线阶段数M micro-batch 数量 # 例4 阶段 × 32 micro-batch → bubble 3/32 ≈ 9.4% # 例4 阶段 × 4 micro-batch → bubble 3/4 75%几乎不可用 # ═══════════════ DeepSpeed Pipeline Parallel ═══════════════ # DeepSpeed 也提供 PP 支持通过配置即可启用 # pipeline: { stages: auto, partition: type:transformer } # 4 GPU 流水线示例8 micro-batches # # GPU0: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU1: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU2: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU3: ▓▓▓▓▓▓▓ F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 ▓▓▓▓▓▓▓ # ↑ bubble ↑ bubble # # F forward, B backward, ▓ bubble (GPU 空闲) # 1F1B 调度GPU0 做完 F1 立即做 B0不需等所有 F 完成六、实战选型决策树终极对比所有方案一览方案显存效率计算效率通信开销实现难度适用规模单卡❌ 最低❌ 最慢—⭐ 1BDPDataParallel❌ 同单卡⚠️ 有限⚠️ 主卡瓶颈⭐不推荐DDP❌ 同单卡✅ ~线性✅ AllReduce⭐⭐ 10BZeRO-1⚠️ 4×✅ 同 DDP✅ 同 DDP⭐⭐1B~10BZeRO-2✅ 8×✅ 同 DDP✅ 同 DDP⭐⭐1B~30BZeRO-3 / FSDP✅ N×⚠️ 略降⚠️ 1.5×⭐⭐⭐10B~100BTP⚠️ 有限⚠️ 通信密集❌ 高⭐⭐⭐⭐单机大模型PP✅ 好⚠️ 有空泡✅ 低⭐⭐⭐⭐跨节点大模型3D 并行✅ 最优⚠️ 复杂调优⚠️ 混合⭐⭐⭐⭐⭐100B~1T 分布式训练深度解析
从单卡到千卡集群理解 DDP 数据并行的本质搞懂 AllReduce 通信原语 看清模型并行TP / PP / ZeRO何时登场。附完整 PyTorch 实现。一、为什么需要分布式训练⏱ 速度瓶颈单卡训练 GPT-3175B需要 355 年。8 卡 DDP 能将 epoch 时间线性缩短至约 1/N。 显存瓶颈单张 A10080GB放不下一个 70B 模型FP16 约 140GB。必须把模型切分到多卡。 数据规模训练数据动辄 TB 级单卡 DataLoader 成为 I/O 瓶颈分布式数据加载是必需项。核心权衡数据并行解决速度问题模型能放进单卡 模型并行解决显存问题模型放不进单卡。 实际训练中两者常混合使用——即3D 并行DP TP PP。二、数据并行DDP核心原理Distributed Data ParallelDDP是 PyTorch 官方推荐的数据并行方案 替代了旧的DataParallelDP。每个 GPU 拥有一份完整的模型副本 但处理不同的 mini-batch 数据切片。Batch (N×B)GPU 0 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₀GPU 1 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₁GPU 2 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₂GPU 3 — 模型副本B/4 样本 · 前向反向 · 本地梯度 ∇₃AllReduceSum / Mean∇ Σ∇ᵢ/4统一梯度 各 GPU 用汇总梯度同步更新参数 → 所有副本保持一致 DDP 数据并行 AllReduce 梯度同步 参数一致性DDP vs DPDataParallel特性DP旧DDP推荐通信方式主卡 Gather → Scatter有瓶颈AllReduce 环状通信无单点瓶颈负载均衡主卡承担额外 scatter/gather 开销所有卡对等负载均匀多进程模型单进程多线程GIL 瓶颈每 GPU 一个独立进程无 GIL性能显著慢于 DDP接近线性加速比使用方式model nn.DataParallel(model)需初始化进程组 包装模型三、AllReduce 通信原语AllReduce 是所有分布式训练框架的核心通信操作将每个进程的本地张量归约reduce后 把结果广播broadcast回所有进程。最常用的实现是Ring AllReduce。Ring AllReduce 算法两步法Ring AllReduce 复杂度通信量 2(N-1)/N × 数据量。 当 N 很大时趋近2×数据量与 GPU 数量无关——这是它能线性扩展的关键。 NCCL 库NVIDIA Collective Communications Library高度优化了这一过程。常用集合通信原语对比原语操作典型用途Broadcast一个进程的数据复制到所有进程分发初始模型权重Reduce所有进程的数据归约到一个进程汇总 loss / 指标到 rank 0AllReduce归约 广播回所有进程DDP 梯度同步最核心AllGather所有进程的数据拼接后广播Ring AllReduce 第二阶段ReduceScatter归约 每进程只接收一个分片Ring AllReduce 第一阶段四、PyTorch DDP 完整代码实现基础DDP#!/usr/bin/env python3 DDP 最小可运行示例 — 数据并行训练 ResNet 风格模型 import os import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, Dataset from torch.utils.data.distributed import DistributedSampler import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP # ═══════════════ Step 1: 初始化进程组 ═══════════════ def setup(rank, world_size): 每个进程调用一次建立通信组 dist.init_process_group( backendnccl, # NVIDIA NCCL 后端GPU init_methodenv://, # 从环境变量读取地址 rankrank, world_sizeworld_size, ) torch.cuda.set_device(rank) # 绑定当前进程到指定 GPU def cleanup(): dist.destroy_process_group() # ═══════════════ Step 2: 准备数据含 DistributedSampler ═══════════════ class DummyDataset(Dataset): def __len__(self): return 10000 def __getitem__(self, idx): return torch.randn(3, 224, 224), torch.randint(0, 10, (1,)) def get_dataloader(rank, world_size, batch_size): dataset DummyDataset() sampler DistributedSampler( dataset, num_replicasworld_size, # 总共 N 个 GPU rankrank, # 当前 GPU 编号 shuffleTrue, drop_lastTrue, ) return DataLoader(dataset, batch_sizebatch_size, samplersampler) # ═══════════════ Step 3: 训练主函数 ═══════════════ def train(rank, world_size, epochs5): setup(rank, world_size) # 模型 → GPU → DDP 包装 model nn.Sequential( nn.Conv2d(3, 64, 3, padding1), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(), nn.Linear(64, 10), ).to(rank) ddp_model DDP(model, device_ids[rank]) dataloader get_dataloader(rank, world_size, batch_size64) criterion nn.CrossEntropyLoss() optimizer optim.SGD(ddp_model.parameters(), lr0.01) for epoch in range(epochs): dataloader.sampler.set_epoch(epoch) # 关键每 epoch 重新 shuffle total_loss 0.0 for batch_idx, (data, target) in enumerate(dataloader): data, target data.to(rank), target.to(rank).squeeze() optimizer.zero_grad() output ddp_model(data) loss criterion(output, target) loss.backward() # DDP 自动触发 AllReduce 同步梯度 optimizer.step() total_loss loss.item() if rank 0: # 仅 rank 0 打印日志 print(fEpoch {epoch} | Loss: {total_loss / len(dataloader):.4f}) cleanup()torchrun启动使用 torchrun 启动 DDP 训练推荐方式 不需要手动设置 rank/world_sizetorchrun 自动注入环境变量 import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def main(): # torchrun 自动设置这些环境变量 local_rank int(os.environ[LOCAL_RANK]) world_size int(os.environ[WORLD_SIZE]) rank int(os.environ[RANK]) dist.init_process_group(backendnccl) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) ddp_model DDP(model, device_ids[local_rank]) # ... 训练循环同上 ... dist.destroy_process_group() if __name__ __main__: main() Shell # 单机 4 卡启动 torchrun --nproc_per_node4 train_ddp.py # 多机多卡2 节点 × 4 GPUs # 在 node0 上 torchrun --nnodes2 --nproc_per_node4 \ --node_rank0 --master_addr192.168.1.10 --master_port29500 \ train_ddp.py # 在 node1 上 torchrun --nnodes2 --nproc_per_node4 \ --node_rank1 --master_addr192.168.1.10 --master_port29500 \ train_ddp.py混合精度DDP混合精度训练 DDPFP16 前向/反向 FP32 主权重 显存减半速度提升利用 Tensor Core import torch from torch.cuda.amp import autocast, GradScaler from torch.nn.parallel import DistributedDataParallel as DDP def train_amp_ddp(rank, world_size): # ... 初始化进程组同上 ... model MyModel().to(rank) ddp_model DDP(model, device_ids[rank]) scaler GradScaler() # 梯度缩放器防止 FP16 下溢 optimizer torch.optim.AdamW(ddp_model.parameters(), lr1e-3) for data, target in dataloader: data, target data.to(rank), target.to(rank) optimizer.zero_grad() with autocast(): # 前向自动转为 FP16 output ddp_model(data) loss criterion(output, target) scaler.scale(loss).backward() # 反向loss × scale → 避免下溢 scaler.step(optimizer) # 更新unscale 梯度 → 正常 step scaler.update() # 动态调整 scale 因子 # ⚠️ DDP AMP 时 backward 中的 AllReduce 在 FP16 精度完成 # 梯度通信量减半进一步加速断电续训DDP 断点续训 / 模型保存的最佳实践 import torch from torch.nn.parallel import DistributedDataParallel as DDP def save_checkpoint(ddp_model, optimizer, epoch, path, rank): 仅 rank 0 保存其他进程跳过 if rank ! 0: return checkpoint { model_state_dict: ddp_model.module.state_dict(), # ↑ 注意 .module 去掉 DDP 包装前缀 optimizer_state_dict: optimizer.state_dict(), epoch: epoch, } torch.save(checkpoint, path) def load_checkpoint(model, optimizer, path, rank): 所有 rank 加载同一份 checkpoint checkpoint torch.load(path, map_locationfcuda:{rank}) model.load_state_dict(checkpoint[model_state_dict]) optimizer.load_state_dict(checkpoint[optimizer_state_dict]) return checkpoint[epoch] # ── 使用示例 ── ddp_model DDP(model, device_ids[rank]) # 保存 save_checkpoint(ddp_model, optimizer, epoch, checkpoint.pt, rank) # 加载在所有进程上调用 start_epoch load_checkpoint(ddp_model, optimizer, checkpoint.pt, rank)关键细节DistributedSampler确保每个 GPU 看到互不重叠的数据分片。sampler.set_epoch(epoch)必须在每个 epoch 调用否则每个 epoch 的 shuffle 结果相同失去随机性。 DDP 在loss.backward()时自动进行梯度 AllReduce——无需手动调用通信操作。常见坑点⚠️ sampler.set_epoch()忘记调用会导致每个 epoch 的 shuffle 顺序完全相同等效于没有 shuffle模型泛化能力下降。⚠️ BatchNorm 同步DDP 默认不同步 BN 统计量。跨卡 BatchNorm 需额外使用SyncBatchNorm小 batch 时尤其重要。⚠️ 随机种子各进程需设置不同种子避免生成相同数据。通常seed base_seed rank。⚠️ loss.backward() 同步DDP 在backward()时进行 AllReduce。如果某进程跳过 backward如条件分支会导致死锁。五、当模型大到放不进单卡模型并行当模型参数量超过单卡显存如 Llama-70B 需要约 140GB FP16而 A100 只有 80GB 数据并行不够用必须将模型本身切分到多张 GPU —— 这就是模型并行。三种主流方案三种模型并行详细对比方案切分粒度通信模式适用场景代表框架Tensor Parallel (TP)单层内矩阵按列/行切分每步 AllReduce通信密集单机多卡NVLink 高带宽Megatron-LMPipeline Parallel (PP)按层切分到不同 GPU仅层边界通信稀疏跨节点低带宽容忍GPipe, PipeDreamZeRO / FSDP优化器状态 / 梯度 / 参数分片等同于 DP 通信量通用性价比最高DeepSpeed, PyTorch FSDP模型并行代码实现DeepSpeed ZeROPipeline Parallel — 按层切分模型到不同 GPU 核心GPU0 持有 Layer 1~8GPU1 持有 Layer 9~16 通过微批次micro-batch流水线掩盖通信延迟 关键概念 - micro-batch将一个大 batch 切分为 M 个 micro-batch - 1F1B 调度一个 forward 一个 backward 交替执行 - bubble流水线空泡GPU 空闲等待的时间比例 (P-1)/M import torch import torch.nn as nn import torch.distributed as dist from queue import Queue # ═══════════════ 原生 PP 实现理解原理 ═══════════════ class PipelineStage(nn.Module): 每个 GPU 上的流水线阶段 def __init__(self, layers, rank, prev_rank, next_rank): super().__init__() self.layers nn.Sequential(*layers) self.rank rank self.prev prev_rank self.next next_rank def forward(self, xNone): if self.prev is not None: # 从上游 GPU 接收激活值 x torch.empty(self.input_shape, deviceself.rank) dist.recv(x, srcself.prev) out self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dstself.next) return out # ═══════════════ 生产环境torch.distributed.pipelining ═══════════════ from torch.distributed.pipelining import ( pipeline, ScheduleGPipe, Schedule1F1B, PipelineStage, SplitPoint, ) def train_pp(): # 1. 定义模型 切分点 model MyBigTransformer() split_spec { model.layers.4: SplitPoint.END, # Layer 0~3 → GPU0 model.layers.8: SplitPoint.END, # Layer 4~7 → GPU1 } # 2. 构建流水线 stages pipeline( modulemodel, num_chunks8, # micro-batch 数量越多 bubble 越小 example_args(example_input,), split_specsplit_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward再全部 backwardbubble 大 # Schedule1F1B: 一次 forward 接一次 backward 交替bubble 小推荐 schedule Schedule1F1B(stages, num_chunks8) for batch in dataloader: # 将 batch 拆成 8 个 micro-batch流水线执行 loss schedule.step(batch) # ═══════════════ Bubble 公式 ═══════════════ # bubble 比例 (P - 1) / M其中 P 流水线阶段数M micro-batch 数量 # 例4 阶段 × 32 micro-batch → bubble 3/32 ≈ 9.4% # 例4 阶段 × 4 micro-batch → bubble 3/4 75%几乎不可用 # ═══════════════ DeepSpeed Pipeline Parallel ═══════════════ # DeepSpeed 也提供 PP 支持通过配置即可启用 # pipeline: { stages: auto, partition: type:transformer } GPU 拓扑示意# 4 GPU 流水线示例8 micro-batches # # GPU0: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU1: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU2: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU3: ▓▓▓▓▓▓▓ F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 ▓▓▓▓▓▓▓ # ↑ bubble ↑ bubble # # F forward, B backward, ▓ bubble (GPU 空闲) # 1F1B 调度GPU0 做完 F1 立即做 B0不需等所有 F 完成Pytorch FSDPPyTorch FSDP — 原生 ZeRO-3 等效实现 (PyTorch ≥ 2.0) 无需安装第三方库API 与 DDP 风格一致 import torch import torch.distributed as dist from torch.distributed.fsdp import ( FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, CPUOffload, BackwardPrefetch, ) from torch.distributed.fsdp.wrap import ( transformer_auto_wrap_policy, size_based_auto_wrap_policy, ) import functools def train_fsdp(rank, world_size): dist.init_process_group(backendnccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) # ═══════════════ 1. FSDP 混合精度配置 ═══════════════ mp_policy MixedPrecision( param_dtypetorch.bfloat16, # 参数存储精度 reduce_dtypetorch.bfloat16, # 梯度 AllReduce 精度 buffer_dtypetorch.bfloat16, # buffer 精度 ) # ═══════════════ 2. Transformer 自动包装策略 ═══════════════ from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained(meta-llama/Llama-2-7b-hf) auto_wrap_policy functools.partial( transformer_auto_wrap_policy, transformer_layer_cls{ type(model.model.layers[0]), # 按 Transformer Block 粒度包装 }, ) # ═══════════════ 3. FSDP 包装模型 ═══════════════ model FSDP( model.to(rank), sharding_strategyShardingStrategy.FULL_SHARD, # ZeRO-3 等效 mixed_precisionmp_policy, auto_wrap_policyauto_wrap_policy, cpu_offloadCPUOffload(offload_paramsTrue), # 可选参数 offload 到 CPU backward_prefetchBackwardPrefetch.BACKWARD_PRE, # 预取下一层参数 device_idrank, ) optimizer torch.optim.AdamW(model.parameters(), lr1e-5) for batch in dataloader: optimizer.zero_grad() loss model(batch).loss loss.backward() optimizer.step() dist.destroy_process_group() # ═══════════════ FSDP 分片策略选择 ═══════════════ # ShardingStrategy.FULL_SHARD → ZeRO-3 等效参数梯度优化器全分片 # ShardingStrategy.SHARD_GRAD_OP → ZeRO-2 等效分片梯度优化器 # ShardingStrategy.NO_SHARD → DDP 等效不分片 # ShardingStrategy.HYBRID_SHARD → 节点内 FULL_SHARD 节点间 NO_SHARDTensor ParallelTensor Parallel — 单层内矩阵切分 核心思想将 Linear 层的权重矩阵按列或行切分到多 GPU 每步需 AllReduce / AllGather 通信适合 NVLink 高带宽场景 原生实现理解原理用→ 生产环境用 Megatron-LM 或 torch.distributed.tensor import torch import torch.nn as nn import torch.distributed as dist from torch.distributed import all_reduce, ReduceOp # ═══════════════ 列并行 LinearColumn Parallel ═══════════════ class ColumnParallelLinear(nn.Module): W: (in_features, out_features) → 按列切到 N 个 GPU 每 GPU 持有 W[:, i*C : (i1)*C]输出需 AllGather 拼接 def __init__(self, in_features, out_features): super().__init__() self.rank dist.get_rank() self.world_size dist.get_world_size() assert out_features % self.world_size 0 self.out_per_gpu out_features // self.world_size self.weight nn.Parameter(torch.empty(in_features, self.out_per_gpu)) nn.init.kaiming_uniform_(self.weight) def forward(self, x): # 每 GPU 计算自己那列 local_out x self.weight # (batch, out_per_gpu) # AllGather 拼接完整输出 gathered [torch.empty_like(local_out) for _ in range(self.world_size)] dist.all_gather(gathered, local_out) return torch.cat(gathered, dim-1) # ═══════════════ 行并行 LinearRow Parallel ═══════════════ class RowParallelLinear(nn.Module): W: (in_features, out_features) → 按行切到 N 个 GPU 每 GPU 持有 W[i*R:(i1)*R, :]输入需提前分片输出 AllReduce 求和 def __init__(self, in_features, out_features): super().__init__() self.rank dist.get_rank() self.world_size dist.get_world_size() assert in_features % self.world_size 0 self.in_per_gpu in_features // self.world_size self.weight nn.Parameter(torch.empty(self.in_per_gpu, out_features)) nn.init.kaiming_uniform_(self.weight) def forward(self, x): # x 已经在前一步被 AllGather 得到完整输入但这里只取自己那行 start self.rank * self.in_per_gpu local_x x[:, start : start self.in_per_gpu] local_out local_x self.weight # AllReduce 求和得到完整输出 dist.all_reduce(local_out, opReduceOp.SUM) return local_out # ═══════════════ 标准 TP MLP Block 组合 ═══════════════ class TP_MLP(nn.Module): ColumnParallel(d → 4d) → GELU → RowParallel(4d → d) → Dropout 列并行输出已 AllGather → 行并行输入直接用 → 输出 AllReduce 求和 一个完整 TP block 只有 1 次 AllReduceRowParallel 末尾 def __init__(self, d_model, d_ff): super().__init__() self.fc1 ColumnParallelLinear(d_model, d_ff) self.fc2 RowParallelLinear(d_ff, d_model) self.act nn.GELU() self.dropout nn.Dropout(0.1) def forward(self, x): return self.dropout(self.fc2(self.act(self.fc1(x)))) # ═══════════════ 生产环境torch.distributed.tensor.parallel ═══════════════ # PyTorch 2.0 内置 TP基于 DTensor一行代码切分模型 from torch.distributed.tensor.parallel import parallelize_module, ColwiseParallel, RowwiseParallel # parallelize_module(model, tp_mesh, { # fc1: ColwiseParallel(), # 列并行 # fc2: RowwiseParallel(), # 行并行自动插入 AllReduce # })Pipeline ParallelPipeline Parallel — 按层切分模型到不同 GPU 核心GPU0 持有 Layer 1~8GPU1 持有 Layer 9~16 通过微批次micro-batch流水线掩盖通信延迟 关键概念 - micro-batch将一个大 batch 切分为 M 个 micro-batch - 1F1B 调度一个 forward 一个 backward 交替执行 - bubble流水线空泡GPU 空闲等待的时间比例 (P-1)/M import torch import torch.nn as nn import torch.distributed as dist from queue import Queue # ═══════════════ 原生 PP 实现理解原理 ═══════════════ class PipelineStage(nn.Module): 每个 GPU 上的流水线阶段 def __init__(self, layers, rank, prev_rank, next_rank): super().__init__() self.layers nn.Sequential(*layers) self.rank rank self.prev prev_rank self.next next_rank def forward(self, xNone): if self.prev is not None: # 从上游 GPU 接收激活值 x torch.empty(self.input_shape, deviceself.rank) dist.recv(x, srcself.prev) out self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dstself.next) return out # ═══════════════ 生产环境torch.distributed.pipelining ═══════════════ from torch.distributed.pipelining import ( pipeline, ScheduleGPipe, Schedule1F1B, PipelineStage, SplitPoint, ) def train_pp(): # 1. 定义模型 切分点 model MyBigTransformer() split_spec { model.layers.4: SplitPoint.END, # Layer 0~3 → GPU0 model.layers.8: SplitPoint.END, # Layer 4~7 → GPU1 } # 2. 构建流水线 stages pipeline( modulemodel, num_chunks8, # micro-batch 数量越多 bubble 越小 example_args(example_input,), split_specsplit_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward再全部 backwardbubble 大 # Schedule1F1B: 一次 forward 接一次 backward 交替bubble 小推荐 schedule Schedule1F1B(stages, num_chunks8) for batch in dataloader: # 将 batch 拆成 8 个 micro-batch流水线执行 loss schedule.step(batch) # ═══════════════ Bubble 公式 ═══════════════ # bubble 比例 (P - 1) / M其中 P 流水线阶段数M micro-batch 数量 # 例4 阶段 × 32 micro-batch → bubble 3/32 ≈ 9.4% # 例4 阶段 × 4 micro-batch → bubble 3/4 75%几乎不可用 # ═══════════════ DeepSpeed Pipeline Parallel ═══════════════ # DeepSpeed 也提供 PP 支持通过配置即可启用 # pipeline: { stages: auto, partition: type:transformer } # 4 GPU 流水线示例8 micro-batches # # GPU0: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU1: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU2: F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 # GPU3: ▓▓▓▓▓▓▓ F0 F1 F2 F3 F4 F5 F6 F7 B0 B1 B2 B3 B4 B5 B6 B7 ▓▓▓▓▓▓▓ # ↑ bubble ↑ bubble # # F forward, B backward, ▓ bubble (GPU 空闲) # 1F1B 调度GPU0 做完 F1 立即做 B0不需等所有 F 完成六、实战选型决策树终极对比所有方案一览方案显存效率计算效率通信开销实现难度适用规模单卡❌ 最低❌ 最慢—⭐ 1BDPDataParallel❌ 同单卡⚠️ 有限⚠️ 主卡瓶颈⭐不推荐DDP❌ 同单卡✅ ~线性✅ AllReduce⭐⭐ 10BZeRO-1⚠️ 4×✅ 同 DDP✅ 同 DDP⭐⭐1B~10BZeRO-2✅ 8×✅ 同 DDP✅ 同 DDP⭐⭐1B~30BZeRO-3 / FSDP✅ N×⚠️ 略降⚠️ 1.5×⭐⭐⭐10B~100BTP⚠️ 有限⚠️ 通信密集❌ 高⭐⭐⭐⭐单机大模型PP✅ 好⚠️ 有空泡✅ 低⭐⭐⭐⭐跨节点大模型3D 并行✅ 最优⚠️ 复杂调优⚠️ 混合⭐⭐⭐⭐⭐100B~1T