经典分布式论文阅读:Parameter Server

4,772 阅读8分钟

本文是李沐大神的Parameter Server论文的学习笔记,李沐大神在OSDI和NIPS上都发过文章,其中OSDI版本偏向于系统设计,而NIPS版本偏向于算法层面,本文显然是OSDI的文章。

本文提出了“参数服务器”的分布式架构来支持分布式机器学习。分布式机器学习主要面临以下挑战:

  • 访问参数消耗大量的网络带宽
  • 很多机器学习算法都是串行的
  • 容错能力非常重要

而本文中的框架给开发者带来两点优势:

  1. 通过分解机器学习的组件,可以让业务代码更加简洁
  2. 能够实现鲁棒、多功能、高性能的分布式机器学习算法

本系统主要有以下五个特点:

  • 高效的通信:使用了异步非阻塞通信模型
  • 灵活的一致性模型:可以允许系统设计人员手动权衡收敛率和系统效率
  • 弹性的规模:能够在运行期间添加新节点
  • 容错和耐久:能从故障中快速恢复,通过向量时钟保证行为的确定性
  • 易用:参数表示为向量和矩阵便于开发机器学习算法

系统的这些特性都是通过选择正确的系统技术,运用在机器学习算法中,以及修改机器学习适应系统来实现。在系统实现过程中,主要面临以下挑战:

  • 通信:通过批量传输参数(向量片段、矩阵的行列而不是单个参数值)来提高通信效率
  • 容错:通过实时备份和热修复和实现

机器学习

机器学习需要从训练数据中学习模型,主要包含三个要素:特征提取目标函数学习。特征提取将原始训练数据转换为特征向量,不在本文赘述。学习的过程就是最小化目标函数从而获得模型。另外,在分布式机器学习任务中,训练数据量也通常是非常巨大的。

风险最小化

监督学习就是风险最小化的过程,例如最小化预测误差。如果有n个训练样本,每个样本的特征向量为x_i,对应的标签为y_i,模型的参数为w,目标函数为

F(w)=\sum_{i=1}^{n} \ell\left(x_{i}, y_{i}, w\right)+\Omega(w)

其中\ell\left(x_{i}, y_{i}, w\right)为损失函数,定义了预测值和真实值之间的误差,\Omega(w)为正则化项,用来防止模型过拟合。在参数服务器框架中可以采用分布式子梯度下降对目标函数进行最小优化

在分布式子梯度下降算法中,每个工作节点只需要计算分配到的参数工作集w对应的梯度,然后由服务节点完成聚合。模型的完整参数w可能会十分巨大,工作节点在使用的时候会面临很大的负担,但是可以通过只保存用到的参数值即可。

生成模型

另外一种机器学习的形式为无监督学习,通常用来学习数据的自身结构。比较典型的就是话题模型:给一些文档,推断出每个文档包含的话题。主题模型的挑战就是:关于当前文档如何生成的参数必须被共享。解决方法就是每个工作节点只保存分配到的文档出现的词有关的参数即可。

架构

参数服务器系统由一个服务节点组和多个工作节点组构成。服务节点之间互相通信来备份和迁移参数,服务管理节点负责维护服务节点元数据之间的一致性。一组工作节点运行一个应用程序,工作节点组中的调度节点负责任务的分配和监控。

参数服务器以命名空间的方式组织参数,模型的参数采用键值的形式保存。不同的应用程序可能会共享命名空间,例如一个应用程序负责模型训练,另一个应用程序负责模型推断。

范围推送和拉取

为提高带宽利用率,系统支持范围推送和拉取。令\mathcal R为键范围,那么

  • w.push(R,dest)w中键范围在\mathcal R中的参数推送到dest
  • w.pull(R,dest)从dest拉去w中键范围在\mathcal R中的参数

服务节点的用户定义函数

服务节点除了从工作节点聚合数据之外,也可以执行用户定义函数。这样一来,用户可以实现可以实现一些更加高级的优化算法。

异步任务和依赖

