Norm重计算

背景与挑战

大模型训练过程中，往往会面临内存不足的问题。

解决方案

类似于激活函数重计算，本特性支持了Norm层的重计算。

运用激活函数重计算特性中的checkpoint机制，对Norm层进行重计算处理，具体细节可参见文献Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism。

使用场景

主要用于训练场景，用户内存不足或要进一步节省内存时。

使用方法

--recompute-norm      # 开启Norm重计算
--recompute-norm-num-layers  N    # N表示Norm重计算的层数

此特性仅支持Mcore分支。

Norm重计算兼容激活函数重计算、全重计算同时开启：

• 同时开启时，仅支持--recompute-method设置为block。
• 同时开启时，将按照指定的全重计算和Norm重计算的层数做各自类型的重计算，即不会有一层既做全重计算又做Norm重计算。

另外，执行优先级是先计算全重计算层，后Norm重计算层。

使用效果

开启后可节省RMSNorm/LayerNorm层的输出激活内存，并且由于Norm计算速度较快，重计算后对整体性能影响较小。对于开启TP及SP的场景，由于该激活内存在TP域内已进行切分，开启后效果不明显，对于未使用TP及SP的模型可考虑使用。