数据并行：DDP 与 FSDP
数据并行是最直观的分布式训练方案：每个 GPU 持有完整模型副本，处理不同数据子集，通过梯度同步保持参数一致。PyTorch 提供 DDP 和 FSDP 两种实现。
DistributedDataParallel (DDP)
DDP 是 PyTorch 经典的数据并行实现，每个 GPU 持有完整模型：
python
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 初始化进程组
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
# 创建模型并移到当前 GPU
model = MyModel().cuda(local_rank)
model = DDP(model, device_ids=[local_rank])
# 训练循环
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for data, target in dataloader:
data, target = data.cuda(local_rank), target.cuda(local_rank)
loss = model(data, target)
loss.backward() # DDP 自动同步梯度
optimizer.step()
optimizer.zero_grad()DDP 梯度同步机制
DDP 在反向传播时自动触发梯度 AllReduce：
- 使用 Bucket 机制将参数分组
- 每个 Bucket 的梯度就绪后立即启动异步 AllReduce
- 与反向传播计算重叠，减少通信等待
FullyShardedDataParallel (FSDP)
FSDP 将模型参数、梯度和优化器状态分片到各 GPU，仅在计算时聚合：
python
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
model = MyModel()
# FSDP 包装
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD, # ZeRO-3
device_id=local_rank,
)
# 训练循环与 DDP 相同
for data, target in dataloader:
loss = model(data.cuda(), target.cuda())
loss.backward()
optimizer.step()
optimizer.zero_grad()DDP vs FSDP 对比
| 维度 | DDP | FSDP |
|---|---|---|
| 参数存储 | 每个 GPU 完整副本 | 分片存储 |
| 显存占用 | O(模型大小) | O(模型大小 / N) |
| 通信量 | 梯度 AllReduce | 参数 AllGather + 梯度 ReduceScatter |
| 适用模型 | 模型能放入单 GPU | 模型超出单 GPU |
| 实现复杂度 | 低 | 中 |
| 性能（小模型） | 更优 | 略低（额外 AllGather） |
| 性能（大模型） | OOM | 可训练 |
选择建议
- 模型能放入单 GPU：使用 DDP（更简单、小模型性能更优）
- 模型超出单 GPU：使用 FSDP（分片存储突破显存限制）
- 极大模型：FSDP + CPU Offload
FSDP 注意事项
FSDP 的 AllGather 通信会增加延迟。对于小模型（<1B），DDP 通常更快。FSDP 在模型 >7B 时优势明显。FSDP 要求模型前向传播中不能跨分片边界共享参数。