论文阅读Optimizing Task Placement and Online Scheduling for Distributed GNN Training Acceleration

提示:文章写完后,目录可以自动生成,如何生成可参考右边的帮助文档


前言

Optimizing Task Placement and Online Scheduling for Distributed GNN Training Acceleration是2022年港大的吴川老师组发表于inforcomm2022的文章,论文传送门https://i.cs.hku.hk/~cwu/papers/xdyi-tpds22.pdf
吴川老师主页传送门https://i.cs.hku.hk/~cwu/publications.html
主要使用一个分层算法进行task placement + 调度的方法,缓解分布式GNN中存在的sampler和worker不好分离问题、数据传输流引发的通信量问题、计算开始后出现的在线调度需求等问题。实验证明该在线调度算法可以在能接收的代价下达到很好的加速效果,有效加速了分布式GNN的训练。
该文章可以针对子图分布式GNN与全图分布式GNN,关于分布式GNN都有什么样的形式,参考我之前的博文https://blog.csdn.net/weixin_43934886/article/details/130141239


一、背景

GNN结构

关于GNN,本质是用神经网络完成图嵌入的过程,分为两种方法:1)Message Passing:将Aggregation过程用MLP等方法替代;2)谱方法:将图嵌入过程用图的拉普拉斯矩阵(D-A)作为卷积核对图结构进行卷积计算。而当输入图规模过大、节点数特征维度等超过万量级时常规GPU无法提供足够的显存支持,必须进行分布式(单机多卡或者多机多卡)计算。

分布式GNN结构

分布式GNN 的标准结构一般可以分为图存储服务器store server(负责存储),节点采样器sampler(负责采样子图或者切分子图),worker(主要训练节点,一般是GPU),Parameter Server(PS服务器,用于同步每个worker的梯度,但是如果使用all-reduce方法的话就可以不需要它)。总的来说,sampler从存储服务器采样子图给各个worker,分布式GNN在计算完各自的梯度后,要把梯度传送给PS服务器求和,求和后得到的梯度传回给每个worker,同步参数更新过程,保证每个woker的参数是相同的

大致流程图如下:
流程图
或者还有一种更细致的流程图:
流程图2

二、问题与挑战

1.挑战

显然任何方法一旦分布式后,必然会带来以下几种问题:

  • 部署在不同节点上的任务的通信代价
  • 带宽限制导致的通信拥塞
  • 程序串行部分与并行部分的比例导致的优化上限
  • 不同任务的调度方案对总时间的影响
  • 在线调度方案的开销以及效果

而对于GNN来说,这个问题就转变成了:

  • 不同机器之间的资源可用性不同,需要调度算法规划任务位置以最小化数据传输流量并最大化资源利用率
  • 分布式GNN训练作业中数据传输和任务执行的优化调度比较复杂,属于强NP-hard多阶段调度问题
  • 图存储和采样器之间的数据传输量随迭代轮数变化,调度问题具有在线性质,需要高效的在线算法设计

2.问题

本篇的比较精华的部分我认为在于这个对上述挑战的建模

首先我们提出一下几点假设
(1)每台机器拥有一个图存储服务器,维护一个分区图(不是用来训练的子图);
(2)每个采样器选择训练节点,从图存储服务器采样节点/边缘特征,并形成子图;
(3)每个worker进行forward和backward计算,在backward计算后将梯度推入参数服务器PS,并从参数服务器中获取参数同步;
(4)一个worker通常与一个或多个采样器相关联,这些采样器专门为worker提供mini-batch;
(5)参数服务器PS聚合所有工人的梯度,更新GNN模型参数,并将更新后的参数分发给所有worker,注意PS可以有多个机器;
(6)上述的图存储服务器、采样器、worker以及PS都可以理解为一个或多个进程,并且在每次调用时可以动态的选择由哪个机器的哪个核心执行(worker的训练过程由CPU唤起GPU执行)
逻辑依赖关系
任务部署示意图如下图:
任务部署示意图
在分布式GNN训练工作中,我们的目标是使总的训练时间最小化,问题的建模包括两个子问题,Task Placement以及Online Execution and Flow Scheduling
首先对Task Placement进行建模
1.令01变量 y j m {y_j}^m yjm 表示任务j是否放置在机器m上。则显然 ∑ m y j m = 1 , y j m = 0 , 1 {\sum_m}{y_j}^m=1,{y_j}^m=0,1 myjm=1,yjm=0,1
2.假设任务j对r型计算任务(存储、采样、训练等)有 w j r {w_j}^r wjr的资源需求,并且机器m上的总资源数为 C m r {C_m}^r Cmr,则显然需要满足部署的任务不能超过机器资源量,即 ∑ m w j r ⋅ y j m ≤ C m r {\sum_m}{w_j}^r·{y_j}^m\le{C_m}^r mwjryjmCmr
这样一来Task Placement部分建模为数学条件式
∑ m y j m = 1 , y j m = 0 , 1 {\sum_m}{y_j}^m=1,{y_j}^m=0,1 myjm=1,yjm=0,1 (1)
∑ m w j r ⋅ y j m ≤ ( 1 + μ ) C m r {\sum_m}{w_j}^r·{y_j}^m\le(1+\mu){C_m}^r mwjryjm(1+μ)Cmr (2) 其中 μ \mu μ表示资源宽松情况

