feat(speculative): DFlash-family loss terms in config; one teacher with an opt-in softcap

CHANGELOG.rst

@@ -18,6 +18,8 @@ Changelog
 *Speculative Decoding*
 
 - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint.
+- Add ``dflash_loss_terms`` to compose the DFlash-family block objective from weighted loss terms (``ce``, ``kd``, ``tvd``, ``tv`` and ``lambda``), e.g. ``{ce: 0.1, tvd: 0.9}``, for any ``projector_type``, online or offline. Left unset, each variant keeps its own objective.
+- Add ``dflash_teacher_softcap`` to cap the distillation teacher with the base model's ``final_logit_softcapping`` (Gemma family), the way the model caps its own logits.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
@@ -37,6 +39,8 @@ 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, 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``.
+- Online DFlash, DFlash2 and LiLiCorr training now distills text and multimodal batches alike from the base model's ``lm_head`` applied to its final hidden state, as offline training and DSpark already did. The teacher therefore leaves out the logit scaling some bases apply at the ``lm_head`` (``logits_scaling``, ``logit_scale``, ``lm_head_multiplier``: Granite, Cohere, Falcon-H1 and others), which conversion now warns about.
+- On a base that caps its logits with ``final_logit_softcapping`` (Gemma family, including the Gemma 3n and Gemma 4 VLMs), online DFlash, DFlash2 and LiLiCorr training now needs ``dflash_teacher_softcap: true`` to keep a capped teacher. A run resumed from an HF-format checkpoint saved before this release restores that checkpoint's config, in which the switch is off, and so continues uncapped; a run resumed from an FSDP-sharded checkpoint takes the switch from the recipe.
 
 **Deprecations**
 
@@ -55,6 +59,8 @@ Changelog
 - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target.
 - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit.
 - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '<tensor>_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT.
+- Fix online DSpark and Domino training on multimodal batches, which ran the base model's language model on the raw token ids and so dropped the image and video features; they now run the base model's multimodal forward, as DFlash's do. Like DFlash's, their Gemma 3 and PaliGemma image batches reach that forward without ``token_type_ids``, so on transformers 5.5 it raises and on 5.12 and later it attends causally over the image.
+- Fix online DFlash-family training on Gemma 4 image and video batches, which failed in the vision tower because their ``image_position_ids`` and ``video_position_ids`` were not passed on.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^

examples/speculative_decoding/doc/dflash.md

@@ -165,6 +165,8 @@ See [`modelopt_recipes/general/speculative_decoding/dflash.yaml`](../../../model
 | `dflash.dflash_loss_objective` | `dpace` | Position weighting: `decay` (static) or `dpace` (dynamic, see below) |
 | `dflash.dflash_dpace_alpha` | 0.5 | D-PACE smoothing factor in (0, 1]; only used when objective is `dpace` |
 | `dflash.dflash_self_logit_distillation` | true | Use target model logits as soft labels (vs hard CE) |
+| `dflash.dflash_loss_terms` | unset | Block objective as `{term: weight}` over `ce`, `kd`, `tvd`, `tv` and `lambda`, e.g. `{ce: 0.1, tvd: 0.9}`; unset keeps the variant's own |
+| `dflash.dflash_teacher_softcap` | false | Cap the teacher with the base's `final_logit_softcapping` (Gemma family) |
 | `dflash.dflash_mask_token_id` | auto | Token ID for masked positions (see note below) |
 | `dflash.dflash_architecture_config.num_hidden_layers` | 5 | Draft decoder layers |
 | `training.answer_only_loss` | false | Mask loss on non-assistant tokens |

modelopt/torch/speculative/config.py

@@ -16,9 +16,9 @@
 """Configurations for speculative decoding modes."""
 
 from copy import deepcopy
-from typing import Literal
+from typing import Annotated, Literal
 
-from pydantic import model_validator
+from pydantic import Field, model_validator
 
 from modelopt.torch.opt.config import ModeloptBaseConfig, ModeloptField
 
@@ -75,6 +75,17 @@ def _get_dflash_default_config():
     },
 }
 
+# The terms a DFlash-family block objective is built from; plugins/dflash_losses.py implements them.
+DFlashLossTerm = Literal["ce", "kd", "tvd", "tv", "lambda"]
+
+# The per-projector_type selectors of the block objective that dflash_loss_terms replaces.
+_LEGACY_OBJECTIVE_FIELDS = (
+    "dflash_self_logit_distillation",
+    "dflash_lk_loss_type",
+    "dflash_ce_loss_alpha",
+    "dflash_l1_loss_alpha",
+)
+
 
 class DFlashConfig(ModeloptBaseConfig):
     """DFlash config for block-wise parallel speculative decoding."""
@@ -98,10 +109,48 @@ class DFlashConfig(ModeloptBaseConfig):
         default=True, description="Whether to freeze base model during DFlash module training."
     )
 
