流水线并行没开Cell共享导致编译时间很长

1 系统环境

硬件环境(Ascend/GPU/CPU): Ascend
MindSpore版本: mindspore=2.2.10
执行模式(PyNative/ Graph):PyNative/ Graph
Python版本: Python=3.8.15
操作系统平台: linux

2 报错信息

使用了流水线并行,模型编译时间很长(几个小时),或者因为编译时间太长误以为是卡住。

3 根因分析

在大模型场景中,编译耗时问题尤为突出,一是大模型的模型结构层次深,节点数多;二是大模型在训练时,由于启用pipeline并行,导致模型规模和节点数进一步加大,如果原来图的规模是O,那开启pipeline并行,单节点图的规模变为(O/X)*Y,其中X为pipeline的stage数量,Y为micro batch的数量。以盘古13B网络为例,计算图中计算节点数量达到13.5万个,单次编译时长可接近3小时。

4 解决方案

使用lazy_inline装饰器优化编译性能
MindSpore 2.2版本前:控制是否使能lazy_inline是通过@lazy_inline加上MS_DEV_CELL_REUSE这个环境变量。
MindSpore 2.2及以后版本:控制lazy_inline是否使能只需通过@lazy_inline。
如果使用的是mindformers:mindformers上为了做到对不同版本的ms,通过同一个环境变量统一控制是否开启lazy_inline,搞了个环境变量ENABLE_CELL_REUSE来控制是否使能。