重计算流水独立调度

背景与挑战

在目前的流水线调度中，重计算由反向计算触发，与反向计算绑定在一起调度，意味着重计算需要等待下一个stage返回梯度才可以开始计算。然而重计算并不需要用到反向计算的梯度，这导致空泡的增多和性能的下降。

解决方案

为了将重计算和反向计算独立调度，需要将重计算的调度修改为由调度器主动触发，并修改调度器，将重计算作为一个调度单元加入到调度器中，这使我们获得了自由地插入或去除部分重计算的能力，进而可以在内存和性能方面做出优化。

通过torch的saved_tensors_hooks实现一种新的重计算方法，在反向计算前合适的时机主动触发或者直接去除部分重计算，从而实现对内存或性能的优化。

使用场景

在虚拟流水线调度中，若用户未开启重计算，则可以利用bubble主动插入重计算，以极小的性能代价换取内存峰值的降低，将需要保留激活值的前向计算块的个数减少到PP * VPP（PP表示流水线并行数，VPP表示虚拟流水线并行数）。

图1 未开启重计算调度图

在虚拟流水线调度中，若用户已开启重计算，则可以通过解除重计算与后一个stage的反向计算间的依赖关系从而提前重计算，以及去除模型最后一层的重计算，实现计算性能的提升。

图2 开启重计算调度图

使用方法

以上两个功能不可同时开启。

使用效果

--recompute-in-bubble可以实现内存节省，--recompute-in-advance可以实现训练性能提升。