diff --git a/nemo_rl/algorithms/loss/loss_functions.py b/nemo_rl/algorithms/loss/loss_functions.py index c615de7199b..9b6ea6ad58f 100755 --- a/nemo_rl/algorithms/loss/loss_functions.py +++ b/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, diff --git a/nemo_rl/distributed/model_utils.py b/nemo_rl/distributed/model_utils.py index 2a6b212b675..4da6b3b1b2a 100644 --- a/nemo_rl/distributed/model_utils.py +++ b/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 diff --git a/tests/unit/algorithms/test_loss_functions.py b/tests/unit/algorithms/test_loss_functions.py index 12aa3f4d4d5..117a01ee85c 100644 --- a/tests/unit/algorithms/test_loss_functions.py +++ b/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 diff --git a/tests/unit/distributed/test_model_utils.py b/tests/unit/distributed/test_model_utils.py index 72f626c001f..cf24a588ba1 100644 --- a/tests/unit/distributed/test_model_utils.py +++ b/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}" + )