分布式训练
随着模型规模的不断增长，单机训练已经无法满足需求。分布式训练通过多机多卡的协同工作，实现了大规模模型的高效训练。
为什么需要分布式训练
模型规模的增长
近年来，模型参数量呈指数级增长：
| 模型 | 参数量 | 发布年份 |
|---|---|---|
| BERT-Large | 3.4亿 | 2018 |
| GPT-2 | 15亿 | 2019 |
| GPT-3 | 1750亿 | 2020 |
| GPT-4 | 约1.8万亿 | 2023 |
一个 1750 亿参数的模型，仅模型参数就需要约 700GB 显存（FP32），远超单卡容量。
训练时间的需求
即使模型可以放入单卡，训练时间也可能难以接受。以 GPT-3 为例：
- 单卡训练时间：约 355 年
- 使用 1024 张 V100：约 1 个月
分布式训练可以将训练时间从年缩短到周甚至天。
分布式训练策略
数据并行
数据并行是最常用的分布式训练策略。
工作原理：
- 每个 GPU 持有完整的模型副本
- 将批次数据分割到各个 GPU
- 各 GPU 独立进行前向和反向传播
- 同步梯度，更新模型参数
┌─────────────────────────────────────────────────────┐
│ 数据并行 │
│ │
│ GPU 0 GPU 1 GPU 2 GPU 3 │
│ ┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐│
│ │模型 │ │模型 │ │模型 │ │模型 ││
│ │副本 │ │副本 │ │副本 │ │副本 ││
│ └─────┘ └─────┘ └─────┘ └─────┘│
│ ↓ ↓ ↓ ↓ │
│ 数据 0 数据 1 数据 2 数据 3 │
│ ↓ ↓ ↓ ↓ │
│ 梯度 0 梯度 1 梯度 2 梯度 3 │
│ └──────────────┴──────────────┴──────────────┘ │
│ ↓ │
│ 梯度同步 │
│ ↓ │
│ 参数更新 │
└─────────────────────────────────────────────────────┘
优点：
- 实现简单，代码改动小
- 适用于模型可以放入单卡显存的场景
缺点：
- 每个GPU需要存储完整模型，显存利用率低
- 通信开销随 GPU 数量增加而增大
模型并行
当模型太大无法放入单卡时，需要使用模型并行。
流水线并行
将模型按层切分到不同 GPU：
GPU 0 GPU 1 GPU 2 GPU 3
┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐
│层 0-5│ ──→ │层 6-11│ ──→ │层 12-17│ ──→ │层 18-23│
└─────┘ └─────┘ └─────┘ └─────┘
特点：
- 各 GPU 顺序处理，存在"气泡"时间
- 需要精心设计 micro-batch 来减少气泡
张量并行
将单个层的参数切分到多个 GPU：
矩阵乘法 Y = XW，将 W 按列切分：
GPU 0: Y₁ = XW₁
GPU 1: Y₂ = XW₂
最终结果: Y = [Y₁, Y₂]（拼接）
特点：
- 通信频繁，需要高带宽互联
- 适合单机多卡场景
混合并行
实际的大模型训练通常结合多种并行策略：
┌─────────────────────────────────────────────────────┐
│ 混合并行 │
│ │
│ 数据并行（跨节点） │
│ ┌─────────────────────────────────────────────┐ │
│ │ 节点 0 │ │
│ │ ┌───────────────────────────────────────┐ │ │
│ │ │ 流水线并行（层间） │ │ │
│ │ │ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ │ │
│ │ │ │ 阶段 0 │ │ 阶段 1 │ │ 阶段 2 │ │ │ │
│ │ │ │ ┌─────┐ │ │ ┌─────┐ │ │ ┌─────┐ │ │ │ │
│ │ │ │ │TP 0 │ │ │ │TP 0 │ │ │ │TP 0 │ │ │ │ │
│ │ │ │ │TP 1 │ │ │ │TP 1 │ │ │ │TP 1 │ │ │ │ │
│ │ │ │ └─────┘ │ │ └─────┘ │ │ └─────┘ │ │ │ │
│ │ │ └─────────┘ └─────────┘ └─────────┘ │ │ │
│ │ └───────────────────────────────────────┘ │ │
│ └─────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────┘
TP = 张量并行（层内）
PyTorch 分布式训练
DistributedDataParallel（DDP）
PyTorch DDP 是最常用的数据并行实现：
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
def setup(rank, world_size):
dist.init_process_group("nccl", rank=rank, world_size=world_size)
def cleanup():
dist.destroy_process_group()
def train(rank, world_size):
setup(rank, world_size)
model = MyModel().to(rank)
ddp_model = DDP(model, device_ids=[rank])
optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
for epoch in range(epochs):
for batch in dataloader:
optimizer.zero_grad()
output = ddp_model(batch)
loss = criterion(output, target)
loss.backward()
optimizer.step()
cleanup()
if __name__ == "__main__":
world_size = torch.cuda.device_count()
torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)
启动方式：
torchrun --nproc_per_node=8 train.py
Fully Sharded Data Parallel（FSDP）
FSDP 是 PyTorch 提供的高级分布式训练方案，通过分片技术大幅降低显存占用：
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy
model = MyModel()
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD,
device_id=torch.cuda.current_device()
)
分片策略：
| 策略 | 说明 | 显存节省 |
|---|---|---|
| FULL_SHARD | 分片参数、梯度、优化器状态 | 最大 |
| SHARD_GRAD_OP | 分片梯度和优化器状态 | 中等 |
| NO_SHARD | 不分片（类似 DDP） | 无 |
DeepSpeed
DeepSpeed 是微软开源的深度学习优化库，提供了强大的显存优化和分布式训练能力。
核心特性
ZeRO 优化：零冗余优化器，通过分片消除冗余：
- ZeRO-1：分片优化器状态
- ZeRO-2：分片优化器状态和梯度
- ZeRO-3：分片优化器状态、梯度和参数
显存对比：
| 配置 | 7B 模型显存 |
|---|---|
| DDP | 112 GB |
| ZeRO-1 | 28 GB |
| ZeRO-2 | 16 GB |
| ZeRO-3 | 8 GB |
使用示例
配置文件 ds_config.json：
{
"train_batch_size": 128,
"gradient_accumulation_steps": 1,
"optimizer": {
"type": "AdamW",
"params": {
"lr": 1e-5
}
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu"
}
}
}
训练代码：
import deepspeed
model_engine, optimizer, _, _ = deepspeed.initialize(
model=model,
model_parameters=model.parameters(),
config="ds_config.json"
)
for batch in dataloader:
outputs = model_engine(batch)
loss = outputs.loss
model_engine.backward(loss)
model_engine.step()
启动方式
deepspeed --num_gpus=8 train.py --deepspeed_config ds_config.json
Megatron-LM
Megatron-LM 是 NVIDIA 开发的大模型训练框架，专注于高效的张量并行和流水线并行。
核心技术
张量并行：将 Transformer 层的计算分散到多个 GPU
from megatron import get_args
from megatron.model import TransformerBlock
# Megatron 自动处理张量并行
model = TransformerBlock(...)
流水线并行：使用 1F1B 调度减少气泡
时间步: 0 1 2 3 4 5 6 7 8 9
GPU 0: F0 F1 F2 F3 F4 B0 B1 B2 B3 B4
GPU 1: F0 F1 F2 F3 F4 B0 B1 B2 B3
GPU 2: F0 F1 F2 F3 F4 B0 B1 B2
GPU 3: F0 F1 F2 F3 F4 B0 B1
F = 前向传播, B = 反向传播
通信优化
梯度累积
减少通信频率，提高计算通信比：
accumulation_steps = 4
for i, batch in enumerate(dataloader):
output = model(batch)
loss = output.loss / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
梯度压缩
减少通信数据量：
- 量化：将 FP32 梯度压缩为 FP16 或 INT8
- 稀疏化：只传输重要的梯度
重叠计算和通信
在反向传播时同步梯度：
# DDP 自动实现计算通信重叠
model = DDP(model, device_ids=[rank])
显存优化技术
混合精度训练
使用 FP16/BF16 进行计算，减少显存占用：
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for batch in dataloader:
optimizer.zero_grad()
with autocast():
output = model(batch)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
梯度检查点
以计算换内存：
from torch.utils.checkpoint import checkpoint
class MyModel(nn.Module):
def forward(self, x):
# 使用检查点保存内存
x = checkpoint(self.layer1, x)
x = checkpoint(self.layer2, x)
return x
激活重计算
DeepSpeed 提供的激活重计算功能：
{
"activation_checkpointing": {
"partition_activations": true,
"cpu_checkpointing": true
}
}
实践建议
选择并行策略
| 场景 | 推荐策略 |
|---|---|
| 模型可放入单卡 | 数据并行（DDP） |
| 模型较大但单机可放 | FSDP 或 ZeRO-2 |
| 超大模型多机训练 | 3D 并行 + ZeRO-3 |
监控和调试
- 监控 GPU 利用率：使用
nvidia-smi或 Prometheus - 分析性能瓶颈：使用 PyTorch Profiler
- 检查梯度同步：确保所有 GPU 梯度一致
常见问题
死锁：确保所有进程执行相同的通信操作
梯度爆炸/消失：检查学习率和梯度裁剪
显存不足：尝试 FSDP、ZeRO 或梯度检查点
小结
分布式训练是大模型时代的必备技能。选择合适的并行策略、优化通信效率、合理管理显存，是构建高效训练系统的关键。在实际应用中，建议优先使用成熟的框架（如 DeepSpeed、Megatron-LM），它们已经解决了大部分工程难题。