[5/8] Add the tvd_ce objective

examples/speculative_decoding/README.md

@@ -165,6 +165,7 @@ torchrun --nproc_per_node 8 main.py \
 | `soft_ce` | cross-entropy against the base's soft distribution (the EAGLE offline objective) |
 | `tvd` | total-variation distance between the two full distributions |
 | `tvd_deploy` | the same, with both sides put through the serving filter (`external_top_k`, then `external_top_p`) first, so the loss is the acceptance the deployed pair sees |
+| `tvd_ce` | `tvd` plus a cross-entropy term on the base's argmax token, weighted by `external_tvd_alpha` / `external_ce_alpha` (default 0.9 / 0.1). TVD is bounded and so ranks weakly among tokens inside the base's support; CE's log penalty does not |
 
 Under rejection sampling the acceptance probability is exactly `sum(min(p, q))`,
 so `tvd_deploy` is `1 - acceptance` — the deployment metric itself rather than a

modelopt/torch/speculative/config.py

@@ -477,16 +477,37 @@ class ExternalDraftConfig(ModeloptBaseConfig):
     hidden states, projected by the base lm_head to obtain teacher logits.
     """
 
-    external_loss: Literal["soft_ce", "tvd", "tvd_deploy"] = ModeloptField(
+    external_loss: Literal["soft_ce", "tvd", "tvd_deploy", "tvd_ce"] = ModeloptField(
         default="tvd",
         description=(
             "Training objective. ``soft_ce`` matches the Eagle offline objective. "
             "``tvd`` is total-variation distance between the two untruncated "
             "distributions. ``tvd_deploy`` puts both through the serving filter first "
             "(``external_top_k`` then ``external_top_p``), making it the acceptance the "
-            "deployed pair will see: ``alpha = 1 - TVD``. Sparse data requires one of the "
-            "TVD objectives, since the dump is already a truncated policy and "
-            "``soft_ce`` needs full base logits."
+            "deployed pair will see: ``alpha = 1 - TVD``. ``tvd_ce`` adds a cross-entropy "
+            "term. Sparse data requires one of the TVD objectives, since the dump is "
+            "already a truncated policy and ``soft_ce`` needs full base logits."
+        ),
+    )
+
+    external_tvd_alpha: float = ModeloptField(
+        default=0.9,
+        ge=0.0,
+        description=(
+            "Weight on the TVD term when ``external_loss='tvd_ce'``. The two weights "
+            "are not normalised: they scale the loss as given, so 0.9/0.1 is a "
+            "convention rather than a constraint, and setting only one of them moves "
+            "the overall gradient magnitude as well as the ratio."
+        ),
+    )
+
+    external_ce_alpha: float = ModeloptField(
+        default=0.1,
+        ge=0.0,
+        description=(
+            "Weight on the cross-entropy term when ``external_loss='tvd_ce'``. CE is "
+            "added because TVD is bounded and so discriminates weakly between tokens "
+            "inside the base's support."
         ),
     )
 

modelopt/torch/speculative/external/external_model.py

@@ -32,12 +32,16 @@ class ExternalDraftModel(DynamicModule):
     def _setup(self):
         self._register_temp_attribute("external_loss", "tvd")
         self._register_temp_attribute("external_report_acc", True)
+        self._register_temp_attribute("external_tvd_alpha", 0.9)
+        self._register_temp_attribute("external_ce_alpha", 0.1)
         self._register_temp_attribute("external_top_k", 20)
         self._register_temp_attribute("external_top_p", 0.95)
 
     def modify(self, config):
         """Base external draft modify function. Child class should implement the details."""
         self.external_loss = config.external_loss
         self.external_report_acc = config.external_report_acc
+        self.external_tvd_alpha = config.external_tvd_alpha
+        self.external_ce_alpha = config.external_ce_alpha
         self.external_top_k = config.external_top_k
         self.external_top_p = config.external_top_p

modelopt/torch/speculative/plugins/hf_external.py

@@ -21,9 +21,10 @@
 projected by the base lm_head to obtain teacher logits.
 
 ``external_loss`` selects the objective: ``soft_ce`` matches the Eagle offline
-objective, ``tvd`` is total-variation distance between the two distributions, and
+objective, ``tvd`` is total-variation distance between the two distributions,
 ``tvd_deploy`` applies the serving filter to both sides first so the loss is the
-acceptance the deployed pair sees.
+acceptance the deployed pair sees, and ``tvd_ce`` adds a cross-entropy ranking
+term to ``tvd``.
 """
 
 import copy
