企业官网建设流程全解析

从单卡到千卡集群：理解 DDP 数据并行的本质，搞懂 AllReduce 通信原语， 看清模型并行（TP / PP / ZeRO）何时登场。附完整 PyTorch 实现。

一、为什么需要分布式训练？

⏱ 速度瓶颈

单卡训练 GPT-3（175B）需要 355 年。8 卡 DDP 能将 epoch 时间线性缩短至约 1/N。

🧠 显存瓶颈

单张 A100（80GB）放不下一个 70B 模型（FP16 约 140GB）。必须把模型切分到多卡。

📈 数据规模

训练数据动辄 TB 级，单卡 DataLoader 成为 I/O 瓶颈，分布式数据加载是必需项。

核心权衡：数据并行解决速度问题（模型能放进单卡）， 模型并行解决显存问题（模型放不进单卡）。 实际训练中两者常混合使用——即3D 并行（DP + TP + PP）。

二、数据并行（DDP）核心原理

Distributed Data Parallel（DDP）是 PyTorch 官方推荐的数据并行方案， 替代了旧的DataParallel（DP）。每个 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 DP（DataParallel）

三、AllReduce 通信原语

AllReduce 是所有分布式训练框架的核心通信操作：将每个进程的本地张量归约（reduce）后， 把结果广播（broadcast）回所有进程。最常用的实现是Ring AllReduce。

Ring AllReduce 算法（两步法）

Ring AllReduce 复杂度：通信量 = 2(N-1)/N × 数据量。 当 N 很大时趋近2×数据量，与 GPU 数量无关——这是它能线性扩展的关键。 NCCL 库（NVIDIA Collective Communications Library）高度优化了这一过程。

常用集合通信原语对比

四、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( backend="nccl", # NVIDIA NCCL 后端（GPU） init_method="env://", # 从环境变量读取地址 rank=rank, world_size=world_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_replicas=world_size, # 总共 N 个 GPU rank=rank, # 当前 GPU 编号 shuffle=True, drop_last=True, ) return DataLoader(dataset, batch_size=batch_size, sampler=sampler) # ═══════════════ Step 3: 训练主函数 ═══════════════ def train(rank, world_size, epochs=5): setup(rank, world_size) # 模型 → GPU → DDP 包装 model = nn.Sequential( nn.Conv2d(3, 64, 3, padding=1), 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_size=64) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(ddp_model.parameters(), lr=0.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(f"Epoch {epoch} | Loss: {total_loss / len(dataloader):.4f}") cleanup()

torchrun启动

"""使用 torchrun 启动 DDP 训练（推荐方式） 不需要手动设置 rank/world_size，torchrun 自动注入环境变量""" 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(backend="nccl") 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_node=4 train_ddp.py # 多机多卡（2 节点 × 4 GPUs） # 在 node0 上： torchrun --nnodes=2 --nproc_per_node=4 \ --node_rank=0 --master_addr="192.168.1.10" --master_port=29500 \ train_ddp.py # 在 node1 上： torchrun --nnodes=2 --nproc_per_node=4 \ --node_rank=1 --master_addr="192.168.1.10" --master_port=29500 \ train_ddp.py

混合精度+DDP

"""混合精度训练 + DDP：FP16 前向/反向 + 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(), lr=1e-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_location=f"cuda:{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 —— 这就是模型并行。三种主流方案：

三种模型并行详细对比

模型并行代码实现

DeepSpeed ZeRO

"""Pipeline Parallel — 按层切分模型到不同 GPU 核心：GPU0 持有 Layer 1~8，GPU1 持有 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, x=None): if self.prev is not None: # 从上游 GPU 接收激活值 x = torch.empty(self.input_shape, device=self.rank) dist.recv(x, src=self.prev) out = self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dst=self.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( module=model, num_chunks=8, # micro-batch 数量（越多 bubble 越小） example_args=(example_input,), split_spec=split_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward，再全部 backward（bubble 大） # Schedule1F1B: 一次 forward 接一次 backward 交替（bubble 小，推荐） schedule = Schedule1F1B(stages, num_chunks=8) 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 FSDP

"""PyTorch 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(backend="nccl", rank=rank, world_size=world_size) torch.cuda.set_device(rank) # ═══════════════ 1. FSDP 混合精度配置 ═══════════════ mp_policy = MixedPrecision( param_dtype=torch.bfloat16, # 参数存储精度 reduce_dtype=torch.bfloat16, # 梯度 AllReduce 精度 buffer_dtype=torch.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_strategy=ShardingStrategy.FULL_SHARD, # ZeRO-3 等效 mixed_precision=mp_policy, auto_wrap_policy=auto_wrap_policy, cpu_offload=CPUOffload(offload_params=True), # 可选：参数 offload 到 CPU backward_prefetch=BackwardPrefetch.BACKWARD_PRE, # 预取下一层参数 device_id=rank, ) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-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_SHARD

Tensor Parallel

"""Tensor 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 # ═══════════════ 列并行 Linear（Column Parallel） ═══════════════ class ColumnParallelLinear(nn.Module): """W: (in_features, out_features) → 按列切到 N 个 GPU 每 GPU 持有 W[:, i*C : (i+1)*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) # ═══════════════ 行并行 Linear（Row Parallel） ═══════════════ class RowParallelLinear(nn.Module): """W: (in_features, out_features) → 按行切到 N 个 GPU 每 GPU 持有 W[i*R:(i+1)*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, op=ReduceOp.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 次 AllReduce（RowParallel 末尾） """ 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 Parallel

"""Pipeline Parallel — 按层切分模型到不同 GPU 核心：GPU0 持有 Layer 1~8，GPU1 持有 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, x=None): if self.prev is not None: # 从上游 GPU 接收激活值 x = torch.empty(self.input_shape, device=self.rank) dist.recv(x, src=self.prev) out = self.layers(x) if self.next is not None: # 发送激活值到下游 GPU dist.send(out, dst=self.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( module=model, num_chunks=8, # micro-batch 数量（越多 bubble 越小） example_args=(example_input,), split_spec=split_spec, ) # 3. 选择调度策略 # ScheduleGPipe: 先全部 forward，再全部 backward（bubble 大） # Schedule1F1B: 一次 forward 接一次 backward 交替（bubble 小，推荐） schedule = Schedule1F1B(stages, num_chunks=8) 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 完成）

六、实战选型决策树

终极对比：所有方案一览