正文
工作负载分片和额外的Kernels
与单GPU推理用例相比,没有引入额外的自定义kernel。对于GEMM、GroupedGEMM和非线性kernel,激活和权重都沿不同维度共享到1/TP,计算/内存开销也共享到1/TP。
如果仅应用张量并行,最后一步应该是AllReduce。或者,如果张量并行与序列并行一起应用,则使用ReduceScatter。
多卡推理的专家并行化
为了启用专家并行化(EP),我们将数据并行维度从路由专家中交换出来,作为路由专家内部的专家并行维度。注意,为了获得更好的GEMM效率,专家并行可以进一步与张量并行交换,但这会增加路由不平衡的风险,我们不会在本博客中介绍这种设计。
如果在token-choice路由中启用了专家并行,由于路由到不同专家组的token数量是动态的,我们必须在使用密集张量或使用静态形状之间做出选择。
-
当优先使用eager模式时,我们使用密集张量和动态形状,以避免运行未Padding的AlltoAll造成的网络流量和内存空间浪费。
-
当优先使用图模式时,我们使用稀疏张量和静态形状,以避免通过运行CUDAGraph导致的CPU启动开销和设备到主机同步产生的GPU气泡。
注意,使用Padding激活的网络流量浪费也可以通过使用自定义AlltoAll实现来避免,但我们不会在本博客中介绍任何关于自定义通信或通信和计算融合kernel的主题。
上图是使用张量并行和专家并行的多主机推理的整体运行时设计。与使用张量并行的单主机推理相比:
Kernel接口和数据流
对于增加的基于专家并行的通信,我们使用3次All2All通信来交换形状和token:
-
第1次A2A:交换设备上关于路由到每个专家的token数量的元数据张量,即
routed_token_counts_per_expert: [E]
,这是由IndexShuffling kernel生成的输出。
-
第2次A2A:将token从基于数据并行转换为基于专家并行,根据路由分发到不同的EP ranks。
-
第3次A2A:将token从基于专家并行转换为基于数据并行,根据路由从不同的EP ranks组合。
此外,我们增加了2个额外的shuffling kernel和1个特殊的scatter kernel:
-
CombineShuffling(密集或Padding)
: 将接收到的token从按rank排序重新排列为按expert排序。后面的T*表示从所有对等节点接收的总token数,可以根据routed_token_counts_per_rank_per_expert张量的形状信息进一步解释为不规则维度。
-
输入:received_tokens: [T*, D](首先按dp ranks排序,然后按专家索引排序); routed_token_counts_per_rank_per_expert: [EP, E // EP];
-
输出:reshuffled_tokens: [T*, D](首先按专家索引排序,然后按dp ranks排序); routed_token_counts_per_expert: [E // EP];
-
SplitShuffling(密集或Padding)
:CombineShuffling的反向过程。将待发送token从专家优先顺序重新排序为rank优先顺序。
-
输入:reshuffuled_tokens: [T*, D](首先按专家索引排序,然后按dp ranks排序); routed_token_counts_per_rank_per_expert: [EP, E // EP];
-
输出:to_send_tokens: [T*, D](首先按dp ranks排序,然后按专家索引排序);
-
ScatterAdd(Padding)
:从Padding张量中scatter adds有效token。
-
输入:共享输出token: [T, D]; 接收到的Padding路由输出token: [EP, K*T, D]; 路由token索引: [K * T]; 每个专家的路由token数量: [E];
-
我们将在"图模式下使用静态形状的Padding通信"部分详细介绍上述kernel。
Eager模式下使用动态形状的非Padding通信
运行时行为的高层示意图。不同组件的实际运行时间可能会根据软件和硬件的不同而变化。
最小化动态形状的使用
由于路由是每个MoE层动态的,所需的最小设备/主机同步次数为每层一次。为了实现这一点,我们延迟了
send_sizes
的D2H复制,并将其与
recv_sizes
连接起来,通过单个D2H复制一起传输。这减少了设备/主机同步次数为每层一次。
最小化动态形状的负面影响
为了进一步隐藏设备/主机同步开销,我们进一步将共享专家分为两部分。
-
我们首先分发第一部分,在路由之后,但在分发A2A之前。然后,当设备/主机同步发生时,设备仍然保持忙碌运行共享专家。
-
我们第二部分在MoE之后,但在组合A2A之前分发。这将进一步帮助重叠第二个A2A。
图模式下使用静态形状的Padding通信
最小化Padding的使用
在无丢弃token选择设计中,路由到任何单个专家的最大可能token数量是T。然而,如果我们通过专家并行分片将多个专家组合在一起并放置在单个GPU上,对于TopK路由:
-
-
-
-
-
路由到K+1个专家的最大token数量仍然是K * T。
-
因此,路由到N个专家组的最大token数量将被限制在min(N, K) * T个token。
对于Top1路由,路由到任意大小的专家组的token数量将始终被限制在T个token,由于有EP个专家组,分配和保存动态token所需的最小内存是EP * T个token。
为了实现最小所需的Padding,我们直接使用AllGather来从不同的EP ranks收集所有活跃token,然后通过自定义kernel在本地拆分和重新排列路由token。激活大小被压缩到1 / (E // EP),这对应于内存和网络流量的减少。
上图展示了Padding设计。每个方框代表一个token,蓝色/绿色表示具有专家分配的有效token,灰色表示Paddingtoken。RiTj表示专家并行组中第i个rank的第j个token。
最小化Padding的负面影响
尽管Padding被减少到最小允许,我们还通过承担设备形状信息
routed_token_counts_per_expert
或
routed_token_counts_per_rank_per_expert
确保Padding只导致内存空间(分配)和网络流量(通信),而不是导致冗余计算(GroupedGEMM / NonLinear),冗余内存带宽(CombineShuffling / SplitShuffling / ScatterAdd)。
激活的概念解释
-
最重要的是,当所有EP ranks上的活跃token总数较小时,这样做很重要,以避免在GroupedGEMM中激活冗余专家并导致额外内存流量。
-
当所有EP ranks上的活跃token总数较大时,这样做也很重要,以避免将GroupedGEMM从memory bound转换为compute bound。
CombineShuffling
: 当前EP rank分配的token被重新排列为从专家优先顺序到rank优先顺序,在AllGather之后。未分配的token不会被复制,并且张量末尾剩余的分配内存空间保持不变。
SplitShuffling
: 当前EP rank分配的token被重新排列为从rank优先顺序到专家优先顺序,在AlltoAll之前。未分配的token不会被复制,并且重新排列的张量具有交错存储的Padding。
ScatterAdd (Padded)
: 每个EP rank最终接收来自所有其他rank计算的激活,它将理解哪些是有效token,哪些是Padding token,然后只读取有效token进行scatter_add。
