作者:Sudarsanan Rajasekaran (1), Manya Ghobadi (1), Aditya Akella (2) ((1) Massachusetts Institute of Technology, (2) UT Austin)
论文摘要原文：We present CASSINI, a network-aware job scheduler for machine learning (ML) clusters. CASSINI introduces a novel geometric abstraction to consider the communication pattern of different jobs while placing them on network links. To do so, CASSINI uses an affinity graph that finds a series of time-shift values to adjust the communication phases of a subset of jobs such that the communication patterns of jobs sharing the same network link are interleaved with each other. Experiments with 13 common ML models on a 24-server testbed demonstrate that compared to the state-of-the-art ML schedulers, CASSINI improves the average and tail completion time of jobs by up to 1.6× and 2.5×,respectively. Moreover, we show that CASSINI reduces the number of ECN marked packets in the cluster by up to 33×.
研究背景
训练大规模模型需要一个高效的GPU集群。
相关研究表明随着GPU数量的增加，分布式训练的通信开销占总时间开销的一大部分。
现有最先进的（state-of-the-art, SOTA）深度学习调度器（Machine Learning Scheduler）在分配任务时忽略了分布式训练任务的通信模式。
研究动机
因此，本文对分布式训练任务在不同并行策略下的通信模式进行了分析，分别是数据并行、流水线并行、张量并行以及3D混合并行，结果如下图所示：
子图(a)-数据并行：无需通信的前反向计算以及最后的AllReduce通信。
子图(b)-流水线并行：使用PipeDream流水线并行方法，前反向过程中有激活值的通信，最后会有一个Embedding的AllReduce通信，开销较大。
子图(c)-张量并行：前反向过程都需要进行通信，在迭代之间，进行数据加载时，没有通信开销。
子图(d)-3D混合并行：在前反向过程中有许多通信操作，这些通信操作对网络带宽的需求不同。
关键发现：
- 网络带宽的需求在迭代之间是重复有规律的。
- 单次迭代过程中存在多个Up和Down阶段（子图(d)中1、3为Down阶段，2、4、5、6为Up阶段），不同Up阶段对网络带宽的需求可能不同。
基于这些发现，本文就提出了CASSINI，通过交错（Interleaving）集群上不同任务Up和Down阶段，避免通信竞争，实现任务训练效率的提升。CASSINI
Single Link Multiple Tasks
Geometric Abstraction
为了解决单链路多任务的问题，CASSINI首先提出了一个几何抽象，对于单个任务，我们使用一个周长为其迭代时间的圆来表示，然后根据其在该链路上通信发生的时间，我们在圆上进行相应的标注， 就得到了一个任务的几何抽象，如下图的左边所示。假设我们有两个相同的任务在该链路上运行，我们就可以通过对几何抽象进行旋转（旋转角度大小对应的是某个任务延迟开始的时间），将两个任务的Up和Down阶段错开，避免通信竞争。
如果是两个不同的任务，其具有不同的迭代时间，我们就取两个任务迭代时间的最小公倍数作为圆的周长，再在圆上进行标注，如下图左边的(a)和(b)所示。对应于Moti中3D并行的情况，如果不同的Up阶段有不同的带宽需求，我们使用颜色的深浅进行区分，如下图右边所示。
至此，我们就可以使用几何抽象统一描述单个链路上多个任务的带宽使用情况（将多个任务使用相同周长的圆表示），然后CASSINI设计了优化函数来为单个链路上的多个任务求解出一个最大的兼容性得分，进而求出降低通信竞争的最优旋转角。形式化表达如下：
至此，我们就可以求解出单链路多任务情况下，为了降低通信竞争，各个任务应该旋转的角度，通过这个角度，我们可以利用下面的公式算出各个任务的开始时间。
Multiple Link Multiple Tasks
接下来，本文将场景扩展到多链路多任务的场景，也就是实际训练集群的场景。在集群中，很容易就出现一个任务与多个任务竞争链路的情况，如下图所示，和竞争，和竞争，如果基于和分别利用单链路多任务的优化函数求解，那么将会有两个开始时间，因此，将单链路多任务扩展到多任务多链路的场景，要解决的就是如何统一多个链路上求解出的开始时间的问题。
为此，CASSINI将任务和链路抽象成了关联图（Affinity graph），如下图左边所示，CASSINI将所有任务放在集合中，将所有链路放在集合中，这样，图中的一条边就表示任务在链路上运行，边的权重为其在链路上的开始时间。
基于关联图，CASSINI基于广度优先搜索算法（Breadth First Search, BFS）设计了算法1来求解各个任务的开始时间，算法1如下所示：
算法1在关联图的各个连通子图上运行，通过传递不同任务在同一链路上开始时间的差值来为各个任务确定开始时间，同时保持各个任务在链路上的通信竞争情况不变。论文在附录中证明了该算法在无环图上的正确性。
Put It All Together
本部分讲述了将CASSINI集成到Themis（NSDI’20的集群调度工作）中实现落地的过程。CASSINI集成到Themis后的总体运行流程如下图所示，相当于在原本Themis的工作流中加入了图中蓝色的Cassini模块。
默认情况下，Themis会直接给出任务运行的GPU数量和对应的设备放置策略，但是由于CASSINI的算法1的正确运行需要无环图的保证，所以这里将Themis原本工作流中为每个任务分配GPU数量和每个任务的设备放置策略解耦，使得Themis能够为每个任务返回N个可能的放置策略，文中也提到不同的放置策略不会影响Themis的评价指标，所以在每一次调度内对Themis没有影响。
基于Themis给出的任务放置策略，CASSINI设计了算法2求解该放置策略下集群的兼容性得分，根据兼容性得分选出最优的放置策略，最后基于最优的放置策略求解出每个任务的开始时间。
最后，配置Themis使用最优的放置策略及各个任务的开始时间即可。
实验结果
实验设置
实验环境
24台 ASUS ESC4000A-E10服务器：一张40GB A100，一张50Gbps Mellanox ConnectX5网卡
使用Ubuntu 18.04 LTS，PyTorch 1.8.0，CUDA 11.1以及NCCL 2.11.4
网络拓扑
搭建了如下图所示的网络拓扑：
模型负载
13种DNN：VGG11、VGG13、VGG19、ResNet50、WideResNet101、BERT、RoBERTa、XLM、CamemBERT、GPT-1、GPT-2、GPT-3和DLRM。
并行策略
数据并行（PyTorch DDP）：VGG、ResNet、BERT
混合并行（Meta’s opensource codebase）：DLRM
混合并行（Microsoft DeepSpeed）：GPT
Traces
- Poission trace: 任务到达时间符合泊松分布，保证集群GPU利用率在80%-100%之间。
- Dynamic trace: 一批DNN任务在集群上运行，另一批DNN任务到达。
- snapshot trace: 集群运行过程中的快照。
参与比较的方法
Themis、Th+CASSNI、Pollux、Po+CASSINI、Ideal、Random
性能提升
图的左半部分展示了集群运行过程中各个任务的迭代时间，除了DRLM其他DNN负载均使用数据并行。可以看到，CASSINI通过消除不同任务在链路上的通信竞争，大大降低了任务运行过程中的迭代时间。不过这里有一个疑惑的点是图中并没有展现出CASSINI决策后延迟任务开始时间的效果，看起来各个任务的开始时间在Themis和Th+CASSINI中是一样的。
图的右半部分展示了迭代时间的CDP曲线，可以看到CASSINI能够将Themis的调度结果提升到接近专用GPU集群训练的效果，相比如Themis，Th+CASSINI将平均的迭代时间提升了1.6倍，将迭代时间的99分位数提高了1.8倍。
模型并行的影响
探究CASSINI在模型并行情况下的效果，使用的模型是DLRM和GPT系列模型。
相比如Themis，Th+CASSINI将平均的迭代时间提升了1.2倍，将迭代时间的99分位数提高了1.6倍。
子图(b)-(e)给出了集群中不同模型每个迭代产生的显式拥塞通知（Explicit Congestion Notification，ECN）数据包的数量，Th+CASSINI方法相较于Themis方法大大降低了ECN数据包数量。
CASSINI降低通信竞争的效果
Dynamic trace：集群上本来运行着VGG16、RoBERTa以及DLRM，此时来了一批DLRM和ResNet50的任务
Pollux和Themis的放置使得新来的DLRM任务与原来的不兼容DLRM的任务共享链路，增加了迭代时间，相比于Themis，Th+CASSINI在平均迭代时间和迭代时间的99分位数上取得了1.2和2.2倍的加速；相比于Pollux，Po+CASSINI在平均迭代时间和迭代时间的99分位数上取得了1.6和2.5倍的加速。
子图(b)-(d)给出了集群中不同模型每个迭代产生的ECN数据包的数量，可以看到Th+CASSINI以及Po+CASSINI方法在不同的模型上相较于原来的方法都可以实现ECN数据包数量的降低。在DLRM模型上，CASSINI相较于Themis和Pollux能够将ECN数据包数量降低27和33倍。
部分兼容情况下的性能
上表展示了五个不同的时刻集群中DNN训练任务在某条链路上的兼容性得分，可以看到，当兼容性得分降低至0.6时，CASSINI带来的性能提升就变得相当有限。为了分析其原因，论文展示了五个不同时刻该链路的使用情况，如下图所示：
可以看到，在前四个时刻，CASSINI很好的将不同模型的通信交错开来，但是当兼容性得分较低时，不同模型的通信同时进行，产生通信竞争，如上图(e)所示。
单机多卡的集群带来的影响
前面的实验都是在单机单卡的集群上进行的，现有的分布式训练集群存在单机多卡的情况，并且偏向于将任务放置在单个节点内，但是随着任务规模的增大，总是会有任务需要跨节点放置，因此会存在不同任务共享链路带宽的情况。在这种情况下，CASSINI的性能收益在跨节点任务上更加明显。
为了验证CASSINI在单机多卡上的性能，通过对实验使用的设备进行手动调整，论文造出了一个单节点两张卡的6节点集群，其网络拓扑如下图左边所示：
在这个集群上使用数据并行和模型并行任务进行实验，实验结果表明相比于Themis，Th+CASSINI在平均迭代时间和迭代时间的99分位数上取得了1.4和1.9倍的加速。这些收益是一些跨节点任务带来的。例如，XLM模型和ResNet50均使用三个GPU进行训练，此时来了一个DLRM的任务，需要超过3个GPU，Themis就会让DLRM与不兼容的XLM模型共享一个节点，带来通信竞争，而Th+CASSINI会让DLRM与兼容的ResNet50模型共享一个节点，提高两个模型的训练效率。
总结
个人觉得本文解决问题的思路很巧妙，乍一想，多链路多任务的场景下，确定降低通信竞争情况下各个任务的开始时间，是一件很困难的事情。为了解决这个问题，本文从单链路多任务出发，首先求解出单链路场景下各个任务的开始时间，扩展到多链路场景时，通过构建关联图，只使用简单的BFS算法就解决了问题。当然这也是长时间沉淀的结果，这篇文章是2022年的一篇Workshop的后续。实验结果也表明，通过错开不同任务在同一链路上的通信操作，本文工作CASSINI能够有效提高各个任务的训练性能。
本文的解法要求能够构建出无环关联图，但是本文并没有讨论构建无环关联图的可能性，如果能够再讨论一下无环关联图构建的可能性，证明可能性较大的话，感觉能进一步提高本文的价值。
