[OSDI'14] Scaling Distributed Machine Learning with the Parameter Server 阅读笔记

发表于 分类于 阅读次数： 本文字数： 4.1k 阅读时长 ≈ 15 分钟

沐神投稿在 OSDI '14 的论文, 提出了第三代 Parameter Server 框架.

沐神在 B 站有亲自讲解这篇论文的视频, 可以参考: 参数服务器 (Parameter Server)逐段精读【论文精读】.

#0. 摘要

本文提出了一个 Parameter Server 框架, 用于解决分布式机器学习的问题. 数据和负载分布在多个计算节点 (Worker Node) 上, 服务节点 (Server Node) 维护全局共享的参数, 参数表示成稠密或稀疏的向量和矩阵. 框架管理节点之间的异步数据通信, 支持灵活的一致性模型, 弹性的可扩展性和持续的容灾.

为了展示该框架的可扩展性, 文中使用了从稀疏 LR 到 LDA 和 Distributed Sketching 等的任务, PB 级别的真实数据、十亿级别的数据样本和参数的规模上的实验结果.

#1. 介绍

分布式优化和推理正在成为解决大规模机器学习问题的先决条件. 真实世界的训练数据量可以达到 1TB 到 1PB. 这使得我们创建一个具有10910^9 (1B) 到101210^12 (100B) 参数量的强大复杂模型. 通常这些模型被所有的计算节点共享, 计算节点还需要经常访问共享参数并进行更新. 共享带来三个挑战:

  • 通信: 所有的计算节点都需要频繁地访问参数, 会导致大量的通信.
  • 性能: 很多机器学习算法是顺序的模型 (算完一个 batch 再继续算下一个), 会引入大量同步机制, 损伤性能.
  • 容灾: 训练任务应该在机器故障、软件错误或者由于人为原因被占用的情况下, 不会停止.

为了更好地说明第三点, 文章收集了一家互联网公司数据中心 (Baidu) 最近三个月的训练作业日志. 统计结果显示, 规模越大的任务的失败率越高. 和实验室环境相比, 工业界环境有资源竞争, 因此必须要考虑容灾问题.

#1.1 贡献

PS 思想最早可以源自于 Alexander Smola 自 2010 年提出的 YahooLDA 框架. YahooLDA 采用了一个分布式的 memcached 服务存储机器学习参数.

第二代 Distbelief 是 Google 提出, 使用独立的 PS 参数存储巨大的深度学习模型的参数.

PS 框架已经在学界和工业界广泛使用, 本文描述了第三代 PS 框架的开源实现, 提供了一个更加通用的设计.

本文重点关注第三代 PS 框架的分布式接口, 对开发者来说, 第三代框架有两个好处:

  1. 把通用的框架代码和特定任务相关代码分离, 保证任务相关代码的简洁性: 例如, 第三代 PS 可以同时支持稀疏 LR、LDA 和 Distributed Sketching 的一系列算法.
  2. 提供了鲁棒、多样化且高性能的实现, 处理各种各样的算法.

服务节点的管理包括节点的添加和删除.

第三代 PS 框架的设计决策基于真实工业界系统计算负载. 作者总结了五个关键特征:

  • 高效通信: 采用了异步的通信模型, 不会阻塞计算, 对通信数据进行压缩;
  • 灵活一致性模型: 核心思想是 trade-off: 舍弃掉机器学习算法部分的一些指标 (收敛速度, 精度等), 换取系统部分更好的性能;
  • 弹性的可扩展性: 在训练时可以动态增加和减少机器;
  • 容灾和持久性: 少量机器挂掉可以在 1s 内恢复, 使用向量时钟确保灾难情况下的行为;
  • 易用性: 当时的主流语言是 C++, 主流库是 Eigen 等, PS 框架把全局参数抽象成向量和矩阵, 支持已有的库;

创新性: 巧妙地调整系统领域和机器学习领域的技术, 使之合理地结合在一起, 得到第一个通用的、能够扩容到企业规模的 ML 系统.

#1.2 工程挑战

解决分布式数据分析问题时, 多个计算节点需要不断地读和写全局参数. PS 框架提供了一个高效的机制来在多个计算节点间汇聚和同步这些参数和一些统计信息. 由于模型的总体参数量很大, 每个服务节点 (Server Node) 都只负责其中的一部分. 计算节点通常会先朝服务节点索要一部分参数, 进行计算, 然后再把计算结果发回服务节点. 构建高性能的 PS 系统的关键挑战在于:

  1. 通信: 传统的 datastore 采用 key-value 模型, 但是这种抽象往往不适用于机器学习场景: 机器学习场景的 value 通常是很小的 floats 或者 integers, 而每个更新都进行一次数据通信的开销很大.

    PS 框架的设计是, 计算节点每次只发送向量或者矩阵的一部分 (vector 的 segment, 或者 matrix 的 row). 这样就可以批量地更新参数, 并允许高效地实现一致性跟踪.

  2. 容灾: PS 框架采用 live replication 的方式, 将参数实时复制到多台机器上, 并且支持 hot failover.

    机器的加入和离开会被视为维修和故障. 服务节点的管理包括节点的添加和删除.

#1.3 相关工作

第一代 PS 包括 YahooLDA, 利用通用的 memcached 数据库作为同步机制, 因此缺乏弹性和性能;

第二代 PS 包括 Distbelief 和 Petuum, 引入了 Bounded delay model, 以及做了一系列工程优化.

对比通用分布式系统, Mahout (基于 Hadoop)和 MLI (基于 Spark)采用迭代式的 MapReduce 框架, 要求同步和迭代式的通信, 不适用于大规模的机器学习任务.

GraphLab 使用图模型, 可扩展性差; Piccolo 缺乏消息压缩、复制和灵活一致性模型.

#2. 机器学习背景知识

相关概念介绍: 特征提取, 目标函数, 学习过程, 损失最小化, 学习率, 特征向量, 生成模型, 主题建模, LDA,

以分布式梯度下降任务为例: 包括任务调度器, 计算节点, 服务节点三个部分.

  • 计算节点: LoadData(), WorkerIterate()
  • 服务节点: ServerIterate()

algo-1

#3. PS 架构

一个 PS 示例包含四大元素: 一个 Server Group, 若干个 Worker Group, 一个 Resource Manager 以及 Training Data.

Worker Group 负责运行一个任务, 包含一个 Task Scheduler 和若干个 Worker Node.

Server Group 负责存储参数, 包含一个 Server Manager 和若干个 Server Node.

分多个组的目的是让系统可以同时执行多个任务, 例如同时训练多个模型或者同时训练和推理 (一个组用于在线服务, 另一个组周期性地更新模型). 多个任务通过命名空间隔离.

#3.1 (Key, Value) 向量

PS 框架中的参数可以表示成一个 (Key, Value) 向量, 其中 Key 是一个整数, Value 是整数、浮点数或向量.

这种表示法使得用户可以使用现有的线性代数库 (例如 BLAS、LAPACK 和 ATLAS) 来提高编程效率.

#3.2 区间 Push 和 Pull

PS 框架允许在 Push 和 Pull 时指定一个参数的区间, 只会传输区间内的参数, 这样可以减少通信开销.

#3.3 用户自定义函数 (UDF)

PS 框架允许用户在服务节点上执行用户定义的函数, 例如更新参数的函数, 这样更加灵活.

#3.4 异步任务和依赖关系

任务通过 RPC 调用来实现.

任务是异步执行的, 发送 RPC 出去后, 发送方不会等待接收方的回复, 而是继续执行下一个运算. 任务的回复可能是 UDF 的返回值, 或者是被请求的 (Key, Value) 对, 或者是一个空的 ACK.

任务之间可以有依赖关系, 例如调用方可以指定一个 execute-after-finished 依赖在任务 A 和任务 B 之间, 表示任务 A 在任务 B 完成后才能执行. 默认情况下, 任务之间是完全并行的.

例如多次 Iteration 可以流水线式的执行, 提高效率.

#3.5 灵活一致性模型

通过依赖关系的不同组合方式可以实现各种不同程度的一致性模型

  • 顺序一致性 (Sequential): 所有任务都有依赖, 即 Bulk Synchronous Processing, BSP
  • 最终一致性 (Eventual): 任务间完全没有依赖, 仅推荐用于算法足够 robust 的情况
  • 有界延迟 (Bounded delay): 第 n 个任务依赖n−τn-\tau 之前的所有任务

所有依赖关系会形成一个有向无环图(DAG). 依赖关系可以是动态的.

#3.6 用户自定义过滤器

用户可以定义过滤器, 来选择性地同步部分参数数据, 达到节省通信的目的.

例如显著修改过滤器, 只推送修改超过阈值的条目.

#4. 实现

本章讲解了 PS 框架的具体实现细节. 服务节点使用一致性哈希来存储参数. 使用链式复制进行多副本. 在数据和向量时钟上进行了压缩.

#4.1 向量时钟 & 4.2 通信消息

消息的格式是一个向量时钟和若干个 (Key, Value) 对.

Server 端会缓存区间的哈希, 当发送相同区间的时候, 可以只发送 Value 部分, 节约通信开销.

由于机器学习场景会有很多零值, PS 框架使用可以删除零值的 Snappy 库来压缩消息.

#4.3 一致性哈希

利用经典的 DHT 技术实现动态地增加和减少服务节点, Key 和节点 ID 都被插入到哈希环中. 为了平衡节点间的负载, 每个服务节点会创建多个虚拟节点.

#4.4 复制和一致性

除了负责自己的参数以外, 每个服务节点额外为哈希环上 k 个逆时针邻居的参数做副本. 该节点称为相应这些键的从节点.

计算节点只会和主节点通信, 对主节点的更新会被同步地复制到从节点 (即 全同步复制).

直接复制需要 k 倍的通信放大, PS 框架允许通过对多个更新先进行一次聚合, 然后再进行复制来减少通信开销.

#4.5 服务节点管理

服务节点的管理包括节点的添加和删除.

当一个新节点加入时, 会进行一个重分配过程:

  1. Server Manager 分配给新节点一段 Key Range. 这可能会导致另一个服务范围从终止的节点中分裂或被删除.
  2. 新节点获取这段 Key Range 的数据, 还有接下来 k 个逆时针邻居的数据.
  3. Server Manager 广播节点更改消息. 消息的接收者可能会缩小自己的数据, 和将未完成的任务重新提交给新节点.

步骤 2, 从源节点 S 获取范围 R 中的数据, 使用了一个两阶段协议. 首先, S 预复制范围 R 内的所有 (key, value) 对及其关联的向量时钟. 这可能会导致向量时钟分裂. 如果新节点在这一阶段失败, 则 S 保持不变. 在第二阶段, S 不再接受影响键范围 R 的消息, 并丢弃未执行和回复的消息 (开启禁写). 同时, S 向新节点发送预复制阶段期间发生在范围 R 的所有更改.

步骤 3, 在收到节点更改消息后, 一个节点 N 首先检查它是否维护着键范围 R. 如果为真, 并且后面不再需要由 N 维护, 则删除 R 中的所有数据和向量时钟. 接下来, N 扫描所有未收到回复的传出消息. 如果一个键范围与 R 相交, 则该消息将被拆分并重新发送.

由于网络延迟、失败和丢失的确认, N 可能会发送两次消息. 由于使用了向量时钟, 原始接收者和新节点都可以拒绝此消息, 并且不会影响正确性.

服务节点的删除 (自愿或由于失败) 和加入过程类似. Server Manager 将新节点分配给离开节点的 Key Range.

Server Manager 通过心跳信号检测节点故障.

与现有集群资源管理器, 例如 Yarn 或 Mesos 的集成将留给未来的工作.

#4.6 计算节点管理

节点 W 加入的过程:

  1. Task Scheduler 给 W 分配一部分训练数据;
  2. W 从 NFS 或者其他计算节点加载这些数据. 训练数据通常是只读的, 所以可以被多个计算节点共享. 然后 W 在从服务节点获取参数;
  3. Task Scheduler 广播这个消息, 其他节点可能会释放一些训练数据.

当删除一个计算节点时, 任务调度程序可以开始替换. 我们给算法设计师提供控制恢复的可选项是出于两个原因:

  • 如果训练数据量很大, 则恢复一个计算节点可能比恢复服务节点更昂贵.
  • 在优化过程中丢失少量训练数据通常只会影响模型很小一部分. 因此, 算法设计师可能宁愿不更换失败的计算节点继续下去. 甚至终止最慢的计算节点也是可以的.

#5. 评估

#5.1 稀疏 LR (1k 台机器)

636 TB 的广告点击预测数据集, 现在看来也非常大了.

#5.2 LDA (Google, 6k 台机器)

#5.3 Sketches (15 台机器)

计算 CountMin Sketches

#参考资料 & FAQ

#1. PS 架构和 TensorFlow, PyTorch 的关系?

PS 架构是一种分布式机器学习的系统架构模式, 核心思想是:

  • Worker: 负责计算 (前向、反向、梯度)
  • Parameter Server (PS): 负责存储和更新模型参数
  • Worker 与 PS 通过网络通信 (pull 参数 / push 梯度)

这是一个通用思想, 不绑定任何具体框架.

TensorFlow 是最早大规模、原生支持 PS 架构的深度学习框架之一.

在 TensorFlow 框架中, 我们可以定义独立的参数节点, 并且将它们放置在独立的 PS 机器上.

在 TensorFlow 1.x (尤其是 1.0~1.15) 中:

  • TensorFlow 的分布式执行模型天然适配 PS
  • 显式区分:
    • ps 节点: 放变量 (tf.Variable), 通过 tf.device("/job:ps/task:i") 放置在 PS 上
    • worker 节点: 放计算图
  • 通信基于 gRPC

在 TensorFlow 2.x 中:

  • 弱化 PS，转向 AllReduce
  • 官方主推 tf.distribute默认策略:
    • MirroredStrategy
    • MultiWorkerMirroredStrategy
  • 核心是 AllReduce (如 NCCL、Ring-AllReduce)

PyTorch 框架原生强支持 Data Parallel + AllReduce 的模式, 原因是动态图 + eager execution:

  • 参数就是本地 Python 对象
  • Autograd 在反向阶段 hook 梯度
  • 梯度 ready -> 立即触发通信

这使得:

  • 梯度通信可以 overlap with backward
  • 实现高效、直观

但是对于 TB 级别 Embedding 的场景, 还是离不开 PS 架构.

#2. 深度学习时代有什么变化?

深度学习的模型通常是稠密的, 参数量小, 计算量大. 因此瓶颈往往在计算上, 而不是通信上. 本文的 PS 框架主要是针对稀疏模型的.

#3. PS 框架的缺点

来自 Horovod 的评论:

  1. 难以确定计算节点和计算节点的正确比例: 如果只有一个计算节点, 它可能会成为通信或者计算的瓶颈. 如果使用多个计算节点, 就变成了 All-to-all 通信, 可能会导致网络饱和,
  2. 应用程序复杂性: 每个应用程序都需要启动计算节点和服务节点, 传递地址和端口号, 等等复杂操作, 占用了实际建模的时间.