fix(loss): count num_valid_samples from the sample mask, not the batch dim

nemo_rl/algorithms/loss/loss_functions.py

@@ -1796,7 +1796,17 @@ def __call__(
 
         metrics = {
             "loss": float(kl_loss.item()) if kl_loss.ndim == 0 else kl_loss,
-            "num_valid_samples": data["input_ids"].shape[0],
+            # From the mask, not the batch dimension: workers gate on
+            # ``num_valid_samples > 0`` to decide whether a microbatch is
+            # recorded at all, so the raw batch size makes a fully-masked
+            # microbatch look like it contributed and dilutes the step's
+            # reported loss with its zero. Falls back to the batch dimension
+            # when there is no mask, where every sample is valid by definition.
+            "num_valid_samples": (
+                torch.count_nonzero(data["sample_mask"]).item()
+                if "sample_mask" in data
+                else data["input_ids"].shape[0]
+            ),
         }
 
         return kl_loss, metrics
@@ -1931,7 +1941,8 @@ def __call__(
             "values_max": values_max,
             "returns_sq_mean": returns_sq_mean,
             "residual_sq_mean": residual_sq_mean,
-            "num_valid_samples": int(values.shape[0]),
+            # See DistillationLossFn: the critic's workers gate on this too.
+            "num_valid_samples": torch.count_nonzero(sample_mask).item(),
         }
 
         return loss, metrics
@@ -2314,7 +2325,8 @@ def __call__(
             "ce_loss": ce_loss.item(),
             "kl_loss_scale": kl_scale.item(),
             "accuracy": accuracy.item(),
-            "num_valid_samples": data["input_ids"].shape[0],
+            # See DistillationLossFn.
+            "num_valid_samples": torch.count_nonzero(data["sample_mask"]).item(),
         }
         metrics.update(per_teacher_metrics)
         return loss, metrics

tests/unit/algorithms/test_num_valid_samples.py

@@ -0,0 +1,218 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""``num_valid_samples`` must come from the sample mask, in every loss.
+
+Workers gate on ``num_valid_samples > 0`` to decide whether a microbatch's
+loss and metrics are recorded at all -- ``dtensor_policy_worker.py``,
+``dtensor_policy_worker_v2.py`` and ``dtensor_value_worker_v2.py`` all do. A
+loss that reports the raw batch dimension makes a fully-masked microbatch look
+like it contributed, and its zero loss then dilutes the step's reported mean.
+
+CPU-only: these construct the loss inputs directly rather than going through
+``prepare_loss_input``, so no GPU and no distributed init.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+import nemo_rl.algorithms.loss.loss_functions as loss_functions_mod
+from nemo_rl.algorithms.loss.loss_functions import (
+    CrossTokenizerDistillationLossConfig,
+    CrossTokenizerDistillationLossFn,
+    DistillationLossConfig,
+    DistillationLossFn,
+    MseValueLossConfig,
+    MseValueLossFn,
+)
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.algorithms.x_token.loss_utils import LocalizedAlignment
+
+B, S = 3, 4
+# One live sample out of three. Everything below asserts on 1.0, not 3.
+LIVE_SAMPLE_MASK = torch.tensor([1.0, 0.0, 0.0])
+FRACTIONAL_SAMPLE_MASK = torch.tensor([0.25, 0.0, 0.75])
+ALL_MASKED = torch.zeros(B)
+
+
+def _value_call(sample_mask: torch.Tensor):
+    values = torch.randn(B, S)
+    data = BatchedDataDict(
+        {
+            "token_mask": torch.ones(B, S),
+            "sample_mask": sample_mask,
+            "returns": torch.randn(B, S),
+            "values": torch.randn(B, S),
+        }
+    )
+    gvs = sample_mask.sum().clamp(min=1.0)
+    gvt = (data["token_mask"] * sample_mask.unsqueeze(-1)).sum().clamp(min=1.0)
+    loss_fn = MseValueLossFn(MseValueLossConfig(scale=1.0, cliprange=None))
+    return loss_fn(values, data, gvs, gvt)
+
+
+def _distillation_call(sample_mask: torch.Tensor, k: int = 5):
+    data = {
+        "input_ids": torch.randint(0, 8, (B, S)),
+        "token_mask": torch.ones(B, S),
+        "sample_mask": sample_mask,
+    }
+    student = torch.randn(B, S - 1, k).log_softmax(-1)
+    teacher = torch.randn(B, S - 1, k).log_softmax(-1)
+    gvs = sample_mask.sum().clamp(min=1.0)
+    gvt = (data["token_mask"] * sample_mask.unsqueeze(-1)).sum().clamp(min=1.0)
+    loss_fn = DistillationLossFn(DistillationLossConfig(kl_type="forward"))
+    return loss_fn(
+        student_topk_logprobs=student,
+        teacher_topk_logprobs=teacher,
+        H_all=None,
+        data=data,
+        global_valid_seqs=gvs,
+        global_valid_toks=gvt,
+    )
+
+
+def _cross_tokenizer_config() -> CrossTokenizerDistillationLossConfig:
+    return {
+        "gold_loss": False,
+        "xtoken_loss": False,
+        "temperature": 1.0,
+        "vocab_topk": 4,
+        "uncommon_topk": 4,
+        "reverse_kl": False,
+        "exact_token_match_only": False,
+        "kl_loss_weight": 1.0,
+        "ce_loss_scale": 1.0,
+        "dynamic_loss_scaling": False,
+        "kd_loss_mode": "sum",
+        "normalize_teacher_by_vocab": False,
+        "alpha": 1.0,
+        "student_vocab_size": 8,
+        "teacher_vocab_sizes": [8],
+        "projection_matrix_paths": [None],
+        "teacher_weights": [1.0],
+        "teacher_gold_loss": [None],
+        "teacher_xtoken_loss": [None],
+    }
+
+
+class TestMseValueLossFn:
+    def test_counts_the_mask_not_the_batch(self):
+        _, metrics = _value_call(LIVE_SAMPLE_MASK)
+        assert metrics["num_valid_samples"] == 1.0
+
+    def test_a_fully_masked_microbatch_reports_zero(self):
+        """The gate is ``num_valid_samples > 0``. Reporting the batch size here
+        lets a microbatch that contributed nothing through, and its zero loss
+        is then averaged into the step."""
+        _, metrics = _value_call(ALL_MASKED)
+        assert metrics["num_valid_samples"] == 0.0
+
+    def test_fractional_loss_weights_are_counted_not_summed(self):
+        _, metrics = _value_call(FRACTIONAL_SAMPLE_MASK)
+        assert metrics["num_valid_samples"] == 2.0
+
+
+class TestDistillationLossFn:
+    def test_counts_the_mask_not_the_batch(self):
+        _, metrics = _distillation_call(LIVE_SAMPLE_MASK)
+        assert metrics["num_valid_samples"] == 1.0
+
+    def test_a_fully_masked_microbatch_reports_zero(self):
+        _, metrics = _distillation_call(ALL_MASKED)
+        assert metrics["num_valid_samples"] == 0.0
+
+    def test_fractional_loss_weights_are_counted_not_summed(self):
+        _, metrics = _distillation_call(FRACTIONAL_SAMPLE_MASK)
+        assert metrics["num_valid_samples"] == 2.0
+
+    def test_without_a_mask_every_sample_is_valid(self):
+        """The unmasked branch has nothing to count, so the batch dimension is
+        the honest answer on current main."""
+        data = {"input_ids": torch.randint(0, 8, (B, S))}
+        student = torch.randn(B, S - 1, 5).log_softmax(-1)
+        teacher = torch.randn(B, S - 1, 5).log_softmax(-1)
+        loss_fn = DistillationLossFn(DistillationLossConfig(kl_type="forward"))
+        _, metrics = loss_fn(
+            student_topk_logprobs=student,
+            teacher_topk_logprobs=teacher,
+            H_all=None,
+            data=data,
+            global_valid_seqs=torch.tensor(1.0),
+            global_valid_toks=torch.tensor(1.0),
+        )
+        assert metrics["num_valid_samples"] == B
+
+
+class TestCrossTokenizerDistillationLossFn:
+    @pytest.mark.parametrize(
+        ("sample_mask", "expected"),
+        [
+            pytest.param(LIVE_SAMPLE_MASK, 1.0, id="one-live-sample"),
+            pytest.param(
+                FRACTIONAL_SAMPLE_MASK,
+                2.0,
+                id="two-fractionally-weighted-samples",
+            ),
+            pytest.param(ALL_MASKED, 0.0, id="all-masked"),
+        ],
+    )
+    def test_call_reports_the_sample_mask_count(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        sample_mask: torch.Tensor,
+        expected: float,
+    ) -> None:
+        """Exercise __call__, where the worker-facing metrics are assembled."""
+        loss_fn = CrossTokenizerDistillationLossFn(_cross_tokenizer_config())
+        monkeypatch.setattr(
+            loss_fn,
+            "_compute_ce",
+            lambda *args, **kwargs: torch.tensor(0.25),
+        )
+        monkeypatch.setattr(
+            loss_fn,
+            "_sum_kd",
+            lambda *args, **kwargs: (torch.tensor(0.75), {}),
+        )
+        monkeypatch.setattr(
+            loss_functions_mod,
+            "next_token_accuracy",
+            lambda *args, **kwargs: torch.tensor(1.0),
+        )
+        data = BatchedDataDict(
+            {
+                "input_ids": torch.zeros((B, S), dtype=torch.long),
+                "token_mask": torch.ones(B, S),
+                "sample_mask": sample_mask,
+            }
+        )
+        align = LocalizedAlignment(
+            sample_mask=sample_mask,
+            student_input_ids=data["input_ids"],
+            student_token_mask=data["token_mask"],
+        )
+
+        _, metrics = loss_fn(
+            data=data,
+            global_valid_seqs=sample_mask.sum().clamp(min=1.0),
+            global_valid_toks=torch.tensor(1.0),
+            logits=torch.empty(B, S, 8),
+            student_logits_contig=torch.empty(B, S, 8),
+            teacher_full_logits_by_idx={0: torch.empty(B, S, 8)},
+            aligns_by_idx={0: align},
+        )
+
+        assert metrics["num_valid_samples"] == expected