任务都是异步执行的:调用者发起一个任务之后,可以马上执行其他运算。为了提高模型收敛率,可以设置某个任务执行完成后运行的依赖关系,设置任务依赖关系可以保证算法的逻辑。

灵活的一致性

工作节点可以并行执行分配的任务,但是可能会对学习算法的收率产生影响。系统效率和一致性之间的权衡关系取决于算法对于不一致的敏感程程度以及系统硬件能力,本框架提供了三种模式供设计者选择:

  • 顺序一致:下一个任务必须在前一个任务完成之后才能执行
  • 最终一致:所有任务一起开始
  • 有界延迟:在\tau时间之前开始的任务全部完成之后才开始任务

用户定义的过滤器

用户可以定义用户定义过滤器选择性地同步部分参数,例如用户可以之推送那些对模型参数有影响的梯度。

实现

向量时钟

为了支持任务依赖图和快速恢复,每个键值对需要一个时钟。如果每个m个参数每个参数都保存一个时间,如果有n个节点,那么一共需要O(nm)空间,更合理的方式是保存范围的时间。

消息

系统中传递的消息有多个在键范围\mathcal R内的键值对以及对应的向量时钟:

\left[\operatorname{vc}(\mathcal{R}),\left(k_{1}, v_{1}\right), \ldots,\left(k_{p}, v_{p}\right)\right] k_{j} \in \mathcal{R} \text { and } j \in\{1, \ldots p\}

消息可能并没有包含范围内全部的键值对,但是那些缺失的键值对的时钟照常更新。

如果每次迭代,工作节点的训练数据没有变化,那么键应该是不变的,那么可以让接收放保存键缓存,而工作节点只需要发送值和键列表的哈希即可。另外,使用用户自定义过滤器可以进一步减少需要发送的键值对数量。

一致哈希

服务节点组中的节点使用分布式哈希表来保存模型参数。为了简化设计,系统使用直接映射,由服务管理节点统一管理。

副本和一致

每个服务节点保存了逆时针方向k个邻居键范围内的参数的副本,作为这些副本的从节点。副本更新的方式可以是

  • 在更新参数的时候,更新消息也会推送给保存副本的从节点
  • 在完成参数参数聚合后推送给从节点

服务节点管理

当一个服务节点加入服务节点组之后:

  1. 服务管理节点分配给新节点一个键范围,新节点将作为这个范围的参数的主节点
  2. 节点获取这个范围内的参数并成为主节点,以及获取k个额外范围的参数作为从节点
  3. 服务管理节点广播更改。其他节点会释放不再需要自己管理的参数,并且叫没完成的任务交给新节点

新节点从某节点\mathcal S拉取范围\mathcal R内的参数的过程可以分为两步:

  1. \mathcal S预先拷贝一份原先的全部键值对数据以及对应的时钟,当新节点下线时可以用来恢复;
  2. \mathcal S不再处理范围\mathcal R内的消息,并且把预拷贝阶段的更改发送给新节点。

当节点N收到节点添加消息后,需要:

  • 删除不再需要自己管理的参数
  • 重新发送未确认的消息,去掉不属于自己管理的内容

在某个服务节点下线后,服务管理节点需要把该节点管理的参数分配给其他节点。

工作节点管理

当一个新的工作节点W被加入工作节点组后:

  1. 任务调度节点分配给W一部分数据
  2. 节点从网络文件系统或者其他工作节点加载分配的训练数据,然后从服务节点拉取的参数
  3. 任务调度节点广播更改,其他工作结果可能需要释放一些重复的训练数据

当一个工作节点离线之后,可以选择重新分配或者无视,系统将这个选项交给设计者,因为:

  1. 当训练数据量非常大的时候,恢复一个工作节点的代价比恢复服务节点的代价大很多
  2. 丢失一小部分训练数据不会对最终模型造成太大影响

参考文献

  1. Li, Mu, et al. "Scaling distributed machine learning with the parameter server." 11th {USENIX} Symposium on Operating Systems Design and Implementation ({OSDI} 14). 2014.
  2. Li, Mu, et al. "Parameter server for distributed machine learning." Big Learning NIPS Workshop. Vol. 6. 2013.