网站突然打不开,中山企业网站推广,免费创建社区平台,网站网站合作建设使用Sum计算Loss和解决梯度累积的Bug
学习 https://unsloth.ai/blog/gradient#xff1a;Bugs in LLM Training - Gradient Accumulation Fix 这篇文章的记录。
在深度学习训练过程中#xff0c;尤其是在大批量#xff08;large batch#xff09;训练中#xff0c;如何高…使用Sum计算Loss和解决梯度累积的Bug
学习 https://unsloth.ai/blog/gradientBugs in LLM Training - Gradient Accumulation Fix 这篇文章的记录。
在深度学习训练过程中尤其是在大批量large batch训练中如何高效地计算损失函数Loss并避免内存溢出一直是一个重要的问题。在许多情况下为了实现更高效的训练我们会使用梯度累积Gradient Accumulation策略模拟大批量训练的效果但同时减少对显存的需求。然而这种方法可能引发一些数值计算错误尤其是在损失计算时。因此理解如何处理梯度累积并确保计算正确性就变得尤为重要。
为什么要使用Sum计算Loss
在代码中reduce_losssum表示我们选择将每个小批量mini-batch的损失值相加而不是取平均。这种做法有几个目的 权重每个样本的影响一致在使用大批量训练时如果直接取每个样本的平均损失可能会使得样本数目更多的批次对模型训练产生不必要的影响。而通过将损失值求和确保了每个样本都以相同的权重参与模型更新。 避免梯度累积带来的偏差如果只使用小批量训练的平均损失进行梯度累积梯度更新随着梯度累积步骤gradient accumulation step增多损失的计算会变得不准确。因为简单地将每个小批量的损失相加可能会导致最终计算的损失过大。
梯度累积中的问题
当我们使用梯度累积时目的是在显存限制下模拟大批量训练的效果。假设每个小批量的计算后我们累积其梯度最后再进行一次反向传播。这种方法可以显著减少显存的使用但也带来了数值误差特别是在损失的计算上。
梯度累积是否和完整批次训练数学上等价
答案是 不完全等价尤其是如果我们不正确地处理梯度的累积。在数学上损失函数通常表示为 L 1 n ∑ i 1 n L i L \frac{1}{n} \sum_{i1}^{n} L_i Ln1i1∑nLi
其中( L i L_i Li ) 是每个样本的损失( n n n ) 是样本总数。我们通常会对每个小批量的损失进行归一化处理使得每个样本对损失的贡献相等。
但是在梯度累积时我们的处理方法是 L accumulated ∑ k 1 G L k L_{\text{accumulated}} \sum_{k1}^{G} L_k Laccumulatedk1∑GLk
其中( G G G ) 是梯度累积的步数 ( L k L_k Lk ) 是每个小批量的损失值。这导致累积损失和原本应计算的损失不同因为没有进行正确的归一化操作。
为什么梯度累积会导致问题
由于梯度累积在计算损失时没有正确处理每个小批量的权重最终的损失值可能比实际应有的损失大特别是当批次大小不一致时。此时我们需要对每个小批量的损失进行缩放以使最终的结果与大批量训练的损失相符。
解决方案修正梯度累积中的损失计算
为了解决梯度累积中损失计算的偏差我们可以在每次计算梯度时对每个小批量的损失进行缩放。具体来说假设我们在训练过程中设置了 ( G G G ) 步梯度累积那么每次计算损失时我们需要将每个小批量的损失除以 ( G G G )以确保最终的梯度更新符合预期。 数学公式如下 L final 1 G ∑ k 1 G L k m k L_{\text{final}} \frac{1}{G} \sum_{k1}^{G} \frac{L_k}{m_k} LfinalG1k1∑GmkLk
其中( m k m_k mk ) 是第 ( k k k ) 个小批量的有效样本数即去除填充后的token数( G G G ) 是梯度累积的步数。这样做的目的是在每次梯度累积时确保损失值的加权平均而不是简单地将其相加。
小结 Loss的求和在使用 reduce_losssum 时我们通过求和的方式来计算损失而不是简单的平均。这可以确保每个样本对损失计算的贡献一致避免了梯度累积中的偏差。 梯度累积的误差梯度累积在计算损失时容易出现数值误差尤其是在不同小批量的损失和有效样本数不一致的情况下。为了解决这一问题我们需要对每个小批量的损失进行缩放处理使得最终的损失与大批量训练的结果一致。
通过正确处理梯度累积的损失计算我们可以在不增加内存消耗的情况下模拟大批量训练的效果从而提高训练效率。
附录代码解析
下面是open instruct框架提供的sum计算 loss的代码它实现了一个损失计算和梯度累积的处理流程。
if args.reduce_loss mean:loss outputs.loss
else:# reduce loss is sum# this ensures that we weight all tokens in the dataset equally,# rather than weighting each overall example equally when# using high amounts of gradient accumulation.# this can result in 5 point improvements in AlpacaEval# see https://github.com/huggingface/transformers/issues/24725 for# more discussion and details.logits outputs.logitslabels batch[labels]# Shift so that tokens n predict nshift_logits logits[..., :-1, :].contiguous()shift_labels labels[..., 1:].contiguous()# Flatten the tokensloss_fct torch.nn.CrossEntropyLoss(reductionsum)shift_logits shift_logits.view(-1, embedding_size)shift_labels shift_labels.view(-1)# Enable model parallelismshift_labels shift_labels.to(shift_logits.device)loss loss_fct(shift_logits, shift_labels)if args.load_balancing_loss:aux_loss args.load_balancing_weight * outputs.aux_lossloss aux_loss
# We keep track of the loss at each logged step
total_loss loss.detach().float()
accelerator.backward(loss)1. Loss的选择reduce_loss参数
代码的第一部分是根据args.reduce_loss参数来决定使用平均mean还是求和sum的损失计算方式。如果选择的是mean则直接使用outputs.loss通常是模型自带的损失。如果选择的是sum则使用后续的处理方式。
2. 处理sum模式的损失计算
如果reduce_loss为sum代码会进入求和模式具体步骤如下 获取logits和labels从模型的输出outputs中获取logits预测值和batch[labels]标签。 Shift操作为了计算交叉熵损失模型的输出logits会做一个移位操作使得每个token的预测目标对应下一个token。这是因为语言模型任务中的预测是基于上一个token的。 shift_logits logits[..., :-1, :].contiguous()
shift_labels labels[..., 1:].contiguous()这里logits[..., :-1, :]表示除去最后一个token的logits而labels[..., 1:]表示去掉第一个token的标签。这样处理的目的是确保每个token的预测值与下一个token的真实标签进行比较。 Flatten操作对处理后的logits和labels进行flatten展平为一维向量以便交叉熵损失函数计算。 shift_logits shift_logits.view(-1, embedding_size)
shift_labels shift_labels.view(-1)这一步是为了将shift_logits和shift_labels转换为适合交叉熵损失函数的格式即shift_logits变为 [batch_size * sequence_length, embedding_size] 维度shift_labels变为 [batch_size * sequence_length] 维度。 交叉熵损失函数使用torch.nn.CrossEntropyLoss计算损失。这里的reductionsum表示对所有token的损失求和而不是取平均。这样可以确保每个token的贡献相同。 loss_fct torch.nn.CrossEntropyLoss(reductionsum)
loss loss_fct(shift_logits, shift_labels)这段代码计算了每个token的损失并将所有token的损失求和。对于大批量训练这样做是为了避免梯度累积时损失计算的不准确。 负载平衡损失Load Balancing Loss如果args.load_balancing_loss为True表示需要加上额外的负载平衡损失。负载平衡损失用于调整模型在不同任务之间的负载以确保训练过程的稳定性。 if args.load_balancing_loss:aux_loss args.load_balancing_weight * outputs.aux_lossloss aux_loss这部分代码将outputs.aux_loss辅助损失加权后加到主损失上以增强训练的稳定性。
3. 反向传播
最后代码通过accelerator.backward(loss)执行反向传播计算梯度。total_loss用于累积损失以便记录每个步骤的损失值。
total_loss loss.detach().float()
accelerator.backward(loss)loss.detach().float()表示将损失从计算图中分离并将其转换为浮动类型以便进行记录和后续的反向传播。
小结 Loss的计算方式如果选择sum会对每个token的损失进行求和而不是取平均以确保每个token对梯度更新的贡献相等。这对于大批量训练非常重要能够避免在梯度累积过程中出现误差。 梯度累积通过逐步计算小批量的梯度并累积能够在不增加显存使用的情况下模拟大批量训练。需要确保在每次累积时正确计算损失并进行适当的缩放。
这段代码的设计和实现考虑到了梯度累积和损失计算中的数值稳定性问题确保在大批量训练时能够正确更新模型参数。
梯度累积公式解释
在梯度累积的公式中( L k L_k Lk ) 表示第 ( k k k ) 步梯度累积中的损失值。具体来说假设我们将训练数据分成 ( G G G ) 个小批量mini-batches进行梯度累积那么在第 ( k k k ) 步即第 ( k k k ) 个小批量中我们计算的损失为 ( L k L_k Lk )。这些损失是通过模型对第 ( k k k ) 小批量的数据进行前向传播后得到的。
梯度累积的背景
梯度累积的目的是在多个小批量上计算梯度并在一定数量的步数之后再进行一次反向传播。这样我们可以模拟更大的批量大小从而提高训练稳定性和效率而不需要一次性处理过大的批量数据。
解释公式
公式中的 ( L k L_k Lk ) 是每个小批量的损失值具体来说它是模型在第 ( k k k ) 步计算得到的损失。公式可以表示为 L accumulated ∑ k 1 G L k L_{\text{accumulated}} \sum_{k1}^{G} L_k Laccumulatedk1∑GLk
这意味着我们将每个小批量即每个 ( k k k )的损失 ( L k L_k Lk ) 累积起来直到进行 ( G G G ) 步梯度累积。此时的累积损失 ( L accumulated L_{\text{accumulated}} Laccumulated ) 还没有进行任何梯度更新操作它只是所有小批量损失的总和。
举个例子
假设我们有一个批量大小为 8 的训练数据集并且我们希望模拟一个批量大小为 32 的训练过程但由于显存限制我们决定使用梯度累积来分 4 个小批量来计算梯度
第一步计算第一小批量batch 1的损失 ( L 1 L_1 L1 )。第二步计算第二小批量batch 2的损失 ( L 2 L_2 L2 )。第三步计算第三小批量batch 3的损失 ( L 3 L_3 L3 )。第四步计算第四小批量batch 4的损失 ( L 4 L_4 L4 )。
此时梯度累积的总损失就可以表示为 L accumulated L 1 L 2 L 3 L 4 L_{\text{accumulated}} L_1 L_2 L_3 L_4 LaccumulatedL1L2L3L4
然后我们对这个累积损失进行一次反向传播更新模型参数。
为什么要进行梯度累积
梯度累积的好处是可以模拟大批量训练而不会导致显存溢出。通过分步计算损失并累积梯度我们可以在不增加显存需求的情况下使用较大的有效批量大小进行训练。
损失值求和的例子
通过一个简单的例子来说明为什么在使用大批量训练时损失值求和的方式比损失值平均更能保证每个样本都以相同的权重参与模型更新。
背景说明
在大批量训练中损失的平均和损失的求和对训练结果的影响不同。具体来说如果我们采用损失平均每个样本对模型更新的贡献会受到批次大小的影响而如果我们采用损失求和每个样本对更新的贡献将保持一致无论批次大小有多大。
例子解释
假设我们有两个不同的批次分别包含不同数量的样本
批次 A包含 4 个样本。批次 B包含 8 个样本。
假设这两个批次的损失分别如下
批次 A 的损失每个样本的损失分别为 2, 3, 1, 4合计为 10。批次 B 的损失每个样本的损失分别为 1, 2, 3, 1, 2, 1, 4, 3合计为 20。
如果使用损失的平均值
批次 A平均损失为 ( 10 4 2.5 \frac{10}{4} 2.5 4102.5 )。批次 B平均损失为 ( 20 8 2.5 \frac{20}{8} 2.5 8202.5 )。
如果我们采用损失平均的方式两个批次的损失平均值相同。因此无论是批次 A 还是批次 B它们对模型的训练贡献是相同的。也就是说每个样本的贡献被平等化了。这个方法在样本数差异较大的情况下可能会引入偏差特别是当批次大小不一致时每个批次的影响可能不成比例。
如果使用损失的求和
批次 A总损失为 10。批次 B总损失为 20。
在这种情况下两个批次的损失对模型的影响是按批次大小加权的。具体来说批次 B 有更多的样本所以它的损失总和较大也就意味着在梯度更新时批次 B 会对模型参数的更新产生更大的影响。如果我们将这两个批次的损失求和最终总损失会更大这会导致模型在训练过程中更多地倾向于批次 B尤其是当使用梯度累积时批次 B 会对模型产生更大的影响。
为什么损失求和更有意义
通过将损失求和确保每个样本的损失对模型更新的影响是等量的。而使用平均损失时由于批次大小不同可能会出现以下问题
样本较多的批次如批次 B会占据更大的权重因为平均值会稀释了损失的影响使得较大的批次对整体损失的影响降低。样本较少的批次如批次 A则可能被忽视因为它对最终的损失计算贡献较小。
梯度累积中的影响
在梯度累积中采用损失求和可以避免不同批次间因批次大小差异而导致的权重不均的问题。这样在进行多次梯度累积时不同批次的损失就可以通过求和确保每个样本都具有相同的影响力。
举个具体的例子
假设我们有 2 个小批量分别为 A 和 B批次大小分别为 4 和 8。每个批次的总损失分别为 10 和 20。如果我们直接进行平均计算得到的损失为 L avg 10 20 4 8 30 12 2.5 L_{\text{avg}} \frac{10 20}{4 8} \frac{30}{12} 2.5 Lavg48102012302.5
而如果我们进行损失求和得到的损失为 L sum 10 20 30 L_{\text{sum}} 10 20 30 Lsum102030
在进行梯度累积时采用求和损失可以避免批次大小的不均衡影响从而确保每个样本的贡献是等量的。
如果batch size都一样的话还会有这个问题吗 如果batch size都一样的话还会有这个问题吗现实中大规模训练batch size开的不一样吗 如果batch size 都一样那么使用损失的平均或者求和不会产生太大的差异。因为在这种情况下每个批次对模型更新的贡献是均衡的无论是计算损失的平均值还是求和最终对模型参数的影响都是一致的。
如果 batch size 相同
平均损失损失的平均值会反映每个样本的损失在训练中的影响因为每个样本的数量是相同的计算平均值不会受到影响。求和损失由于批次大小相同求和后的损失也是一个与批次大小成正比的值而在训练过程中更新的频率和力度会基于每个批次的大小进而影响训练的稳定性和收敛性。
所以如果每个批次大小一样通常来说不会出现上述问题。两种方式平均损失和求和损失的训练效果是差不多的。 现实中大规模训练的情况
在现实中的大规模训练中批次大小batch size通常是不同的这取决于多个因素例如 显存限制不同的硬件设备如 GPU、TPU的显存大小不同支持的 batch size 也不同。在显存较小的设备上可能无法一次性处理大量数据因此会选择较小的批次进行梯度累积而在显存较大的设备上则可以使用较大的 batch size 进行训练。 分布式训练在多卡multi-GPU或者多机multi-node训练的情况下为了充分利用计算资源批次大小也会有所不同。通常跨多个 GPU 或节点分布训练时为了平衡每个卡上的负载可能会采用不同的 batch size。 性能优化有时为了加速训练训练人员可能会根据任务的需要进行优化选择不同的批次大小以提升计算效率。例如在某些任务中增大批次大小可以提高并行效率但也可能导致训练不稳定这时就需要进行梯度累积。 模型设计有些模型架构比如需要大输入的模型可能会需要更大的批次以适应内存需求或者为了更高效地训练。 在大批量训练中为什么使用梯度累积
即便batch size 不同如果我们需要模拟一个大批量训练的效果并且又不能增加显存的使用我们通常会使用梯度累积来应对这个问题。 模拟大批量效果如果显存不足无法处理大批量数据我们可以通过梯度累积模拟大批量训练的效果。通过在多个小批量上计算损失并累积梯度我们可以达到类似大批量训练的效果避免了显存溢出的问题。 不同批次大小的影响如果不同批次的大小不一致在计算损失时损失求和的方式就显得更为重要了因为批次大小较大的批次对总损失的贡献较大而平均损失的计算可能会削弱这种差异。而损失求和方式则会保证较大批次对总损失的影响相对较大从而保持了每个样本的影响力一致。 总结
如果batch size 相同那么损失求和和损失平均不会有太大差异。现实中大规模训练确实可能会有不同的 batch size尤其是在显存和计算资源有限的情况下。为了应对这个问题梯度累积是一个常见的解决方案用于模拟大批量训练的效果。在使用不同大小的 batch size 时损失求和能够保证每个样本在训练中的权重一致而损失平均可能会导致较大批次的样本对训练过程的影响被稀释。
为什么要除以 m k m_k mk直接用 L k L_k Lk不行吗?
这是一个非常重要且值得深入思考的问题。让我们仔细分析为什么在梯度累积时每个小批量的损失还需要除以 ( m k m_k mk)即有效token数。
1. 有效token数量的影响
首先需要理解的是损失 ( L k L_k Lk) 通常是基于每个小批量的有效token计算的。每个小批量中的有效token数量去除填充后的部分可能是不同的因此直接用 ( L k L_k Lk) 来累积会导致训练时某些小批量包含更多有效token的批次对最终梯度的贡献过大。
例如
假设有两个批次批次 1 包含 3 个有效token批次 2 包含 6 个有效token。如果我们直接将 ( L 1 L_1 L1) 和 ( L 2 L_2 L2) 累加批次 2 的损失会占据过多权重因为它包含更多有效token。这样会导致更新偏向于包含更多token的批次从而影响梯度的准确性。
2. 损失求和和平均的逻辑
不除以 ( m k m_k mk) 的问题
如果不除以有效token数 ( m k m_k mk)我们实际上在做损失累积时忽视了每个批次的“贡献”大小特别是当批次中有效token数目差异较大时。这样会使得损失计算不均衡。损失 ( L k L_k Lk) 并不直接等于每个token的损失之和它实际上是对有效token的损失的总和。所以批次大小的不同会导致损失的总和不成比例地反映出每个批次的有效学习。
除以 ( m k m_k mk) 的好处
通过除以每个小批量的有效token数 ( m k m_k mk)我们实际上是在对每个小批量的损失进行“加权平均”。这样做的目的就是保证每个token对最终梯度更新的贡献是公平的不会因为某个小批量有更多有效token而对最终梯度产生不合理的影响。
3. 为什么需要这样做
假设你使用了梯度累积目标是模拟大批量训练的效果。如果每个小批量中有效token的数量不同直接将损失加起来会导致更大的小批量含有更多有效token的批次对最终梯度更新产生更大的影响。为了避免这种偏差我们需要对每个小批量的损失进行标准化处理即通过除以 ( m k m_k mk) 来校准每个小批量的损失。
数学解释 公式 L final 1 G ∑ k 1 G L k m k L_{\text{final}} \frac{1}{G} \sum_{k1}^{G} \frac{L_k}{m_k} LfinalG1k1∑GmkLk 这里的 ( L k m k \frac{L_k}{m_k} mkLk ) 表示对每个小批量的损失进行标准化确保每个小批量在梯度累积中的贡献是按有效token数量来加权的而不是简单地将所有小批量的损失直接加起来。
举个简单的例子
假设我们有两个批次
批次 13 个有效token损失 ( L 1 1.5 L_1 1.5 L11.5)批次 26 个有效token损失 ( L 2 3.0 L_2 3.0 L23.0)
如果我们直接累加损失 L accumulated L 1 L 2 1.5 3.0 4.5 L_{\text{accumulated}} L_1 L_2 1.5 3.0 4.5 LaccumulatedL1L21.53.04.5
但是如果我们将损失除以每个批次的有效token数 L final 1 2 ( L 1 m 1 L 2 m 2 ) 1 2 ( 1.5 3 3.0 6 ) 1 2 ( 0.5 0.5 ) 0.5 L_{\text{final}} \frac{1}{2} \left( \frac{L_1}{m_1} \frac{L_2}{m_2} \right) \frac{1}{2} \left( \frac{1.5}{3} \frac{3.0}{6} \right) \frac{1}{2} \left( 0.5 0.5 \right) 0.5 Lfinal21(m1L1m2L2)21(31.563.0)21(0.50.5)0.5
这样做的目的是确保每个批次无论它的大小如何对损失的贡献是公平的。如果没有这种标准化损失较大的批次会导致最终的梯度更新不准确偏向于更多token的批次。
总结
除以 ( m k m_k mk) 的目的是为了对损失进行标准化处理确保每个小批量的有效token对梯度累积的贡献是均衡的特别是在批次之间有效token数量差异较大的情况下。这样可以避免某些批次因包含更多有效token而在梯度更新中占据不成比例的权重从而使训练更加稳定和准确。
后记
2025年1月18日15点45分于上海 在OpenAI o1大模型辅助下完成。
