From a3e154543d74e201e35881ed7228c57dc3485c86 Mon Sep 17 00:00:00 2001 From: zhangwenhe1007 Date: Wed, 8 Jul 2026 00:34:27 +0000 Subject: [PATCH] [DFlash/DSpark] Optional direct-acceptance (LK) loss Adds a differentiable direct-acceptance objective that trains the drafter to maximise expected accepted length E[tau]=sum_k prod_{j<=k} alpha_j directly, instead of only the per-token cross-entropy proxy. It makes the acceptance quantity already logged as the tau_probabilistic metric trainable. - deepspec/modeling/dspark/loss.py: accept_loss term (gated on accept_loss_alpha>0) - deepspec/trainer/dspark_trainer.py: wire accept_loss_alpha (getattr default 0.0) - config/dflash/dflash_qwen3_8b.py: documented accept_loss_alpha knob (default 0.0) - scripts/test_accept_loss.py: CPU unit test (zero-regression, differentiable, rewards matching) Implementation prepared with AI assistance. --- README.md | 2 + config/dflash/dflash_qwen3_8b.py | 5 ++ deepspec/modeling/dspark/loss.py | 44 +++++++++++++- deepspec/trainer/dspark_trainer.py | 1 + scripts/test_accept_loss.py | 95 ++++++++++++++++++++++++++++++ 5 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 scripts/test_accept_loss.py diff --git a/README.md b/README.md index dbb79990..02f43df4 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,8 @@ bash scripts/train/train.sh Hardware: the default configs and scripts assume a single node with 8 GPUs. For fewer GPUs, reduce `CUDA_VISIBLE_DEVICES`. +**Direct-acceptance (LK) loss (optional).** Set `accept_loss_alpha > 0` in a DFlash/DSpark config's `model` block to add a differentiable objective that maximises the drafter's expected accepted length E[tau] = sum_k prod_{j<=k} alpha_j directly, rather than relying only on the per-token cross-entropy proxy. It makes the acceptance quantity the trainer already logs as `tau_probabilistic` trainable; `0.0` (default) keeps prior behaviour. + ## Evaluation diff --git a/config/dflash/dflash_qwen3_8b.py b/config/dflash/dflash_qwen3_8b.py index 979b1d74..34177751 100644 --- a/config/dflash/dflash_qwen3_8b.py +++ b/config/dflash/dflash_qwen3_8b.py @@ -24,6 +24,11 @@ loss_decay_gamma=4.0, ce_loss_alpha=1.0, l1_loss_alpha=0.0, + + # Optional direct-acceptance (LK) loss. >0 trains the drafter to maximise + # expected accepted length E[tau]=sum_k prod_{j<=k} alpha_j directly, + # rather than only the per-token cross-entropy proxy. 0.0 keeps prior behaviour. + accept_loss_alpha=0.0, ) train = dict( diff --git a/deepspec/modeling/dspark/loss.py b/deepspec/modeling/dspark/loss.py index 88dd56ba..3f7b1f38 100644 --- a/deepspec/modeling/dspark/loss.py +++ b/deepspec/modeling/dspark/loss.py @@ -14,7 +14,7 @@ def _all_reduce_loss_denominators( world_size: int, ) -> dict[str, torch.Tensor]: denominators = {} - for key in ("ce_loss_den", "l1_loss_den", "confidence_loss_den"): + for key in ("ce_loss_den", "l1_loss_den", "confidence_loss_den", "accept_loss_den"): tensor = loss_terms[key].detach().clone() if world_size > 1: dist.all_reduce(tensor, op=dist.ReduceOp.SUM) @@ -92,6 +92,7 @@ def _collect_local_terms( outputs: DSparkForwardOutput, loss_decay_gamma: Optional[float], l1_loss_alpha: float, + accept_loss_alpha: float, ) -> tuple[dict[str, torch.Tensor], bool]: draft_logits = outputs.draft_logits target_ids = outputs.target_ids @@ -143,6 +144,21 @@ def _collect_local_terms( valid_block_weights=valid_block_weights, ) + accept_loss_num = zero + accept_loss_den = zero + if accept_loss_alpha > 0: + assert accept_rate_3d is not None, ( + "aligned_target_logits is required when accept_loss_alpha > 0." + ) + # Trainable form of the ``tau_probabilistic`` metric: maximise the + # expected accepted draft length E[accepted] = sum_k prod_{j<=k} alpha_j, + # with alpha the per-position acceptance probability. Negated so that + # minimising the loss maximises accepted length. + valid_accept_rate = accept_rate_3d * eval_mask.to(torch.float32) + expected_draft_accepted = valid_accept_rate.cumprod(dim=-1).sum(dim=-1) + accept_loss_num = -(expected_draft_accepted * valid_block_weights).sum() + accept_loss_den = valid_block_weights.sum() + has_confidence = outputs.confidence_pred is not None confidence_loss_num = zero confidence_loss_den = zero @@ -187,6 +203,8 @@ def _collect_local_terms( "l1_loss_den": l1_loss_den, "confidence_loss_num": confidence_loss_num, "confidence_loss_den": confidence_loss_den, + "accept_loss_num": accept_loss_num, + "accept_loss_den": accept_loss_den, } for pos_idx in range(block_size): @@ -231,6 +249,7 @@ def _build_loss( ce_loss_alpha: float, l1_loss_alpha: float, confidence_head_alpha: float, + accept_loss_alpha: float, has_confidence: bool, world_size: int, ) -> torch.Tensor: @@ -245,10 +264,16 @@ def _build_loss( confidence_loss = loss_terms["confidence_loss_num"] / ( global_denominators["confidence_loss_den"] + 1e-6 ) + accept_loss = ce_loss.new_zeros(()) + if global_denominators["accept_loss_den"].item() > 0: + accept_loss = loss_terms["accept_loss_num"] / ( + global_denominators["accept_loss_den"] + 1e-6 + ) return ( ce_loss_alpha * ce_loss + l1_loss_alpha * l1_loss + confidence_head_alpha * confidence_loss + + accept_loss_alpha * accept_loss ) * world_size @@ -259,11 +284,13 @@ def compute_dspark_loss( ce_loss_alpha: float, l1_loss_alpha: float, confidence_head_alpha: float, + accept_loss_alpha: float = 0.0, ): loss_terms, has_confidence = _collect_local_terms( outputs=outputs, loss_decay_gamma=loss_decay_gamma, l1_loss_alpha=float(l1_loss_alpha), + accept_loss_alpha=float(accept_loss_alpha), ) world_size = dist.get_world_size() global_denominators = _all_reduce_loss_denominators( @@ -273,6 +300,7 @@ def compute_dspark_loss( ce_loss_alpha = float(ce_loss_alpha) l1_loss_alpha = float(l1_loss_alpha) confidence_head_alpha = float(confidence_head_alpha) + accept_loss_alpha = float(accept_loss_alpha) local_ce_loss = loss_terms["ce_loss_num"] / (loss_terms["ce_loss_den"] + 1e-6) local_l1_loss = local_ce_loss.new_zeros(()) @@ -285,10 +313,16 @@ def compute_dspark_loss( local_confidence_loss = loss_terms["confidence_loss_num"] / ( loss_terms["confidence_loss_den"] + 1e-6 ) + local_accept_loss = local_ce_loss.new_zeros(()) + if loss_terms["accept_loss_den"].item() > 0: + local_accept_loss = loss_terms["accept_loss_num"] / ( + loss_terms["accept_loss_den"] + 1e-6 + ) local_loss = ( ce_loss_alpha * local_ce_loss + l1_loss_alpha * local_l1_loss + confidence_head_alpha * local_confidence_loss + + accept_loss_alpha * local_accept_loss ) add_metric( @@ -311,6 +345,13 @@ def compute_dspark_loss( den=loss_terms["confidence_loss_den"], tag="train", ) + if loss_terms["accept_loss_den"].item() > 0: + add_metric( + "accept_loss", + loss_terms["accept_loss_num"], + den=loss_terms["accept_loss_den"], + tag="train", + ) add_metric( "loss", local_loss, @@ -323,6 +364,7 @@ def compute_dspark_loss( ce_loss_alpha=ce_loss_alpha, l1_loss_alpha=l1_loss_alpha, confidence_head_alpha=confidence_head_alpha, + accept_loss_alpha=accept_loss_alpha, has_confidence=has_confidence, world_size=world_size, ) diff --git a/deepspec/trainer/dspark_trainer.py b/deepspec/trainer/dspark_trainer.py index 487e99a2..9e04d158 100644 --- a/deepspec/trainer/dspark_trainer.py +++ b/deepspec/trainer/dspark_trainer.py @@ -35,6 +35,7 @@ def run_batch(self, batch): ce_loss_alpha=float(self.args.model.ce_loss_alpha), l1_loss_alpha=float(self.args.model.l1_loss_alpha), confidence_head_alpha=float(self.args.model.confidence_head_alpha), + accept_loss_alpha=float(getattr(self.args.model, "accept_loss_alpha", 0.0)), ) return loss diff --git a/scripts/test_accept_loss.py b/scripts/test_accept_loss.py new file mode 100644 index 00000000..b2d8f965 --- /dev/null +++ b/scripts/test_accept_loss.py @@ -0,0 +1,95 @@ +"""CPU unit test for the optional direct-acceptance (LK) loss in DSpark. + +Runs without a GPU or a real model: it mocks a ``DSparkForwardOutput`` and checks + (1) accept_loss_alpha=0 is exactly the pre-existing CE(+L1+confidence) loss (zero regression); + (2) the accept term is differentiable (gradient reaches draft_logits); + (3) it rewards a draft that matches the target (lower loss when draft==target logits). + +Usage: python scripts/test_accept_loss.py +""" +import os +import sys + +import torch +import torch.distributed as dist + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +# metrics may accumulate to global rank state; make it a no-op for the unit test. +import deepspec.utils.metrics as _metrics # noqa: E402 +_metrics.add_metric = lambda *a, **k: None +import deepspec.modeling.dspark.loss as L # noqa: E402 +L.add_metric = lambda *a, **k: None +from deepspec.modeling.dspark.common import DSparkForwardOutput # noqa: E402 + + +def _init_dist(): + if not dist.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29591") + dist.init_process_group("gloo", rank=0, world_size=1) + + +def _mock(B=1, A=2, block=4, V=32, draft_eq_target=False, seed=0): + g = torch.Generator().manual_seed(seed) + target_logits = torch.randn(B, A, block, V, generator=g) + draft = target_logits.clone() if draft_eq_target else torch.randn(B, A, block, V, generator=g) + draft = draft.detach().requires_grad_(True) + target_ids = target_logits.argmax(dim=-1) + return draft, DSparkForwardOutput( + draft_logits=draft, + target_ids=target_ids, + eval_mask=torch.ones(B, A, block, dtype=torch.bool), + block_keep_mask=torch.ones(B, A, dtype=torch.bool), + confidence_pred=None, + aligned_target_logits=target_logits, + ) + + +def _loss(out, accept_alpha): + return L.compute_dspark_loss( + outputs=out, loss_decay_gamma=None, + ce_loss_alpha=1.0, l1_loss_alpha=0.0, confidence_head_alpha=0.0, + accept_loss_alpha=accept_alpha, + ) + + +def main(): + _init_dist() + + # (1) zero regression: accept_loss_alpha=0 == the CE-only loss to the bit + _, o0 = _mock(seed=1) + base = _loss(o0, 0.0) + _, o0b = _mock(seed=1) + off = _loss(o0b, 0.0) + assert torch.allclose(base, off), (base, off) + print(f"(1) zero-regression OK (CE-only loss={base.item():.4f})") + + # (2) differentiable: grad reaches draft_logits, and the accept term changes the loss + draft, o1 = _mock(seed=2) + withacc = _loss(o1, 0.5) + _, o1b = _mock(seed=2) + ceonly = _loss(o1b, 0.0) + assert not torch.allclose(withacc, ceonly), "accept term did not change the loss" + withacc.backward() + assert draft.grad is not None and torch.isfinite(draft.grad).all() and draft.grad.abs().sum() > 0 + print(f"(2) differentiable OK (loss +accept={withacc.item():.4f} vs CE-only={ceonly.item():.4f}; grad flows)") + + # (3) rewards matching: a draft equal to target has a LOWER (more negative) accept term + dm, om = _mock(seed=3, draft_eq_target=True) # draft == target + dr, orr = _mock(seed=3, draft_eq_target=False) # random draft + # isolate the accept term (ce_alpha=0, accept_alpha=1) + acc_match = L.compute_dspark_loss(outputs=om, loss_decay_gamma=None, + ce_loss_alpha=0.0, l1_loss_alpha=0.0, + confidence_head_alpha=0.0, accept_loss_alpha=1.0) + acc_rand = L.compute_dspark_loss(outputs=orr, loss_decay_gamma=None, + ce_loss_alpha=0.0, l1_loss_alpha=0.0, + confidence_head_alpha=0.0, accept_loss_alpha=1.0) + assert acc_match.item() < acc_rand.item(), (acc_match.item(), acc_rand.item()) + print(f"(3) rewards matching OK (accept-loss match={acc_match.item():.4f} < random={acc_rand.item():.4f})") + + print("\nALL PASS") + + +if __name__ == "__main__": + main()