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

发布时间：2026/11/5 1:30:45
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:int=1,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_mode=="token-mean":ifbatch_num_tokensisNone:batch_num_tokens=loss_mask.sum()loss=verl_F.masked_sum(loss_mat,loss_mask)/batch_num_tokens*dp_sizeelifloss_agg_mode=="seq-mean-token-sum":seq_losses=torch.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_size=seq_mask.sum()loss=verl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size# seq-meanelifloss_agg_mode=="seq-mean-token-mean":seq_mask=torch.sum(loss_mask,dim=-1)# per-sequence token countseq_losses=torch.sum(loss_mat*loss_mask,dim=-1)/(seq_mask+1e-8)# token-meanseq_mask=(seq_mask>0).float()# exclude fully masked sequencesifglobal_batch_sizeisNone:global_batch_size=seq_mask.sum()loss=verl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size# seq-meanelifloss_agg_mode=="seq-mean-token-sum-norm":seq_losses=torch.sum(loss_mat*loss_mask,dim=-1)ifloss_scale_factorisNone:loss_scale_factor=loss_mask.shape[-1]loss=torch.sum(seq_losses)/loss_scale_factorelse:raiseValueError(f"Invalid loss_agg_mode:{loss_agg_mode}")returnloss

1 输入参数回顾

  • 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.1"token-mean"：所有有效 token 的平均损失

loss=verl_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.2"seq-mean-token-sum"：先对每个序列内 token 损失求和，再对所有序列求平均

seq_losses=torch.sum(loss_mat*loss_mask,dim=-1)# 每个序列的总损失seq_mask=(torch.sum(loss_mask,dim=-1)>0).float()# 哪些序列至少有一个有效 tokenloss=verl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size
  • 直观理解：
    • 第一步：对每个回答，把里面所有有效字的损失加起来，得到“这个回答的总损失”。
    • 第二步：把这些回答的总损失加起来，再除以回答的数量，得到“平均每个回答的总损失”。
    • 把每个回答当成一个整体，先看这个回答整句话写得有多差（总损失），然后看看所有回答的平均整句损失。这样，一个很长的回答虽然字多，但它的总损失可能会很大，所以它对最终结果的影响也比较大。
  • 计算：
    1. 对每个序列，将有效 token 的损失相加，得到序列级总损失seq_losses。
    2. 生成序列掩码seq_mask（有效序列为 1，全填充序列为 0）。
    3. 对有效序列的总损失求和，除以全局序列数global_batch_size，得到序列平均总损失。最后乘以dp_size。
  • 作用：这种聚合方式将每个序列视为一个样本，不论序列长度如何，每个序列对损失的贡献权重相同。因此，长序列的 token 级损失会被隐式地“压缩”成序列级损失，可能更适合某些任务。

2.3"seq-mean-token-mean"：先对每个序列内 token 损失求平均，再对所有序列求平均

seq_mask=torch.sum(loss_mask,dim=-1)# 每个序列的有效 token 数seq_losses=torch.sum(loss_mat*loss_mask,dim=-1)/(seq_mask+1e-8)# 每个序列的平均损失seq_mask=(seq_mask>0).float()# 有效序列掩码loss=verl_F.masked_sum(seq_losses,seq_mask)/global_batch_size*dp_size
  • 直观理解：
    • 第一步：对每个回答，把里面所有有效字的损失加起来，除以这个回答的有效字数，得到“这个回答每个字的平均损失”。
    • 第二步：把这些回答的平均字损失加起来，除以回答的数量，得到“平均每个回答的平均字损失”。
    • 先看每个回答里平均每个字写得有多差，然后把这些平均值再平均。这样，不管一个回答是长是短，它对最终结果的贡献只取决于它内部字的平均质量，而不是总损失。这样可以避免长句子因为字多而主导训练。
  • 计算：
    1. 计算每个序列的有效 token 数seq_mask（即序列长度）。
    2. 对每个序列，计算有效 token 损失的平均值seq_losses。
    3. 对有效序列的平均损失求和，除以全局序列数，得到序列平均损失的平均。
  • 作用：这种模式是原始 GRPO 论文中使用的聚合方式。它将每个序列的损失归一化到相同尺度（每个序列的平均损失），然后再平均。这样做的动机是让每个序列对梯度的贡献相等，避免长序列因 token 多而主导梯度。但在长思维链场景中，每个序列的平均损失可能波动较大，导致训练不稳定，因此 verl 默认未采用此模式。

2.4"seq-mean-token-sum-norm"：所有序列的总损失之和除以一个固定常数

seq_losses=torch.sum(loss_mat*loss_mask,dim=-1)# 每个序列的总损失ifloss_scale_factorisNone:loss_scale_factor=loss_mask.shape[-1]# 默认使用固定的响应长度loss=torch.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通常由框架在全局范围内计算并传入，确保聚合时使用的是全局统计数据，而不是局部统计。