[4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss

modelopt/torch/speculative/plugins/hf_dflash2.py

@@ -29,9 +29,11 @@
 - Take the backbone's top-k candidates per block position.
 - Score each candidate against its *teacher-forced* predecessor token, so the
   positions train in parallel exactly as the backbone does.
-- When the gold token is missing from the top-k, substitute it into the last
-  candidate slot. Without this the selector sees no positive class on the hard
-  positions and never learns those edges.
+- Supervise only the strict top-k, the candidate set serving builds. A gold token
+  outside it is a backbone recall miss, not a selector example, and carries no
+  selector gradient.
+- Weight each position as the backbone loss does (decay or D-PACE) and divide by the
+  backbone's weight sum, misses included.
 """
 
 import torch
@@ -106,15 +108,20 @@ def get_exporter(self):
 
         return DFlash2Exporter(self)
 
-    def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask):
+    def _selector_loss(
+        self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum
+    ):
         """Cross-entropy over the selector's candidate set, and its top-1 accuracy.
 
         Args:
             logits: Backbone logits per block position ``[B, N, block_size, V]``.
             target_ids: Gold token ids ``[B, N, block_size]``.
             hidden: Backbone hidden states ``[B, N, block_size, H]``.
             predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``.
-            weight_mask: Per-position loss weights ``[B, N, block_size]``.
+            supervised_mask: The positions the backbone loss supervises, unweighted
+                ``[B, N, block_size]``.
+            loss_weights: The backbone loss's per-position weights ``[B, N, block_size]``.
+            weight_sum: The backbone loss's denominator, the sum of ``loss_weights``.
 
         Returns:
             ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised
@@ -130,7 +137,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
         # token the backbone did not propose is a backbone recall failure, not a
         # selector classification example: substituting it in would teach the selector
         # to override the unary ranking on a set it will never be shown. Those
-        # positions carry no selector gradient and leave the denominator instead.
+        # positions carry no selector gradient.
         gold_matches = candidate_ids == target_ids.unsqueeze(-1)
         gold_in_topk = gold_matches.any(dim=-1)
         gold_slot = gold_matches.long().argmax(dim=-1)
@@ -139,26 +146,27 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
             candidate_ids, unary_logits, hidden, predecessor_ids
         )
 
-        covered = weight_mask * gold_in_topk.to(weight_mask.dtype)
-        flat_weights = covered.reshape(-1)
-        denominator = flat_weights.sum() + 1e-6
+        in_topk = gold_in_topk.to(loss_weights.dtype)
         per_token = F.cross_entropy(
             selector_logits.float().reshape(-1, top_k),
             gold_slot.reshape(-1),
             reduction="none",
         )
-        loss = (per_token * flat_weights).sum() / denominator
+        # The backbone's denominator, misses included: dividing by the covered positions
+        # alone would scale the term up whenever backbone recall drops.
+        loss = (per_token * (loss_weights * in_topk).reshape(-1)).sum() / weight_sum
 
         with torch.no_grad():
+            covered = (supervised_mask * in_topk).reshape(-1)
             chosen = selector_logits.argmax(dim=-1).reshape(-1)
-            accuracy = (
-                (chosen == gold_slot.reshape(-1)).float() * flat_weights
-            ).sum() / denominator
+            accuracy = ((chosen == gold_slot.reshape(-1)).float() * covered).sum() / (
+                covered.sum() + 1e-6
+            )
             # Coverage keeps the full supervised mask as its denominator: it measures
             # how often the selector was given a solvable problem at all.
             # clamp, not an epsilon: a fully covered batch must read exactly 1.0.
-            supervised = weight_mask.reshape(-1).sum()
-            coverage = flat_weights.sum() / supervised.clamp(min=1.0)
+            supervised = supervised_mask.reshape(-1).sum()
+            coverage = covered.sum() / supervised.clamp(min=1.0)
         # Detached tensors, not Python scalars: .item() would force a CPU-GPU sync on
         # every training step. The trainer converts them at the logging boundary.
         return loss, accuracy.detach(), coverage.detach()
