训练框架对比：DDP vs DeepSpeed vs Megatron

分布式训练不用从零写。现成的框架很多，但各有特点。

这篇对比最常用的几个框架，帮你选对工具。

框架概览

框架开发者主要特点适合场景
PyTorch DDPPyTorch简单，官方支持数据并行入门
PyTorch FSDPPyTorchZeRO 类实现中等模型
DeepSpeed微软功能全，ZeRO大模型训练
Megatron-LMNVIDIA性能极致超大模型
ColossalAI国产易用入门大模型

PyTorch DDP

是什么

Distributed Data Parallel，PyTorch 官方的分布式数据并行。

原理

每张卡放完整模型，各自处理不同数据，AllReduce 同步梯度。

代码示例

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'])
torch.cuda.set_device(local_rank)

# 模型
model = MyModel().cuda()
model = DDP(model, device_ids=[local_rank])

# 数据加载器
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler)

# 训练
for epoch in range(epochs):
    sampler.set_epoch(epoch)  # 打乱数据
    for batch in dataloader:
        loss = model(batch)
        loss.backward()
        optimizer.step()

启动方式

# torchrun（推荐）
torchrun --nproc_per_node=8 train.py

# 或者 python -m torch.distributed.launch
python -m torch.distributed.launch --nproc_per_node=8 train.py

优点

  • 官方支持，稳定
  • 简单易用
  • 文档完善
  • 和 PyTorch 无缝集成

缺点

  • 只支持数据并行
  • 每张卡要存完整模型
  • 不适合超大模型

适合场景

  • 模型能放进单卡
  • 入门分布式训练
  • 不想引入额外依赖

PyTorch FSDP

是什么

Fully Sharded Data Parallel，PyTorch 官方实现的 ZeRO 类技术。

原理

类似 ZeRO Stage 3，把参数、梯度、优化器状态分片到各 GPU。

代码示例

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

# 自动分片策略
auto_wrap_policy = functools.partial(
    size_based_auto_wrap_policy,
    min_num_params=1e6
)

# 包装模型
model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    ),
)

优点

  • 官方支持，持续更新
  • 显存效率高
  • 支持混合精度
  • 和 PyTorch 生态兼容

缺点

  • 配置复杂
  • 性能不如 DeepSpeed 优化得好
  • 相对较新，坑较多

适合场景

  • 想用官方方案
  • 中等规模模型（7B-30B）
  • 不想引入第三方库

DeepSpeed

是什么

微软开源的分布式训练框架，功能最全。

核心特性

ZeRO 优化：三个 Stage，逐级降低显存占用。

ZeRO-Offload：把数据卸载到 CPU 内存或 NVMe。

混合精度：FP16/BF16 训练。

梯度累积：等效更大 batch size。

稀疏注意力：长序列优化。

代码示例

import deepspeed

# 配置文件 ds_config.json
{
    "train_batch_size": 32,
    "gradient_accumulation_steps": 4,
    "fp16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu"
        }
    }
}

# 初始化
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config="ds_config.json"
)

# 训练
for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

ZeRO 各 Stage 对比

Stage切分内容显存节省通信开销
Stage 1优化器状态4x低
Stage 2+ 梯度8x中
Stage 3+ 参数N x高

启动方式

deepspeed --num_gpus=8 train.py --deepspeed_config ds_config.json

优点

  • 功能最全
  • ZeRO 效果好
  • 支持 Offload（显存不够时救命）
  • 社区活跃
  • 和 HuggingFace 集成好

缺点

  • 配置复杂
  • 有学习成本
  • 和 PyTorch 版本兼容问题偶发
  • 性能极致场景不如 Megatron

适合场景

  • 大模型训练
  • 显存紧张（需要 ZeRO/Offload）
  • 用 HuggingFace 生态

Megatron-LM

是什么

NVIDIA 开源的大模型训练框架，性能极致。

核心特性

张量并行：高效的 Transformer 层切分。

流水线并行：优化的调度策略。

序列并行：长序列优化。

混合精度：FP16/BF16/FP8。

代码示例

Megatron 不是「包装」你的代码，而是要用它的代码结构：

# 启动训练
python pretrain_gpt.py \
    --tensor-model-parallel-size 8 \
    --pipeline-model-parallel-size 4 \
    --num-layers 96 \
    --hidden-size 12288 \
    --num-attention-heads 96 \
    --micro-batch-size 1 \
    --global-batch-size 1024 \
    ...

优点

  • 性能最强（NVIDIA 优化）
  • 3D 并行支持完善
  • FP8 支持
  • 大规模训练验证过

缺点

  • 侵入性强，要改代码结构
  • 学习曲线陡
  • 灵活性差
  • 文档相对少

适合场景

  • 超大模型训练（100B+）
  • 追求极致性能
  • 有专业团队

ColossalAI

是什么

国产分布式训练框架，主打易用。

特点

  • 封装了 ZeRO、张量并行等
  • API 更简单
  • 中文文档

代码示例

import colossalai
from colossalai.booster import Booster
from colossalai.booster.plugin import GeminiPlugin

colossalai.launch_from_torch()

plugin = GeminiPlugin()
booster = Booster(plugin=plugin)

model, optimizer, _, dataloader, _ = booster.boost(
    model, optimizer, dataloader=dataloader
)

for batch in dataloader:
    outputs = model(**batch)
    loss = outputs.loss
    booster.backward(loss, optimizer)
    optimizer.step()

优点

  • 易用
  • 中文社区
  • 更新快

缺点

  • 稳定性不如前几个
  • 大规模验证少
  • 生态不如 DeepSpeed

适合场景

  • 入门大模型训练
  • 国内团队
  • 快速实验

怎么选

按模型规模

模型规模推荐框架
< 1BDDP
1B - 10BDDP + 梯度累积，或 FSDP/DeepSpeed
10B - 100BDeepSpeed ZeRO-2/3
> 100BMegatron-LM 或 DeepSpeed + Megatron

按使用场景

场景推荐
快速实验DDP 或 ColossalAI
生产训练DeepSpeed
追求极致性能Megatron-LM
用 HuggingFaceDeepSpeed（集成好）
显存不够DeepSpeed ZeRO-3 + Offload

按团队情况

团队情况推荐
个人/小团队DeepSpeed（文档全）
有 infra 团队Megatron-LM
入门学习DDP → DeepSpeed

组合使用

实际上，大规模训练经常组合使用：

DeepSpeed + Megatron

微软的 DeepSpeed-Megatron：

  • 用 Megatron 的张量/流水线并行
  • 用 DeepSpeed 的 ZeRO

FSDP + 其他

PyTorch 2.0 后 FSDP 可以和张量并行组合。

HuggingFace Accelerate

HuggingFace 的 Accelerate 库可以统一调用 DDP/FSDP/DeepSpeed：

from accelerate import Accelerator

accelerator = Accelerator()
model, optimizer, dataloader = accelerator.prepare(
    model, optimizer, dataloader
)

一套代码，切换后端。

小结

框架选择速查：

框架一句话
DDP简单场景，官方首选
FSDP官方 ZeRO，中等模型
DeepSpeed功能全，大模型首选
Megatron性能极致，超大模型
ColossalAI易用，入门友好

建议路径：

  1. 先用 DDP 跑通
  2. 模型变大上 DeepSpeed
  3. 追求极致上 Megatron

训练篇到此结束。下一部分讲推理。