fix(training): sanitize masked sentinel and non-finite inputs in the DSpark loss

Padded teacher-forcing layouts fill unsupervised positions with sentinel target
ids (-1 / -100) and non-finite logits. A sentinel id indexes out of bounds in the
gather, and zero-weighting a non-finite row does not save the step: 0 * NaN is
NaN in the forward, and the backward of a non-finite row poisons the gradient
regardless of its weight.

Replace the inputs at unsupervised positions instead of masking the outputs.
torch.where on draft_logits selects exact zeros into those slots in the backward
rather than multiplying by zero, so no NaN reaches draft_logits.grad; target_ids
are filled with a valid row; the reduced L1 and the confidence logits are
scrubbed for non-finite teacher rows without copying the full-vocab teacher.

Keying off the combined supervision mask means dropped anchors get the same
protection as eval-masked positions. A poisoned batch now produces bit-identical
losses, metrics, and gradients to the equivalent clean one, which the tests
assert term by term.

Note in the head's docstring that the embedding lookup stays strict: it is the
producer's job to replace sentinels, and clamping there would silently train on
row 0 rather than fail.

Signed-off-by: khazic <khazzz1c@gmail.com>

Author: khazic

src/megatron/bridge/training/post_training/dspark/heads.py

@@ -54,7 +54,15 @@ def __init__(self, *, vocab_size: int, markov_rank: int):
         self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False)
 
     def get_prev_embeddings(self, token_ids: torch.Tensor) -> torch.Tensor:
-        """Look up ``markov_w1`` for ``token_ids`` (any shape) -> ``[..., markov_rank]``."""
+        """Look up ``markov_w1`` for ``token_ids`` (any shape) -> ``[..., markov_rank]``.
+
+        ``token_ids`` must already be valid rows in ``[0, vocab_size)``. The
+        teacher-forcing producer owns replacing sentinel pads (``-1`` / ``-100`` at
+        positions outside the supervision mask) with any valid id before calling,
+        mirroring how :func:`~.loss.dspark_loss` sanitizes ``target_ids``. The
+        lookup stays strict on purpose: a sentinel that reaches it is a producer
+        bug, and clamping here would silently train on row 0 instead of failing.
+        """
         return self.markov_w1(token_ids.long())
 
     def project_bias(self, latent_states: torch.Tensor) -> torch.Tensor:

src/megatron/bridge/training/post_training/dspark/loss.py

@@ -83,7 +83,9 @@ class DSparkForwardOutput:
         target_ids: Teacher-forced next token per position
             ``[batch, num_blocks, block_size]``.
         eval_mask: Bool/float supervision mask ``[batch, num_blocks, block_size]``
-            (a block is a contiguous, in-bounds, loss-enabled prefix).
+            (a block is a contiguous, in-bounds, loss-enabled prefix). Positions
+            outside the mask may carry sentinel ids or non-finite padded values;
+            :func:`dspark_loss` sanitizes them before computing any term.
         block_keep_mask: Kept-anchor mask ``[batch, num_blocks]``. Dropped anchors
             are excluded from every loss term, from ``num_tokens``, and from the
             acceptance metrics, independently of what ``eval_mask`` says.
