浅谈后向传递的计算量大约是前向传递的两倍

【点击】 加入大模型技术交流群

1. 前言

训练神经网络的一次迭代分为三步：（1）前向传递计算损失函数；（2）后向传递计算梯度；（3）优化器更新模型参数。在实验中，我们观察到一个现象： 后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时几乎可以忽略 。要解释这个现象，我们要从前向传递、后向传递和优化器参数更新的浮点数计算次数入手。

上图表示一次训练迭代中各个环节（前向传递、后向传递、通信环节、优化器更新）的耗时占比，来自于《PyTorch Distributed: Experiences on Accelerating Data Parallel Training》。上图中，纵轴表示耗时占比，FWD表示一次训练迭代中前向传递的耗时占比，BWD则表示一次训练迭代中后向传递的耗时占比，OPT表示一次训练迭代中优化器更新模型参数的耗时占比。从上图中可以看到，一次训练迭代中，后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时占比很小，几乎可以忽略。

上图表示GPipe流水线并行的调度策略，来自于《Efficient large-scale language model training on gpu clusters using megatron-lm》。上图中，横轴表示耗时，一个蓝色小块表示一个微批次的前向传递，一个绿色小块表示一个微批次的后向传递，黑色竖线表示一次流水线刷新，也就是优化器更新模型参数。从上图中可以看到，一个绿色小块的宽度大约是蓝色小块的二倍，一次训练迭代中，后向传递的耗时几乎是前向传递的两倍，相比之下，优化器更新的耗时占比很小，几乎可以忽略。

从上表中，我们可以观察到，对于多层前馈神经网络模型，有以下结论：

1. 相比于线性层，激活函数ReLU和损失函数的浮点数运算量可以忽略。

2. 对于第一层，后向传递-前向传递的FLOPs比率是1:1。

3. 对于其他层，后向传递-前向传递的FLOPs比率是2:1。

4. 采用随机梯度下降作为优化器，权重更新的FLOPs是模型参数规模的2倍。

3.2.1 第一层与其他层的区别

对于多层前馈神经网络模型，采用随机梯度下降作为优化器。

第一层的后向传递-前向传递的FLOPs比率是1:1，其他层的后向传递-前向传递的FLOPs比率是2:1。模型参数更新的FLOPs是模型参数规模的2倍。

3.2.2 batch_size的影响

前向传递和后向传递的计算量FLOPs与batch_size成正比，即随着batch_size增大而线性增长。

优化器更新模型参数的FLOPs与batch_size无关，只与模型参数规模和优化器类型有关。

随着batch_size增大，前向传递和后向传递的FLOPs线性增长，而权重更新的FLOPs保持不变，这使得权重更新的FLOPs变得逐渐可以忽略不计了。