Janus: Data-Centric MoE 通讯成本分析（2）

文章链接：Janus: A Unified Distributed Training Framework for Sparse Mixture-of-Experts Models

发表会议: ACM SIGCOMM 2023 (计算机网络顶会)

系统学习：Janus: 逆向思维，以数据为中心的MoE训练范式（1）

前言

在之前的blog中，我们学习了Janus的理论基础和模型搭建。基于专家的规模小于数据规模的假设，作者得到了Data-Centric的思想灵感，并验证了算法的有效性。
通过Data-Centric这一范式思路和巧妙的读取测略，Janus极大的减少了算法的通讯成本。这篇blog将从数学的角度定量解读Janus是如何降低通讯开销的。

通讯成本分析

1. Expert Parallelism and all-to-all

一个MoE模型的尺寸可能会大到超过单个gpu的能力。为了在gpu上训练一个大规模的MoE模型，专家并行（Expert Parallelism，EP）被提出并得到广泛应用。

右图显示了EP的概念。专家并行是将专家层划分为几个部分，并分配给GPU。每个GPU中都有专家层的专家，不同的GPU中有不同的专家。对于MoE模型的其他部分(如attention layer和Gate)，每个GPU持有一个独立的副本。

目前EP的实现默认是以专家为中心的。图(a)说明了以专家为中心的训练过程。当MoE块处理token序列时，gate需要为每个token分配专家，而token由gate分发到承载分配专家的gpu。

这种token到gpu的分发是由一个all-to-all通信原语完成的，因为由gpu生成的令牌的目标gpu很可能包括所有的gpu。token被指定的专家处理后，需要将结果发送回其原始的gpu，这需要再次进行all-to-all通信。由于MoE模型通常具有多个MoE块，因此MoE模型的训练可能涉及到多次的all-to-all通信操作。

---

2. Traffic Comparison between D-C. and E-C.

模型 参数	MoE-Bert	MoE-GPT	MoE-Transformer-xl
Batch size $B$	256	256	64
Seq-lenth $S$	128	64	512
Top-K in Gate	2	4	2
Expert dim. $H$	768	768	256
MoE Block	4	1	12
Total block	12	12	12
Expert num	16    32	16    32	16    32
GPU num	16    32	16    32	16    32
Model Size(B)	0.42   0.73	0.23   0.31	0.11   0.21
E-C. Traffic(GB)	6   9	1.5   2.25	6   9
D-C. Traffic(GB)	0.56   1.69	0.14   0.42	0.19   0.56

上表展示了以专家为中心的范式(E-C.)和数据中心范式(D-C.)的模型配置和通讯成本大小。
可以看到，在相同的模型配置下，D-C.的通讯量明显低于E-C.

---

3. Communication Efficiency Analysis

由于整个系统的通信瓶颈在于节点间通信而非节点内通信，故将节点间通信量作为衡量训练系统潜在通信效率的指标。

记号	含义
$n$	设备数量
$m$	每个设备的工作进程
$E$	每个工作进程的专家数量
$H$	向专家输入的token的维度
$T$	工作进程生成的token数
$k$	门参数 top-K
$B$	每个工作进程的训练任务的批量大小
$S$	训练任务的序列长度

A. Forward Phase

在MoE模型中，每个Expert通常是一个由两个线性层组成的前馈网络(FFN)。对于一个FFN模块，第一层包括一个形状为 $H\ast 4H$ 的矩阵，第二层包括一个形状为 $4H\ast H$ 的矩阵。因此，一个FFN模块的大小为 $8H^2$。

每个工作进程有 $E$ 个专家，每台设备有 $m$ 个专家。由于每台设备都需要将这些专家广播给其他 $n-1$ 台设备，因此在MoE块的训练过程中，D-C.的通讯量为:
$$Com{m}_{DC}=8H^{2}Em(n-1)$$

在E-C.的训练系统中，专家之间的token分配通常是不平衡的。完成通信所需的时间取决于发送/接收数据量最大的设备。显然，非平衡分布下的通信时间几乎总是比平衡分布下的通信时间长。

现在计算E-C.的训练系统中要传输的token大小，即节点间通信量。每个工作进程生成 $T$ 个token，然后一个 m-worker (GPU)设备可以生成 $mT$ 个token。

在token均衡分配的假设下，token被发送到其他设备的百分比为$\frac{n-1}{n}$。在E-C.中，MoE型块需要在前向计算阶段执行两个all-to-all通信操作。因此，MoE块中以E-C.的通信量为:
$$Com{m}_{EC}=2mHT\cdot \frac{n-1}{n}$$

---

B. Backward Phase

对于后向阶段，在E-C.中，系统需要传输生成梯度所需的所有中间结果，并且这个量等于它在前向阶段中发送的token的量。

在D-C.中，系统可以重用在前向阶段提取和缓存的专家。在专家模块计算出梯度后，梯度应该被发送回原始工作进程。梯度的大小与拉取的专家模型相同，通信方向相反。

此外，同一个专家在一台设备上的多个梯度被减少和合并，然后被送回。因此，在D-C.中，后向阶段中的通讯量也等于前向阶段中的通讯量。

---

C. Ratio measure

在分析了两种模式的通讯量后，文章定义了一个度量 $R$ 来评估 D-C. 的理论增益，它是两种范式下节点间通信量的比率:
$$R=\frac{Com{m}_{EC}}{Com{m}_{DC}}=\frac{T}{4nHE}$$

给定训练参数可以计算出token的数量 $T$，包括批大小 $B$、序列长度 $S$ 以及与相关的Gate参数 top-K：$k$ 。$T=BSk$ 于是有：
$$R=\frac{Com{m}_{EC}}{Com{m}_{DC}}=\frac{BSk}{4nHE}$$

$R>1$ 表明D-C.效率优于E-C.效率。

---

Be Janus

Janus是一个以专家为中心的范式和以数据为中心的范式的统一框架。Janus在混合专家模型模型的训练开始前进行评估。

对于其中 $R\le 1$ 的MoE块，Janus默认使用专家中心范式。
对于其中 $R\ge 1$ 的MoE块，Janus将使用以数据为中心的范式。