@@ -272,26 +274,38 @@ def dspark_loss(
         raise ValueError("aligned_target_logits is required for the L1 loss or the confidence head.")
     needs_teacher = l1_alpha > 0 or has_confidence
 
+    # Unsupervised positions may carry sentinel target ids (-1 / -100) or
+    # non-finite padded logits. Zero-weighting them afterwards is not enough:
+    # 0 * NaN is NaN, and a non-finite forward row still produces NaN gradients.
+    # Replacing the inputs with torch.where is what makes this safe in both
+    # directions, because where's backward *selects* zeros into the masked slots
+    # rather than multiplying by zero, so no NaN reaches draft_logits.grad.
+    supervised_bool = supervised > 0
+    draft_logits = torch.where(supervised_bool.unsqueeze(-1), draft_logits, 0.0)
+    target_ids = outputs.target_ids.long().masked_fill(~supervised_bool, 0)
+
     # One FP32 normalization of the draft distribution serves both terms.
     ce_per_token, l1_dist = _ce_and_l1_per_token(
-        draft_logits, outputs.target_ids, target_logits if needs_teacher else None, chunk_tokens
+        draft_logits, target_ids, target_logits if needs_teacher else None, chunk_tokens
     )
     ce_num = (ce_per_token * weight).sum()
 
     zero = ce_num.new_zeros(())
     l1_num = conf_num = zero
     accept_rate = None
     if needs_teacher:
+        # A non-finite teacher row at an unsupervised position surfaces as NaN in
+        # the reduced distance; scrub it there rather than copying the full-vocab
+        # teacher tensor just to mask a few rows.
+        l1_dist = torch.where(supervised_bool, l1_dist, 0.0)
         accept_rate = (1.0 - 0.5 * l1_dist).clamp(0.0, 1.0)
         if l1_alpha > 0:
             # Raw weighted L1 (paper Eq. 10). The 1/2 lives in accept_rate above.
             l1_num = (l1_dist * weight).sum()
         if has_confidence:
+            confidence_pred = torch.where(supervised_bool, outputs.confidence_pred.float(), 0.0)
             conf_num = (
-                F.binary_cross_entropy_with_logits(
-                    outputs.confidence_pred.float(), accept_rate.detach(), reduction="none"
-                )
-                * weight
+                F.binary_cross_entropy_with_logits(confidence_pred, accept_rate.detach(), reduction="none") * weight
             ).sum()
 
     loss = ce_alpha * ce_num + l1_alpha * l1_num + confidence_alpha * conf_num

tests/unit_tests/training/post_training/test_dspark_loss.py

@@ -242,6 +242,52 @@ def test_accept_rate_and_tau_exclude_dropped_blocks():
     assert _den(report, "dspark tau") == pytest.approx(B * (N - 1))
 
 
+def _masked_outputs(*, poison: bool):
+    out = _outputs(with_target=True, with_confidence=True)
+    out.eval_mask[:, :, -1] = False  # trailing position unsupervised
+    if poison:
+        with torch.no_grad():
+            out.draft_logits[:, :, -1] = float("-inf")
+        out.aligned_target_logits[:, :, -1] = float("nan")
+        out.target_ids[:, :, -1] = -1
+        out.confidence_pred[:, :, -1] = float("nan")
+    return out
+
+
+def test_masked_positions_tolerate_sentinel_ids_and_nonfinite_padding():
+    # Padded layouts fill unsupervised positions with sentinel ids (-1 / -100) and
+    # non-finite logits; they must neither crash nor leak NaN into loss or grads.
+    clean = _masked_outputs(poison=False)
+    ref_loss, _, ref_report = dspark_loss(clean, ce_alpha=0.1, l1_alpha=0.9, confidence_alpha=1.0)
+    poisoned = _masked_outputs(poison=True)
+    loss, _, report = dspark_loss(poisoned, ce_alpha=0.1, l1_alpha=0.9, confidence_alpha=1.0)
+    assert torch.isfinite(loss)
+    assert loss.item() == pytest.approx(ref_loss.item(), rel=1e-6)
+    for key in ref_report:
+        assert _num(report, key) == pytest.approx(_num(ref_report, key), rel=1e-6), key
+        assert _den(report, key) == pytest.approx(_den(ref_report, key), rel=1e-6), key
+    loss.backward()
+    assert torch.isfinite(poisoned.draft_logits.grad).all()
+    assert torch.all(poisoned.draft_logits.grad[:, :, -1] == 0)
+
+
+def test_dropped_anchors_also_tolerate_poisoned_padding():
+    # The sanitizer keys off the combined mask, so a poisoned position inside a
+    # dropped anchor is covered even when its eval_mask says supervised.
+    out = _outputs(with_target=True, with_confidence=True)
+    out.block_keep_mask[:, 0] = False
+    with torch.no_grad():
+        out.draft_logits[:, 0] = float("-inf")
+    out.aligned_target_logits[:, 0] = float("nan")
+    out.target_ids[:, 0] = -1
+    out.confidence_pred[:, 0] = float("nan")
+    loss, _, _ = dspark_loss(out, ce_alpha=0.1, l1_alpha=0.9, confidence_alpha=1.0)
+    assert torch.isfinite(loss)
+    loss.backward()
+    assert torch.isfinite(out.draft_logits.grad).all()
+    assert torch.all(out.draft_logits.grad[:, 0] == 0)
+
+
 def test_eval_mask_zero_positions_do_not_contribute():
     out = _outputs(with_target=True, with_confidence=False)
     out.eval_mask = torch.zeros(B, N, L, dtype=torch.bool)