长文本场景下,LLM训练中梯度累加存在的BUG
Unsloth的梯度累积修复确保训练过程和损失计算能够准确和正确地执行。梯度累积的目标是在减少显存(VRAM)使用量的同时模拟完整批次训练。由于梯度累积也用于DDP和多GPU设置中,因此这个问题同样影响着大规模训练。 from unsloth import unsloth_train # trainer_stats = trainer.train() << 存在bug的梯度累积 trainer_stats = unsloth_train(trainer) 复现问题 理论上,梯度累积在数学上应该等效于完整批次训练。我们使用有效批次大小16进行训练,因此批次大小(bsz) * 梯度累积步数(ga)应保持恒定。我们测试了bsz=1,2,4,8和16,发现使用较大梯度累积步数的训练损失始终更高。 什么是梯度累积? 在训练或微调过程中,每一步都会从训练数据集中选择一定数量的随机样本来更新模型权重。但应该选多少样本呢?对于非常大的预训练任务,批次大小可能达到数百万,就像在Llama 3.1中那样,这有助于减少过拟合并提高模型的泛化能力。而对于像Unsloth的Llama 3.2笔记本中的微调任务,批次大小可能只有较小的32。 问题在于大批次的内存使用量非常大。如果1个批次使用1单位内存,那么100万大小的批次将需要100万单位内存。我们如何模拟大批次训练但又不消耗大量内存呢? 这就是梯度累积的用武之地!我们通过在每次新的小批次到来时即时创建梯度,然后将所有小梯度加起来,进行适当缩放,从而获得最终的大批次梯度。 可能的解释 一种流行的理论认为梯度累积在累积步骤中存在数值误差。但研究人员发现,即使在float32中进行累积也会产生相同的问题。我们的研究表明,确实存在一些微小的累积误差。 第二种理论是损失计算中存在bug,我们确认了这一点。 数学上是否等价? 梯度累积和完整批次训练在数学上是否等价?遗憾的是,如果简单地将梯度加起来,答案是否定的!我们首先注意到交叉熵损失是通过以下方式计算的: $$ \frac{1}{\sum \mathbb{I}\{y_i \neq -100\}} \sum L_i $$ 注意分母计算的是非填充或非忽略的token数量 - 即它通过每个文本片段中有效训练token的数量来归一化损失。指示函数实际上是未填充token的总和,也就是所有序列长度的总和,即: $$ \mathbb{I}\{y_i \neq -100\} = \sum m_i $$ 因此我们得到最终方程为: $$ \frac{\sum L_i}{\sum m_i} $$ 然后我们在分子和分母中同时添加 $\frac{1}{n}$ - 这是允许的,因为两者可以相互抵消: $$ \frac{\frac{1}{n}\sum L_i}{\frac{1}{n}\sum m_i} $$ 这意味着最终损失是平均损失值除以所有未填充序列长度的平均值: $$ \frac{\bar{L}}{\bar{m}} $$ 在进行梯度累积时,我们需要分别计算每个小批次的损失,然后将它们加起来得到最终损失。我们首先利用每个分区的平均损失和平均序列长度。 但我们发现,最终总和不等于原始的完整批次损失 - 实际上它比原来大$G$倍(其中$G$是梯度累积步骤的数量)。 ...