2012 年 Google 的 Jeff Dean 发表了《Large Scale Distributed Deep Networks》一文。这篇文章不仅提出了 DistBelief 框架(TensorFlow 前身),还同时基于参数服务器架构(Parameter Server)提出了 Downpour SGD 纯异步模式。在取得性能大幅提升的同时,也有非常不错的效果收益。
2014 年,李沐所在的 DMLC 组发表了论文《Scaling Distributed Machine Learning with the Parameter Server》,提出了第三代参数服务器架构,相比之前的 DistBelief 提供了更加通用的架构:在设计上包含一个Server Group 和若干个 Worker Group,Server Group 用来做参数服务器,每个Server Node 存放一个参数分片,由 Server Manager 管理整个 Server Group,维持整个 Server Group 的元数据的一致性视图,以及参数分片情况。每个 Worker Group 运行一个应用,Worker Node 只跟 Server Node 通信用来更新参数,Worker Node 之间没有任何交互。每个 Worker Group 内有一个调度器,负责给 Worker Nodes 分配任务以及监控,如果有Worker Node 挂掉或者新加入,调度器负责重新调度剩余的任务。PS-Lite 针对网络带宽的优化主要是针对 Server之间的参数复制提供了采用先聚合再复制的方式:
Replication after Aggregation
Server 之间复制主要是为容错考虑,因此 Worker 和 Server 之间的数据传输仍依赖参数服务器本身异步机制带来的带宽节省:在应用于深度学习时,主要借助于称作 Delayed Block Proximal Gradient 的方法(包括收敛性证明):每次迭代只更新一个 block 的参数;Worker 节点计算梯度的同时还需要计算特定坐标的学习速率,即该 block 的二阶偏导数的对角线。在数据传输时,PS-Lite 还会引入部分定义的 Filter 来避免传输对模型影响不大的参数,例如 Random Skip或者KKT Filter,引入这些Filter可以让传输的模型参数量减少十倍以上。
在 Parameter Server 中,每个 PServer 实际上只负责分到的部分参数(PServer 共同维持一个全局的共享参数,每个 Worker 也只分到部分数据和处理任务)它具有以下优势:
Ease of Use:全局共享的参数可以被表示成各种形式:vector、matrices 或是 sparse类型,同时框架还提供对线性代数类型提供高性能的多线程计算库。
它的具体实现:
Vector Clock
使用 range vector clock 来记录每个节点参数的时间戳,用来跟踪数据状态或避免数据重复发送。由于参数都是 Range Push/Range Pull,因此同一个 key range 里的参数可以共享一个时间戳,相较于传统的 vector clock 进行了压缩,降低了内存和网络带宽开销。
Messages
节点之间通信发送的 message 由 range vector clock 和 对组成。
由于频繁更新模型参数,需要对 message 进行压缩以减少网络带宽开销,Parameter Server 采用两种方法来压缩 message:
value 的压缩:有些参数更新并非对最终优化有价值,因此用户可以自定义过滤规则来过滤掉一些不必要的参数。例如对于梯度下降,大量 value 为 0 或者很小的梯度是低效的,可以过滤。
在数据一致性上使用一致性哈希算法,参数 key 和 server id 被插入到哈希环中,有两种方式保证了主节点和备份节点之间的数据一致性:
Chain Replication
Repiication after Aggregation
Server 管理
添加 server:server manager 给新 server 分配 key range,其他 server 的 key range 做出相应更改;新 server 获取作为主 server 维护的 key range 和作为从 server 维护的 k 个 key range;server manager 广播节点的更改;
删除 server:当 server manager 通过心跳信号发现 server 出现故障后,会将该 server 的 key range 分配给新的 server,并删除该 server
该系统包含若干个 Server Group 和 Worker Group。每个 Server Group 维护完整的模型参数,相邻的 Server Group 定期同步模型参数。每个 Worker Group 负责计算,在数据并行中,每个 Worker Group 负责一个训练数据的分片,所有的 Worker Group 节点跟Server Group之间异步通信,然而在 Worker Group 内部,则是同步处理SINGA值得称道之处在于同时支持数据并行和模型并行,以及混合并行(两种并行的组合),在模型并行时,每个Worker Group 只负责更新部分模型参数。
Facebook 在 2019 年发表了一篇论文讲DLRM模型《Deep Learning Recommendation Model for Personalization and Recommendation Systems》。它基于 pytorch 和 caffe2 实现。在 embedding table的模型并行模式上实现了一种特殊的并行范式(butterfly shuffle)。同时,这篇论文将这种并行模式运行在Facebook 的超级 AI 计算机 Big Basin 平台上。
蝶形 shuffle
2019 年阿里推出 PAI Tensorflow(阿里云优化的商业版 Tensorflow),它支持的功能包括:静态特征回填(新上线特征);实时样本拼接;embedding variable 功能。这个功能提供动态的弹性特征的能力。每个新的特征会新增加一个 slot。并支持特征淘汰,比如说下架一个商品,对应的特征就会被删掉;实时训练模型校正;模型回退及样本回放。
2019 年字节跳动开源 BytePS,它继承了许多加速技术,比如分层策略、流水线、张量分区、NUMA 感知本地通信、基于优先级的调度等等。还提出了 BytePS ring-allreduce 方案:一是引入 CPU 节点,用于 ReduceScatter 操作;二是 ReduceScatter 和 AllGather 异步执行,从而缩短了数据传输时间开销。本质上还是通过引入 CPU Server,相当于额外增加一张网卡。
BytePS 架构
2020 年 360 公司开源了 TensorNet,该系统是针对 tensorflow 开发的轻量级 PS,支持 tf2.2 及以后的版本,TensorNet 能支持到百亿特征规模,它通过引入 virtual sparse feature,由原来直接从 parameter server 查找稀疏特征变为二级查找。TensorNet将每个batch 的特征 ID 从 0 开始重新编排作为输入,这样输入特征的 index 分布在 [0, 1024) 之间;同时根据原始的特征 ID 从 server 拉取对应的 embedding 向量,填充到 tensorflow 的 embedding 矩阵中,这样每个特征 field/slot 的 embedding 矩阵大小就只有 1024 x embedding_size.
[4] Poseidon: A System Architecture for Efficient GPU-based Deep Learning on Multiple Machines, Zhang, Hao and Hu, Zhiting and Wei, Jinliang and Xie, Pengtao and Kim, Gunhee and Ho, Qirong and Xing, Eric, arXiv preprint arXiv:1512.06216
[5] GeePS: Scalable deep learning on distributed GPUs with a GPU-specialized parameter server, Cui, Henggang and Zhang, Hao and Ganger, Gregory R and Gibbons, Phillip B and Xing, Eric P, Proceedings of the Eleventh European Conference on Computer Systems, 2016
[6] https://github.com/yahoo/CaffeOnSpark
[7] Deep Learning At Scale and At Ease, Wang, Wei and Chen, Gang and Chen, Haibo and Dinh, Tien Tuan Anh and Gao, Jinyang and Ooi, Beng Chin and Tan, Kian-Lee and Wang, Sheng, arXiv preprint arXiv:1603.07846
[9] Deep learning with Elastic Averaging SGD, Sixin Zhang, Anna Choromanska, Yann LeCun, arXiv preprint arXiv:1412.06651
[10] Communication Efficient Distributed Machine Learning with the Parameter Server, Li, Mu and Andersen, David G and Smola, Alex J and Yu, Kai, NIPS 2014
[11] 1-bit stochastic gradient descent and its application to data-parallel distributed training of speech DNNs, Seide, Frank and Fu, Hao and Droppo, Jasha and Li, Gang and Yu, Dong, INTERSPEECH 2014
[12] SINGA:SINGA Architecture
[13] NuPS: A Parameter Server for Machine Learning with Non-Uniform Parameter Access: https://arxiv.org/pdf/2104.00501.pdf
[14] 顶级 AI 学者邢波教授:机器学习缺乏清晰理论与工程框架,需重新思考评估方法及目标: 新闻公告-中国科学院自动化研究所图书馆
[15] Simplifying and Automating Parallel Machine Learning via a Programmable and Composable Parallel ML System: https://sites.google.com/view/aaai-2021-tutorial-ah9/home
[17] Distributed Hierarchical GPU Parameter Server for Massive Scale Deep Learning Ads Systems: https://mlsys.org/media/Slides/mlsys/2020/balla(02-08-00)-02-09-40-1408-distributed_hie.pdf)
[29] A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters:https://www.usenix.org/system/files/osdi20-jiang.pdf
[30] Deep Learning Recommendation Model for Personalization and Recommendation Systems: https://arxiv.org/pdf/1906.00091.pdf