Reduce time cost of BuildOpHappensBefore #50137
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
PR types
Performance optimization
PR changes
Others
Describe
【相关背景】
PR #41471 为新执行器在预分析阶段引入了存储依赖图上结点可达关系的矩阵
op_happens_before_
,用于辅助对存储OP依赖关系的op_downstream_map_
进行剪枝。op_happens_before_
的构建需要对网络做n次BFS遍历,时间复杂度为O(n^2+nm)),其中n为OP数量,m为未经剪枝的数据依赖边数量。在PR #48809 支持跨step多流同步后,多流的依赖同步需要拼接两个step的网络结构进行分析,这使得模型中的OP和依赖边数量都成倍增加,进一步增大了新执行器预分析阶段的时间开销。
虽然实际训练时图依赖关系的构建只需要在预分析阶段执行一次,但在网络较大时,较长的预处理时间仍是不可接受的。当前自动并行场景下GPT 10B模型新执行器进行一次预分析所需时间超过10小时;手动数据并行
transformer_base_bs4096_amp_fp16
单机八卡训练切换成新执行器后,预分析阶段op_happens_before_的计算也耗费了极长的时间,导致CE监控触发模型运行超时失败报警。【具体工作】
本PR将新执行器依赖剪枝逻辑从构建完整的数据依赖边 -> 根据数据依赖边n次BFS构建可达矩阵 -> 根据可达矩阵裁剪数据依赖边更改为在构建数据依赖时实时维护和更新可达矩阵,并在添加依赖时提前判断并跳过冗余依赖。
在具体的可达矩阵计算方法上,不会再通过n次BFS来计算
op_happens_before
,而是在每次加依赖边时进行实时更新。当添加依赖边a->b时:查询所有b的可达结点c,将其添加到a的可达列表中
查询所有可达a的结点c,将b添加到c的可达列表中
为了高效查询以上两种情况中的c,额外维护两个邻接表
ops_before_
和op_behind_
,分别存储可达结点i的结点列表
以及结点i的可达结点列表
,通过空间换时间的方式降低代码复杂度。基于新的算法,添加依赖边时实时跳过冗余依赖,只在加实际依赖边时更新邻接表,总的时间复杂度为O(m'),其中m’为剪枝后的数据依赖边数量。每个结点i最坏情况下会被通过所有
可达i的结点列表
以及i的可达结点列表
各查询一次,这两个结点列表的总数不超过网络总的结点数n,假设平均可达结点数为k,则总的查询复杂度为O(kn)。整个算法计算可达关系的复杂度为O(kn + m + m'),通常情况下m >> m',因而时间复杂度为O(kn + m)。另外,由于剪枝操作的存在,k通常情况下是远小于n的。【优化效果】
基于本PR的修改,计算可达矩阵的时间复杂度从O(n^2+nm))降低为O(kn + m)*,同时依赖关系
op_downstream_map_
的构建也省去了先添加再剪掉冗余边的操作开销**。对于transformer_base_bs4096_amp_fp16模型,在跨step多流拼接两倍网络的情况下:
n = 6391
m = 16407
m' = 7772
优化前最大计算量:n^2+n*m = 145702018
优化后即使k按打满n计,最大计算量:n^2 + m= 40861288
从亿级降低到千万级,由于k在实际网络中至少能比n小一个数量级,因而实际的计算量至少可以优化到百万级。
实测本PR优化后transformer_base_bs4096_amp_fp16模型可在感知较小的时间开销内完成一轮预分析,且前后依赖分析结果逐一对比完全一致。
对于自动并行GPT 10B模型,优化效果未实测,后续由@JZ-LIANG 进行验证。