perf(loss): defer the fp32 upcast in the distillation top-k path

nemo_rl/algorithms/loss/loss_functions.py

@@ -1197,7 +1197,16 @@ def __call__(
                 global_normalization_factor=global_valid_toks,
             )
         else:
-            kl_loss = per_token_kl.mean()
+            # Enforced rather than assumed: the non-distributed top-k gather in
+            # ``get_distillation_topk_logprobs_from_logits`` accumulates its
+            # backward on the duplicate indices that padded positions carry, and
+            # is only equivalent to the pre-gather-upcast formulation because
+            # those positions reduce to exactly 0.0. An unmasked mean would give
+            # them nonzero gradient and quietly break that equivalence.
+            raise ValueError(
+                "DistillationLossFn requires 'token_mask' and 'sample_mask' in "
+                "data; got keys: " + str(sorted(data.keys()))
+            )
 
         metrics = {
             "loss": float(kl_loss.item()) if kl_loss.ndim == 0 else kl_loss,

nemo_rl/distributed/model_utils.py

@@ -2018,8 +2018,6 @@ def get_distillation_topk_logprobs_from_logits(
             "topk=0 is not supported as it would result in empty tensor operations."
         )
 
-    # Ensure float32 for stability
-    student_logits = student_logits.to(torch.float32)
     # Move teacher topk indices to the same device as student logits
     teacher_topk_indices = teacher_topk_indices.to(student_logits.device)
 
@@ -2069,6 +2067,24 @@ def get_distillation_topk_logprobs_from_logits(
         student_logits = student_logits
         parallel_group = None
 
+    # Two different reasons to keep the full [B, S, V] fp32 copy here:
+    #
+    #   * ``zero_outside_topk`` takes a log_softmax over the whole vocabulary,
+    #     so it genuinely reads every column;
+    #   * the TP/CP paths reach ``ChunkedDistributedGatherLogprob``, whose
+    #     ``backward`` builds its grad as
+    #     ``torch.zeros_like(vocab_parallel_logits, dtype=torch.float32)`` and
+    #     returns it. Hand it a bf16 input and autograd rejects the fp32
+    #     gradient. (The forward would be fine either way -- it upcasts each
+    #     chunk itself -- so this is a backward-side constraint, and narrowing
+    #     the guard further would mean changing that autograd.Function.)
+    #
+    # The remaining path reads K columns, so it upcasts the gathered [B, S, K]
+    # instead -- gather-then-cast is equivalent to cast-then-gather there, and
+    # avoids materializing the full-vocab fp32 tensor.
+    if zero_outside_topk or parallel_group is not None or cp_size > 1:
+        student_logits = student_logits.to(torch.float32)
+
     # Automodel owns the sequence layout: shard the teacher indices into the
     # model's local layout. The legacy CP state was neutralized above so its
     # load-balanced relayout stays out of the way. Gather back to canonical order
@@ -2157,9 +2173,25 @@ def get_distillation_topk_logprobs_from_logits(
 
         # Non-distributed processing
         else:
+            # Gathering K columns and then widening is bitwise equal to
+            # widening the whole vocab axis first: bf16 -> fp32 is exact and
+            # gather is pure selection.
+            #
+            # Duplicate indices do occur -- padded sequence positions carry K
+            # copies of index 0, because dtensor_policy_worker.py pads
+            # topk_indices with value=0 -- so the gather backward accumulates
+            # there rather than scattering. The equality holds anyway, because
+            # DistillationLossFn masks those positions to exactly 0.0 before
+            # reduction, so the accumulated gradient is zero either way.
+            # That is enforced, not assumed: DistillationLossFn rejects data
+            # without ``token_mask``/``sample_mask`` rather than falling back to
+            # an unmasked mean, which would give those positions real gradient.
+            # ``cp_sharder.shard_token_tensor(..., fill=0)`` above is a second
+            # source of duplicate indices, on the padded tail; the same masking
+            # covers them.
             student_topk_logits = student_logits.gather(
                 dim=-1, index=indices_for_logits
-            )
+            ).to(torch.float32)
 
         student_topk_logprobs = torch.nn.functional.log_softmax(
             student_topk_logits, dim=-1

tests/unit/algorithms/test_loss_functions.py

@@ -2264,6 +2264,48 @@ def test_distillation_loss_topk_filtering(k, zero_outside_topk):
         assert loss.item() != 0.0  # Should have some meaningful loss
 
 
+def test_distillation_loss_requires_masks():
+    """Missing masks must raise, not silently fall back to an unmasked mean.
+
+    The non-distributed top-k gather in
+    ``get_distillation_topk_logprobs_from_logits`` accumulates its backward on
+    the duplicate indices that padded positions carry, and matches the
+    pre-gather-upcast formulation only because those positions reduce to
+    exactly 0.0. An unmasked mean gives them real gradient and breaks that.
+    """
+    batch, seq_minus_one, k = 2, 3, 4
+    student = torch.log_softmax(torch.randn(batch, seq_minus_one, k), dim=-1)
+    teacher = torch.log_softmax(torch.randn(batch, seq_minus_one, k), dim=-1)
+    loss_fn = DistillationLossFn(
+        DistillationLossConfig(
+            kl_type="forward",
+            mixed_kl_weight=0.5,
+            zero_outside_topk=False,
+        )
+    )
+    full = {
+        "input_ids": torch.zeros(batch, seq_minus_one + 1, dtype=torch.long),
+        "token_mask": torch.ones(batch, seq_minus_one + 1),
+        "sample_mask": torch.ones(batch),
+    }
+    global_seqs = torch.tensor(float(batch))
+    global_toks = torch.tensor(float(batch * seq_minus_one))
+
+    # Both masks present: reduces normally.
+    loss, _ = loss_fn(
+        student, teacher, None, BatchedDataDict(full), global_seqs, global_toks
+    )
+    assert torch.isfinite(loss)
+
+    # Either one missing: refuse rather than fall back.
+    for missing in ("token_mask", "sample_mask"):
+        stripped = BatchedDataDict(
+            {key: value for key, value in full.items() if key != missing}
+        )
+        with pytest.raises(ValueError, match="requires 'token_mask'"):
+            loss_fn(student, teacher, None, stripped, global_seqs, global_toks)
+
+
 def test_distillation_loss_invalid_k_zero():
     """Test that k=0 should raise a ValueError."""
     # Test with k=0 which should be invalid

tests/unit/distributed/test_model_utils.py

@@ -32,6 +32,7 @@
     from_parallel_logits_to_logprobs,
     from_parallel_logits_to_logprobs_packed_sequences,
     gather_logits_at_global_indices,
+    get_distillation_topk_logprobs_from_logits,
 )
 from nemo_rl.distributed.named_sharding import NamedSharding
 from nemo_rl.distributed.ray_actor_environment_registry import (
@@ -1583,3 +1584,176 @@ def test_distributed_vocab_topk_ops(
         worker_group.shutdown(force=True)
     finally:
         cluster.shutdown()
+
+
+def test_distillation_topk_non_tp_defers_fp32_upcast():
+    """The non-TP top-k gather must not need a full-vocab fp32 copy.
+
+    Only K columns are read on this path, and ``gather`` then ``to(float32)``
+    is equivalent to ``to(float32)`` then ``gather``. Pin that the returned
+    log-probs and the gradient w.r.t. the student logits stay bitwise identical
+    to the previous full-vocab-upcast formulation, and that the result is still
+    fp32.
+    """
+    torch.manual_seed(0)
+    batch, seq, vocab, k = 1, 8, 512, 4
+    logits = torch.randn(batch, seq, vocab, dtype=torch.bfloat16)
+    teacher_topk_logits = torch.randn(batch, seq, k)
+    # distinct per position, as torch.topk returns: with duplicate indices the
+    # gather backward would accumulate into the (now bf16) input rather than an
+    # fp32 copy, which is a different computation
+    teacher_topk_indices = torch.stack(
+        [
+            torch.stack([torch.randperm(vocab)[:k] for _ in range(seq)])
+            for _ in range(batch)
+        ]
+    )
+
+    student = logits.clone().requires_grad_(True)
+    topk_logprobs, _, h_all = get_distillation_topk_logprobs_from_logits(
+        student_logits=student,
+        teacher_topk_logits=teacher_topk_logits,
+        teacher_topk_indices=teacher_topk_indices,
+        zero_outside_topk=False,
+        calculate_entropy=False,
+    )
+    topk_logprobs.sum().backward()
+
+    reference = logits.clone().requires_grad_(True)
+    expected = torch.nn.functional.log_softmax(
+        reference.to(torch.float32).gather(dim=-1, index=teacher_topk_indices),
+        dim=-1,
+    )[:, :-1, :]
+    expected.sum().backward()
+
+    assert h_all is None
+    assert topk_logprobs.dtype == torch.float32
+    assert torch.equal(topk_logprobs, expected)
+    assert torch.equal(student.grad, reference.grad)
+
+
+def _topk_grad_with(indices, upstream, logits):
+    """Gradient w.r.t. ``logits`` under the deferred-upcast implementation."""
+    student = logits.clone().requires_grad_(True)
+    topk_logprobs, _, _ = get_distillation_topk_logprobs_from_logits(
+        student_logits=student,
+        teacher_topk_logits=torch.zeros_like(upstream),
+        teacher_topk_indices=indices,
+        zero_outside_topk=False,
+        calculate_entropy=False,
+    )
+    (topk_logprobs * upstream[:, :-1, :]).sum().backward()
+    return student.grad
+
+
+def _topk_grad_reference(indices, upstream, logits):
+    """Gradient under the previous full-vocab-upcast formulation."""
+    reference = logits.clone().requires_grad_(True)
+    expected = torch.nn.functional.log_softmax(
+        reference.to(torch.float32).gather(dim=-1, index=indices),
+        dim=-1,
+    )[:, :-1, :]
+    (expected * upstream[:, :-1, :]).sum().backward()
+    return reference.grad
+
+
+def test_deferred_upcast_matches_only_because_padded_positions_are_masked():
+    """The equivalence is conditional, and this pins the condition.
+
+    ``gather``'s backward is a scatter-add, and padded positions carry K copies
+    of index 0 -- ``dtensor_policy_worker.py`` pads the sequence axis with
+    ``value=0``. There the old formulation accumulates in fp32 and this one
+    accumulates in bf16, which is a different number.
+
+    They agree because ``DistillationLossFn`` masks those positions to exactly
+    0.0, so the incoming gradient is zero and there is nothing to accumulate.
+    Both directions are checked: masked agrees, unmasked does not. The second
+    half is why the loss now raises on missing masks instead of falling back to
+    an unmasked mean.
+
+    Two things would make this test pass without testing anything:
+      - top-k indices are distinct within a position, so only the padded tail
+        has duplicates at all -- ``torch.randperm`` everywhere reproduces
+        nothing;
+      - a uniform upstream gradient makes the unmasked case agree too, since
+        ``log_softmax`` over K duplicated logits is uniform and its gradient
+        for an all-ones upstream is exactly zero.
+    """
+    torch.manual_seed(0)
+    batch, seq, vocab, k = 2, 8, 512, 4
+    pad_from = 6
+
+    logits = torch.randn(batch, seq, vocab, dtype=torch.bfloat16)
+    indices = torch.stack(
+        [
+            torch.stack([torch.randperm(vocab)[:k] for _ in range(seq)])
+            for _ in range(batch)
+        ]
+    )
+    indices[:, pad_from:, :] = 0  # the padded tail, as the worker writes it
+
+    upstream = torch.randn(batch, seq, k)  # non-uniform, deliberately
+    mask = torch.ones(batch, seq, 1)
+    mask[:, pad_from:, :] = 0.0
+
+    assert torch.equal(
+        _topk_grad_with(indices, upstream * mask, logits),
+        _topk_grad_reference(indices, upstream * mask, logits),
+    ), "masked: the deferred upcast must be bitwise equal"
+
+    assert not torch.equal(
+        _topk_grad_with(indices, upstream, logits),
+        _topk_grad_reference(indices, upstream, logits),
+    ), (
+        "unmasked: the two must DIFFER -- if they do not, this test has stopped "
+        "reproducing the duplicate-index accumulation and the ValueError in "
+        "DistillationLossFn is no longer pinned by anything"
+    )
+
+
+def test_distillation_topk_gathers_before_upcasting(monkeypatch):
+    """The full-vocab tensor must still be bf16 when ``gather`` runs.
+
+    The equivalence test above cannot catch a revert: its reference is the
+    old ``to(float32).gather(...)`` formulation, so it holds whether or not
+    the upcast was deferred. This one observes the deferral itself.
+
+    It spies on ``Tensor.gather``, so it is coupled to that specific call. An
+    equivalent rewrite to ``torch.gather(...)`` or ``index_select`` would keep
+    the deferral but stop tripping the spy, failing here as a false alarm --
+    update the spy in that case rather than reading it as a regression.
+    """
+    torch.manual_seed(0)
+    batch, seq, vocab, k = 1, 8, 512, 4
+    logits = torch.randn(batch, seq, vocab, dtype=torch.bfloat16)
+    teacher_topk_logits = torch.randn(batch, seq, k)
+    teacher_topk_indices = torch.stack(
+        [
+            torch.stack([torch.randperm(vocab)[:k] for _ in range(seq)])
+            for _ in range(batch)
+        ]
+    )
+
+    seen_dtypes = []
+    real_gather = torch.Tensor.gather
+
+    def spy(self, dim, index):
+        # only the full-vocab gather is interesting; K-wide ones are downstream
+        if self.shape[-1] == vocab:
+            seen_dtypes.append(self.dtype)
+        return real_gather(self, dim, index)
+
+    monkeypatch.setattr(torch.Tensor, "gather", spy, raising=True)
+
+    get_distillation_topk_logprobs_from_logits(
+        student_logits=logits,
+        teacher_topk_logits=teacher_topk_logits,
+        teacher_topk_indices=teacher_topk_indices,
+        zero_outside_topk=False,
+        calculate_entropy=False,
+    )
+
+    assert seen_dtypes, "expected a gather over the vocab axis"
+    assert all(d == torch.bfloat16 for d in seen_dtypes), (
+        f"the vocab-wide tensor was upcast before the gather: {seen_dtypes}"
+    )