@@ -44,7 +45,7 @@
 # On the sparse path the base's truncation is baked into the dump, so both TVD
 # variants are computable there; they differ only in whether the *draft* is put
 # through the serving filter too. soft_ce needs full base logits and is dense-only.
-SPARSE_CAPABLE_LOSSES = ("tvd", "tvd_deploy")
+SPARSE_CAPABLE_LOSSES = ("tvd", "tvd_deploy", "tvd_ce")
 
 __all__ = ["SPARSE_CAPABLE_LOSSES", "HFExternalDraftModel"]
 
@@ -308,6 +309,14 @@ def _top1_accuracy(teacher_top1, draft_logits, mask):
             correct = (teacher_top1 == draft_logits.detach().argmax(dim=-1)) & valid
             return correct.sum().float() / valid.sum().clamp_min(1).float()
 
+    def _mix_ce(self, tvd, draft_logits, target, mask):
+        """``tvd_ce``: add CE on ``target``, since bounded TVD ranks in-support tokens weakly."""
+        ce = torch.nn.functional.cross_entropy(
+            draft_logits.float().flatten(0, 1), target.flatten(), reduction="none"
+        ).view_as(mask)
+        ce = (ce * mask).sum() / (mask.sum() + 1e-5)
+        return self.external_tvd_alpha * tvd + self.external_ce_alpha * ce
+
     def _tvd_sparse(self, draft_logits, topk_tok, topk_prob):
         """Plain TVD against the stored teacher policy, over the whole vocabulary.
 
@@ -345,12 +354,16 @@ def compute_sparse_loss(self, draft_logits, topk_tok, topk_prob, loss_mask):
         tok, prob = topk_tok[:, 1:seq], topk_prob[:, 1:seq]
         mask = loss_mask[:, 1:seq].to(draft_logits.dtype)
         denom = mask.sum() + 1e-5
-        tvd = self._tvd_sparse(draft, tok, prob)
-        loss = (tvd * mask).sum() / denom
+        loss = (self._tvd_sparse(draft, tok, prob) * mask).sum() / denom
+        # A row the dump left empty has no teacher top-1: its argmax is slot 0, i.e.
+        # whatever the zero-fill holds. CE and accuracy both skip those rows.
+        teacher_top1 = tok.gather(-1, prob.argmax(-1, keepdim=True)).squeeze(-1)
+        has_mass = mask * (prob.sum(-1) > 0)
+        if self.external_loss == "tvd_ce":
+            loss = self._mix_ce(loss, draft, teacher_top1, has_mass)
         accuracy = None
         if self.external_report_acc:
-            teacher_top1 = tok.gather(-1, prob.argmax(-1, keepdim=True)).squeeze(-1)
-            accuracy = self._top1_accuracy(teacher_top1, draft, mask)
+            accuracy = self._top1_accuracy(teacher_top1, draft, has_mass)
         return loss, accuracy
 
     def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=None):
@@ -372,7 +385,7 @@ def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=No
         mask = loss_mask[:, :seq_len].to(draft_logits.dtype)
         denom = mask.sum() + 1e-5
 
-        if self.external_loss in ("tvd", "tvd_deploy"):
+        if self.external_loss in ("tvd", "tvd_deploy", "tvd_ce"):
             if self.external_loss == "tvd_deploy":
                 # Both sides through the serving filter, so the loss is the
                 # acceptance the deployed pair will actually see.
@@ -381,6 +394,8 @@ def compute_loss(self, draft_logits, base_logits, loss_mask, base_predict_tok=No
                 p = torch.softmax(base_logits.float(), dim=-1)
             q = self._draft_probs(draft_logits)
             loss = (0.5 * (p - q).abs().sum(-1) * mask).sum() / denom
+            if self.external_loss == "tvd_ce":
+                loss = self._mix_ce(loss, draft_logits, base_logits.argmax(-1), mask)
         else:
             # Identical in form to hf_eagle._eagle_loss so switching modes does not
             # silently change the objective.

modelopt_recipes/general/speculative_decoding/external_draft.yaml

@@ -26,7 +26,7 @@ data:
   data_path:
   offline_data_path:
   # Dumped top-k teacher deployment policy. Selects data.mode=sparse, which needs no
-  # base lm_head at train time and requires a TVD objective (tvd or tvd_deploy).
+  # base lm_head at train time and requires a TVD objective (tvd, tvd_deploy or tvd_ce).
   #
   # The first run over a dump builds a memmapped cache beside the shards, which takes
   # hours on a large corpus. Under torchrun only rank 0 builds it and the others wait,
@@ -70,7 +70,8 @@ training:
 external:
   # soft_ce matches the EAGLE offline objective. tvd is plain total-variation
   # distance; tvd_deploy filters both sides to the serving policy first and is the
-  # acceptance the deployed pair sees.
+  # acceptance the deployed pair sees. tvd_ce adds a cross-entropy term on the base's
+  # argmax, weighted by external_tvd_alpha / external_ce_alpha.
   #
   # How much of the base's distribution the loss can see depends on the data:
   #   offline_data_path  hidden states -> base lm_head -> the FULL distribution,

tests/unit/torch/speculative/test_external_draft_offline.py

@@ -146,6 +146,24 @@ def test_dense_tvd_uses_full_vocabulary_and_the_filter_changes_it():
     assert not torch.allclose(a, b), "the serving filter must change the loss"
 
 
+def test_dense_tvd_ce_is_weighted_tvd_plus_ce_on_the_base_top1():
+    """One reference check pins the CE target, the weights and the dense composition.
+
+    Going through `compute_loss` directly rather than a forward pass: the forward is
+    shared with the other objectives and covered by `test_training_steps_reduce_loss`,
+    while the target choice is what a wrong CE would get silently wrong here.
+    """
+    draft, _ = _external_draft(external_loss="tvd_ce")
+    torch.manual_seed(0)
+    d, b = torch.randn(BATCH, SEQ, VOCAB), torch.randn(BATCH, SEQ, VOCAB)
+    mask = torch.ones(BATCH, SEQ)
+    mixed, _ = draft.compute_loss(d, b, mask)
+    draft.external_loss = "tvd"
+    tvd, _ = draft.compute_loss(d, b, mask)
+    ce = torch.nn.functional.cross_entropy(d.flatten(0, 1), b.argmax(-1).flatten())
+    assert torch.allclose(mixed, 0.9 * tvd + 0.1 * ce, atol=1e-5)
+
+
 def test_prenorm_without_a_norm_raises_rather_than_corrupting():
     """A producer declaring a pre-norm hidden with no norm attached must fail loudly.
 

tests/unit/torch/speculative/test_external_draft_sparse.py

@@ -106,10 +106,9 @@ def test_dataset_masks_response_positions_and_pads_to_train_len(tmp_path):
     assert batch["teacher_topk_prob"].shape == (3, 16, TOPK)
 
 
-# The pure-TVD objectives, named rather than taken from SPARSE_CAPABLE_LOSSES: the
-# reference below is the TVD formula, and a later objective that adds a term to it
-# would silently join this parametrization and fail against the wrong reference.
-@pytest.mark.parametrize("loss_fn", ["tvd", "tvd_deploy"])
+# Named rather than taken from SPARSE_CAPABLE_LOSSES, so a later objective cannot
+# silently join this parametrization and be scored against the wrong reference.
+@pytest.mark.parametrize("loss_fn", ["tvd", "tvd_deploy", "tvd_ce"])
 def test_sparse_tvd_matches_one_minus_acceptance(loss_fn):
     """The sparse path must compute the same acceptance objective as the dense one.
 
@@ -140,8 +139,13 @@ def test_sparse_tvd_matches_one_minus_acceptance(loss_fn):
         else torch.softmax(draft_logits[:, :-1].float(), dim=-1)
     )
     q = q_full.gather(-1, tok[:, 1:])
-    expected = 0.5 * ((p - q).abs().sum(-1) + (1.0 - q.sum(-1)).clamp_min(0.0))
-    assert torch.allclose(loss, expected.mean(), atol=1e-5)
+    expected = 0.5 * ((p - q).abs().sum(-1) + (1.0 - q.sum(-1)).clamp_min(0.0)).mean()
+    if loss_fn == "tvd_ce":
+        # one assertion pins the CE target, the shift and both weights
+        top1 = tok[:, 1:].gather(-1, prob[:, 1:].argmax(-1, keepdim=True)).squeeze(-1)
+        ce = torch.nn.functional.cross_entropy(draft_logits[0, :-1], top1[0])
+        expected = model.external_tvd_alpha * expected + model.external_ce_alpha * ce
+    assert torch.allclose(loss, expected, atol=1e-5)
 
 
 def test_every_sparse_capable_loss_runs_and_backprops(tmp_path):
@@ -405,6 +409,38 @@ def test_warns_when_stored_policy_is_a_truncated_nucleus(tmp_path):
         SparsePolicyDataset([str(p)], max_length=32)
 
 
+def test_tvd_ce_ignores_rows_the_dump_left_empty():
+    """An empty policy row has no argmax; CE must not read that as a target of id 0.
+
+    The ragged fallback zero-fills a short row while ``__getitem__`` still scores the
+    position, and TVD only degrades to its bounded out-of-support term there. CE's
+    penalty is unbounded, so an unguarded target would drag the draft toward token 0.
+    """
+    torch.manual_seed(0)
+    seq = 4
+    tok = torch.randint(1, VOCAB, (1, seq, TOPK))
+    prob = torch.rand(1, seq, TOPK).softmax(-1)
+    prob[0, 2] = 0.0  # the row the dump never filled
+    tok[0, 2] = 0
+    mask = torch.ones(1, seq)
+    m = _draft(external_loss="tvd_ce")
+
+    # neither the loss nor the accuracy may depend on what the draft predicts there.
+    # b spikes on a different id: with both uniform their argmax is also id 0, so the
+    # accuracies would agree even with the guard removed.
+    logits = torch.randn(1, seq, VOCAB)
+    a = logits.clone()
+    a[0, 1] = torch.full((VOCAB,), -30.0).scatter(-1, torch.tensor([0]), 30.0)
+    b = logits.clone()
+    b[0, 1] = torch.full((VOCAB,), -30.0).scatter(-1, torch.tensor([1]), 30.0)
+    # shift: position 2 of the policy is scored against draft position 1
+    la, acc_a = m.compute_sparse_loss(a, tok, prob, mask)
+    lb, acc_b = m.compute_sparse_loss(b, tok, prob, mask)
+    assert torch.allclose(la, lb), f"empty row still steers the CE term: {la} vs {lb}"
+    assert acc_a == acc_b, "empty row still counts toward accuracy"
+    assert torch.isfinite(la)
+
+
 def _write_sourced(path, sources):
     """One record per entry in ``sources``, tagged with that source."""
     torch.manual_seed(0)