对Online Execution and Flow Scheduling进行建模
假设GNN模型训练需要N次迭代才能收敛
在一个给定的Task Placement下,调度工作变成了决定每个任务(包括计算与通信)的开始时间
设01变量 x j , n t {x_{j,n}}^t xj,nt为第n轮的任务j是否在时间t开始执行(开始为1否则为0);
k ( j , t ) → ( j ′ , t ′ ) {k_{(j,t)\rightarrow(j',t')}} k(j,t)(j,t)为第n轮的任务j向第𝑛’轮的任务𝑗′的通信量,这里可以有𝑛=𝑛′,那么通信时间在知道网络带宽的情况下理论上可以计算出来;
设succ()集为某个任务的继承任务,如worker从采样器通信到数据后紧跟着训练;
设任务j的执行时间为 p j {p_j} pj

根据这些设定,可以获得如下的约束条件:
1.任何任务在整个运行过程中只部署/开始一次: ∑ t x j , n t = 1 {\sum_t}{x_{j,n}}^t=1 txj,nt=1
2.任务j的后一个任务需要在完成后才能执行: m i n { t ∣ k ( j , n ) → ( j ′ , n ′ ) t > 0 } ≥ ∑ t t ⋅ x j , n t + p j , ( j ′ , n ′ ) ∈ s u c c ( j , n ) min\{t|k^t_{(j,n)\rightarrow(j',n')}>0 \}\ge{\sum_t}t·{x_{j,n}}^t+{p_j}, (j',n')\in succ(j,n) min{tk(j,n)(j,n)t>0}ttxj,nt+pj,(j,n)succ(j,n)
3.同理任务j的开始时间也不能迟于后继任务开始时间: m a x { t ∣ k ( j , n ) → ( j ′ , n ′ ) t > 0 } < ∑ t t ⋅ x j , n t , ( j ′ , n ′ ) ∈ s u c c ( j , n ) max\{t|k^t_{(j,n)\rightarrow(j',n')}>0 \}<{\sum_t}t·{x_{j,n}}^t, (j',n')\in succ(j,n) max{tk(j,n)(j,n)t>0}<ttxj,nt,(j,n)succ(j,n)
4. 必须满足任务j完成后再进行后继任务: ∑ t t ⋅ x j , n t + p j < ∑ t t ⋅ x j ′ , n ′ t {\sum_t}t·{x_{j,n}}^t+{p_j}<{\sum_t}t·{x_{j',n'}}^t ttxj,nt+pj<ttxj,nt
5. 同4,同一任务不同轮也需满足先后关系: ∑ t t ⋅ x j , n t + p j < ∑ t t ⋅ x j , n + 1 t {\sum_t}t·{x_{j,n}}^t+{p_j}<{\sum_t}t·{x_{j,n+1}}^t ttxj,nt+pj<ttxj,n+1t
6. 下轮的数据传输必须在此轮的数据传输之后: m a x { t ∣ k ( j , n ) → ( j ′ , n ′ ) t > 0 } ≥ m i n { t ∣ k ( j , n + 1 ) → ( j ′ , n ′ + 1 ) t > 0 } , ( j ′ , n ′ ) ∈ s u c c ( j , n ) max\{t|k^t_{(j,n)\rightarrow(j',n')}>0 \}\ge min\{t|k^t_{(j,n+1)\rightarrow(j',n'+1)}>0 \},(j',n')\in succ(j,n) max{tk(j,n)(j,n)t>0}min{tk(j,n+1)(j,n+1)t>0},(j,n)succ(j,n)

变量表如下:
变量表

