2、常见分布式框架对比：PyTorch DDP、Horovod、DeepSpeed、Megatron-LM 的选型与适用场景

做分布式训练这几年，我接触过的框架少说也有七八个。说实话，没有哪个框架是「万能钥匙」。选型这件事，说白了就是看你的场景、你的卡、你的模型规模。

今天我就把四个最主流的框架——PyTorch DDP、Horovod、DeepSpeed、Megatron-LM——掰开揉碎了讲。每个框架我都踩过坑，有些经验是用加班换来的。

2.1 先看一张全景图

我习惯在选型前先画一张图，把框架的定位理清楚。下面这张图是我自己总结的，你可以把它当作选型时的「导航地图」。

2.2 PyTorch DDP：入门首选，但别贪心

PyTorch DDP（Distributed Data Parallel）是我用得最多的框架。为什么？因为它简单。你只需要把模型用 DistributedDataParallel 包一下，再改改数据加载器，基本就完事了。

我记得有一次，团队里新来的同学问我：「为什么我用 DataParallel 跑8卡，速度还不如单卡？」我笑了笑说：「你换成 DistributedDataParallel 试试。」结果速度直接翻了6倍。这里面的门道是——DataParallel 用单线程管理多卡，而 DDP 每个进程管一张卡，没有GIL锁的问题。

# 最简单的 DDP 用法
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
model = MyModel().cuda()
model = DDP(model)

# 数据加载也要改
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)

2.3 Horovod：多框架混用的救星

Horovod 是 Uber 开源的分布式训练框架。说实话，我一开始对它没什么好感——觉得多此一举。直到有一次，我需要在一个项目里同时用 TensorFlow 和 PyTorch，Horovod 才真正派上用场。

Horovod 最大的优势是「框架无关」。它底层用 NCCL 或 MPI 做通信，上层统一了 API。你想想看，如果你的团队既有 TensorFlow 用户，又有 PyTorch 用户，Horovod 就是那个「最大公约数」。

# Horovod 的 PyTorch 用法
import horovod.torch as hvd

hvd.init()
model = MyModel().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01 * hvd.size())

# 关键一步：广播初始参数
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

# 分布式优化器
optimizer = hvd.DistributedOptimizer(optimizer)

2.4 DeepSpeed：大模型训练的「显存魔术师」

DeepSpeed 是微软开源的框架，它的 ZeRO 优化器是我见过最聪明的显存管理方案。说白了，ZeRO 就是把模型状态（优化器状态、梯度、参数）分散到各个GPU上，用的时候再取回来。

我习惯把 ZeRO 分成三个阶段：

ZeRO阶段	优化内容	显存节省	通信开销
ZeRO-1	优化器状态分片	约4倍	低
ZeRO-2	优化器状态 + 梯度分片	约8倍	中
ZeRO-3	优化器状态 + 梯度 + 参数分片	约16倍	高

我在项目中遇到过这样一个场景：用 4 张 A100 跑 20B 参数的模型，单卡显存 80GB。如果用 DDP，一张卡都放不下。但用 DeepSpeed 的 ZeRO-3，居然跑起来了。虽然速度慢了点，但至少能跑。

# DeepSpeed 配置示例（ds_config.json）
{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu"
    }
  }
}

2.5 Megatron-LM：为超大模型而生

Megatron-LM 是 NVIDIA 开源的框架，专门为训练超大规模语言模型设计。它支持三种并行策略：数据并行、张量并行、流水线并行。这三种策略组合起来，就是所谓的「3D并行」。

说实话，Megatron-LM 的学习曲线是最陡的。我第一次用的时候，光配置并行策略就花了两天。但它的效果也是最好的——我见过用 Megatron 在 1024 张 GPU 上训练 530B 参数的模型，效率还能保持在 50% 以上。

# Megatron 的并行配置（伪代码）
model = MegatronModel(
    tensor_model_parallel_size=4,   # 张量并行
    pipeline_model_parallel_size=4, # 流水线并行
    data_parallel_size=8            # 数据并行
)
# 总 GPU 数 = 4 × 4 × 8 = 128

2.6 选型决策树

说了这么多，到底怎么选？我总结了一个简单的决策流程：

1. 模型能塞进单卡吗？ → 能，用 DDP；不能，往下看
2. 模型在 10B 参数以内？ → 用 DeepSpeed ZeRO-2/3
3. 模型超过 10B 参数？ → 用 Megatron-LM 的 3D 并行
4. 团队同时用多个深度学习框架？ → 考虑 Horovod
5. 想快速验证想法？ → 先用 DDP 跑通，再优化

想要更多资料，联系可可