跳转到内容
Jason's Study

Parameter Server — 多机训练前 AllReduce 时代的工业标准

是什么

Parameter Server（参数服务器，简称 PS） 是一套让几百台机器协作训练同一个机器学习模型的架构。日常类比：班里 100 个同学要合作写一本书，但一本书太厚没法每人都抄一份；于是把书拆成 10 章，10 个图书管理员各管一章；同学想改第 3 章就找 3 号管理员，改完同步回去。

把”书”换成模型参数，“同学”换成worker（算梯度的机器），“管理员”换成server（存参数的机器），就是 PS。

2014 年 Mu Li 在 OSDI 发的这篇论文，把这套架构第一次系统化讲清楚——之前 Google DistBelief（2012）已经在用类似思路，但 PS 把接口、容错、一致性都标准化了。

为什么重要

不理解 PS，下面这些事都没法解释：

  • 为什么 2015-2017 年所有”分布式深度学习”教程都先讲 PS（TensorFlow 1.x 默认就是 PS 模式）
  • 为什么后来 Horovod / PyTorch DDP 火了之后，PS 反而变成”过时方案”——但推荐系统团队还在用
  • 为什么今天 LLM 训练里 ZeRO 的参数分片和 PS 的 server 分片几乎是同一个思想
  • 为什么联邦学习的拓扑（中心服务器聚合梯度、客户端本地训练）跟 PS 一模一样

核心要点

PS 的设计可以拆成 三个角色 + 两个动作：

三个角色

  1. server（参数服务器）：存模型参数。参数太大一台存不下，按 key（参数 ID）切分到多台 server，类比图书管理员每人管一架书。
  2. worker（计算节点）：拿一批数据，算出梯度，推给 server。类比写作小组成员，每人负责改一章。
  3. scheduler（调度器）：监控 worker 健康、分配任务、做容错。类比班长。

两个动作

  • push(key, gradient)：worker 算完梯度，推给负责这个 key 的 server
  • pull(key)：worker 开始下一轮前，从 server 拉最新的参数

整个训练循环就是：每个 worker 反复 pull → 算梯度 → push。

关键设计：异步更新 + Bounded Delay

最朴素的做法是同步：所有 worker 算完一轮才一起更新。但只要有一台慢（straggler，掉队者），全班等它。

PS 用异步：worker 不等别人，自己 push 完就拉新参数继续算。问题：可能拿到的参数太旧（别人已经更新了 50 轮，你还在用 100 轮前的版本），训练发散。

折中方案叫 Bounded Delay：允许异步，但落后超过 τ 轮就必须等。τ=0 是同步，τ=∞ 是完全异步，工业上常取 τ=4。

容错：server 挂了怎么办

PS 用一致性哈希 + 链式复制：每个参数 key 在 3 台 server 上有副本，主副本挂了从副本顶上。worker 看到的接口不变。

实践案例

案例 1：训练千亿参数 LR 模型

百度 2014 年用 PS 训练点击率预测的 LR 模型（逻辑回归）：

  • 参数量：1000 亿（每个广告位×特征组合一个权重）
  • 数据量：PB 级日志
  • 机器：6000 worker + 600 server
  • 关键：参数稀疏，每条数据只激活几百个 key，所以 push/pull 只走少量 key，带宽吃得起

这是 AllReduce 完全做不到的——AllReduce 要每台机都存全量参数，1000 亿权重单机塞不下。

案例 2：PyTorch RPC 里的 PS 实现

PyTorch 1.4 之后官方支持 torch.distributed.rpc，可以手写 PS：

# server 端：存参数
class ParamServer:
    def __init__(self):
        self.params = torch.zeros(1000)
    def push(self, grad):
        self.params -= 0.01 * grad
    def pull(self):
        return self.params


# worker 端：算梯度
def worker_loop():
    while True:
        params = rpc.rpc_sync("ps", ParamServer.pull)
        grad = compute_grad(params, data)
        rpc.rpc_sync("ps", ParamServer.push, args=(grad,))

这就是 PS 的最简骨架。真实工业版本要加分片、副本、压缩。

案例 3：今天还在用 PS 的地方

推荐系统的 sparse embedding 表：

  • 一个 user × item 矩阵分解模型，user embedding 表 10 亿行 × 64 维 = 256 GB
  • 单机塞不下，但每个 batch 只读几千个 user → PS 完美匹配
  • 阿里 XDL、腾讯 Angel、ByteDance BytePS 都用 PS 做 embedding，dense 部分用 AllReduce

联邦学习：中心服务器 = server，客户端 = worker，push 梯度、pull 全局模型，拓扑完全一样。

踩过的坑

  1. 异步训练收敛慢：τ 设太大模型发散，设太小退化成同步。τ 是个调参变量，不是越大越好。

  2. server 是瓶颈：所有 worker 都往 server 推，server 网卡先打满。后来工程上用梯度压缩（FP16 / 1-bit SGD）缓解。

  3. PS 不适合 dense 大模型：BERT/GPT 这种参数密集且每步全量更新的模型，PS 的”按 key 拉一部分”优势没了，AllReduce 反而更快。

  4. 代码复杂度高：worker 和 server 是两套不同代码，调试痛苦。这也是 DDP（一份代码跑所有节点）后来流行的原因。

适用 vs 不适用场景

适用：

  • 参数稀疏（LR / LDA / 协同过滤 / 推荐系统 embedding）
  • 模型太大单机塞不下，但每步只更新一小部分
  • 需要异步训练（容忍 straggler 的场景）
  • 联邦学习（天生中心化拓扑）

不适用：

  • dense 模型（CNN / Transformer），每步全量更新 → AllReduce 更优
  • GPU 间带宽足够（NVLink / IB），不需要中心聚合点
  • 强同步需求（金融模型、复现实验）→ 同步 AllReduce 更稳

历史小故事（可跳过）

  • 2012 年：Google DistBelief（Jeff Dean 团队）用 PS 训练 ImageNet，证明可行
  • 2014 年：Mu Li 在 CMU 读博，跟 Alex Smola、David Andersen 把 PS 系统化，OSDI 发文
  • 2015 年：Mu Li 主导写 MXNet，PS 是默认通信后端
  • 2017 年：Uber 开源 Horovod，AllReduce 在 GPU 训练上替代 PS
  • 2020+：BytePS 把 PS 和 AllReduce 混用，证明两者不是对立而是互补

学到什么

  1. 架构选择跟硬件演化绑死：2014 年 CPU 集群 + 千兆网，PS 最优；2017 年 GPU + IB 网，AllReduce 翻盘。没有永恒的最佳架构，只有当下硬件下的最优解。
  2. 异步 vs 同步是经典权衡：吞吐 vs 收敛性。Bounded Delay 是这道题的工程解。
  3. 稀疏 vs 稠密决定通信模式：参数稀疏→中心化 PS，参数稠密→点对点 AllReduce。
  4. 工业系统的论文价值：这篇 OSDI 论文不是新算法，而是把”已经在跑的东西”标准化。这种工作价值不亚于纯理论。

延伸阅读

关联

  • pytorch —— DDP（AllReduce）是 PS 的后继者
  • alpa-2022 —— 现代分布式训练统一框架
  • paxos —— PS 的 server 复制借鉴了类似一致性思路

反向链接

  • alpa-2022 —— Alpa — 把张量/流水/数据并行统一成一道搜索题
  • paxos —— Paxos — 分布式共识算法
  • pytorch —— PyTorch — 深度学习主流框架