PyTorch中为何需将学习率乘以梯度累积步数？

PyTorch中损失函数使用"mean"归约方式时，无论batch size（批次大小）如何，模型梯度的量级大致相同。因此当增大批次大小时，需要按比例提高学习率——因为梯度并不会随批次增大而变大。

但PyTorch的梯度累积机制中，梯度会被累加N次（N是调用step()前调用backward()的次数）。按直觉这会增大梯度的量级，应该降低学习率才对，至少不该提高。不过看到两处代码都将学习率乘以了梯度累积步数：

第一处代码：

if args.scale_lr:
    args.learning_rate = (
        args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
    )

第二处代码：

model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr

我能理解按批次大小缩放学习率的原因，但搞不懂为什么要乘以梯度累积步数，有两个疑问：

• 是不是他们会把损失除以N来降低梯度量级？不然为什么要把学习率乘以梯度累积步数？
• 多GPU间的梯度是怎么累积的？取均值还是求和？如果是求和，为什么要把学习率乘以GPU数量？

问题1解答 #

核心是模拟大batch训练的效果。当用梯度累积N次时，本质是把N个小batch当成一个完整的大batch来训练。

如果损失用"mean"归约，每个小batch的梯度是该batch内样本的均值梯度。累积N次后，总梯度是N个小batch梯度的和（因为每次backward()都是累加操作）。而如果直接用一个大小为N*小batch_size的大batch训练，损失是所有样本的均值，对应的梯度是所有样本梯度的均值——也就是N个小batch梯度之和再除以N。

两者的梯度差了N倍：累积N次的梯度是大batch梯度的N倍。要让梯度更新的效果和大batch训练完全一致，有两种方案：

1. 每次backward()前把损失除以N，这样累积后的梯度就和大batch的梯度量级一致；
2. 不修改损失，而是把学习率乘以N，这样学习率*梯度的结果就和大batch场景下（大batch学习率*大batch梯度）的更新量对齐。

那些代码选择的是第二种方案，所以要乘以梯度累积步数，目的就是让梯度累积的训练效果和直接用大batch训练对齐。

问题2解答 #

多GPU训练（比如用DDP）时，梯度默认是取均值的：每个GPU计算自身batch的梯度，之后所有GPU的梯度会被收集起来求平均，再把均值广播给每个GPU。

乘以GPU数量的原因还是为了对齐大batch效果。假设用K个GPU，每个GPU的batch size是B，总等效batch size就是K*B。如果不用梯度累积，直接训练的话，DDP输出的梯度是总batch的均值梯度，和单GPU用K*B大小batch训练的梯度量级一致。而单GPU大batch训练时，学习率需要按总batch size缩放，所以这里乘以K就是为了匹配总batch size对应的学习率。

如果再加上梯度累积N次，总等效batch size就是N*K*B，所以学习率要同时乘以这三个参数，和代码里的逻辑完全对应。

内容的提问来源于stack exchange，提问作者offchan