- 作者：Tianyu Zhang, Kaige Liu, Jack Kosaian, Juncheng Yang, Rashmi Vinayak
- 机构：Carnegie Mellon University
- 备注：The paper was accepted by the Very Large Data Bases Conference(VLDB‘23).
- 链接：https://dl.acm.org/doi/abs/10.14778/3611479.3611514
研究背景
深度学习推荐模型(Deep-learning-based recommendation models, DLRMs)由参数量占比较大的嵌入表(embedding table)和参数量占比较小神经网络(Neural Network, NN)组成，采用分布式训练，将embedding table和NN的参数分布式存储在server上，然后让worker以数据并行的方式进行异步模型训练，训练过程中先从server取参数，计算完成后再将梯度发送给server进行参数更新，如下图所示：
现阶段，DLRM的规模不断增大
-> DLRM的训练需要大量的资源并且训练时间长
-> DLRM训练过程中容易发生故障，本文关注于server故障，因为server保存了模型参数和状态，相比worker更为重要
-> 检查点是现有的容错方法，需要周期性的暂停训练，将参数和优化器状态写入非易失性存储中，Facebook报告指出检查点存储的时间占了DLRM训练时间的12%
DLRM容错面临的问题与挑战
DLRM的特性：
- 参数量占比较小的NN，参数量占比较大的embedding table
- 每个训练样例的计算需要部分embedding table中的条目(embedding table entry)以及完整的NN
- DLRM的参数量在不断增加，并且增加的主要是embedding table entry的数量
检查点容错的不足：
- 训练时，需要周期性暂停训练，进行检查点存储
- 利用检查点进行恢复时，需要暂停训练，将模型回滚到最新检查点的状态，并重做最新检查点到故障出现这段时间的训练过程
基于检查点的容错方案在训练过程中的开销以及恢复过程中的开销存在一个tradeoff， 即在降低检查点在训练过程中的存储频率时， 会增加恢复过程中的开销；而增加检查点训练过程中的存储频率时，会降低恢复过程中的开销
降低检查点存储开销方法的不足：
- 近似检查点(approximate checkpoints)：收敛性得不到保证
- 异步检查点(asynchronous checkpoint)：无法保证模型状态的一致性，收敛性得不到保证
- 日志(Logging)：梯度产生的速度快于日志存储的速度，无法应用于DLRM容错
理想的DLRM容错方法应该具有：
- 在训练和故障恢复过程中不会暂停训练
- 有较低的内存开销
- 不会影响到模型的收敛性
- 能够扩展到更大的模型
研究内容
本文基于纠删码以及DLRM的特点，设计了DLRM容错系统ECRec。在保证相同精度的情况下，ECRec的容错方法在训练过程中不需要暂停训练、恢复过程中不需要进行回滚
Overview
传统的DLRM训练系统：
ECRec的DLRM训练系统：
不同之处：
- Embedding table以及对应的优化器状态的纠删码块
- NN以及对应的优化器状态的副本
- 冗余参数的更新（图中的虚线部分）
Erasure-coded embedding table entries
将k个server上的embedding table及其优化器状态使用纠删码方法产生r个纠删码块，本文中只关注r=1的情况，因为单点故障时数据中心中最常见的情况，图中的p就是r=1的情况下产生的纠删码块
因为纠删码块的计算涉及了k个server上的embedding table，所以这k个embedding table中的任何一个的更新，对应的纠删码块都需要进行更新，也就是说纠删码块的更新频率比正常的embedding table的频率要高，为了平衡各个server的负载，ECRec对纠删码块使用了循环放置的方法，让每个server上的纠删码块数量相等，其实也就是将k+1个server上的embedding table及其优化器状态均匀划分为k份，然后组织成一个对角线为空的k+1阶矩阵，然后将计算出来的纠删码块填充到对角线上
纠删码块的更新
比较直接的想法就是将worker计算出的梯度直接传给存储对应纠删码块的server，如下图所示：
但是对于带状态的优化器，这样的更新不能保证结果的正确性，以Adam为例，纠删码块中的动量和方差是多个server上动量和方差的和，用这样的动量和方差去更新参数显然是错误的
为了保证纠删码块的正确更新，本文提出了“difference propagation”的方法，不直接将梯度传输给包含纠删码块的server，而是将embedding table以及优化器状态进行梯度更新后的差值传输给含有纠删码块的server，从而正确更新纠删码块
Replicated neural network parameters
NN及其优化器状态使用复制副本方法来进行容错，其原因是：
- worker上只有NN副本，没有对应的优化器状态副本，不能使用worker上的NN副本进行容错
- NN每次迭代都要更新，正常训练过程中NN更新的通信开销在Avazu数据集下占DLRM通信开销的80%-90%，而纠删码方法需要通信参数和优化器状态的差值，会对性能产生严重影响
- NN更新的通信开销较大，但是参数量较小，占比小于DLRM参数总量的1%
因此本文使用了复制副本的方法来对NN及其优化器状态进行容错，在参数更新时，NN的梯度会被wokrer发送给包含NN及其优化器状态副本的server，各个server接收到梯度以后本地对参数进行更新，相比于纠删码方法，降低了对通信开销
Pause-free recovery from failure
假设总共有k+1个server，当某个server出现故障时，不同部分的故障恢复方式：
-
Embedding table:通过剩余的k个server上的embedding table以及纠删码块计算出来
-
NN：从另一存有NN副本的server读取
但是由于embedding table使用了纠删码方法，故障恢复的通信开销较大，如果在故障恢复时暂停训练，对性能影响比较大，ECRec就使用了故障恢复时继续训练的方法，这就带来了一个问题，如果某个server上的embedding table是纠删码方法计算所需要的，但是server收到了embedding table的梯度，该如何处理
如果对embedding table进行更新，那么最后使用纠删码计算出来的结果就会不对，不更新的话训练无法继续进行。为此，本文提出了embedding table entry级别的锁机制，当embedding table entry是纠删码计算所需要的，且该entry对应的部分还未计算完成，就对该entry加锁，同时将其缓存到一个buffer，worker对该entry的读取以及更新操作均在buffer内进行，直到该entry在纠删码中的计算完成，释放锁，将buffer中的entry覆盖embedding table中的原entry
Recovering a consistent DLRM
最后，本文阐述了ECRec如何保持异步DLRM训练的一致性，针对一种特定的情况：embedding table参数更新过程中，在server更新完embedding table之后，将差值传输给纠删码块所在server之前，server出现故障
为了应对这种情况，ECRec提出了两阶段提交（two-phase commit, 2PC）的方法，在第一个阶段，参数更新（包括原始参数和冗余参数）会被计算并且缓存起来，在第二个阶段，利用缓存的值进行参数更新操作
下图展示了一个k=2, r=1的情况下2PC的例子：
第一阶段：
- worker发送梯度给对应的server
- server利用收到的梯度以及优化器状态计算出更新后的参数值，并缓存起来
- server发送差值给存储纠删码块的server
- 收到差值后，server返回一个确认信号给原server
- 原server发送一个确认信号给发送梯度的worker
当worker收到所有被它发送梯度的server的确认信号时， 第一阶段完成
第二阶段：
- worker向所有它发送了梯度的server发送一个提交信号
- 收到提交信号的server将缓存的更新操作应用
- 更新完的server向worker发送一个确认信号
如果第一阶段发生server故障，worker重新开始第一阶段的步骤
如果第二阶段发生server故障，worker只从没有故障的server接受确认信号，由于本文只考虑单点故障，所以剩下的server的参数是能够保证故障server的正确恢复的
实验评估
Evaluation setup
ECRec基于阿里的XDL实现；使用动量SGD优化器
Models
基础模型Criteo-Original，embedding table加上优化器状态是220G
通过增加embedding table entry数量产生的系列模型：Criteo-2S（440G）、Criteo-4S（880G）、Criteo-8S（1760G）
Baselines
Ckpt-30：每隔30分钟存储一次检查点到HDFS（Hadoop Distributed File System）
Ckpt-60：每隔60分钟存储一次检查点到HDFS
No FT：无容错机制
Cluster setup
server：5个AWS实例，256GB内存，25Gbps带宽，对于大于440GB、大于880GB的DLRM会使用更大内存的实例
worker：15个V100，10Gbps带宽
persist storage： 15个配备NVMe SSD的HDFS节点，25Gbps带宽
Performance during normal operation
首先比较了各个容错方法与No FT方法的端到端训练时间，ECRec相比于检查点方法，随着DLRM模型的增大，端到端训练时间增加相对平缓。例如在k=4的情况下，模型从Criteo-Original增加为Criteo-8S，ECRec的训练时间变为原来的1.2倍，而Ckpt-30变为原来的7.2倍，Ckpt-60变为原来的7倍。因此，相比于检查点容错方法，ECRec可以显著的减少模型的端到端训练时间
虽然在小模型上检查点方法的效果比ECRec好，但是在DLRM模型大小不断增加的大背景下，ECRec更加适用于未来的DLRM模型，并且相比于暂停训练的检查点方法，ECRec在故障恢复时性能更好
接下来比较了各个方法在正常训练时的吞吐量，检查点方法由于在进行检查点存储时需要暂停训练，所以检查点方法的平均吞吐量小于ECRec的吞吐量
Q：k=2和k=4情况下，ECRec的吞吐量完全一样
Performance during recovery
文章比较了Ckpt-30以及Ckpt-60与ECRec在故障恢复时的吞吐量，使用的模型是Criteo-4S（880G），故障恢复期间，吞吐量仅仅降低了7%-13%
然后比较了各个方法的故障恢复时间，实验结果如下图，ECRec在k=4时与Ckpt-30和Ckpt-60相比，取得了1.2-6.7倍、0.8-3.5倍的加速，并且在故障出现时可以在30秒内通过容错方法保证训练的继续进行。图中也表明了当k值较大时，ECRec的故障恢复时间较长，另一方面，ECRec的故障恢复时间随着DLRM大小的增加而增加，且增加的速度超过检查点方法，但是ECRec在故障恢复期间可以继续训练并且保持较高的吞吐量，所以这些对DLRM训练时间的影响不大
总结
本文提出了DLRM容错系统ECRec， 基于DLRM的特性，提出了对Embedding table采用纠删码、NN采用复制副本的in-memory冗余方法。ECRec满足了理想DLRM容错系统的四个方面: 1)在训练和故障恢复过程中不会暂停训练; 2)有较低的内存开销; 3)不会影响到模型的收敛性; 4)能够扩展到更大的模型。实验结果表明ECRec相比于检查点方法能够将训练时间减少66%，故障恢复时速度提高9.8倍
论文写作上，本文的Intro部分故事还是比较完整的，递进的突出了检查点方法带来的挑战以及引入纠删码方法带来的挑战；方法部分受创新点数量的限制，比较考验写作，写了两段实际上可有可无的话
对于Embedding table应用纠删码的思路可以考虑应用到MoE层，都是模型中参数量占比大的稀疏部分