本文主要是介绍Janus: Data-Centric MoE 通讯成本分析（2），希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！
文章链接：Janus: A Unified Distributed Training Framework for Sparse Mixture-of-Experts Models
发表会议: ACM SIGCOMM 2023 (计算机网络顶会)
系统学习：Janus: 逆向思维，以数据为中心的MoE训练范式（1）
目录
- 前言
- 通讯成本分析
- 1. Expert Parallelism and all-to-all
- 2. Traffic Comparison between D-C. and E-C.
- 3. Communication Efficiency Analysis
- A. Forward Phase
- B. Backward Phase
- C. Ratio measure
- Be Janus
前言
在之前的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 B B | 256 | 256 | 64 |
| Seq-lenth S S S | 128 | 64 | 512 |
| Top-K in Gate | 2 | 4 | 2 |
| Expert dim. H H 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 n n | 设备数量 |
| m m m | 每个设备的工作进程 |
| E E E | 每个工作进程的专家数量 |
| H H H | 向专家输入的token的维度 |
| T T T | 工作进程生成的token数 |
| k k k | 门参数 top-K |
| B B B | 每个工作进程的训练任务的批量大小 |
| S S S | 训练任务的序列长度 |
A. Forward Phase
在MoE模型中，每个Expert通常是一个由两个线性层组成的前馈网络(FFN)。对于一个FFN模块，第一层包括一个形状为 H ∗ 4 H H*4H H∗4H 的矩阵，第二层包括一个形状为 4 H ∗ H 4H*H 4H∗H 的矩阵。因此，一个FFN模块的大小为 8 H 2 8H^{2} 8H2。
每个工作进程都有专家，每台设备就有专家。由于每台设备都需要将这些专家广播给其他 n − 1 n−1 n−1 台设备，因此在MoE块的训练过程中，D-C.的通讯量为:
C o m m D C = 8 H 2 E m ( n − 1 ) Comm_{DC}=8H^{2}Em(n-1) CommDC=8H2Em(n−1)
在E-C.的训练系统中，专家之间的token分配通常是不平衡的。完成通信所需的时间取决于发送/接收数据量最大的设备。显然，非平衡分布下的通信时间几乎总是比平衡分布下的通信时间长。
现在计算E-C.的训练系统中要传输的token大小，即节点间通信量。每个工作进程生成token，然后一个 m-worker (GPU)设备可以生成 m T mT mT 个token。
在token均衡分配的假设下，token被发送到其他设备的百分比为 n − 1 n \frac{n-1}{n} nn−1。在E-C.中，MoE型块需要在前向计算阶段执行两个all-to-all通信操作。因此，MoE块中以E-C.的通信量为:
C o m m E C = 2 m H T ⋅ n − 1 n Comm_{EC}=2mHT\cdot \frac{n-1}{n} CommEC=2mHT⋅nn−1
B. Backward Phase
对于后向阶段，在E-C.中，系统需要传输生成梯度所需的所有中间结果，并且这个量等于它在前向阶段中发送的token的量。
在D-C.中，系统可以重用在前向阶段提取和缓存的专家。在专家模块计算出梯度后，梯度应该被发送回原始工作进程。梯度的大小与拉取的专家模型相同，通信方向相反。
此外，同一个专家在一台设备上的多个梯度被减少和合并，然后被送回。因此，在D-C.中，后向阶段中的通讯量也等于前向阶段中的通讯量。
C. Ratio measure
在分析了两种模式的通讯量后，文章定义了一个度量 R R R 来评估 D-C. 的理论增益，它是两种范式下节点间通信量的比率:
R = C o m m E C C o m m D C = T 4 n H E R=\frac{Comm_{EC}}{Comm_{DC}}=\frac{T}{4nHE} R=CommDCCommEC=4nHET
给定训练参数可以计算出token的数量 T T T，包括批大小 B B B、序列长度 S S S 以及与相关的Gate参数 top-K： k k k 。 T = B S k T=BSk T=BSk 于是有：
R = C o m m E C C o m m D C = B S k 4 n H E R=\frac{Comm_{EC}}{Comm_{DC}}=\frac{BSk}{4nHE} R=CommDCCommEC=4nHEBSk
R > 1 R>1 R>1 表明D-C.效率优于E-C.效率。
Be Janus
Janus是一个以专家为中心的范式和以数据为中心的范式的统一框架。Janus在混合专家模型模型的训练开始前进行评估。
对于其中 R ≤ 1 R\le 1 R≤1 的MoE块，Janus默认使用专家中心范式。
对于其中 R ≥ 1 R\ge 1 R≥1 的MoE块，Janus将使用以数据为中心的范式。
这篇关于Janus: Data-Centric MoE 通讯成本分析（2）的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！
