基于深度学习的任务放置在具有异质工作负载的分布式机器学习集群中的应用（论文笔记）

问题

随着机器集群上分布式机器学习任务负载的增加，集群调度器面临了一个重要的挑战，即如何高效地把不同的ML任务放到服务器上，以实现较高的资源利用效率和训练吞吐。

同一机器上的不同任务可能会共享机器的资源，如CPU、GPU、内存、CPU缓存、磁盘I/O、网络I/O以及总线等。因此，任务之间可能会互相干扰。减少不同任务之间的干扰能够优化训练性能。

本文提出了一种基于深度学习的调度器，Harmony，隐式地将工作负载的干扰编码入神经网络中，将集群和任务状态映射成任务的放置决定。

系统概述

对于一个提交的ML任务，用户需要提交以下信息：1) 运行worker和PS的资源需求，2) 使用的worker和PS的数量，3) 训练数据集的回合数。

时间被划分成小的时间片，一个时间片是一个调度间隔。在每个间隔中，将新到达的任务视为一个批次。Harmony决定整个批次的作业布置，即，一个作业中的每个worker和每个参数服务器应该在哪个服务器上运行（作为虚拟机或容器）。然后根据布置决策相应地部署作业，Harmony运行这些作业直到完成，即，每个作业的worker和参数服务器的布置在整个训练过程中都不会改变。

离线训练

使用纯粹的Online训练在刚开始的时候效果很差，因此需要提前进行离散训练。此文的离线训练分成两步。

1. Reward model training. 使用历史工作轨迹作为训练数据，Harmony利用监督学习训练奖励预测神经网络。输入包括作业集信息和不同服务器上的放置信息，标签是每个并发作业的奖励（训练速度）。该模型为任务集的放置决定提供了快速的奖励估计。
2. DRL model training. DRL NN的输入是作业集、现存的放置和集群资源可用性，输出是所有新作业的放置决定。此文根据奖励模型估计奖励，并将此作为DRL训练的奖励。 有奖励预测模型，作者能够高效地扩展可用的轨迹集合，为DRL离线训练产生足够的训练样本。

在线推断和模型更新

然后将离线训练的模型用于在线决策。在每个调度间隔中，Harmony通过DRL神经网络的推理来决定新作业批次的放置，并观察与放置决策对应的实际奖励。通过这种方式，我们可以收集到新的奖励样本，这些样本可能从未出现过，因此，在离线阶段没有被训练过。

我们使用在线收集的样本定期重新训练DRL神经网络和奖励神经网络，随着时间的推移不断改进决策。

DRL

DRL用于学习任务的放置策略，最大化任务训练速度。训练架构如下：

状态空间

输入状态是一个长为N的序列，N是目前并发任务的数量，包括新到达的任务和未完成的任务。包括已经被放置过的任务的原因是，允许DRL模型学习潜在的新任务和老任务之间的干扰。

序列中的一个元素$s_n$包括以下信息：

1. $x_n$，是任务n模型的独热编码。具有相同架构和批量大小的DNN被认为是同一种模型。
2. $r_n$，是一个2(1+K)维的向量，编码worker和PS的资源需求。其中K是组成一个worker或PS的资源类型的数量。第一个值代表任务n需要的worker的数量，下K个值表示每个worker的K种资源的需求，剩余的1+K个值表示PS的数量及其资源需求。对于All-Reduce架构来说，将PS的资源需求置为0即可。
3. $w_n(p_n)$，一个整数，表示分配给任务n的worker和PS的数量。
4. $\text{v}$，一个M*K的矩阵，表示每个服务器上每种资源的可用数量，M是物理服务器的数量。
5. $d_n$，一个M*2大小的向量，编码任务n的worker和PS在服务器上的放置。任务n在服务器m上的worker的数量和PS的数量分别在$2m^{th}$和$2m-1^{th}$位置上。
6. $u_n$，一个整数，指示任务n是新到达的任务，或者是在之前的调度间隔上已经放置的任务。

动作空间

对每个新到达的任务，一个接着一个地产生放置决定，产生一个放置序列。

具体来说，动作空间包括2M个动作，如下：

1. (0, m)，表示将一个新到达的任务的一个worker放到服务器m上
2. (1, m)，表示将一个新到达的任务的一个PS放到服务器m上

在每个调度间隔中，我们允许对神经网络进行多次推理，每次推理为一个新到达的作业选择一个动作。对于每个DNN推断，Harmony决定作业的PS或worker的放置位置，然后相应地更新状态（包括放置和剩余的可用资源）。然后在接下来的推理中，Harmony将根据更新后的输入状态进行放置决策，避免与干扰密集的作业放在同一位置。

在多次推理之后，我们为新作业中的所有worker和PS提出了一组完整的放置决策（或者当没有足够的资源放置任何额外的PS/worker时，推理停止）。这样，我们通过对所有新到达的作业逐个决定(0, m)和(1,m)来有效地减少动作空间。

奖励

奖励r是一个调度间隔内所有并发任务的标准训练速度之和：

其中$c_n$是任务n在这个间隔内的训练回合数，$e_n$是任务n设定的总训练回合数。

表示网络

表示网络在每个调度间隔将所有并发作业的状态和服务器状态作为输入，提取特征，并生成一个表示（即尺寸较小的向量），然后在解码器网络中使用该表示来产生调度决策。

由于输入的任务序列是可变的，此文采用了Transformer中的编码部分来编码任务和服务器的信息，结果是一个固定大小的矩阵的序列。注意力机制被用于捕获输入序列中的不同任务之间的关系。Transformer的encoder输出被返回作为状态的表示。

编码器网络

编码器网络被用来分析表示网络得出的编码序列，一个接一个地为新到达的任务产生放置决定。

模型训练

目标是最大化累计折扣奖励，即${\sum }_{i\ge 0}{\gamma }^i{r}_i$。

奖励预测模型

给定任务和集群状态，模型预测奖励。输入是任务信息的序列，其中任务信息包括任务的模型类型、分配的worker和PS的数量以及在每个服务器上的放置。输出是输入任务的预测的训练速度。模型的架构同样使用了多头注意力机制，以获取并发任务之间的相互关系。

使用历史轨迹数据进行训练。