接着对Online Execution and Flow Scheduling进行建模
对通信部分进行建模
设第n轮的任务j对第𝑛’轮的任务𝑗′的总通信量为 d ( j , n ) → ( j ′ , n ′ ) d_{(j,n)\rightarrow (j',n')} d(j,n)(j,n)
∑ t k ( j , n ) → ( j ′ , n ′ ) t = d ( j , n ) → ( j ′ , n ′ ) , ( j ′ , n ′ ) ∈ s u c c ( j , n ) {\sum_t}k^t_{(j,n)\rightarrow(j',n')}=d_{(j,n)\rightarrow (j',n')},(j',n')\in succ(j,n) tk(j,n)(j,n)t=d(j,n)(j,n),(j,n)succ(j,n)

每个t时刻,机器m上输出(输入)的总流量不应超过其可用带宽
∑ n ∑ j : y j m = 1 ∑ ( j ′ , n ′ ) ∈ s u c c ( j , n ) : y j ′ m = 0 k ( j , n ) → ( j ′ , n ′ ) t < B o u t m {\sum_n}{\sum_{j:{y_j}^m=1}}{\sum_{(j',n')\in succ(j,n):{y_j'}^m=0}k^t_{(j,n)\rightarrow(j',n')}<{B^{m}_{out}}} nj:yjm=1(j,n)succ(j,n):yjm=0k(j,n)(j,n)t<Boutm 表述输出
以及 ∑ n ∑ j : y j m = 0 ∑ ( j ′ , n ′ ) ∈ s u c c ( j , n ) : y j ′ m = 1 k ( j , n ) → ( j ′ , n ′ ) t > B i n m {\sum_n}{\sum_{j:{y_j}^m=0}}{\sum_{(j',n')\in succ(j,n):{y_j'}^m=1}k^t_{(j,n)\rightarrow(j',n')}>{B^{m}_{in}}} nj:yjm=0(j,n)succ(j,n):yjm=1k(j,n)(j,n)t>Binm 表述输出

在满足以上所有条件式以后,需要满足最优化目标,即离线执行调度问题的最优解,即最后一次任务的实际结束时间要取得最小值(这么设计目标可以把优化工作看做一个贪心算法,即每次都搜索前一次的最后一个任务的最小值)
目标公式

三、调度与任务部署方法

1.在线调度算法OES

流程图
算法

1.初始化运行集 F a c t F_{act} Fact与排队集 F p e n d F_{pend} Fpend,并将所有 x j , n 1 {x_{j,n}}^1 xj,n1置;
以t为循环对象执行以下2-5循环
2.如果任务j的 ( j , n − 1 ) → ( j ′ , n ′ − 1 ) ∈ F a c t ∪ F p e n d (j,n-1)\rightarrow (j',n'-1)\in F_{act}\cup F_{pend} (j,n1)(j,n1)FactFpend则将一轮的通信放入 F p e n d F_{pend} Fpend,否则放入 F a c t F_{act} Fact)(说明上一轮通信结束可以进行下一轮通信了);
3.对t-1时刻完成的 ( j , n ) → ( j ′ , n ′ ) (j,n)\rightarrow (j',n') (j,n)(j,n),如果在 F p e n d F_{pend} Fpend中的有 ( j , n + 1 ) → ( j ′ , n ′ + 1 ) (j,n+1)\rightarrow (j',n'+1) (j,n+1)(j,n+1),则把该通信任务从 F p e n d F_{pend} Fpend中转移到 F a c t F_{act} Fact中;
4.计算每个机器的 Δ m i n {\Delta^m}_{in} Δmin Δ m o u t {\Delta^m}_{out} Δmout(流入与流出度);
5.限制每个通信的通信量 k ( j , n ) → ( j ′ , n ′ ) t k^t_{(j,n)\rightarrow(j',n')} k(j,n)(j,n)t
6.输出 x j , n {x_{j,n}} xj,n, k ( j , n ) → ( j ′ , n ′ ) t k^t_{(j,n)\rightarrow(j',n')} k(j,n)(j,n)t, T O E S T_{OES} TOES

OES算法效果理论推导
F o n e _ i t e r F_{one\_iter} Fone_iter为一次训练迭代中所有机器间流的集合,包括在此迭代中计算从PS转移到worker的更新参数的通信
定义一个迭代中通信入流度为 Δ i n m ^ \widehat{{\Delta}^m_{in}} Δinm ,流出度为 Δ o u t m ^ \widehat{{\Delta}^m_{out}} Δoutm
Δ i n m ^ = ∣ ( j ′ , n ′ ) → ( j , n ) ∣ ( j ′ , n ′ ) → ( j , n ) ∈ F o n e _ i t e r , y j m = 1 \widehat{{\Delta}^m_{in}}=|(j',n')\rightarrow(j,n)|(j',n')\rightarrow(j,n)\in F_{one\_iter},{{y_j}^m}=1 Δinm =(j,n)(j,n)(j,n)(j,n)Fone_iter,yjm=1
Δ o u t m ^ = ∣ ( j , n ) → ( j ′ , n ′ ) ∣ ( j , n ) → ( j ′ , n ′ ) ∈ F o n e _ i t e r , y j m = 1 \widehat{{\Delta}^m_{out}}=|(j,n)\rightarrow(j',n')|(j,n)\rightarrow(j',n')\in F_{one\_iter},{{y_j}^m}=1 Δoutm =(j,n)(j,n)(j,n)(j,n)Fone_iter,yjm=1

引理1 在任何时刻任何一台机器上, Δ i n m {{\Delta}^m_{in}} Δinm Δ o u t m {{\Delta}^m_{out}} Δoutm都不会大于 Δ i n m ^ \widehat{{\Delta}^m_{in}} Δinm Δ o u t m ^ \widehat{{\Delta}^m_{out}} Δoutm

定理1 算法OES实现的整体训练开销不大于 Δ \Delta Δ乘以离线执行调度问题的最优目标值 T ∗ T^* T,即算法OES中在线算法的竞争比为 Δ \Delta Δ
竞争比:这里有一个重要的概念在于:对于任意给定的S,算法A产生的cost最多是最优算法OPT的α 倍。

2.探索性部署IFS算法

bushu

算法采用马尔科夫链搜索框架,首先构建一个可行的初始布局解 Y 0 = { y j m } Y_0=\{{{y_j}^m}\} Y0={yjm},然后生成布局序列𝑌_1, 𝑌_2 ,…,直到时间预算 I I I耗尽
算法伪代码
算法概述
定理2 IFS算法可以在多项式复杂度内部署完任务

3.搜索更好的部署计划ETP

执行在线调度是由于抽样图数据的大小变化,应该根据预期中的通信量变化确定一个最佳部署。为此,需要通过对一些迭代运行GNN训练来分析任务执行时间和任务间通信量,并计算出它们的分布。
使用算法OES模拟在placement为Y下的训练过程,从分布中提取时间和通信量,并推导出训练完成时间 T Y ′ {T'_Y} TY
c o s t ( Y ) = T Y ′ ( 1 + ∑ m , r m a x { ∑ j w j r ⋅ y j m − C m r C m r } ) cost(Y)={T'_Y}(1+\sum_{m,r}max\{\frac{{\sum_j}{w^r_j}·{y^m_j}-{C^r_m}}{{C^r_m}}\}) cost(Y)=TY(1+m,rmax{CmrjwjryjmCmr})
其中, T Y ′ {T'_Y} TY乘以1加上资源违规的惩罚 (计算为所有类型的资源和所有机器的容量违规百分比之和)。算法搜索通过从一个位置 Y Z Y_Z YZ转移到另一个位置 Y Z + 1 Y_Z+1 YZ+1来探索Placement方案,总共转移 I I I次(时间预算)。
算法过程
伪代码

4.整体流程

整体过程
伪代码
整体流程是,1.先用IFS获得一个初始部署,2.接着用OES进行调度,3.之后再执行探索算法ETP找寻是否有更好部署,每次重新部署后都要执行OES算法,重复3和2总共 I I I次,4.获得最终部署与调度方案

四.实验结果

这里就不细看了,感兴趣的童鞋仔细研究研究把
通过试验台实验和模拟研究对DGTP进行评价
测试平台由4台GPU服务器组成,通过一台戴尔Z9100-ON交换机连接,任何两台服务器之间的峰值带宽为50Gbps。
每台服务器配置1块50GbE网卡、1个8核Intel E5-1660 CPU、2块GTX 1080Ti gpu和48GB DDR4 RAM。为了模拟资源异构性,使用tc将两台服务器的带宽容量限制为10Gbps。
benchmark
Benchmark选用OGB的product和paper,以及reddit论坛的社交网络数据集,主要采用DistDGL库部署,将两个数据集上的迷你批处理大小设置为2000(子图)。

实验结果

结果2


总结

文章有个小问题,那就是对于分布式GNN中不同的划分图策略其实并没有很好的适配,换句话说,其实该方法除了在分布式GNN上,任何一个分布式深度学习其实都挺适用的,顶多去掉那个sampler,而算法层面其实并没有针对sample做特别的处理,所以其实个人感觉主要就是建模做的很好。可能本人水平有限,欢迎大家补充!

  • 2
    点赞
  • 2
    收藏
    觉得还不错? 一键收藏
  • 3
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 3
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值