- 在解决的是什么问题?第一个自动做模型并行的编译器,结合了数据、算子和模型并行,而且也支持 MoE 上的并行策略
- 为何成功,标志/准是什么?比速度,而且适配的模型类型跟多(不限于transformer),不需要手工指定并行计划
- 在前人基础上的关键创新是什么?把并行抽象为两层:算子间和算子内,这样提供了新的层次化(2层)的解空间。而已有的方案要么需要手工创建并行策略,要么只能从有限的
- 关键结果有哪些?transformer模型上跟 Megatron 差不多,MoE 模型上比 DeepSpeed 要好,还能支持其他人都不支持的 Wide-ResNet
- 有哪些局限性?如何优化?只支持能拿到整个计算图,而且在编译期就知道所有的tensor shapes的情况
- 这个工作可能有什么深远的影响?
为了解决上述问题,我们设计和实现了 Alpa,是第一个自动产生并行执行计划,覆盖了所有数据、算子、流水并行方法。给定模型的描述和集群配置(只是说几台机器?),Alpa 就可以把集群划分为几个 device meshes,然后把模型的计算图划分到不同的阶段(stages)。然后把 stages 分配到 device meshes,在同一个device mesh内部自动编排算子内并行,在 device meshes 间执行算子间并行(data、pipeline)
我们的核心观察是:可以把不同的并行技术组织为层次化的空间,正好映射到计算机集群的两层层次化结构上:同一台主机上的设备通信带宽很高,而不同主机间带宽低。
我们根据是否需要在算子粒度进行划分,分为两种:
- 算子间(inter-op)并行: 流水并行,只需要在邻接的两个 stage 上通信,切分好了通信很小,但是优于调度限制会有空闲时间。所以正好放在机器间
- 算子内(intra-op)并行: 包含数据和算子并行,有更高的设备利用率,但是在每个iteration粒度上有split和merge上通信开销。正好放在机器内
其次这样划分的好处是可以把两层各自当作一个可解决的子问题来得到近似最优解,尽管联合起来的执行计划并不能保证全局最优,但是已经证明性能在很多大模型上很高了
主要的贡献如下:
- 构建了一个两级的并行执行计划空间
- 设计可解的优化算法来继承每一级上的执行计划
- 实现了 alpa,给GPU集群上的分布式DL使用的编译系统
它主要特性:
- 一堆编译器的pass,利用层次化的优化算法来产生执行计划
- 一个新的运行时架构,可以在不同主机(device mesh)间编排算子间并行
- 一堆系统优化技术,可以提高编译速度,解决跨mesh通信的问题
几种 pipeline 并行的调度策略:
Gpipe: fill-drain
PipeDream:
dynamic schedules: lookahead, registers/back-pressure
synchronous 1F1B: 本论文里选用的这个,因为有同样的pipeline latency,但是跟 Gpipe 比有更低的峰值显存使用
虽然有一些同期工作已经提出了类似的分类,但是 Alpa 是第一个端到端自动产生执行计划,而且支持完整的空间
inter-op 的优化依赖于知道由intra-op优化提供的当前每个子图的执行速度。
inter-op 先根据 dp 来划分出不同的 stage,然后向 intra-op optimization 查询这些stage 的速度
intra-op 里根据 ILP 公式建模执行速度(跟ILP很像),汇报分析出的最低代价。这样不断请求每个 stage 组合的情况,inter-op就利用DP来最小化 inter-op 并行执行的延迟,最终得到最佳计划
R: replicated
S: 在某个维度上 shard
Reshard: 为了在不同op间输入tensor形状不同时做layaout转换
是针对 XLA 的 HLO 里的各种op(大约80种)做了各种算子内并行的实现(那可以看看conv这里的?)。当并行后,比原来速度快,就可以搞
优化的目标,分为两部分:第一项是计算和通信的代价,第二块是 reshard 的代价
为什么 dv 部分是0?
通过噢把计算开销小的算子合并到其他算子来减少图里的节点
在ILP之前,有一波优化,比如把所有 allreduce 替换为 reduce-scatter 和 all-gather,因为后者的replicated tensor更少,计算少,但是通信量相同?我们能用吗?
两种优化:
- 剪枝
- operator分组。ReLU 等非计算密集型,对执行开销影响很小
公式里 FLOP 部分是做什么的?公式里复杂度怎么看?K是什么?
- 是否支持 checkpoint
- 与 DAPPLE 的差异?
- 与 Unity 的差异? Unity 是基于flexflow,使用的人少
- 真正执行前,编译一次需要多久?半小时到一小时这种粒度,跟模型大小和机器数量有关
- 为什么这样划分后,能得到全局最优?不能保证,但效果挺好的
- 看看跨mesh通信的代码生成
- 看看 Wide-ResNet 上4 nodes 上80% 线性扩展效率是啥意思?我们以后也可以这样评估?
- 我们可以在ILP求解时,把 checkpoint 也考虑进去?类似 monet
- 怎么没在论文里看到描述 dp 粒度?是给定各种并行度,找一个最优的划分?
- 看看 wideresnet 的13B 模型里,它是的策略,在哪里需要用到TP?