@@ -238,24 +246,12 @@ def _compute_loss(
 
         offsets = torch.arange(block_size, device=device).view(1, 1, -1)
         label_indices = anchor_positions.unsqueeze(-1) + offsets
-        valid_label = label_indices < seq_len
         safe_label_indices = label_indices.clamp(max=seq_len - 1)
         expanded_ids = input_ids.unsqueeze(1).expand(-1, n_blocks, -1)
         target_ids = torch.gather(expanded_ids, 2, safe_label_indices)
 
-        # Same supervision mask as the backbone loss: valid block, in bounds, not the
-        # anchor slot, and inside the answer span. Position weighting (decay/D-PACE) is
-        # deliberately not applied — it shapes *where* the backbone spends capacity,
-        # while the selector should learn every position's transition equally.
-        weight_mask = block_keep_mask.unsqueeze(-1).expand(-1, -1, block_size).float()
-        weight_mask = weight_mask * valid_label.float()
-        weight_mask = weight_mask * (offsets > 0).float()
-        weight_mask = weight_mask * torch.gather(
-            loss_mask.unsqueeze(1).expand(-1, n_blocks, -1), 2, safe_label_indices
-        )
-
         # Teacher-forced predecessor of block position k is the real token at anchor+k-1.
-        # Only offsets 1..block_size-1 are supervised (weight_mask zeroes slot 0), so the
+        # Only offsets 1..block_size-1 are supervised (the backbone's mask zeroes slot 0), so the
         # first supervised position, k=1, has the anchor's own token as its predecessor.
         # Slot 0's entry here resolves to anchor-1 and never contributes.
         #
@@ -264,12 +260,16 @@ def _compute_loss(
         # that hands greedy_path the full 0..block_size-1 candidate set is off by one.
         predecessor_ids = torch.gather(expanded_ids, 2, (safe_label_indices - 1).clamp(min=0))
 
+        # The backbone's own mask and position weights: a block position's choice only
+        # matters if every earlier one was accepted, which decay / D-PACE model.
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
             logits.reshape(bsz, n_blocks, block_size, -1),
             target_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
-            weight_mask,
+            terms.supervised_mask.view_as(target_ids),
+            terms.weights.view_as(target_ids),
+            terms.weight_sum,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,

tests/unit/torch/speculative/plugins/test_hf_dflash2.py

@@ -397,9 +397,11 @@ def _selector(self, vocab=16, rank=4, top_k=3, hidden=8):
             torch.nn.init.normal_(selector.successor_codebook, std=0.5)
         return selector
 
-    def _selector_terms(self, selector, target_ids, candidate_ids):
+    def _selector_terms(self, selector, target_ids, candidate_ids, loss_weights=None):
         """Drive HFDFlash2Model._selector_loss with a hand-built candidate set."""
         b, length, k = candidate_ids.shape
+        supervised = torch.ones(b, length, dtype=torch.double)
+        weights = supervised if loss_weights is None else loss_weights
         model = HFDFlash2Model.__new__(HFDFlash2Model)
         model.dflash_module = type("_M", (), {"candidate_selector": selector})()
         logits = torch.zeros(b, length, 32, dtype=torch.double)
@@ -410,9 +412,13 @@ def _selector_terms(self, selector, target_ids, candidate_ids):
             model,
             logits,
             target_ids,
-            torch.randn(b, length, 8, dtype=torch.double),
+            torch.randn(
+                b, length, 8, generator=torch.Generator().manual_seed(1), dtype=torch.double
+            ),
             torch.zeros(b, length, dtype=torch.long),
-            torch.ones(b, length, dtype=torch.double),
+            supervised,
+            weights,
+            weights.sum(),
         )
 
     def test_a_miss_carries_no_selector_gradient(self):
@@ -453,6 +459,26 @@ def test_a_covered_position_does_train_the_selector(self):
         loss.backward()
         assert selector.successor_codebook.grad.abs().sum() > 0.0
 
+    def test_positions_weigh_like_the_backbone_and_misses_stay_in_the_denominator(self):
+        """The term is the backbone-weighted mean over every supervised position.
+
+        A miss adds nothing to the numerator but keeps its weight in the denominator, so
+        the term does not grow as backbone recall drops.
+        """
+        selector = self._selector(vocab=32, top_k=2)
+        candidates = torch.tensor([[[5, 6], [5, 6], [5, 6]]])
+        targets = torch.tensor([[5, 6, 9]])  # two covered positions, one miss
+        weights = torch.tensor([[1.0, 3.0, 2.0]], dtype=torch.double)
+
+        loss, _, _ = self._selector_terms(selector, targets, candidates, weights)
+        per_position = [
+            self._selector_terms(selector, targets, candidates, torch.eye(3)[i : i + 1].double())[0]
+            for i in range(3)
+        ]
+        assert float(per_position[2]) == 0.0
+        expected = sum(w * term for w, term in zip(weights[0], per_position)) / weights.sum()
+        torch.testing.assert_close(loss, expected)
+
     def test_greedy_path_position_zero_is_seeded_by_the_anchor(self):
         """``greedy_path`` position 0 scores against the anchor, i.e. block offset 1.