3. 分布式训练框架

说实话，每次有新人问我「分布式训练该选哪个框架」，我都觉得这个问题挺有意思。因为答案从来不是「A比B好」，而是「你的场景适合哪个」。我自己这些年三个框架都用过，踩过坑也填过坑，今天就把我的真实感受摊开来聊聊。

3.1 三个框架的「出身」与「性格」

先说说它们各自是怎么来的，这决定了它们擅长什么。

• Horovod：Uber 出品，2017 年开源。它的核心思路是「让单机训练代码几乎不改就能跑多机」。我最早接触它的时候，觉得这玩意儿太神奇了——加几行代码，单卡变多卡，简直傻瓜式操作。
• PyTorch DDP：PyTorch 官方自带的分布式模块，从 1.0 版本开始逐步完善。说白了就是「亲儿子」，跟 PyTorch 的生态无缝衔接。我现在大部分项目都优先考虑它。
• DeepSpeed：微软出品，主打「大模型训练」。它的看家本领是 ZeRO 优化器，能把显存压到极致。我去年训一个 70B 的模型，全靠它才在 8 卡上跑起来。

你看，这三个框架的定位其实很清晰：Horovod 强在兼容性，DDP 强在原生集成，DeepSpeed 强在显存优化。

3.2 核心机制对比：通信、同步与显存

嗯，这里要深入一点了。我们搞分布式训练，最关心的无非三件事：通信效率、同步方式、显存占用。

维度	Horovod	PyTorch DDP	DeepSpeed
通信后端	MPI + NCCL	NCCL / Gloo	NCCL + 自定义
同步方式	AllReduce（Ring）	AllReduce（Bucket）	AllReduce + ZeRO 分段
显存优化	无	无	ZeRO-1/2/3 + Offload
易用性	高（代码侵入小）	中（需改 DataLoader）	低（配置复杂）
适用规模	小-中（1-64 卡）	中（1-256 卡）	大（64-1024 卡）

为什么会这样？我一个个说。

Horovod 的 Ring-AllReduce：它把梯度通信拆成一个个环，每个节点只跟邻居通信。好处是带宽利用率高，坏处是延迟敏感。我在 8 卡机器上测过，小模型（<1B 参数）Horovod 比 DDP 慢 10% 左右，因为通信开销占比太大。

PyTorch DDP 的 Bucket AllReduce：它把梯度先攒成桶（bucket），然后一次性通信。这招很聪明——减少了通信次数，对小 batch 特别友好。我习惯用 DDP 做日常实验，因为改起来快，调试也方便。

DeepSpeed 的 ZeRO：这才是真正的黑科技。ZeRO-1 只分片优化器状态，ZeRO-2 加上梯度分片，ZeRO-3 连模型参数都分片了。说白了就是用通信换显存。我曾经用 ZeRO-3 在 4 张 V100 上训了一个 13B 的模型，换成 DDP 的话，至少需要 16 张卡。

3.3 代码实战：同一个模型，三种写法

光说不练假把式。我们拿一个简单的 ResNet-50 训练任务，看看三个框架的代码差异。

3.3.1 Horovod 版本

import horovod.torch as hvd
hvd.init()
torch.cuda.set_device(hvd.local_rank())

model = ResNet50().cuda()
optimizer = optim.SGD(model.parameters(), lr=0.01)
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters())
hvd.broadcast_parameters(model.state_dict(), root_rank=0)

for data, target in dataloader:
    output = model(data.cuda())
    loss = F.nll_loss(output, target.cuda())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

你看，核心改动就三行：hvd.init()、DistributedOptimizer、broadcast_parameters。我最早用 Horovod 就是被这个简洁性吸引的。但注意，dataloader 需要自己处理 shuffle 和 sampler，否则每个节点会看到相同的数据。

3.3.2 PyTorch DDP 版本

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

dist.init_process_group(backend='nccl')
torch.cuda.set_device(local_rank)

model = ResNet50().cuda()
model = DDP(model, device_ids=[local_rank])
optimizer = optim.SGD(model.parameters(), lr=0.01)

train_sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=train_sampler)

for epoch in range(epochs):
    train_sampler.set_epoch(epoch)
    for data, target in dataloader:
        optimizer.zero_grad()
        output = model(data.cuda())
        loss = F.nll_loss(output, target.cuda())
        loss.backward()
        optimizer.step()

DDP 的写法更「PyTorch 原生」。唯一要注意的是 DistributedSampler 必须设置 epoch，否则每个 epoch 的数据顺序都一样。我刚开始用的时候忘了这茬，模型训了 50 个 epoch 精度死活上不去，排查了两天才发现是数据没打乱。

3.3.3 DeepSpeed 版本

import deepspeed

model = ResNet50()
ds_config = {
    "train_batch_size": 256,
    "gradient_accumulation_steps": 2,
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {"device": "cpu"}
    },
    "fp16": {"enabled": True}
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params=ds_config
)

for data, target in dataloader:
    output = model_engine(data.cuda())
    loss = F.nll_loss(output, target.cuda())
    model_engine.backward(loss)
    model_engine.step()

DeepSpeed 的配置全在 ds_config 里。我个人觉得它的 API 设计有点「重」，但好处是调参方便——改个 stage 就能切换 ZeRO 等级。我曾经在 ZeRO-2 和 ZeRO-3 之间反复横跳，就为了找到显存和通信开销的平衡点。

3.4 选型决策树：一张图说清楚

我画了一张决策图，帮你快速定位该用哪个框架。

这张图是我自己总结的，不一定适合所有场景，但至少能帮你少走弯路。举个例子，我有个朋友非要用 Horovod 训一个 7B 的模型，结果显存爆了，换成 DeepSpeed 之后才跑起来。

3.5 我的最终建议

说了这么多，其实就三句话：

• 日常实验、小模型、快速迭代：无脑 PyTorch DDP。它跟 PyTorch 的生态绑得最紧，torch.compile、torch.fx 这些新特性都能直接用。
• 显存不够、卡数有限、模型较大：上 DeepSpeed。ZeRO 的显存优化是实打实的，而且微软一直在更新，现在 ZeRO-3 已经支持 CPU/NVMe Offload 了。
• 多机多卡、异构环境、需要兼容 TensorFlow：考虑 Horovod。它的 MPI 后端在跨机通信上很成熟，而且支持 TensorFlow、Keras 甚至 MXNet。

好了，框架选型这块就聊到这儿。记住，没有银弹，只有最适合你当前场景的工具。