4、模型训练平台：分布式训练框架（TensorFlow/PyTorch），参数服务器与AllReduce架构
好，咱们今天聊聊分布式训练。说实话，这是推荐系统里最“硬核”的部分之一。很多同学把模型搭好了，数据也准备好了，结果一上多机训练就崩。嗯，我当年也踩过不少坑。
分布式训练的核心就一句话：把计算和通信拆开，让多张卡、多台机器像一个人一样思考。但怎么拆、怎么通信，这里面的门道可不少。
4.1 为什么推荐系统需要分布式训练？
你想想看，一个典型的推荐模型，Embedding表可能有几十GB，甚至上百GB。单卡显存才多大？A100也就80GB。更别说每天几亿条样本，单机训练一个迭代要跑好几天。
我做过一个项目，模型参数有50亿。单机训练？别想了。后来上了分布式，训练时间从一周压缩到6小时。说白了，分布式训练不是“锦上添花”，而是“雪中送炭”。
核心痛点：
- 单卡显存装不下大Embedding
- 单机算力跟不上数据增长
- 训练时间过长，无法快速迭代
4.2 两大主流框架：TensorFlow vs PyTorch
现在工业界基本就是TensorFlow和PyTorch二分天下。我个人习惯用PyTorch做研究，用TensorFlow做生产。为什么？
| 维度 | TensorFlow | PyTorch |
|---|---|---|
| 分布式原生支持 | tf.distribute.Strategy 很成熟 | torch.distributed 后来居上 |
| 动态图 | Eager模式（后来加的） | 原生动态图，调试方便 |
| 生产部署 | TF Serving 生态完善 | TorchServe 相对年轻 |
| 推荐场景 | Google系（Wide&Deep等） | Meta系（DLRM等） |
我记得有一次，团队用PyTorch写了个模型，训练没问题，但上线时发现TF Serving不支持。最后不得不转成ONNX，折腾了两周。所以我的建议是：如果团队偏工程，选TF；如果偏研究，选PyTorch。
4.3 参数服务器架构（Parameter Server）
参数服务器，说白了就是“主从模式”。有一组机器专门存参数，其他机器只负责计算。这个架构在推荐系统里特别常见。
为什么？因为推荐模型的Embedding太大了。你想想看，几亿个用户ID，每个ID对应一个128维的向量。如果每个worker都存一份，内存直接爆炸。参数服务器把参数集中管理，worker只拉取自己需要的部分。
我的经验：参数服务器架构适合“大Embedding + 小DNN”的场景。比如CTR预估模型，Embedding占了90%的参数，但计算量不大。这时候用PS模式，通信开销可控。
具体流程是这样的：
- Worker从PS拉取当前batch需要的参数
- Worker计算梯度
- Worker把梯度推送给PS
- PS更新参数，再分发给下一个batch
但这里有个坑：通信瓶颈。我曾经遇到一个项目，100个worker同时向10个PS发梯度，结果PS的网络被打满了。后来我们做了梯度压缩，把32位浮点压成8位，通信量直接降到1/4。
避坑指南：我曾经在参数服务器上踩过一个雷——同步更新。所有worker必须等最慢的那个算完才能更新参数。如果某个worker卡住了，整个训练就停了。后来我改成了异步更新，虽然收敛慢一点，但至少不会死锁。
4.4 AllReduce架构
AllReduce是另一种思路。它没有中心节点，所有机器既是计算节点，也是通信节点。说白了，大家“平起平坐”。
最经典的算法是Ring AllReduce。想象一下，4台机器围成一个环，每台机器只跟左右邻居通信。数据分成4份，每台机器先把自己的一份传给邻居，同时接收邻居的一份。这样转一圈，每台机器都拿到了所有数据。
为什么AllReduce适合小模型？因为通信量跟模型大小成正比，跟机器数量无关。比如一个1GB的模型，4台机器做AllReduce，每台机器只传1GB数据。而参数服务器呢？每台机器要传4GB（因为要跟所有PS通信）。
| 对比项 | 参数服务器 | AllReduce |
|---|---|---|
| 通信模式 | 中心化（PS） | 去中心化（Ring） |
| 适合场景 | 大Embedding模型 | 小参数模型（如ResNet） |
| 通信量 | O(N*M) N=worker数, M=PS数 | O(P) P=模型大小 |
| 容错性 | PS单点故障 | 任意节点故障影响全局 |
我个人的建议是：如果你的模型Embedding占比超过70%，用参数服务器；如果模型小而深，用AllReduce。当然，现在也有混合方案，比如Horovod就支持两种模式混用。
4.5 实战：如何选择分布式策略？
好，理论说完了，咱们聊聊实战。假设你现在要训练一个CTR预估模型，Embedding表有50GB，DNN部分只有500MB。你怎么选？
我的做法是这样的：
- 第一步：看模型大小。50GB的Embedding，单卡肯定装不下。所以必须用参数服务器。
- 第二步：看数据量。每天1亿条样本，单机处理要3天。所以必须多机并行。
- 第三步：看网络。如果机器都在同一个机架，万兆网络，那PS模式没问题。如果跨机房，延迟高，那就得考虑梯度压缩。
我记得有一次，我们用了32台机器做PS模式训练。结果发现网络利用率只有30%。排查了半天，发现是PS的线程数没调好。默认是4个线程，但我们的机器有64核。改成32个线程后，网络利用率直接飙到90%。
核心原则：
- 大Embedding + 小DNN → 参数服务器
- 小Embedding + 大DNN → AllReduce
- 混合场景 → 分层策略（Embedding用PS，DNN用AllReduce）
4.6 代码示例：PyTorch分布式训练
最后，给一段PyTorch的分布式训练代码。这是我在项目中常用的模板。
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
# 初始化进程组
dist.init_process_group(backend='nccl')
local_rank = dist.get_rank()
# 模型放到对应GPU
model = MyModel().cuda(local_rank)
model = DistributedDataParallel(model, device_ids=[local_rank])
# 数据加载器（每个进程只取一部分）
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = DataLoader(train_dataset, batch_size=32, sampler=train_sampler)
# 训练循环
for epoch in range(10):
train_sampler.set_epoch(epoch) # 打乱数据
for data, target in train_loader:
data, target = data.cuda(local_rank), target.cuda(local_rank)
output = model(data)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
optimizer.zero_grad()
这段代码看起来简单，但有几个关键点：
DistributedSampler保证每个进程拿到的数据不重复nccl是NVIDIA的通信库，比gloo快很多set_epoch必须在每个epoch开始时调用，否则数据顺序不变
小技巧：我习惯在训练脚本里加一个--local_rank参数，这样可以用torch.distributed.launch启动。比如：python -m torch.distributed.launch --nproc_per_node=4 train.py
好了，分布式训练的核心内容就这些。说白了，选对架构比调参更重要。你想想看，如果架构选错了，后面再怎么优化也是事倍功半。下一章咱们聊聊模型评估和A/B测试，那又是另一个坑了。