[2/2] refactor(speculative): train the DFlash family through one forward with shared loss terms

CHANGELOG.rst

@@ -36,7 +36,7 @@ Changelog
 
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
-- The DFlash2 and LiLiCorr training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics`` and ``lilicorr_metrics``.
+- The DFlash2, LiLiCorr, DSpark and Domino training forwards now return their metrics in one ``draft_metrics`` dict of detached 0-dim tensors (call ``.item()`` for a Python float), replacing ``selector_metrics``, ``lilicorr_metrics``, ``dspark_metrics`` and ``domino_metrics``. DSpark's ``ce_loss``/``l1_loss``/``base_accuracy`` are now ``loss_ce``/``loss_tvd``/``backbone_accuracy``, and Domino's ``final_loss``/``base_loss``/``base_accuracy`` are now ``loss_final``/``loss_backbone``/``backbone_accuracy``.
 
 **Deprecations**
 

modelopt/torch/speculative/plugins/dflash_losses.py

@@ -41,6 +41,10 @@
 A training step's supervision: which token each block position is trained on and where
 (:func:`_block_targets`), how the positions are weighted (:func:`_block_position_weights`),
 and :class:`LossContext`, which carries them with the draft's logits to the loss.
+
+The divergences a block objective is built from live in :data:`DIVERGENCES`. An objective
+is an ordered list of ``(name, weight, term)`` entries, summed left to right by
+:func:`objective_loss`.
 """
 
 import functools
@@ -50,7 +54,14 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 
-__all__ = ["DraftLogits", "LossContext", "LossOutput"]
+__all__ = [
+    "DIVERGENCES",
+    "DraftLogits",
+    "LossContext",
+    "LossOutput",
+    "objective_loss",
+    "objective_term",
+]
 
 
 def _block_positions(anchor_positions, block_size, offset, seq_len):
@@ -228,7 +239,7 @@ class LossContext:
     objective backpropagates through it; otherwise it runs under ``no_grad``.
     """
 
-    def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad):
+    def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad, tvd_chunk):
         """Weigh the step's block positions.
 
         Args:
@@ -238,13 +249,16 @@ def __init__(self, batch, logits, targets, *, weighting, teacher, ce_grad):
             weighting: ``_block_position_weights``' scheme, first, decay_factor and dpace_alpha.
             teacher: Returns the base model's logits for every label position.
             ce_grad: Whether the objective backpropagates through ``ce``.
+            tvd_chunk: The per-chunk kernel ``tvd`` runs, possibly compiled.
         """
         self.batch = batch
         self.logits = logits
         self.targets = targets
         self._teacher_fn = teacher
         self._ce_grad = ce_grad and torch.is_grad_enabled()
+        self._tvd_chunk = tvd_chunk
         self._ce = {}