+    dflash_loss_terms: (
+        dict[DFlashLossTerm, Annotated[float, Field(ge=0.0, allow_inf_nan=False)]] | None
+    ) = ModeloptField(
+        default=None,
+        description=(
+            "The block objective as ``{term: weight}``, summed in the order given, e.g. "
+            "``{ce: 0.1, tvd: 0.9}``. Terms: ``ce`` (cross-entropy on the labels), ``kd`` "
+            "(cross-entropy against the base model's distribution), ``tvd`` (total variation "
+            "to the base model's distribution), ``tv`` (``1 - q(gold)``, total variation to "
+            "the label) and ``lambda`` (``ce`` and ``tv`` blended by acceptance, see "
+            "dflash_lk_ce_scale). Any term works with any projector_type, online and offline; "
+            "``kd`` and ``tvd`` read the base model's distribution, which offline batches must "
+            "then carry. Weights are >= 0 with at least one > 0; a term of weight 0 is "
+            "reported without being trained on. None (default) keeps each projector_type's "
+            "objective: ``kd`` or ``ce`` per dflash_self_logit_distillation for DFlash and "
+            "LiLiCorr, the same or dflash_lk_loss_type's term for DFlash2, "
+            "dflash_ce_loss_alpha * ``ce`` + dflash_l1_loss_alpha * ``tvd`` for DSpark, and "
+            "``ce`` for Domino. Setting it requires those selectors to keep their defaults. "
+            "The terms a variant adds itself (the DFlash2 selector, the LiLiCorr lattice, the "
+            "DSpark confidence head) and Domino's blend of the objective on its corrected and "
+            "backbone logits keep their own settings."
+        ),
+    )
+
     dflash_self_logit_distillation: bool = ModeloptField(
         default=True, description="Whether to use logit distillation from base model."
     )
 
+    dflash_teacher_softcap: bool = ModeloptField(
+        default=False,
+        description=(
+            "Cap the teacher's logits the way the base model caps its own, "
+            "``cap * tanh(logits / cap)`` with ``cap`` the base config's "
+            "``final_logit_softcapping`` (Gemma family). The teacher is the base's ``lm_head`` "
+            "applied to its final hidden state for every batch, text or multimodal, online or "
+            "offline, which leaves that cap out unless this is set. Other scaling a base applies "
+            "at its ``lm_head``, such as ``logits_scaling`` or ``logit_scale``, stays out; "
+            "conversion warns about it. Requires a base config that sets "
+            "``final_logit_softcapping``. The draft's own logits are never capped."
+        ),
+    )
+
     dflash_loss_decay_factor: float = ModeloptField(
         default=0.0,
         description="Gamma for exponential loss decay weighting (paper Eq.4). "
@@ -148,7 +197,7 @@ class DFlashConfig(ModeloptBaseConfig):
     dflash_use_torch_compile: bool = ModeloptField(
         default=False,
         description=(
-            "Whether to torch.compile the draft decoder stack and the DSpark TVD chunk in "
+            "Whether to torch.compile the draft decoder stack and the TVD term's chunk in "
             "training. Compiles once, on the first step, and needs a working Inductor "
             "toolchain."
         ),
@@ -424,18 +473,18 @@ class DFlashConfig(ModeloptBaseConfig):
         default=1.0,
         ge=0.0,
         description=(
-            "DFlash2 only: scale of the CE share in the dflash_lk_loss_type='lambda' "
-            "blend. 1.0 starts the run as pure CE. Ignored for other loss types."
+            "Scale of the CE share in the 'lambda' blend (dflash_lk_loss_type='lambda', or the "
+            "'lambda' term of dflash_loss_terms). 1.0 starts the run as pure CE. Ignored "
+            "without that term."
         ),
     )
 
     dflash_lk_ce_decay: float = ModeloptField(
         default=1.0,
         ge=0.0,
         description=(
-            "DFlash2 only: how fast the CE share decays as acceptance rises in the "
-            "dflash_lk_loss_type='lambda' blend. 0 pins the blend at dflash_lk_ce_scale. "
-            "Ignored for other loss types."
+            "How fast the CE share decays as acceptance rises in the 'lambda' blend. 0 pins "
+            "the blend at dflash_lk_ce_scale. Ignored without that term."
         ),
     )
 
@@ -466,6 +515,28 @@ def _check_dpace_alpha(self) -> "DFlashConfig":
                 )
         return self
 
+    @model_validator(mode="after")
+    def _check_loss_terms(self) -> "DFlashConfig":
+        if self.dflash_loss_terms is None:
+            return self
+        if not any(self.dflash_loss_terms.values()):
+            raise ValueError(
+                f"dflash_loss_terms needs a term of weight > 0, got {self.dflash_loss_terms}."
+            )
+        # Against the defaults, not model_fields_set: main.py converts from model_dump() and
+        # modelopt_state restores from it, which marks every field as set.
+        moved = [
+            name
+            for name in _LEGACY_OBJECTIVE_FIELDS
+            if getattr(self, name) != type(self).model_fields[name].default
+        ]
+        if moved:
+            raise ValueError(
+                f"dflash_loss_terms sets the block objective, so {moved} must keep their "
+                "defaults; express the objective in dflash_loss_terms alone."
+            )
+        return self
+
 
 class MedusaConfig(ModeloptBaseConfig):
     """Medusa config."""

modelopt/torch/speculative/plugins/dflash_losses.py

@@ -42,9 +42,9 @@
 (: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`.
+The divergences a block objective is built from live in :data:`DIVERGENCES`, under the
+names ``dflash_loss_terms`` takes. An objective is an ordered list of
+``(name, weight, term)`` entries, summed left to right by :func:`objective_loss`.
 """
 
 import functools
@@ -54,6 +54,8 @@
 import torch.nn.functional as F
 import torch.utils.checkpoint
 
+from ..config import DFlashLossTerm
+
 __all__ = [
     "DIVERGENCES",
     "DraftLogits",
@@ -387,7 +389,7 @@ def __call__(self, ctx, source):
 
 
 # The closed set of terms an objective is built from, by the name its metric reports.
-DIVERGENCES = {
+DIVERGENCES: dict[DFlashLossTerm, type[_Divergence]] = {
     "ce": _CrossEntropy,
     "kd": _SoftCrossEntropy,
     "tvd": _TeacherTotalVariation,