集合通信原语学习

在阅读《From ATOP to ZCube》写作理解《From ATOP to ZCube》中的ZCube网络拓扑的过程中,我发现自己对于集合通信原语的理解还不够熟悉,所以在此小记一下。

最基础的

发送send、接收receive、复制copy、组内进程栅障同步Barrier以及节点间进程同步(signal+wait)

其中的Barrier相当于一个特殊的节点。所有进程都必须调用同一个 Barrier。每个进程执行到 Barrier 时会阻塞,直到组内所有进程都到达了这一屏障点,之后所有进程才会同时被释放,继续向后执行。

关于signal和wait,可以这么理解:

A: 正常运行 B: 正常运行

A: 正常运行 B: 调用wait, B停下

A: 正常运行,调用signal B: 收到signal B继续 运行

可以认为Barrier是全局锁,而signal和sync是局部定义的阻塞与等待。

Broadcast

一个发送方,多个接收方,通信完成后所有节点内容相同。

Scatter

一对多,但是是把发送方的数据切成多个等长的块发给每个接收方

Gather

多对一,最终唯一的接收方得到了所有的数据

All Gather

先多对一Gather到一处,再由刚才的接收方Broadcast到全部。可以理解为 All Gather = Gather + Broadcast。

AlltoAll

这个似乎复杂一点,n个节点中每一个节点将目前已经有的数据切为n份,第n份发送到第n个节点。这个通信原语似乎在MoE模型中用的比较多,具体的集合通讯原语与训练和推理过程的关系在后面再学习吧

Reduce

Reduce将多个节点的数据规约运算到一个节点上。

什么是规约计算呢?就是把很多的数据变成少的数据,背后是求和,取最大值等等操作。通信库往往就此做一些优化。

ReduceScatter

Reduce之后Scatter一下,每个节点拿到Reduce后的彩虹条(见上图)中的一部分,分发的方式和Scatter一样。

AllReduce

有两种理解:

  1. Reduce + Broadcast 先reduce到了一个节点,那个节点再broadcast出去了
  2. ReduceScatter + AllGather

这里第二种看起来比第一种更长更复杂,我一开始好奇为什么要像第二种那样实现。研究了实际情形下的负载才知道:Reduce操作需要选一个节点作为主节点,其他的节点需要把数据都全量传过来,这个主节点进行规约计算,然后这个主节点还要向其他所有节点发送数据。这并没有利用到多节点的计算能力,并且在主节点形成了大量的收与发的流量,这是性能低下的。

而ReduceScatter + AllGather是这样的:

  1. 每个节点将所有数据中的第i部分发送到第i个节点,注意到这里不是发送所有数据,并且数据是从不同的节点过来的,能充分利用网络路径

  2. 每个节点将自己拿到的数据进行规约计算(这里利用到了每个节点的计算能力)

  3. 每个节点将自己规约计算后的结果发送给其他所有节点,注意,在所有节点将第i部分发给节点i之后,节点i实际上是有n部分数据的,也就是原先一整块数据的大小,但是在那里进行了规约计算,计算后的产物和第i部分一样大。这个第i个节点向其他所有节点发送的是1/n大小的数据量,所以流量也不是很大。