+        self._tvd = {}
         self._teacher = None
         self.weights = _block_position_weights(
             targets.mask,
@@ -277,3 +291,121 @@ def teacher(self):
         if self._teacher is None:
             self._teacher = self._teacher_fn()
         return self._teacher
+
+    def tvd(self, source):
+        """Per-token ``||softmax(logits.<source>) - softmax(teacher)||_1``: [B * N * bs]."""
+        if source not in self._tvd:
+            logits = self.flat_logits(source)
+            self._tvd[source] = _tvd_per_token(
+                logits, self.teacher().reshape(logits.shape), chunk_fn=self._tvd_chunk
+            )
+        return self._tvd[source]
+
+
+class _Divergence:
+    """A loss term on one logits source: ``term(ctx, source)`` -> 0-dim loss.
+
+    Each term reduces under the step's position weights itself. ``needs_teacher``: it reads
+    the base model's distribution. ``needs_ce_grad``: it backpropagates through ``ctx.ce``.
+    """
+
+    needs_teacher = False
+    needs_ce_grad = False
+
+
+class _CrossEntropy(_Divergence):
+    """``-log q(gold)``: cross-entropy on the labels."""
+
+    needs_ce_grad = True
+
+    def __call__(self, ctx, source):
+        return (ctx.ce(source) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _SoftCrossEntropy(_Divergence):
+    """Cross-entropy against the base model's distribution (logit distillation)."""
+
+    needs_teacher = True
+
+    def __call__(self, ctx, source):
+        logits = ctx.flat_logits(source)
+        target_soft = torch.softmax(ctx.teacher().reshape(logits.shape), dim=-1)
+        draft_logsoft = torch.log_softmax(logits, dim=-1)
+        per_token = -(target_soft * draft_logsoft).sum(dim=-1)
+        return (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _TeacherTotalVariation(_Divergence):
+    """``||softmax(draft) - softmax(teacher)||_1``: the L1 distance to the base model's distribution."""
+
+    needs_teacher = True
+
+    def __call__(self, ctx, source):
+        return (ctx.tvd(source) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _LabelTotalVariation(_Divergence):
+    """``1 - q(gold)``: total variation to the one-hot label, the per-position rejection rate."""
+
+    needs_ce_grad = True
+
+    def __call__(self, ctx, source):
+        gold_probability = torch.exp(-ctx.ce(source))
+        return ((1.0 - gold_probability) * ctx.weights.view(-1)).sum() / ctx.weight_sum
+
+
+class _AcceptanceBlend(_Divergence):
+    """Cross-entropy and label total variation, blended by how often the draft is accepted.
+
+    ::
+
+        q     = exp(-ce)                 draft probability of the gold token
+        L_ce  = <ce>_w                   the cross-entropy term
+        L_tv  = <1 - q>_w                the label total-variation term
+        a     = <q>_mask                 mean acceptance over supervised positions
+        L     = s*exp(-d*a) * L_ce + (1 - s*exp(-d*a)) * L_tv
+
+    ``<.>_w`` averages under the position weights, ``<.>_mask`` under the unweighted
+    supervision mask, as the reported accuracy is. The blend weight is detached, so it
+    reshapes the objective without adding a gradient path of its own.
+    """
+
+    needs_ce_grad = True
+
+    def __init__(self, scale, decay):
+        """Blend with ``s = scale`` and ``d = decay``."""
+        self.scale = scale
+        self.decay = decay
+
+    def __call__(self, ctx, source):
+        tv_loss = _LabelTotalVariation()(ctx, source)
+        ce_loss = _CrossEntropy()(ctx, source)
+        mask = ctx.targets.mask.view(-1)
+        acceptance = (torch.exp(-ctx.ce(source)).detach() * mask).sum() / (mask.sum() + 1e-6)
+        ce_share = self.scale * torch.exp(-self.decay * acceptance)
+        return ce_share * ce_loss + (1.0 - ce_share) * tv_loss
+
+
+# The closed set of terms an objective is built from, by the name its metric reports.
+DIVERGENCES = {
+    "ce": _CrossEntropy,
+    "kd": _SoftCrossEntropy,
+    "tvd": _TeacherTotalVariation,
+    "tv": _LabelTotalVariation,
+    "lambda": _AcceptanceBlend,
+}
+
+
+def objective_term(name, weight, **params):
+    """One objective entry, ``(name, weight, term)``, with ``params`` bound to the term."""
+    return name, weight, DIVERGENCES[name](**params)
+
+
+def objective_loss(objective, ctx, source):
+    """The objective on ``logits.<source>``, and each term's value as ``loss_<name>``.
+
+    Summed left to right from 0, so a single term of weight 1.0 is that term exactly.
+    """
+    values = [(name, weight, term(ctx, source)) for name, weight, term in objective]
+    loss = sum(weight * value for _, weight, value in values)
+    return loss, {f"loss_{name}": value.detach() for name, _, value in values}

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -85,7 +85,16 @@
 from ..dflash.conversion import DFlashDMRegistry
 from ..dflash.dflash_model import DFlashModel
 from .dflash_flex_attention import build_draft_block_mask
-from .dflash_losses import DraftLogits, LossContext, LossOutput, _block_targets, _masked_top1
+from .dflash_losses import (
+    DraftLogits,
+    LossContext,
+    LossOutput,
+    _block_targets,
+    _masked_top1,
+    _tvd_chunk,
+    objective_loss,
+    objective_term,
+)
 from .modeling_dflash import (  # noqa: F401
     DFlashAttention,
     DFlashBaseModelOutput,
@@ -220,7 +229,7 @@ def _base_model_norm(self):
     @property
     def _needs_base_logits(self) -> bool:
         """Whether the loss reads the target's logits; offline training then reconstructs them."""
-        return bool(self.dflash_self_logit_distillation)
+        return any(term.needs_teacher for _, _, term in self._objective)
 
     @property
     def _base_llm_config(self):
@@ -450,10 +459,22 @@ def modify(self, config):
         if self.dflash_use_torch_compile:
             self.dflash_module.compile_body()
 
+        self._objective = self._build_objective(config)
+        # Compiling fuses the TVD chunk's six vocab-wide elementwise passes.
+        self._tvd_chunk_fn = (
+            torch.compile(_tvd_chunk, dynamic=False, fullgraph=True)
+            if self.dflash_use_torch_compile
+            else _tvd_chunk
+        )
+
     def _build_draft_module(self, dflash_config):
         """Build the draft module. Subclasses override to use an augmented module."""
         return DFlashModule(dflash_config)
 
+    def _build_objective(self, config):
+        """The block objective, ``[(name, weight, term)]``: KD from the base model, or CE."""
+        return [objective_term("kd" if config.dflash_self_logit_distillation else "ce", 1.0)]
+
     def _base_device(self):
         """Device the frozen base model sits on. ``meta`` during a from_pretrained restore."""
         if self.dflash_offline:
@@ -900,33 +921,36 @@ def _loss_context(self, batch):
                 "dpace_alpha": self.dflash_dpace_alpha,
             },
             teacher=functools.partial(self._teacher_logits, batch.base_outputs, targets.positions),
-            # The KD path optimizes the soft target; its D-PACE confidences need no gradient.
-            ce_grad=not self.dflash_self_logit_distillation,
+            ce_grad=any(term.needs_ce_grad for _, _, term in self._objective),
+            tvd_chunk=self._tvd_chunk_fn,
         )
 
     def _compute_loss(self, ctx):
-        """The weighted cross-entropy (or KD) loss on the served logits, and its accuracy.
+        """The objective on the served logits and its accuracy.
 
-        A step whose position weights leave nothing to train returns a zero loss and
-        accuracy 0.0.
+        Metrics: each objective term, and the backbone's accuracy when a head sits on it.
+        A step whose position weights leave nothing to train gets ``_zero_output``.
         """
-        final = ctx.logits.final
+        final, backbone = ctx.logits
         if not ctx.trainable:
             # Through the draft's output, so an FSDP unit around the draft still runs its backward.
-            return LossOutput(final.sum() * 0.0, 0.0, {})
-
-        if self.dflash_self_logit_distillation:
-            flat_logits = ctx.flat_logits("final")
-            target_soft = torch.softmax(ctx.teacher().reshape(flat_logits.shape), dim=-1)
-            draft_logsoft = torch.log_softmax(flat_logits, dim=-1)
-            per_token = -(target_soft * draft_logsoft).sum(dim=-1)
-        else:
-            per_token = ctx.ce("final")
-        loss = (per_token * ctx.weights.view(-1)).sum() / ctx.weight_sum
-
+            return self._zero_output(ctx, final.sum() * 0.0)
+        loss, metrics = objective_loss(self._objective, ctx, "final")
         with torch.no_grad():
             accuracy = _masked_top1(final, ctx.targets.token_ids, ctx.targets.mask)
-        return LossOutput(loss, accuracy.item(), {})
+            if backbone is not final:
+                metrics["backbone_accuracy"] = _masked_top1(
+                    backbone, ctx.targets.token_ids, ctx.targets.mask
+                )
+        return LossOutput(loss, accuracy.item(), metrics)
+
+    def _zero_output(self, ctx, loss):
+        """``loss``, a zero, for a step with nothing to train; accuracy and metrics read 0."""
+        zero = torch.zeros((), device=loss.device)
+        names = [f"loss_{name}" for name, _, _ in self._objective]
+        if ctx.logits.backbone is not ctx.logits.final:
+            names.append("backbone_accuracy")
+        return LossOutput(loss, 0.0, dict.fromkeys(names, zero))
 
     def forward(
         self,

modelopt/torch/speculative/plugins/hf_dflash2.py

@@ -18,7 +18,7 @@
 DFlash2 differs from DFlash only in the draft module (grouped dynamic convolutions
 around every sublayer, plus a candidate selector) and in one extra loss term, so
 this wrapper reuses ``HFDFlashModel``'s forward wholesale and overrides just
-:meth:`_compute_loss`.
+:meth:`_build_objective` and :meth:`_compute_loss`.
 
 The convolutions need no supervision of their own: they sit inside the backbone
 and are trained by the backbone loss. The selector does, because at serving time
@@ -39,7 +39,7 @@
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DFlash2DMRegistry
-from .dflash_losses import LossOutput, _block_tokens
+from .dflash_losses import LossOutput, _block_tokens, objective_term
 from .hf_dflash import HFDFlashModel
 from .modeling_dflash2 import DFlash2Module
 
@@ -74,12 +74,9 @@ def modify(self, config):
             )
         super().modify(config)
         self.dflash_selector_loss_alpha = config.dflash_selector_loss_alpha
-        self.dflash_lk_loss_type = config.dflash_lk_loss_type
-        self.dflash_lk_ce_scale = config.dflash_lk_ce_scale
-        self.dflash_lk_ce_decay = config.dflash_lk_ce_decay
-        if self.dflash_lk_loss_type != "ce" and self.dflash_self_logit_distillation:
+        if config.dflash_lk_loss_type != "ce" and self.dflash_self_logit_distillation:
             raise ValueError(
-                f"dflash_lk_loss_type={self.dflash_lk_loss_type!r} needs the draft's "
+                f"dflash_lk_loss_type={config.dflash_lk_loss_type!r} needs the draft's "
                 "probability of the gold token, which the KD path never forms -- it "
                 "optimizes a soft target instead. Set dflash_self_logit_distillation=false, "
                 "or dflash_lk_loss_type='ce' to keep distillation."
@@ -148,35 +145,17 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
         # every training step. The trainer converts them at the logging boundary.
         return loss, accuracy.detach(), coverage.detach()
 
-    def _lk_loss(self, ctx):
-        """Re-weight the block objective between cross-entropy and acceptance.
-
-        Both terms are read off the same per-position cross-entropy the backbone loss
-        already produced, so the target alignment and position weighting are shared::
-
-            q     = exp(-ce)                 draft probability of the gold token
-            L_ce  = <ce>_w                   today's objective
-            L_tv  = <1 - q>_w                total variation to the one-hot target,
-                                             i.e. the per-position acceptance loss
-            a     = <q>_mask                 mean acceptance over supervised positions
-            L     = s*exp(-d*a) * L_ce + (1 - s*exp(-d*a)) * L_tv
-
-        ``<.>_w`` averages under the position weighting, ``<.>_mask`` under the
-        unweighted supervised mask, matching how the reported accuracy is normalized.
-        The blend weight is detached, so it reshapes the objective without adding a
-        gradient path of its own.
-        """
-        ce, weights, weight_sum = ctx.ce("final"), ctx.weights.view(-1), ctx.weight_sum
-        gold_probability = torch.exp(-ce)
-        tv_loss = ((1.0 - gold_probability) * weights).sum() / weight_sum
-        if self.dflash_lk_loss_type == "tv":
-            return tv_loss
-
-        ce_loss = (ce * weights).sum() / weight_sum
-        mask = ctx.targets.mask.view(-1)
-        acceptance = (gold_probability.detach() * mask).sum() / (mask.sum() + 1e-6)
-        ce_share = self.dflash_lk_ce_scale * torch.exp(-self.dflash_lk_ce_decay * acceptance)
-        return ce_share * ce_loss + (1.0 - ce_share) * tv_loss
+    def _build_objective(self, config):
+        """DFlash's objective, or the label TV / CE-TV blend ``dflash_lk_loss_type`` selects."""
+        if config.dflash_lk_loss_type == "tv":
+            return [objective_term("tv", 1.0)]
+        if config.dflash_lk_loss_type == "lambda":
+            return [
+                objective_term(
+                    "lambda", 1.0, scale=config.dflash_lk_ce_scale, decay=config.dflash_lk_ce_decay
+                )
+            ]
+        return super()._build_objective(config)
 
     def _compute_loss(self, ctx):
         """Backbone DFlash loss plus the candidate-selector cross-entropy.
@@ -185,11 +164,8 @@ def _compute_loss(self, ctx):
         comparable; the selector's own accuracy and coverage are added to the metrics.
         """
         out = super()._compute_loss(ctx)
-        loss = out.loss
-        if self.dflash_lk_loss_type != "ce":
-            loss = self._lk_loss(ctx)
         if self.dflash_selector_loss_alpha <= 0:
-            return out._replace(loss=loss)
+            return out
 
         batch, targets = ctx.batch, ctx.targets
         bsz, n_blocks = batch.anchors.shape
@@ -223,5 +199,5 @@ def _compute_loss(self, ctx):
             "selector_coverage": selector_coverage,
         }
         return LossOutput(
-            loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics
+            out.loss + self.dflash_selector_loss_alpha * selector_loss, out.accuracy, metrics
         )