fix(loss): exclude -inf logprobs from the actor reduction, not the KL

nemo_rl/algorithms/grpo.py

@@ -59,6 +59,7 @@
 from nemo_rl.algorithms.utils import (
     WALL_CLOCK_EFFICIENCY_CATEGORIES,
     calculate_baseline_and_std_per_prompt,
+    finalize_actor_token_metrics,
     calculate_trivial_reward_distributions,
     get_gdpo_reward_component_keys,
     log_generation_metrics,
@@ -3838,6 +3839,7 @@ def _grpo_train_impl(
                         metrics[k] = np.sum(v).item()
                     else:
                         print(f"Skipping aggregation for {k} ({type(v)})")
+                finalize_actor_token_metrics(metrics)
 
                 metrics.update(rollout_metrics)
                 metrics["generation_logger_metrics"] = generation_logger_metrics
@@ -5739,6 +5741,7 @@ def _flush_collector_telemetry() -> None:
                         metrics[k] = np.mean(v).item()
                     else:
                         metrics[k] = np.sum(v).item()
+                finalize_actor_token_metrics(metrics)
                 metrics.update(rollout_metrics)
                 if generation_logger_metrics is not None:
                     metrics["generation_logger_metrics"] = generation_logger_metrics

nemo_rl/algorithms/grpo_sync.py

@@ -70,6 +70,7 @@
 from nemo_rl.algorithms.reward_functions import apply_reward_shaping
 from nemo_rl.algorithms.utils import (
     calculate_baseline_and_std_per_prompt,
+    finalize_actor_token_metrics,
     calculate_trivial_reward_distributions,
     get_gdpo_reward_component_keys,
     log_generation_metrics,
@@ -1218,6 +1219,7 @@ def grpo_train_sync(
                         metrics[k] = np.sum(v).item()
                     else:
                         print(f"Skipping aggregation for {k} ({type(v)})")
+                finalize_actor_token_metrics(metrics)
 
                 metrics.update(rollout_metrics)
                 metrics["generation_logger_metrics"] = generation_logger_metrics

nemo_rl/algorithms/loss/interfaces.py

@@ -31,7 +31,7 @@ class MetricNormalizer(enum.Enum):
     Losses reduce most metrics with ``masked_mean(...,
     global_normalization_factor=...)`` where the factor is the global valid
     *token* count, the global valid *sequence* count, or absent entirely (raw
-    counts, per-microbatch means, extrema). Split-API trainers run each
+    counts/numerator fragments, per-microbatch means, extrema). Split-API trainers run each
     microbatch with placeholder ``global_valid_*=1`` (collecting raw sums) and
     rescale once per optimizer step — to do that they must know, per metric,
     which denominator applies.

nemo_rl/algorithms/loss/loss_functions.py

@@ -28,7 +28,12 @@
     LossType,
     MetricNormalizer,
 )
-from nemo_rl.algorithms.utils import calculate_kl, masked_mean
+from nemo_rl.algorithms.utils import (
+    ACTOR_TOKEN_COUNT_METRIC,
+    ACTOR_TOKEN_MEAN_METRICS,
+    calculate_kl,
+    masked_mean,
+)
 from nemo_rl.algorithms.x_token.loss_utils import (
     LocalizedAlignment,
     build_exact_token_map,
@@ -179,6 +184,9 @@ class ClippedPGLossDataDict(TypedDict):
     reference_policy_logprobs: torch.Tensor
     token_mask: torch.Tensor
     sample_mask: torch.Tensor
+    # Set by ``prepare_loss_input`` when top-k/top-p filtering leaves -inf
+    # logprobs; shape [B, S-1], already in the next-token frame.
+    curr_logprobs_keep_mask: NotRequired[torch.Tensor]
     __extra__: Any
 
 
@@ -353,14 +361,19 @@ def __init__(
             # Normalized like the gradient (loss_type-dependent).
             "loss": grad_normalizer,
             "kl_penalty": grad_normalizer,
-            # Token-normalized diagnostics, independent of loss_type.
-            "probs_ratio": MetricNormalizer.TOKENS,
-            "probs_ratio_clamped": MetricNormalizer.TOKENS,
+            # Actor-mask means are emitted as raw numerator fragments and
+            # normalized once their matching retained-token counts have been
+            # summed across microbatches and ranks.
+            **{
+                metric_name: MetricNormalizer.NONE
+                for metric_name in ACTOR_TOKEN_MEAN_METRICS
+            },
+            ACTOR_TOKEN_COUNT_METRIC: MetricNormalizer.NONE,
+            # Full-token diagnostics, independent of loss_type.
             "token_mult_prob_error": MetricNormalizer.TOKENS,
             "gen_kl_error": MetricNormalizer.TOKENS,
             "policy_kl_error": MetricNormalizer.TOKENS,
             "js_divergence_error": MetricNormalizer.TOKENS,
-            "approx_entropy": MetricNormalizer.TOKENS,
             # Keyed on sequence_level_importance_ratios, NOT loss_type.
             "sampling_importance_ratio": (
                 MetricNormalizer.SEQUENCES
@@ -472,6 +485,29 @@ def __call__(
 
         mask = token_mask * sample_mask.unsqueeze(-1)
 
+        # Positions whose ``curr_logprobs`` were filtered out carry a
+        # substituted value, so every quantity derived from ``curr_logprobs``
+        # reduces over ``actor_mask`` instead.
+        #
+        # The reference-policy KL keeps the full mask and is genuinely clean
+        # there: it reads ``curr_logprobs_unfiltered``, which is finite at
+        # exactly those positions.
+        #
+        # The mismatch diagnostics keep the full mask for a weaker reason, and
+        # it is worth being precise about it. They read ``prev_logprobs``,
+        # which carries its OWN substitutions -- the same helper runs on the
+        # logprob-inference pass (see ``automodel/train.py``,
+        # ``megatron/train.py`` and ``dtensor_policy_worker.py``). So they are
+        # not clean; they simply cannot be narrowed here, because the keep mask
+        # published this step identifies the positions *this* forward filtered,
+        # not the ones the earlier forward did, and narrowing with it would
+        # drop the wrong ones. ``force_on_policy_ratio`` is the single case
+        # where the two sets coincide, since prev is then an alias of curr.
+        keep_mask = data.get("curr_logprobs_keep_mask")
+        actor_token_mask = token_mask if keep_mask is None else token_mask * keep_mask
+        actor_mask = actor_token_mask * sample_mask.unsqueeze(-1)
+        num_valid_actor_tokens = actor_mask.sum()
+
         # For truly on-policy training, use curr_logprobs as prev_logprobs
         # This avoids computing prev_logprobs upstream
         if self.force_on_policy_ratio:
@@ -604,7 +640,7 @@ def __call__(
             if self.sequence_level_importance_ratios:
                 seq_log_ratio_mean = masked_mean(
                     log_ratios,
-                    token_mask,
+                    actor_token_mask,
                     dim=-1,
                 ).unsqueeze(-1)
                 seq_ratio = seq_log_ratio_mean.exp()
@@ -754,14 +790,14 @@ def __call__(
         if self.loss_type == LossType.TOKEN_LEVEL:
             actor_loss = masked_mean(
                 importance_weights_to_use * clip_loss,
-                mask,
+                actor_mask,
                 global_normalization_factor=global_valid_toks,
             )
         else:
             actor_loss = masked_mean(
                 masked_mean(
                     importance_weights_to_use * clip_loss,
-                    token_mask,
+                    actor_token_mask,
                     dim=-1,
                 ),
                 sample_mask,
@@ -786,10 +822,13 @@ def __call__(
         # 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,
-                global_normalization_factor=global_valid_toks,
+            # Keep this as a numerator fragment. The narrowed token count is
+            # only known after every microbatch/rank has run, so downstream
+            # aggregation divides the global sum by the global actor-token sum.
+            seq_entropy_approx = -torch.sum(
+                torch.exp(curr_logprobs - generation_logprobs)
+                * curr_logprobs
+                * actor_mask
             )
 
         # -----------------------------------------------------------------
@@ -799,7 +838,7 @@ def __call__(
         nll_loss = torch.tensor(0.0, device=mask.device)
         if self.positive_example_nll_weight > 0 and "rewards" in data:
             correct_sample_mask = (data["rewards"] > 0).float()  # [batch]
-            correct_mask = mask * correct_sample_mask.unsqueeze(-1)
+            correct_mask = actor_mask * correct_sample_mask.unsqueeze(-1)
             correct_valid_toks = correct_mask.sum()
             if correct_valid_toks > 0:
                 nll_loss = masked_mean(
@@ -810,20 +849,14 @@ def __call__(
 
         loss = actor_loss + kl + self.positive_example_nll_weight * nll_loss
         with torch.no_grad():
-            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()
+            # These are numerator fragments, not per-microbatch means. A local
+            # mean cannot be added across packed sequences or uneven DP shards.
+            probs_ratio = torch.sum(ratios.detach() * actor_mask).item()
+            probs_ratio_clamped = torch.sum(ratios_clamped.detach() * actor_mask).item()
 
             # Calculate min/max values for ratios (only for valid tokens)
-            masked_ratios = ratios.detach()[mask.bool()]
-            masked_ratios_clamped = ratios_clamped.detach()[mask.bool()]
+            masked_ratios = ratios.detach()[actor_mask.bool()]
+            masked_ratios_clamped = ratios_clamped.detach()[actor_mask.bool()]
 
             # Handle edge case where there might be no valid tokens
             if masked_ratios.numel() > 0:
@@ -837,9 +870,10 @@ def __call__(
                 probs_ratio_clamped_min = float("inf")
                 probs_ratio_clamped_max = float("-inf")
 
-        # 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.
+        # Most metrics are fragments normalized by global_valid_{seqs/toks} and
+        # become the true step value after summing microbatches. Actor-mask means
+        # are raw numerator fragments paired with num_valid_actor_tokens instead;
+        # callers sum both globally and finalize that ratio once per step.
         return (
             loss,
             {
@@ -857,6 +891,7 @@ def __call__(
                 "js_divergence_error": js_divergence_error,
                 "sampling_importance_ratio": sample_importance_ratio.item(),
                 "num_valid_samples": sample_mask.sum().item(),
+                ACTOR_TOKEN_COUNT_METRIC: num_valid_actor_tokens.item(),
                 "approx_entropy": seq_entropy_approx.item(),
                 **_is_filter_metrics,
                 "positive_nll_loss": nll_loss.item(),

nemo_rl/algorithms/loss/utils.py

@@ -571,7 +571,14 @@ def prepare_loss_input(
             # mask out negative infinity logprobs
             # prev_logprobs is already masked out in the previous step
             mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-            logprobs = mask_out_neg_inf_logprobs(logprobs, mask[:, 1:], "curr_logprobs")
+            logprobs, keep = mask_out_neg_inf_logprobs(
+                logprobs, mask[:, 1:], "curr_logprobs"
+            )
+            # Publish the narrowing instead of folding it into ``token_mask``:
+            # the actor term must drop these positions, but the
+            # reference-policy KL reduces with ``token_mask`` and is computed
+            # from ``curr_logprobs_unfiltered``, which is finite here.
+            data["curr_logprobs_keep_mask"] = keep
 
             # compute unfiltered logprobs for reference policy KL penalty
             if (
@@ -891,7 +898,14 @@ def prepare_packed_loss_input(
     # use filtered curr_logprobs for actor loss, but keep unfiltered values for KL.
     if need_top_k_or_top_p_filtering(sampling_params):
         mask = data["token_mask"] * data["sample_mask"].unsqueeze(-1)
-        logprobs = mask_out_neg_inf_logprobs(logprobs, mask[:, 1:], "curr_logprobs")
+        logprobs, keep = mask_out_neg_inf_logprobs(
+            logprobs, mask[:, 1:], "curr_logprobs"
+        )
+        # Publish the narrowing instead of folding it into ``token_mask``:
+        # the actor term must drop these positions, but the
+        # reference-policy KL reduces with ``token_mask`` and is computed
+        # from ``curr_logprobs_unfiltered``, which is finite here.
+        data["curr_logprobs_keep_mask"] = keep
 
         if (
             hasattr(loss_fn, "reference_policy_kl_penalty")

nemo_rl/algorithms/ppo.py

@@ -51,6 +51,7 @@
     apply_reward_shaping,
 )
 from nemo_rl.algorithms.utils import (
+    finalize_actor_token_metrics,
     print_efficiency_summary,
     print_performance_metrics,
     set_seed,
@@ -1916,6 +1917,7 @@ def ppo_train(
                         metrics[k] = np.mean(v).item()
                     elif isinstance(v, (np.ndarray, list)):
                         metrics[k] = np.sum(v).item()
+                finalize_actor_token_metrics(metrics)
 
                 metrics.update(rollout_metrics)
                 metrics["generation_logger_metrics"] = generation_logger_metrics
@@ -2907,6 +2909,7 @@ def _raise_if_collector_stopped(waiting_for: str) -> None:
                         metrics[k] = np.mean(v).item()
                     elif isinstance(v, (np.ndarray, list)):
                         metrics[k] = np.sum(v).item()
+                finalize_actor_token_metrics(metrics)
 
                 metrics.update(rollout_metrics)
                 if generation_logger_metrics is not None:

nemo_rl/algorithms/single_controller_utils/utils.py

@@ -22,6 +22,7 @@
 import torch
 from tensordict import TensorDict
 
+from nemo_rl.algorithms.utils import finalize_actor_token_metrics
 from nemo_rl.data_plane import KVBatchMeta
 
 # Reduction rules for all_mb_metrics. Mirror grpo.py / grpo_sync.py.
@@ -92,6 +93,7 @@ def aggregate_step_metrics(train_result: dict[str, Any]) -> dict[str, Any]:
             metrics[k] = float(np.mean(v))
         else:
             metrics[k] = float(np.sum(v))
+    finalize_actor_token_metrics(metrics)
 
     # Deferred-mode counterpart to rollout_reassembler's direct-mode
     # routed_experts_row_coverage/routed_experts_sentinel_token_fraction:

nemo_rl/algorithms/utils.py

@@ -39,6 +39,46 @@
     from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 
 
+ACTOR_TOKEN_COUNT_METRIC = "num_valid_actor_tokens"
+ACTOR_TOKEN_MEAN_METRICS = frozenset(
+    {"probs_ratio", "probs_ratio_clamped", "approx_entropy"}
+)
+
+
+def finalize_actor_token_metrics(metrics: dict[str, Any]) -> None:
+    """Normalize actor-mask metric numerators after global aggregation.
+
+    ``ClippedPGLossFn`` emits raw numerator fragments for metrics reduced over
+    its narrowed actor mask, together with the matching valid-token count.
+    Callers must first sum both across every microbatch and data-parallel rank,
+    then call this helper exactly once. Normalizing fragments locally would
+    weight sparse shards and packed sequences equally instead of by their
+    retained-token counts.
+
+    Args:
+        metrics: Step metrics whose loss fragments have already been summed.
+            The mapping is updated in place.
+    """
+    if ACTOR_TOKEN_COUNT_METRIC not in metrics:
+        return
+
+    num_valid_actor_tokens = float(metrics[ACTOR_TOKEN_COUNT_METRIC])
+    if num_valid_actor_tokens < 0:
+        raise ValueError(
+            f"{ACTOR_TOKEN_COUNT_METRIC} must be non-negative, "
+            f"got {num_valid_actor_tokens}"
+        )
+
+    for metric_name in ACTOR_TOKEN_MEAN_METRICS:
+        if metric_name not in metrics:
+            continue
+        metrics[metric_name] = (
+            float(metrics[metric_name]) / num_valid_actor_tokens
+            if num_valid_actor_tokens > 0
+            else 0.0
+        )
+
+
 def get_gdpo_reward_component_keys(batch) -> list[str]:
     """Return batch keys that are named reward components (e.g. reward/correctness) in sorted order."""
     return sorted(
@@ -277,7 +317,7 @@ def masked_mean(
 
 def mask_out_neg_inf_logprobs(
     logprobs: torch.Tensor, mask: torch.Tensor, logprobs_name: str
-) -> torch.Tensor:
+) -> tuple[torch.Tensor, torch.Tensor]:
     """Mask out negative infinity log probabilities.
 
     Handling sampling mask mismatch:
@@ -286,13 +326,22 @@ def mask_out_neg_inf_logprobs(
     token X may fall outside the training policy's top-k/p set -> curr_logprobs[X] = -inf, prev_logprobs[X] = -inf
     Detect positions with -inf in any logprobs (generation_logprobs is always finite for valid tokens)
 
+    The substituted ``0.0`` is *not* a neutral value: it means log p = 0, i.e.
+    p = 1, while ``generation_logprobs`` at the same position is finite and
+    typically around -5. Anything that compares the two -- the importance ratio,
+    ``token_mult_prob_error``, the sequence-level weights -- reads a large
+    fabricated difference unless the position is also dropped from the reduction
+    mask. Callers must therefore fold ``keep`` into their own mask; returning it
+    is what makes the warning above true.
+
     Args:
         logprobs: Log probabilities.
         mask: Mask.
         logprobs_name: Name of the logprobs tensor. Used for printing warning messages.
 
     Returns:
-        Masked log probabilities.
+        ``(logprobs, keep)``, where ``keep`` is 0 at positions whose logprob was
+        -inf and 1 elsewhere, with the same shape as ``mask``.
     """
     is_neginf = torch.isinf(logprobs)
     neginf_count = (is_neginf & mask.bool()).sum().item()
@@ -302,10 +351,11 @@ def mask_out_neg_inf_logprobs(
             "(policy top-k/top-p mismatch). Masking out these positions."
         )
 
-    mask = mask * (~is_neginf).float()
+    keep = (~is_neginf).to(mask.dtype)
+    mask = mask * keep
     logprobs = torch.where(mask.bool(), logprobs, 0.0)
 
-    return logprobs
+    return logprobs, keep
 
 
 def masked_var(