论文

反向状态策略是学习算法的一部分

Backward-State Policy Is Part of the Learning Algorithm

模型训练预训练

摘要

低精度训练会对张量舍入,而反向传播可能为了计算多个梯度而再次读取这些张量;每次使用可读取前向传播时的舍入值、原始值,或重新随机舍入的值。这种反向状态策略看起来只是内存与精度细节,似乎可由复制准确性和最终损失判定。我们认为它是学习算法的一部分,而这两项检查都不能证明其正确。复制准确性并不决定结果:在三组配对的390M模型训练、采用模拟FP8反向传播时,注意力反向传播复用前向舍入输出会导致训练失败,而从同一分布重新舍入则成功。按照我们的参照,即实际执行的前向传播的梯度、且梯度穿过舍入时保持不变,连最准确的原始值也可能是错误的。例如,低精度保存的归一化输出用于两个梯度:增益参数的梯度需要原始值,而下一层权重的梯度需要该层实际相乘的舍入值。最终损失也不能排除两处都读取原始值的错误:这种错误持续存在于采用该存储方式训练的模型中,而预先设计的损失比较仍处于提前设定的容差内。因此,我们依据该参照推导每次使用必须读取的值,或在固定前向传播条件下能产生相同平均梯度的替代值,并无需训练、在单个算子上检查这些逐次使用要求。在采用PyTorch与Transformer Engine的三项测试中,这些要求提前预测了复用是否会使反向计算的平均结果偏离独立副本,所有预测均成立。因此,反向状态策略是学习算法的一部分,应按每次使用分别规定和检查,不能只由复制准确性及最终损失决定。