refactor: update loss aggregation mode in GSPO functions to seq-mean-token-mean

Author: 0x404

verl/trainer/ppo/core_algos.py

@@ -893,17 +893,17 @@ def compute_policy_loss_gspo(
     log_prob: torch.Tensor,
     advantages: torch.Tensor,
     response_mask: torch.Tensor,
-    loss_agg_mode: str = "token-mean",
+    loss_agg_mode: str = "seq-mean-token-mean",
     config: Optional[DictConfig | ActorConfig] = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """
     Compute the clipped policy objective for vanilla GSPO (sequence-level optimization).
     See https://arxiv.org/pdf/2507.18071 for more details.
-    
+
     Implements equations 6-8 from the paper:
     J_GSPO(θ) = E[x~D, {y_i}_{i=1}^G ~ π_θold(·|x)] [1/G Σ_{i=1}^G min(s_i(θ)Â_i, clip(s_i(θ), 1-ε, 1+ε)Â_i)]
     where s_i(θ) = (π_θ(y_i|x)/π_θold(y_i|x))^(1/|y_i|)
-    
+
     Args:
         old_log_prob (torch.Tensor):
             Log-probabilities of actions under the old policy, shape (batch_size, response_length).
@@ -914,8 +914,8 @@ def compute_policy_loss_gspo(
         response_mask (torch.Tensor):
             Mask indicating which tokens to include in the loss, shape (batch_size, response_length).
         loss_agg_mode (str, optional):
-            Aggregation mode for `agg_loss`. Defaults to "token-mean".
-        config (Optional[DictConfig | ActorConfig]): 
+            Aggregation mode for `agg_loss`. For GSPO, it is recommended to use "seq-mean-token-mean".
+        config (Optional[DictConfig | ActorConfig]):
             Configuration parameters
     """
 
@@ -929,7 +929,9 @@ def compute_policy_loss_gspo(
     # Clamp negative_approx_kl for numerical stability
     negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)
 
-    # compute sequence-level importance ratio: si(θ) = (π_θ(yi|x)/π_θold(yi|x))^(1/|yi|) = exp [(1/|y_i|) * Σ_t log(π_θ(y_i,t|x,y_i,<t)/π_θold(y_i,t|x,y_i,<t))]
+    # compute sequence-level importance ratio:
+    # si(θ) = (π_θ(yi|x)/π_θold(yi|x))^(1/|yi|)
+    # = exp [(1/|y_i|) * Σ_t log(π_θ(y_i,t|x,y_i,<t)/π_θold(y_i,t|x,y_i,<t))]
     seq_lengths = torch.sum(response_mask, dim=-1).clamp(min=1)
     negative_approx_kl_seq = torch.sum(negative_approx_kl * response_mask, dim=-1) / seq_lengths
 
@@ -940,27 +942,26 @@ def compute_policy_loss_gspo(
     seq_ratio_expanded = seq_importance_ratio.unsqueeze(-1).expand_as(log_prob)
 
     pg_losses1 = -advantages * seq_ratio_expanded
-    pg_losses2 = -advantages * torch.clamp(
-        seq_ratio_expanded, 1 - clip_ratio_low, 1 + clip_ratio_high
-    )
+    pg_losses2 = -advantages * torch.clamp(seq_ratio_expanded, 1 - clip_ratio_low, 1 + clip_ratio_high)
     pg_losses = torch.maximum(pg_losses1, pg_losses2)
-    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
+    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode="seq-mean-token-mean")
 
     # For compatibility, return zero for pg_clipfrac_lower (not used in standard GSPO)
     pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)
     pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)
-    
+
     ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)
 
     return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower
 
+
 @register_policy_loss("gspo_token")
 def compute_policy_loss_gspo_token(
     old_log_prob: torch.Tensor,
     log_prob: torch.Tensor,
     advantages: torch.Tensor,
     response_mask: torch.Tensor,
-    loss_agg_mode: str = "token-mean",
+    loss_agg_mode: str = "seq-mean-token-mean",
     config: Optional[DictConfig | ActorConfig] = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """
@@ -978,7 +979,7 @@ def compute_policy_loss_gspo_token(
         response_mask (torch.Tensor):
             Mask indicating which tokens to include in the loss, shape (batch_size, response_length).
         loss_agg_mode (str, optional):
-            Aggregation mode for `agg_loss`. Defaults to "token-mean".
+            Aggregation mode for `agg_loss`. For GSPO, it is recommended to use "seq-mean-token-mean".
     """
 
     assert config is not None
@@ -991,28 +992,29 @@ def compute_policy_loss_gspo_token(
     # Clamp negative_approx_kl for numerical stability
     negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)
 
-    # compute sequence-level importance ratio: si(θ) = (π_θ(yi|x)/π_θold(yi|x))^(1/|yi|) = exp [(1/|y_i|) * Σ_t log(π_θ(y_i,t|x,y_i,<t)/π_θold(y_i,t|x,y_i,<t))]
+    # compute sequence-level importance ratio:
+    # si(θ) = (π_θ(yi|x)/π_θold(yi|x))^(1/|yi|) =
+    # exp [(1/|y_i|) * Σ_t log(π_θ(y_i,t|x,y_i,<t)/π_θold(y_i,t|x,y_i,<t))]
     seq_lengths = torch.sum(response_mask, dim=-1).clamp(min=1)
     negative_approx_kl_seq = torch.sum(negative_approx_kl * response_mask, dim=-1) / seq_lengths
 
-    # Combined ratio at token level: s_i,t(θ) = sg[s_i(θ)] · π_θ(y_i,t|x, y_i,<t) / sg[π_θ(y_i,t|x, y_i,<t)]
+    # Combined ratio at token level:
+    # s_i,t(θ) = sg[s_i(θ)] · π_θ(y_i,t|x, y_i,<t) / sg[π_θ(y_i,t|x, y_i,<t)]
     # In log space: log(s_i,t(θ)) = sg[log(s_i(θ))] + log_prob - sg[log_prob]
     log_seq_importance_ratio = negative_approx_kl_seq.detach().unsqueeze(-1) + log_prob - log_prob.detach()
 
     # finaly exp() to remove log
     seq_importance_ratio = torch.exp(log_seq_importance_ratio)
 
     pg_losses1 = -advantages * seq_importance_ratio
-    pg_losses2 = -advantages * torch.clamp(
-        seq_importance_ratio, 1 - clip_ratio_low, 1 + clip_ratio_high
-    )
+    pg_losses2 = -advantages * torch.clamp(seq_importance_ratio, 1 - clip_ratio_low, 1 + clip_ratio_high)
     pg_losses = torch.maximum(pg_losses1, pg_losses2)
-    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
+    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode="seq-mean-token-mean")
 
     # For compatibility, return zero for pg_clipfrac_lower (not used in standard GSPO)
     pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)
     pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)
-    
+
     ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)
 
     return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower