feat: Fixed metric calculation and made all grpo metrics token-level (#373)

Signed-off-by: Sahil Jain <sahilj@nvidia.com>

Author: SahilJain314

nemo_rl/algorithms/dpo.py

@@ -302,7 +302,7 @@ def validate(
 
             else:
                 for k, v in val_results["all_mb_metrics"].items():
-                    if k in {"lr", "normalization_factor"}:
+                    if k in {"lr", "global_valid_seqs", "global_valid_toks"}:
                         val_metrics[k] += np.mean(v).item()
                     else:
                         val_metrics[k] += np.sum(v).item()
@@ -491,7 +491,7 @@ def dpo_train(
             }
             metrics.update(train_results["all_mb_metrics"])
             for k, v in metrics.items():
-                if k in {"lr", "normalization_factor"}:
+                if k in {"lr", "global_valid_seqs", "global_valid_toks"}:
                     metrics[k] = np.mean(v).item()
                 else:
                     metrics[k] = np.sum(v).item()

nemo_rl/algorithms/grpo.py

@@ -570,7 +570,7 @@ def grpo_train(
         }
         metrics.update(train_results["all_mb_metrics"])
         for k, v in metrics.items():
-            if k in {"lr", "reward", "normalization_factor"}:
+            if k in {"lr", "reward", "global_valid_seqs", "global_valid_toks"}:
                 metrics[k] = np.mean(v).item()
             else:
                 metrics[k] = np.sum(v).item()

nemo_rl/algorithms/interfaces.py

@@ -27,7 +27,11 @@ class LossFunction(Protocol):
     """
 
     def __call__(
-        self, next_token_logits: torch.Tensor, data: BatchedDataDict
+        self,
+        next_token_logits: torch.Tensor,
+        data: BatchedDataDict,
+        global_valid_seqs: torch.Tensor,
+        global_valid_toks: torch.Tensor,
     ) -> Tuple[torch.Tensor, Dict[str, Any]]:
         """Compute loss and metrics from logprobs and other data.
 
@@ -40,6 +44,14 @@ def __call__(
             data: Dictionary containing all relevant data for loss computation
                   such as rewards, values, actions, advantages, masks, and other
                   algorithm-specific information needed for the particular loss calculation.
+            global_valid_seqs: torch.Tensor
+                this tensor should contain the number of valid sequences in the microbatch.
+                It's used for global normalization for losses/metrics that are computed at the sequence level
+                and needs to be aggregated across all microbatches.
+            global_valid_toks: torch.Tensor
+                This tensor should contain the number of valid tokens in the microbatch.
+                It's used for global normalization for losses/metrics that are computed at the token level
+                and needs to be aggregated across all microbatches.
 
         Returns:
             tuple: (loss, metrics)

nemo_rl/algorithms/loss_functions.py

@@ -113,7 +113,8 @@ def __call__(
         self,
         next_token_logits: torch.Tensor,
         data: BatchedDataDict[ClippedPGLossDataDict],
-        total_valid_tokens_or_seqs: torch.Tensor,
+        global_valid_seqs: torch.Tensor,
+        global_valid_toks: torch.Tensor,
     ) -> Tuple[torch.Tensor, dict]:
         """Clipped Policy Gradient RL loss function."""
         token_mask = data["token_mask"][:, 1:]
@@ -128,21 +129,12 @@ def __call__(
         # token_mult_prob_error
         # See more details and other metrics in docs/guides/grpo.md#metrics
         lp_error = torch.abs(generation_logprobs - prev_logprobs)  # noqa: F841  (precommit ignore for now)
-        if self.loss_type == LossType.TOKEN_LEVEL:
-            # average over all tokens in the microbatch
-            mult_prob_error = masked_mean(
-                torch.exp(lp_error * mask),
-                mask,
-                global_normalization_factor=total_valid_tokens_or_seqs,
-            ).item()
-        else:
-            # first average over tokens per sample, then average over samples
-            # multiply lp_error by mask before exp to prevent inf for large lp_error values on masked tokens
-            mult_prob_error = masked_mean(
-                masked_mean(torch.exp(lp_error) * token_mask, token_mask, dim=-1),
-                sample_mask,
-                global_normalization_factor=total_valid_tokens_or_seqs,
-            ).item()
+        # average over all tokens in the microbatch
+        mult_prob_error = masked_mean(
+            torch.exp(lp_error * mask),
+            mask,
+            global_normalization_factor=global_valid_toks,
+        ).item()
 
         next_token_logits = next_token_logits.to(torch.float32)
 
@@ -184,13 +176,13 @@ def __call__(
             )
             if self.loss_type == LossType.TOKEN_LEVEL:
                 kl = masked_mean(
-                    kl, mask, global_normalization_factor=total_valid_tokens_or_seqs
+                    kl, mask, global_normalization_factor=global_valid_toks
                 )
             else:
                 kl = masked_mean(
                     masked_mean(kl, token_mask, dim=-1),
                     sample_mask,
-                    global_normalization_factor=total_valid_tokens_or_seqs,
+                    global_normalization_factor=global_valid_seqs,
                 )
         else:
             kl = 0
@@ -235,7 +227,7 @@ def __call__(
             actor_loss = masked_mean(
                 importance_weights_to_use * clip_loss,
                 mask,
-                global_normalization_factor=total_valid_tokens_or_seqs,
+                global_normalization_factor=global_valid_toks,
             )
         else:
             actor_loss = masked_mean(
@@ -245,41 +237,41 @@ def __call__(
                     dim=-1,
                 ),
                 sample_mask,
-                global_normalization_factor=total_valid_tokens_or_seqs,
+                global_normalization_factor=global_valid_seqs,
             )
 
+        # See: docs/guides/grpo.md#sampling-importance-ratio
+        sample_importance_ratio = masked_mean(
+            actor_importance_weights,
+            mask,
+            global_normalization_factor=global_valid_toks,
+        )
+
         # Approximating entropy as E_{s ~ \pi_{gen}(s)}[-(\pi_{curr}/\pi_{gen})log(\pi_{curr}(s))]
         # See more details and other metrics in docs/guides/grpo.md#metrics
         with torch.no_grad():
             seq_entropy_approx = -masked_mean(
-                torch.exp(curr_logprobs - generation_logprobs) * curr_logprobs, mask
+                torch.exp(curr_logprobs - generation_logprobs) * curr_logprobs,
+                mask,
+                global_normalization_factor=global_valid_toks,
             )
 
         loss = actor_loss + kl
         with torch.no_grad():
-            if self.loss_type == LossType.TOKEN_LEVEL:
-                probs_ratio = masked_mean(
-                    ratios.detach(),
-                    mask,
-                    global_normalization_factor=total_valid_tokens_or_seqs,
-                ).item()
-                probs_ratio_clamped = masked_mean(
-                    ratios_clamped.detach(),
-                    mask,
-                    global_normalization_factor=total_valid_tokens_or_seqs,
-                ).item()
-            else:
-                probs_ratio = masked_mean(
-                    masked_mean(ratios.detach(), token_mask, dim=-1),
-                    sample_mask,
-                    global_normalization_factor=total_valid_tokens_or_seqs,
-                ).item()
-                probs_ratio_clamped = masked_mean(
-                    masked_mean(ratios_clamped.detach(), token_mask, dim=-1),
-                    sample_mask,
-                    global_normalization_factor=total_valid_tokens_or_seqs,
-                ).item()
+            probs_ratio = masked_mean(
+                ratios.detach(),
+                mask,
+                global_normalization_factor=global_valid_toks,
+            ).item()
+            probs_ratio_clamped = masked_mean(
+                ratios_clamped.detach(),
+                mask,
+                global_normalization_factor=global_valid_toks,
+            ).item()
 
+        # If you provided a global_valid_{seqs/toks}, all metrics here are globally normalized
+        # by either sequence or token count, depending on particular metric.
+        # To get the true metric, you'll need to sum over the microbatch.
         return (
             loss,
             {
@@ -288,9 +280,7 @@ def __call__(
                 "probs_ratio_clamped": probs_ratio_clamped,
                 "kl_penalty": kl.item() / self.reference_policy_kl_penalty if kl else 0,
                 "token_mult_prob_error": mult_prob_error,
-                "sampling_importance_ratio": masked_mean(
-                    actor_importance_weights, mask
-                ).item(),
+                "sampling_importance_ratio": sample_importance_ratio.item(),
                 "num_valid_samples": sample_mask.sum().item(),
                 "approx_entropy": seq_entropy_approx.item(),
             },
@@ -306,7 +296,8 @@ def __call__(
         self,
         next_token_logits: torch.Tensor,
         data: BatchedDataDict,
-        total_valid_tokens_or_seqs: torch.Tensor,
+        global_valid_seqs: torch.Tensor | None,
+        global_valid_toks: torch.Tensor,
         dpo_loss: bool = False,
         dpo_average_log_probs: bool = False,
     ) -> Tuple[torch.Tensor, dict]:
@@ -346,7 +337,7 @@ def __call__(
             loss = -masked_mean(
                 token_logprobs,
                 mask,
-                global_normalization_factor=total_valid_tokens_or_seqs,
+                global_normalization_factor=global_valid_toks,
             )
 
         return loss, {
@@ -446,7 +437,7 @@ def preference_loss(
         self,
         next_token_logits: torch.Tensor,
         data: BatchedDataDict[DPOLossDataDict],
-        total_valid_tokens_or_seqs: torch.Tensor,
+        global_valid_seqs: torch.Tensor,
     ) -> torch.Tensor:
         ## TODO(@ashors): there's some duplicate code here with the NLLLoss function. We should refactor
         token_mask = data["token_mask"][:, 1:]
@@ -490,53 +481,58 @@ def preference_loss(
             masked_mean(
                 per_sample_loss,
                 sample_mask[::2],
-                global_normalization_factor=total_valid_tokens_or_seqs / 2,
+                global_normalization_factor=global_valid_seqs / 2,
             ),
             masked_mean(
                 rewards_chosen > rewards_rejected,
                 sample_mask[::2],
-                global_normalization_factor=total_valid_tokens_or_seqs / 2,
+                global_normalization_factor=global_valid_seqs / 2,
             ),
             masked_mean(
                 rewards_chosen,
                 sample_mask[::2],
-                global_normalization_factor=total_valid_tokens_or_seqs / 2,
+                global_normalization_factor=global_valid_seqs / 2,
             ),
             masked_mean(
                 rewards_rejected,
                 sample_mask[1::2],
-                global_normalization_factor=total_valid_tokens_or_seqs / 2,
+                global_normalization_factor=global_valid_seqs / 2,
             ),
         )
 
     def __call__(
         self,
         next_token_logits: torch.Tensor,
         data: BatchedDataDict[DPOLossDataDict],
-        total_valid_tokens_or_seqs: torch.Tensor,
+        global_valid_seqs: torch.Tensor,
+        global_valid_toks: torch.Tensor | None,
     ) -> Tuple[torch.Tensor, dict]:
         sft_loss_chosen = torch.tensor(0.0)
         if self.sft_loss_weight > 0:
+            assert global_valid_toks is not None, (
+                "global_valid_toks must be provided for SFT loss"
+            )
             sft_loss, _ = self.sft_loss(
                 next_token_logits,
                 data,
-                total_valid_tokens_or_seqs=total_valid_tokens_or_seqs,  ## unused because sft loss returned is at the sample level
+                global_valid_seqs=global_valid_seqs,
+                global_valid_toks=global_valid_toks,  ## unused because sft loss returned is at the sample level
                 dpo_loss=True,
                 dpo_average_log_probs=self.sft_average_log_probs,
             )
             sft_loss_chosen, sft_loss_rejected = self.split_output_tensor(sft_loss)
             sft_loss_chosen = masked_mean(
                 sft_loss_chosen,
                 data["sample_mask"][::2],
-                global_normalization_factor=total_valid_tokens_or_seqs / 2,
+                global_normalization_factor=global_valid_seqs / 2,
             )
 
         (
             preference_loss,
             accuracy,
             rewards_chosen_mean,
             rewards_rejected_mean,
-        ) = self.preference_loss(next_token_logits, data, total_valid_tokens_or_seqs)
+        ) = self.preference_loss(next_token_logits, data, global_valid_seqs)
 
         dpo_loss = (
             self.sft_loss_weight * sft_loss_chosen

nemo_rl/algorithms/sft.py

@@ -490,7 +490,7 @@ def sft_train(
             }
             metrics.update(train_results["all_mb_metrics"])
             for k, v in metrics.items():
-                if k in {"lr", "normalization_factor"}:
+                if k in {"lr", "global_valid_seqs", "global_valid_toks"}:
                     metrics[k] = np.mean(v).item()
                 else:
                     metrics[k] = np.sum(v).item()

nemo_rl/algorithms/utils.py

@@ -123,7 +123,7 @@ def masked_mean(
     values,
     mask,
     dim: Optional[int] = None,
-    global_normalization_factor: Optional[torch.Tensor] = None,
+    global_normalization_factor: Optional[torch.Tensor | float] = None,
 ):
     """Computes the mean of a microbatch, using a global statistic as the normalization factor."""
     normalization_factor = (

nemo_rl/models/policy/dtensor_policy_worker.py

@@ -315,26 +315,29 @@ def train(
                     "sample_mask must be present in the data!"
                 )
                 ## get the normalization factor for the loss
-                if loss_fn.loss_type == LossType.TOKEN_LEVEL:
-                    assert "token_mask" in global_batch, (
-                        "token_mask must be present in the data when using token-level loss"
+                local_valid_seqs = torch.sum(global_batch["sample_mask"])
+
+                if not "token_mask" in global_batch:
+                    local_valid_toks = (
+                        local_valid_seqs * global_batch["input_ids"].shape[1]
                     )
-                    ## get number of tokens in the global batch
-                    total_valid_tokens_or_seqs = torch.sum(
+                else:
+                    local_valid_toks = torch.sum(
                         global_batch["token_mask"][:, 1:]
                         * global_batch["sample_mask"].unsqueeze(-1)
                     )
-                    torch.distributed.all_reduce(
-                        total_valid_tokens_or_seqs, group=self.dp_mesh.get_group()
-                    )
-                elif loss_fn.loss_type == LossType.SEQUENCE_LEVEL:
-                    ## get number of valid samples in the global batch
-                    total_valid_tokens_or_seqs = torch.sum(global_batch["sample_mask"])
-                    torch.distributed.all_reduce(
-                        total_valid_tokens_or_seqs, group=self.dp_mesh.get_group()
+
+                to_reduce = torch.tensor([local_valid_seqs, local_valid_toks]).cuda()
+                torch.distributed.all_reduce(to_reduce, group=self.dp_mesh.get_group())
+                global_valid_seqs, global_valid_toks = to_reduce[0], to_reduce[1]
+
+                if (
+                    hasattr(loss_fn, "loss_type")
+                    and loss_fn.loss_type == LossType.TOKEN_LEVEL
+                ):
+                    assert "token_mask" in global_batch, (
+                        "token_mask must be present in the data when using token-level loss"
                     )
-                else:
-                    raise ValueError(f"Unknown loss type: {loss_fn.loss_type}")
 
                 self.optimizer.zero_grad()
                 mb_losses = []
@@ -386,16 +389,17 @@ def train(
                     if "generation" in self.cfg and self.cfg["generation"] is not None:
                         logits.div_(self.cfg["generation"]["temperature"])
 
-                    loss, loss_metrics = loss_fn(logits, mb, total_valid_tokens_or_seqs)
+                    loss, loss_metrics = loss_fn(
+                        logits, mb, global_valid_seqs, global_valid_toks
+                    )
                     ## scale by the number of global batches so we get the correct
                     ## value when summing metrics across all microbatches
                     for k in loss_metrics.keys():
                         loss_metrics[k] /= num_global_batches
                     num_valid_samples = loss_metrics["num_valid_samples"]
                     loss_metrics["lr"] = self.optimizer.param_groups[0]["lr"]
-                    loss_metrics["normalization_factor"] = (
-                        total_valid_tokens_or_seqs.cpu()
-                    )
+                    loss_metrics["global_valid_seqs"] = global_valid_seqs.item()
+                    loss_metrics["global_valid_toks"] = global_valid_toks.item()
 
                     # Backward pass
                     if not eval_mode: