diff --git a/examples/speculative_decoding/README.md b/examples/speculative_decoding/README.md index 943cc6e005f..4d4076c2bef 100644 --- a/examples/speculative_decoding/README.md +++ b/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 diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py index eaaf37b62cb..d6b42f4e1b0 100644 --- a/modelopt/torch/speculative/config.py +++ b/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." ), ) diff --git a/modelopt/torch/speculative/external/external_model.py b/modelopt/torch/speculative/external/external_model.py index 29a8355771d..a0750d72646 100644 --- a/modelopt/torch/speculative/external/external_model.py +++ b/modelopt/torch/speculative/external/external_model.py @@ -32,6 +32,8 @@ 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) @@ -39,5 +41,7 @@ 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 diff --git a/modelopt/torch/speculative/plugins/hf_external.py b/modelopt/torch/speculative/plugins/hf_external.py index e9095a640f8..77f27a4027f 100644 --- a/modelopt/torch/speculative/plugins/hf_external.py +++ b/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. diff --git a/modelopt_recipes/general/speculative_decoding/external_draft.yaml b/modelopt_recipes/general/speculative_decoding/external_draft.yaml index c7415a9faf2..235afe4f18b 100644 --- a/modelopt_recipes/general/speculative_decoding/external_draft.yaml +++ b/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, diff --git a/tests/unit/torch/speculative/test_external_draft_offline.py b/tests/unit/torch/speculative/test_external_draft_offline.py index e0cf35fec0f..752edb75f3f 100644 --- a/tests/unit/torch/speculative/test_external_draft_offline.py +++ b/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. diff --git a/tests/unit/torch/speculative/test_external_draft_sparse.py b/tests/unit/torch/speculative/test_external_draft_sparse.py index c91d9619cc6..e87e09a618b 100644 --- a/tests/unit/torch/speculative/test_external_draft_sparse.py +++ b/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)