VeRL 框架 RL 微调大语言模型，token-level loss 聚合方式详解

考虑这个场景在 RL 微调大模型的过程中我们让模型生成了一批回答每个回答是一串文字然后我们给每个回答里的每个字token算了一个“损失值”可以理解为这个字写得“好不好”的评分损失越大说明越不好。同时我们还有一个“掩码”mask标记哪些字是有效的比如忽略填充的无意义字。此时为了最小化损失值我们需要把这堆零散的字级别的损失汇总成一个总的损失值然后用这个总损失去更新模型。这个汇总方式可以有很多种就像我们计算平均分可以用“所有字的总分除以字数”也可以用“每句话的平均分再平均”等等不同的汇总方式会影响模型学习的方向。VeRL 将这一部分封装为agg_loss函数以下为该函数的具体实现# 位于 verl/trainer/ppo/core_algos.pydefagg_loss(loss_mat:torch.Tensor,loss_mask:torch.Tensor,loss_agg_mode:str,dp_size:int1,batch_num_tokens:Optional[int]None,global_batch_size:Optional[int]None,loss_scale_factor:Optional[int]None,): Aggregate the loss across global batch to ensure the loss is invariant to fsdp/megatron parallelism. NOTE: The returned loss has different behaviors for different backend: - FSDP: the loss is directly used for backward. - Megatron: the loss should be scaled by num_microbatches and cp_size for pp schedule. Args: loss_mat: micro batch loss matrix, (bs, response_length) loss_mask: micro batch loss mask, (bs, response_length) loss_agg_mode: method to aggregate the loss matrix into a scalar dp_size: data parallel size batch_num_tokens: number of valid tokens in global batch global_batch_size: global batch size loss_scale_factor: scale factor for seq-mean-token-sum-norm mode. If None, uses loss_mask.shape[-1]. Set this to a constant value to ensure consistent normalization throughout training. Returns: loss: a scalar torch.Tensor aggregated loss ifloss_agg_modetoken-mean:ifbatch_num_tokensisNone:batch_num_tokensloss_mask.sum()lossverl_F.masked_sum(loss_mat,loss_mask)/batch_num_tokens*dp_sizeelifloss_agg_modeseq-mean-token-sum:seq_lossestorch.sum(loss_mat*loss_mask,dim-1)# token-sumseq_mask(torch.sum(loss_mask,dim-1)0).float()# exclude fully masked sequencesifglobal_batch_sizeisNone:global_batch_sizeseq_mask.sum()lossverl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size# seq-meanelifloss_agg_modeseq-mean-token-mean:seq_masktorch.sum(loss_mask,dim-1)# per-sequence token countseq_lossestorch.sum(loss_mat*loss_mask,dim-1)/(seq_mask1e-8)# token-meanseq_mask(seq_mask0).float()# exclude fully masked sequencesifglobal_batch_sizeisNone:global_batch_sizeseq_mask.sum()lossverl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size# seq-meanelifloss_agg_modeseq-mean-token-sum-norm:seq_lossestorch.sum(loss_mat*loss_mask,dim-1)ifloss_scale_factorisNone:loss_scale_factorloss_mask.shape[-1]losstorch.sum(seq_losses)/loss_scale_factorelse:raiseValueError(fInvalid loss_agg_mode:{loss_agg_mode})returnloss1 输入参数回顾loss_mat: 形状(bs, response_length)每个 token 的损失值例如策略梯度损失或 KL 散度。loss_mask: 相同形状标记哪些 token 是有效的1 表示有效0 表示 padding 或忽略。loss_agg_mode: 聚合方式可选值有token-mean,seq-mean-token-sum,seq-mean-token-mean,seq-mean-token-sum-norm。dp_size: 数据并行度分布式训练时使用的 GPU 数量用于将局部损失缩放为全局等价损失。batch_num_tokens: 全局批次中有效 token 的总数跨所有 GPU。如果未提供会在token-mean模式中自动计算局部和。global_batch_size: 全局批次中有效序列的数量。如果未提供会在某些模式中自动使用局部有效序列数。loss_scale_factor: 仅用于seq-mean-token-sum-norm模式作为分母的固定缩放因子。如果未提供默认使用loss_mask.shape[-1]即固定的响应长度。2 聚合模式详解2.1token-mean所有有效 token 的平均损失lossverl_F.masked_sum(loss_mat,loss_mask)/batch_num_tokens*dp_size直观理解把每个回答里所有有效字的损失加起来然后除以所有有效字的总数。最后得到一个“平均每个字有多差”的值。计算先计算所有有效 token 的损失总和masked_sum相当于(loss_mat * loss_mask).sum()然后除以有效 token 总数batch_num_tokens得到每个 token 的平均损失。最后乘以dp_size使结果相当于全局 token 平均损失在分布式训练中每个 GPU 只计算局部平均乘以dp_size是为了统一到全局视角。作用这是最直观的平均方式每个有效 token 对损失的贡献权重相等。verl 的 GRPO 默认使用此模式因为它简单稳定尤其适合长文本生成任务。2.2seq-mean-token-sum先对每个序列内 token 损失求和再对所有序列求平均seq_lossestorch.sum(loss_mat*loss_mask,dim-1)# 每个序列的总损失seq_mask(torch.sum(loss_mask,dim-1)0).float()# 哪些序列至少有一个有效 tokenlossverl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size直观理解第一步对每个回答把里面所有有效字的损失加起来得到“这个回答的总损失”。第二步把这些回答的总损失加起来再除以回答的数量得到“平均每个回答的总损失”。把每个回答当成一个整体先看这个回答整句话写得有多差总损失然后看看所有回答的平均整句损失。这样一个很长的回答虽然字多但它的总损失可能会很大所以它对最终结果的影响也比较大。计算对每个序列将有效 token 的损失相加得到序列级总损失seq_losses。生成序列掩码seq_mask有效序列为 1全填充序列为 0。对有效序列的总损失求和除以全局序列数global_batch_size得到序列平均总损失。最后乘以dp_size。作用这种聚合方式将每个序列视为一个样本不论序列长度如何每个序列对损失的贡献权重相同。因此长序列的 token 级损失会被隐式地“压缩”成序列级损失可能更适合某些任务。2.3seq-mean-token-mean先对每个序列内 token 损失求平均再对所有序列求平均seq_masktorch.sum(loss_mask,dim-1)# 每个序列的有效 token 数seq_lossestorch.sum(loss_mat*loss_mask,dim-1)/(seq_mask1e-8)# 每个序列的平均损失seq_mask(seq_mask0).float()# 有效序列掩码lossverl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size直观理解第一步对每个回答把里面所有有效字的损失加起来除以这个回答的有效字数得到“这个回答每个字的平均损失”。第二步把这些回答的平均字损失加起来除以回答的数量得到“平均每个回答的平均字损失”。先看每个回答里平均每个字写得有多差然后把这些平均值再平均。这样不管一个回答是长是短它对最终结果的贡献只取决于它内部字的平均质量而不是总损失。这样可以避免长句子因为字多而主导训练。计算计算每个序列的有效 token 数seq_mask即序列长度。对每个序列计算有效 token 损失的平均值seq_losses。对有效序列的平均损失求和除以全局序列数得到序列平均损失的平均。作用这种模式是原始 GRPO 论文中使用的聚合方式。它将每个序列的损失归一化到相同尺度每个序列的平均损失然后再平均。这样做的动机是让每个序列对梯度的贡献相等避免长序列因 token 多而主导梯度。但在长思维链场景中每个序列的平均损失可能波动较大导致训练不稳定因此 verl 默认未采用此模式。2.4seq-mean-token-sum-norm所有序列的总损失之和除以一个固定常数seq_lossestorch.sum(loss_mat*loss_mask,dim-1)# 每个序列的总损失ifloss_scale_factorisNone:loss_scale_factorloss_mask.shape[-1]# 默认使用固定的响应长度losstorch.sum(seq_losses)/loss_scale_factor直观理解第一步把每个回答的总损失算出来和模式2一样。第二步把所有回答的总损失加起来然后除以一个固定的数字比如你预设的最大回答长度或者直接使用固定的响应长度。这个固定数字在训练过程中不变。通俗理解把所有回答的“总损失”加起来然后除以一个固定的大数。这个固定的大数不随当前批次回答的长度变化所以最终损失的大小不受句子长短的影响。这样可以避免模型为了降低损失而故意把句子写得很长因为长句子总损失更大除以固定数后可能仍然大但不会因为短句子分母小而被放大。计算先计算每个序列的总损失然后对所有序列的总损失求和最后除以一个固定的缩放因子loss_scale_factor可以是一个常数如最大响应长度。注意这里没有乘以dp_size也没有使用global_batch_size。作用这是 DrGRPO 中引入的模式旨在消除长度偏差。在标准 GRPO 中组内奖励归一化可能导致模型偏向生成长响应尤其是错误答案时而此模式通过使用固定的分母例如预定义的最大响应长度来保证损失大小不受序列实际长度影响从而缓解优化偏差。如果loss_scale_factor设置为一个全局常数如 2048那么不论当前批次序列长度如何损失的计算都会除以这个常数使得梯度更新步长与长度解耦。3 为什么需要这么多聚合模式在强化学习训练 LLM 时损失的计算方式会影响梯度更新的方向和幅度。不同的聚合策略对应着不同的假设token-mean每个 token 同等重要适合希望模型在 token 级别精细调整的场景。seq-mean-token-sum每个序列同等重要但保留序列内 token 损失的原始和适合希望模型优化序列整体回报的任务。seq-mean-token-mean每个序列同等重要且每个序列的损失被归一化到 token 平均适合希望平衡长短序列贡献的场景。seq-mean-token-sum-norm消除长度对损失的影响专门用于纠正 GRPO 可能产生的长度偏差。4 分布式训练的注意事项参数dp_size用于将局部损失单个 GPU 上的聚合结果放大到全局等效值。因为在分布式训练中每个 GPU 只处理一部分数据计算出的损失通常是局部平均但反向传播时需要保证梯度与全局 batch 大小匹配。乘以dp_size是一种常见的做法相当于将局部损失视为全局损失的一部分然后通过 all-reduce 平均梯度时能得到正确结果。batch_num_tokens和global_batch_size通常由框架在全局范围内计算并传入确保聚合时使用的是全局统计数据而不是局部统计。