计算通信重叠
多机多卡的专家并行会引入比较大的通信开销,所以我们使用了双 batch 重叠来掩盖通信开销,提高整体吞吐。
对于 prefill 阶段,两个 batch 的计算和通信交错进行,一个 batch 在进行计算的时候可以去掩盖另一个 batch 的通信开销;
Prefill 阶段的双 batch 重叠
对于 decode 阶段,不同阶段的执行时间有所差别,所以我们把 attention 部分拆成了两个 stage,共计 5 个 stage 的流水线来实现计算和通信的重叠。
Decode 阶段
的双 ba
tch 重叠
尽可能地负载均衡
由于采用了很大规模的并行(包括数据并行和专家并行),如果某个 GPU 的计算或通信负载过重,将成为性能瓶颈,拖慢整个系统;同时其他 GPU 因为等待而空转,造成整体利用率下降。因此我们需要尽可能地为每个 GPU 分配均衡的计算负载、通信负载。
1. Prefill Load Balancer
1)核心问题:
不同数据并行(DP)实例上的请求个数、长度不同,导致 core-attention 计算量、dispatch 发送量也不同
2)优化目标:
各 GPU 的计算量尽量相同(core-attention 计算负载均衡)、输入的 token 数量也尽量相同(dispatch 发送量负载均衡),避免部分 GPU 处理时间过长
2. Decode Load Balancer
1)核心问题:
不同数据并行(DP)实例上的请求数量、长度不同,导致 core-attention 计算量(与 KVCache 占用量相关)、dispatch 发送量不同
2)优化目标:
各 GPU 的 KVCache 占用量尽量相同(core-attention 计算负载均衡)、请求数量尽量相同(dispatch 发送量负载均衡)
3. Expert-Parallel Load Balancer
1)核心问题:
对于给定 MoE 模型,存在一些天然的高负载专家(expert),导致不同 GPU 的专家计算负载不均衡
2)优化目标:
每个 GPU 上的专家计算量均衡(即最小化所有 GPU 的 dispatch 接收量的最大值)
参考架构图
线上系统的实际统计数据