From 04b7ce3763d5a1b26177d3c62be58703b39cf2a8 Mon Sep 17 00:00:00 2001
From: mrusanovsky <mrusanovsky@nvidia.com>
Date: Tue, 29 Sep 2026 08:18:57 -0700
Subject: [PATCH 1/4] feat(speculative): LiLiCorr offline/streaming training
 and an optional calibration loss

- LiLiCorr no longer refuses offline/streaming training. HFDFlashModel gains a
  _needs_base_logits property; LiLiCorr overrides it when a term reads the target's
  logits, and the offline path reconstructs them from the captured final hidden state.
- Add dflash_lilicorr_w_cal (default 0.0, off): a cross-entropy from the target's
  distribution renormalized over each slot's candidates onto the reranker's candidate
  distribution, usable in place of the distractor penalty. The penalty's candidate-logit
  gather moves into a helper both terms share.
- Fix DFlash conversion on a target whose config declares rope_theta but leaves it None
  (e.g. a NoPE model): the draft keeps its own RoPE base instead of inheriting None.

Signed-off-by: mrusanovsky <mrusanovsky@nvidia.com>
---
 CHANGELOG.rst                                 |   3 +
 modelopt/torch/speculative/config.py          |  22 +++-
 .../torch/speculative/plugins/hf_dflash.py    |  16 ++-
 .../torch/speculative/plugins/hf_lilicorr.py  | 123 +++++++++++++-----
 .../speculative_decoding/lilicorr.yaml        |   5 +-
 .../speculative/plugins/test_hf_dflash.py     |  13 ++
 .../speculative/plugins/test_hf_lilicorr.py   |  67 +++++++++-
 7 files changed, 204 insertions(+), 45 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 76f31c28d90..6c5e5d28cd1 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -38,6 +38,8 @@ Changelog
 *Speculative Decoding*
 
 - Add the DFlash2 draft variant, selected with ``dflash_architecture_config.projector_type="dflash2"``: DFlash's one-pass parallel backbone plus a grouped dynamic convolution around every attention/MLP sublayer (``conv_kernel_size`` / ``conv_group_size``) and a low-rank candidate selector (``selector_rank`` / ``selector_top_k``, weighted by ``dflash_selector_loss_alpha``). Exported checkpoints declare ``DFlash2DraftModel`` and load in the SGLang/vLLM DFlash2 serving path.
+- LiLiCorr now trains offline and in streaming mode, not only online: the target logits its distractor penalty reads are reconstructed from the captured final hidden state, the same path self-logit distillation uses.
+- Add ``dflash_lilicorr_w_cal`` (default ``0.0``, off): an optional LiLiCorr proposal-calibration term, a cross-entropy from the target's distribution over each slot's candidates onto the reranker's candidate distribution, usable in place of the distractor penalty.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
@@ -105,6 +107,7 @@ Changelog
 **Bug Fixes**
 
 - Fix Megatron unified HF export of MoE models with grouped-GEMM experts when only the experts are quantized (e.g. ``nvfp4_experts_only-*`` recipes): ``hf_quant_config.json`` and the ``quantization_config`` in ``config.json`` were not written, so the quantized experts were served as unquantized weights. Re-export such checkpoints.
+- Fix DFlash conversion on a NoPE target whose config declares ``rope_theta`` without setting it: the draft keeps its own RoPE base instead of inheriting ``None``, which failed when the draft's rotary embedding was built.
 - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled.
 - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and
   projector names instead of restoring the hub layout; deployment loaders could skip those weights.
diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index b8ecf0aa9c9..e5db62d2cde 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -322,7 +322,8 @@ class DFlashConfig(ModeloptBaseConfig):
         description=(
             "LiLiCorr only: absolute weight of the cross-entropy term on the reranker's "
             "per-slot conditional. The objective is "
-            "loss = dflash_loss + w_ce*CE + w_margin*hinge + w_pen*penalty. The weights are "
+            "loss = dflash_loss + w_ce*CE + w_margin*hinge + w_pen*penalty "
+            "(+ w_cal*calibration, see dflash_lilicorr_w_cal). The weights are "
             "absolute — there is no outer multiplier scaling the three terms as a group — so "
             "each is the coefficient with which its term enters the total, and "
             "`loss == origin_loss + lilicorr_loss` holds exactly. Both halves and all three "
@@ -369,8 +370,23 @@ class DFlashConfig(ModeloptBaseConfig):
             "reranker's expected target-rejection over its own candidate distribution. Each "
             "competing candidate is weighted by the target model's logit gap to the ground "
             "truth, so candidates the target finds plausible are penalized lightly and "
-            "confident wrong ones hard. Requires the target's logits, hence online training "
-            "(dflash_offline=False). Both shipped variants use 0.25. "
+            "confident wrong ones hard. Requires the target's logits: online training reads "
+            "them directly, offline training reconstructs them from the captured final hidden "
+            "state. Both shipped variants use 0.25. "
+            "Ignored unless dflash_architecture_config.projector_type == 'lilicorr'."
+        ),
+    )
+
+    dflash_lilicorr_w_cal: float = ModeloptField(
+        default=0.0,
+        ge=0.0,
+        allow_inf_nan=False,
+        description=(
+            "LiLiCorr only: absolute weight of the optional calibration term, a cross-entropy "
+            "from the target's distribution renormalized over each slot's k candidates onto "
+            "the reranker's. An alternative to w_pen whose gradient does not scale with the "
+            "target's logit gap. Requires the target's logits, like w_pen. 0.0 (default) is a "
+            "no-op, and it is not part of the three weights' all-or-nothing validation. "
             "Ignored unless dflash_architecture_config.projector_type == 'lilicorr'."
         ),
     )
diff --git a/modelopt/torch/speculative/plugins/hf_dflash.py b/modelopt/torch/speculative/plugins/hf_dflash.py
index 084040a5e06..629275a76b6 100644
--- a/modelopt/torch/speculative/plugins/hf_dflash.py
+++ b/modelopt/torch/speculative/plugins/hf_dflash.py
@@ -241,6 +241,11 @@ def _base_model_norm(self):
         path = getattr(self, "base_model_norm_path", None)
         return self.get_submodule(path) if path else None
 
+    @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)
+
     @property
     def _base_llm_config(self):
         return (
@@ -458,6 +463,11 @@ def modify(self, config):
                 base_val = getattr(base_config, attr)
             else:
                 continue
+            # A NoPE target's config class can declare rope_theta without the checkpoint
+            # setting it. There is no target rotary to align to, so the draft keeps its own
+            # rather than inheriting None as the RoPE base.
+            if base_val is None:
+                continue
             user_val = getattr(self.dflash_config, attr, None)
             if user_val is not None and user_val != base_val:
                 logger.warning(
@@ -1027,14 +1037,14 @@ def forward(
         # 1. Run base model → extract target hidden states
         if self.dflash_offline:
             assert "base_model_outputs" in kwargs
-            # For self-logit-distillation, from_offline_dict reconstructs base logits from the
-            # captured hidden (final norm re-applied as needed) when the producer didn't supply
+            # When the loss needs them (see _needs_base_logits), from_offline_dict reconstructs
+            # base logits from the captured hidden (final norm re-applied as needed) when the producer didn't supply
             # them, and raises if anything needed for that is missing.
             base_outputs = DFlashBaseModelOutput.from_offline_dict(
                 kwargs["base_model_outputs"],
                 self._base_model_norm,
                 self._base_model_lm_head,
-                need_logits=self.dflash_self_logit_distillation,
+                need_logits=self._needs_base_logits,
             )
             target_hidden = base_outputs.target_hidden
         else:
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index 6733fd3896a..13a171c65c4 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -47,9 +47,9 @@
 between them, then commits a path greedily. This module trains that scorer jointly
 with the backbone.
 
-The objective adds three weighted terms to the DFlash loss::
+The objective adds weighted terms to the DFlash loss::
 
-    loss = dflash_loss + w_ce * L_ce + w_margin * L_margin + w_pen * L_pen
+    loss = dflash_loss + w_ce * L_ce + w_margin * L_margin + w_pen * L_pen + w_cal * L_cal
 
 - ``L_ce``: a softmax over each slot's ``k`` candidate scores, pushing the
   ground-truth candidate up and the competing ones down.
@@ -59,6 +59,9 @@
 - ``L_pen``: the head's own probability mass on the competing candidates, each
   weighted by the target model's logit gap to the ground truth — so a candidate the
   target finds plausible is penalized lightly and a confident wrong one hard.
+- ``L_cal`` (optional, off by default): cross-entropy from the target's distribution
+  renormalized over the ``k`` candidates onto the head's. An alternative to ``L_pen``
+  whose gradient does not scale with the target's logit gap.
 
 The weights are absolute, with no outer multiplier, so
 ``loss == origin_loss + lilicorr_loss`` holds exactly. Two variants ship and differ
@@ -105,13 +108,6 @@ def _build_draft_module(self, dflash_config):
 
     def modify(self, config):
         """Initialize the LiLiCorr draft module and resolve the objective weights."""
-        if config.dflash_offline:
-            raise ValueError(
-                "LiLiCorr (projector_type='lilicorr') requires online training. Its "
-                "distractor penalty weights every competing candidate by the target "
-                "model's own logit gap to the ground truth, and offline mode has no "
-                "target model to read that from. Use data.mode=online."
-            )
         # The grouped convolution is OPTIONAL here (unlike DFlash2, which requires it),
         # but the two geometry keys are all-or-nothing. Setting one alone would build a
         # draft with no convolutions at all and report nothing, and the serving loader
@@ -141,11 +137,12 @@ def modify(self, config):
         self.dflash_lilicorr_w_ce = weights["dflash_lilicorr_w_ce"]
         self.dflash_lilicorr_w_margin = weights["dflash_lilicorr_w_margin"]
         self.dflash_lilicorr_w_pen = weights["dflash_lilicorr_w_pen"]
-        if not any(weights.values()):
+        self.dflash_lilicorr_w_cal = float(getattr(config, "dflash_lilicorr_w_cal", 0.0))
+        if not any(weights.values()) and self.dflash_lilicorr_w_cal == 0.0:
             raise ValueError(
-                "All three LiLiCorr objective weights are 0, so the reranker would never "
+                "All LiLiCorr objective weights are 0, so the reranker would never "
                 "receive a gradient and would be exported randomly initialized. Set at "
-                "least one of dflash_lilicorr_w_ce / _w_margin / _w_pen above 0."
+                "least one of dflash_lilicorr_w_ce / _w_margin / _w_pen / _w_cal above 0."
             )
 
         self.dflash_lilicorr_margin = float(getattr(config, "dflash_lilicorr_margin", -1.0))
@@ -161,9 +158,9 @@ def modify(self, config):
         logger.info(
             "LiLiCorr enabled: candidate_topk=%d, factor_dim=%d, num_layers=%d, "
             "num_heads=%d, logit_scale=%.3g, vector_eps=%.3g | objective: "
-            "w_ce=%.4g, w_margin=%.4g, w_pen=%.4g, margin=%.4g. Total loss is "
-            "dflash_loss + w_ce*CE + w_margin*hinge + w_pen*penalty, with no outer "
-            "multiplier.",
+            "w_ce=%.4g, w_margin=%.4g, w_pen=%.4g, w_cal=%.4g, margin=%.4g. Total loss is "
+            "dflash_loss + w_ce*CE + w_margin*hinge + w_pen*penalty + w_cal*calibration, "
+            "with no outer multiplier.",
             head.candidate_topk,
             head.factor_dim,
             len(head.layers),
@@ -173,10 +170,20 @@ def modify(self, config):
             self.dflash_lilicorr_w_ce,
             self.dflash_lilicorr_w_margin,
             self.dflash_lilicorr_w_pen,
+            self.dflash_lilicorr_w_cal,
             self.dflash_lilicorr_margin,
         )
         self._lilicorr_metrics = None
 
+    @property
+    def _needs_base_logits(self) -> bool:
+        """The distractor penalty and the calibration term both read the target's logits."""
+        return (
+            super()._needs_base_logits
+            or self.dflash_lilicorr_w_pen > 0.0
+            or self.dflash_lilicorr_w_cal > 0.0
+        )
+
     def get_exporter(self):
         """Get the exporter for the LiLiCorr draft model."""
         from modelopt.torch.export.plugins.hf_spec_export import LiLiCorrExporter
@@ -241,7 +248,7 @@ def _compute_loss(
         accuracy, which is reported as ``origin_accuracy`` in the metrics instead.
 
         The two tensors this needs beyond the shared signature -- the target-layer hidden
-        states it anchors on, and the target's logits its distractor penalty weights by --
+        states it anchors on, and the target's logits its penalty and calibration terms read --
         both come from ``base_outputs``, the container ``HFDFlashModel.forward`` already
         builds for them.
         """
@@ -298,7 +305,7 @@ def _compute_lilicorr_loss(
         target_hidden,
         target_logits,
     ):
-        """Score the candidate lattice and evaluate the three-term objective.
+        """Score the candidate lattice and evaluate the LiLiCorr objective.
 
         Returns ``(loss, accuracy, metrics)``. ``loss`` already carries the absolute
         term weights, ``accuracy`` is the greedy path-prefix fraction, and ``metrics``
@@ -441,6 +448,18 @@ def _compute_lilicorr_loss(
             )
             loss = loss + self.dflash_lilicorr_w_pen * penalty_loss
 
+        cal_loss = ce_loss.new_zeros(())
+        if self.dflash_lilicorr_w_cal > 0.0:
+            cal_loss = self._calibration_loss(
+                node_potentials=node_potentials,
+                candidate_ids=candidate_ids,
+                anchor_positions=anchor_positions,
+                target_logits=target_logits,
+                supervised=supervised,
+                denominator=denominator,
+            )
+            loss = loss + self.dflash_lilicorr_w_cal * cal_loss
+
         raw_metrics = self._lattice_metrics(
             log_start=log_start,
             log_pair=log_pair,
@@ -459,6 +478,7 @@ def _compute_lilicorr_loss(
         raw_metrics["lilicorr_ce"] = ce_loss
         raw_metrics["lilicorr_margin"] = margin_loss
         raw_metrics["lilicorr_penalty"] = penalty_loss
+        raw_metrics["lilicorr_calibration"] = cal_loss
 
         # One device sync for every scalar, rather than one per metric.
         names = list(raw_metrics)
@@ -469,6 +489,7 @@ def _compute_lilicorr_loss(
         metrics["lilicorr_w_ce"] = self.dflash_lilicorr_w_ce
         metrics["lilicorr_w_margin"] = self.dflash_lilicorr_w_margin
         metrics["lilicorr_w_pen"] = self.dflash_lilicorr_w_pen
+        metrics["lilicorr_w_cal"] = self.dflash_lilicorr_w_cal
         accuracy = metrics["lilicorr_selected_prefix"] / float(num_slots)
         return loss, accuracy, metrics
 
@@ -492,20 +513,40 @@ def _distractor_penalty(
         ``p(j) * (w_j - E[w])`` pushes the ground truth up, the target-rejected
         confuser down, and defers on genuine ties.
         """
+        # p_head(j | slot): the same per-slot distribution the CE term scores, reused here
+        # as the weights of the expectation.
+        potentials = torch.stack(node_potentials, dim=1)
+        head_probs = F.log_softmax(potentials.float(), dim=-1).exp()
+
+        candidate_target_logits = self._candidate_target_logits(
+            candidate_ids=candidate_ids,
+            anchor_positions=anchor_positions,
+            target_logits=target_logits,
+            requested_by="dflash_lilicorr_w_pen",
+        )
+        # w_j is 0 at the ground truth itself, and 0 for any candidate the target scores
+        # at least as highly — those are not distractors, they are alternatives.
+        gt_target_logit = candidate_target_logits.gather(2, gt_indices.unsqueeze(-1))
+        rejection_weight = (gt_target_logit - candidate_target_logits).clamp_min(0.0)
+
+        # E_{p_head}[w] per slot, averaged over the supervised slots only.
+        per_slot = (head_probs * rejection_weight).sum(dim=-1)
+        return (per_slot * supervised).sum() / denominator
+
+    def _candidate_target_logits(
+        self, *, candidate_ids, anchor_positions, target_logits, requested_by
+    ):
+        """The target's logits for exactly the candidates in play: ``[blocks, slots, k]``."""
         if target_logits is None:
             raise ValueError(
-                "dflash_lilicorr_w_pen > 0 requires the target model's logits, which are "
-                "only available in online training (dflash_offline=False)."
+                f"{requested_by} > 0 requires the target model's logits. Online training "
+                "reads them from the target; offline training reconstructs them from the "
+                "captured final hidden state."
             )
         batch_blocks, num_slots, topk = candidate_ids.shape
         bsz, n_blocks = anchor_positions.shape
         device = candidate_ids.device
 
-        # p_head(j | slot): the same per-slot distribution the CE term scores, reused here
-        # as the weights of the expectation.
-        potentials = torch.stack(node_potentials, dim=1)
-        head_probs = F.log_softmax(potentials.float(), dim=-1).exp()
-
         target_seq_len = target_logits.shape[1]
         # The target's next-token logits that predict the token at anchor+1+s sit at
         # position anchor+s. `sample_index` maps each flattened block back to its batch
@@ -523,16 +564,34 @@ def _distractor_penalty(
         )
         # Advanced indexing reads only the [blocks, slots, k] scalars in play rather
         # than gathering whole vocabulary rows. Teacher weights, hence detached.
-        candidate_target_logits = (
-            target_logits[sample_expanded, positions_expanded, candidate_ids].detach().float()
+        return target_logits[sample_expanded, positions_expanded, candidate_ids].detach().float()
+
+    def _calibration_loss(
+        self,
+        *,
+        node_potentials,
+        candidate_ids,
+        anchor_positions,
+        target_logits,
+        supervised,
+        denominator,
+    ):
+        """Cross-entropy ``H(p|K, q)`` from the target's candidate distribution to the head's.
+
+        ``p|K`` is the target renormalized over the ``k`` candidates, so the gathered
+        candidate logits suffice and the full-vocabulary normalizer is not needed.
+        """
+        candidate_target_logits = self._candidate_target_logits(
+            candidate_ids=candidate_ids,
+            anchor_positions=anchor_positions,
+            target_logits=target_logits,
+            requested_by="dflash_lilicorr_w_cal",
         )
-        # w_j is 0 at the ground truth itself, and 0 for any candidate the target scores
-        # at least as highly — those are not distractors, they are alternatives.
-        gt_target_logit = candidate_target_logits.gather(2, gt_indices.unsqueeze(-1))
-        rejection_weight = (gt_target_logit - candidate_target_logits).clamp_min(0.0)
+        target_probs = F.softmax(candidate_target_logits, dim=-1)
+        potentials = torch.stack(node_potentials, dim=1)
+        head_log_probs = F.log_softmax(potentials.float(), dim=-1)
 
-        # E_{p_head}[w] per slot, averaged over the supervised slots only.
-        per_slot = (head_probs * rejection_weight).sum(dim=-1)
+        per_slot = -(target_probs * head_log_probs).sum(dim=-1)
         return (per_slot * supervised).sum() / denominator
 
     @torch.no_grad()
diff --git a/modelopt_recipes/general/speculative_decoding/lilicorr.yaml b/modelopt_recipes/general/speculative_decoding/lilicorr.yaml
index f96fa44eeef..d5c67c5f058 100644
--- a/modelopt_recipes/general/speculative_decoding/lilicorr.yaml
+++ b/modelopt_recipes/general/speculative_decoding/lilicorr.yaml
@@ -11,7 +11,8 @@
 # The objective is the DFlash block loss plus three weighted terms:
 #   loss = dflash_loss + w_ce*CE + w_margin*hinge + w_pen*penalty
 # No outer multiplier, so `loss == origin_loss + lilicorr_loss` holds exactly.
-# Online training is required (the penalty reads the target model's logits).
+# The penalty reads the target model's logits: online training has them directly,
+# offline/streaming training reconstructs them from the captured final hidden state.
 #
 # This file is the published `base` variant. The `margin` variant differs only in
 # the composition of the cross-entropy weight — see the two-line override below.
@@ -29,8 +30,6 @@ model:
 
 # maps to DataArguments (main.py)
 data:
-  # Online only: the distractor penalty weights each candidate by the target
-  # model's own logit gap, so there has to be a target model in the process.
   mode: online
   data_path:
   # Jinja chat template with {% generation %} tags for answer_only_loss.
diff --git a/tests/unit/torch/speculative/plugins/test_hf_dflash.py b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
index 76050c0e568..47db8871e11 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_dflash.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_dflash.py
@@ -111,6 +111,19 @@ def test_convert_sets_mask_token_id(self):
         assert hasattr(model, "mask_token_id")
         assert model.mask_token_id == 0
 
+    def test_convert_nope_target_keeps_the_draft_rope_theta(self):
+        """A target config that declares rope_theta but leaves it unset is not inherited."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        model.config.rope_theta = None
+        if isinstance(getattr(model.config, "rope_parameters", None), dict):
+            model.config.rope_parameters.pop("rope_theta", None)
+        mtsp.convert(model, [("dflash", get_dflash_config())])
+        draft = model.dflash_config
+        draft_rope_params = getattr(draft, "rope_parameters", None) or {}
+        assert draft_rope_params.get("rope_theta", getattr(draft, "rope_theta", None)) is not None
+        model.train()
+        assert torch.isfinite(model(**_dflash_batch(model.config.vocab_size)).loss)
+
 
 def test_qwen3_vl_transformers_530_position_ids_expand_video_grid(monkeypatch):
     """Only mRoPE receives a per-frame video grid on Transformers 5.3.0."""
diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
index a21a8bdecf4..ab1bb6380e1 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
@@ -222,10 +222,22 @@ def test_hinge_without_margin_raises(self):
         with pytest.raises(ValueError, match="dflash_lilicorr_margin"):
             _converted(variant="margin", margin=0.0)
 
-    def test_offline_rejected(self):
-        """The distractor penalty needs a target model, which offline mode does not have."""
-        with pytest.raises(ValueError, match="requires online training"):
-            _converted(dflash_offline=True)
+    @pytest.mark.parametrize(
+        ("weights", "needs_logits"),
+        [
+            ({}, True),
+            ({"dflash_lilicorr_w_pen": 0.0}, False),
+            ({"dflash_lilicorr_w_pen": 0.0, "dflash_lilicorr_w_cal": 0.5}, True),
+        ],
+        ids=["penalty", "no_target_terms", "calibration_only"],
+    )
+    def test_offline_reconstructs_target_logits_only_when_needed(self, weights, needs_logits):
+        """Offline converts; target logits are rebuilt only for terms that read them."""
+        model = get_tiny_llama(num_hidden_layers=4)
+        # Set by the offline loader, which a directly constructed model never went through.
+        model.config.num_orig_hidden_layers = model.config.num_hidden_layers
+        mtsp.convert(model, [("dflash", _get_lilicorr_config(dflash_offline=True, **weights))])
+        assert model._needs_base_logits is needs_logits
 
     def test_topk_above_vocab_rejected(self):
         """The lattice cannot be wider than the vocabulary."""
@@ -305,6 +317,53 @@ def test_penalty_term_is_active(self):
         metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics
         assert metrics["lilicorr_penalty"] > 0.0
 
+    @pytest.mark.parametrize("w_cal", [0.0, 0.5])
+    def test_calibration_enters_the_loss_with_its_weight(self, w_cal):
+        """Off by default and an exact no-op at 0; otherwise one more absolute term."""
+        model = _converted(dflash_lilicorr_w_cal=w_cal)
+        model.train()
+        metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics
+        assert metrics["lilicorr_w_cal"] == w_cal
+        assert (metrics["lilicorr_calibration"] > 0.0) is (w_cal > 0.0)
+        weights = VARIANTS["base"]
+        assert metrics["lilicorr_loss"] == pytest.approx(
+            weights["w_ce"] * metrics["lilicorr_ce"]
+            + weights["w_pen"] * metrics["lilicorr_penalty"]
+            + w_cal * metrics["lilicorr_calibration"],
+            abs=1e-5,
+        )
+
+    def test_calibration_loss_is_minimal_when_the_head_matches_the_target(self):
+        """Zero gradient when the reranker's candidate distribution already matches the target's."""
+        model = _converted(dflash_lilicorr_w_cal=0.5)
+        torch.manual_seed(0)
+        num_slots = BLOCK_SIZE - 1
+        target_logits = torch.randn(1, SEQ_LEN, VOCAB_SIZE)
+        candidate_ids = torch.stack(
+            [torch.randperm(VOCAB_SIZE)[:CANDIDATE_TOPK] for _ in range(2 * num_slots)]
+        ).view(2, num_slots, CANDIDATE_TOPK)
+        inputs = {
+            "candidate_ids": candidate_ids,
+            "anchor_positions": torch.tensor([[0, BLOCK_SIZE]]),
+            "target_logits": target_logits,
+        }
+        at_target = model._candidate_target_logits(**inputs, requested_by="test")
+
+        def grad_at(potentials):
+            potentials = potentials.clone().requires_grad_(True)
+            supervised = torch.ones(2, num_slots)
+            model._calibration_loss(
+                node_potentials=list(potentials.unbind(1)),
+                supervised=supervised,
+                denominator=supervised.sum(),
+                **inputs,
+            ).backward()
+            return potentials.grad
+
+        # A per-slot shift leaves the softmax, and so the optimum, unchanged.
+        assert grad_at(at_target + 3.0).abs().max() < 1e-6
+        assert grad_at(torch.zeros_like(at_target)).abs().max() > 1e-3
+
     @pytest.mark.parametrize(
         ("kwargs", "degenerate"),
         [

From 3a4caf6cd7c34df2b351e9fd2792bd3784f3b3dc Mon Sep 17 00:00:00 2001
From: mrusanovsky <mrusanovsky@nvidia.com>
Date: Tue, 29 Sep 2026 08:52:15 -0700
Subject: [PATCH 2/4] review: state the NoPE fix by outcome; cover every
 candidate in the calibration loss test

Signed-off-by: mrusanovsky <mrusanovsky@nvidia.com>
---
 CHANGELOG.rst                                            | 2 +-
 tests/unit/torch/speculative/plugins/test_hf_lilicorr.py | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 6c5e5d28cd1..2e4514a9cc7 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -107,7 +107,7 @@ Changelog
 **Bug Fixes**
 
 - Fix Megatron unified HF export of MoE models with grouped-GEMM experts when only the experts are quantized (e.g. ``nvfp4_experts_only-*`` recipes): ``hf_quant_config.json`` and the ``quantization_config`` in ``config.json`` were not written, so the quantized experts were served as unquantized weights. Re-export such checkpoints.
-- Fix DFlash conversion on a NoPE target whose config declares ``rope_theta`` without setting it: the draft keeps its own RoPE base instead of inheriting ``None``, which failed when the draft's rotary embedding was built.
+- Fix DFlash conversion on NoPE targets whose config leaves ``rope_theta`` unset.
 - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled.
 - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and
   projector names instead of restoring the hub layout; deployment loaders could skip those weights.
diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
index ab1bb6380e1..bf2fc8c256d 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
@@ -320,7 +320,9 @@ def test_penalty_term_is_active(self):
     @pytest.mark.parametrize("w_cal", [0.0, 0.5])
     def test_calibration_enters_the_loss_with_its_weight(self, w_cal):
         """Off by default and an exact no-op at 0; otherwise one more absolute term."""
-        model = _converted(dflash_lilicorr_w_cal=w_cal)
+        # Every token is a candidate, so the ground truth is always in the lattice and the
+        # term has supervised slots to be positive on.
+        model = _converted(topk=VOCAB_SIZE, dflash_lilicorr_w_cal=w_cal)
         model.train()
         metrics = model(**_make_batch(model.dflash_config.vocab_size)).lilicorr_metrics
         assert metrics["lilicorr_w_cal"] == w_cal

From 246ae90788221962b92c26360e583ad4017d38ec Mon Sep 17 00:00:00 2001
From: mrusanovsky <mrusanovsky@nvidia.com>
Date: Wed, 30 Sep 2026 12:07:39 -0700
Subject: [PATCH 3/4] review: offline logit reconstruction casts to the lm_head
 dtype; log calibration as a KL

Cast the captured final hidden state to the lm_head weight dtype before reconstructing target logits offline, cover the offline path with a forward test that matches the online loss, and report the calibration term as KL(p|K || q), whose gradient is identical to the cross-entropy's.

Signed-off-by: mrusanovsky <mrusanovsky@nvidia.com>
---
 CHANGELOG.rst                                 |  3 +-
 modelopt/torch/speculative/config.py          |  4 +-
 .../torch/speculative/plugins/hf_lilicorr.py  | 13 ++--
 .../speculative/plugins/modeling_dflash.py    |  3 +-
 .../speculative/plugins/test_hf_lilicorr.py   | 65 +++++++++++++++++--
 5 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 2e4514a9cc7..644fd7a8040 100755
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -39,7 +39,7 @@ Changelog
 
 - Add the DFlash2 draft variant, selected with ``dflash_architecture_config.projector_type="dflash2"``: DFlash's one-pass parallel backbone plus a grouped dynamic convolution around every attention/MLP sublayer (``conv_kernel_size`` / ``conv_group_size``) and a low-rank candidate selector (``selector_rank`` / ``selector_top_k``, weighted by ``dflash_selector_loss_alpha``). Exported checkpoints declare ``DFlash2DraftModel`` and load in the SGLang/vLLM DFlash2 serving path.
 - LiLiCorr now trains offline and in streaming mode, not only online: the target logits its distractor penalty reads are reconstructed from the captured final hidden state, the same path self-logit distillation uses.
-- Add ``dflash_lilicorr_w_cal`` (default ``0.0``, off): an optional LiLiCorr proposal-calibration term, a cross-entropy from the target's distribution over each slot's candidates onto the reranker's candidate distribution, usable in place of the distractor penalty.
+- Add ``dflash_lilicorr_w_cal`` (default ``0.0``, off): an optional LiLiCorr proposal-calibration term, a KL divergence from the target's distribution over each slot's candidates to the reranker's candidate distribution, usable in place of the distractor penalty.
 
 *Megatron Framework (M-LM / M-Bridge)*
 
@@ -108,6 +108,7 @@ Changelog
 
 - Fix Megatron unified HF export of MoE models with grouped-GEMM experts when only the experts are quantized (e.g. ``nvfp4_experts_only-*`` recipes): ``hf_quant_config.json`` and the ``quantization_config`` in ``config.json`` were not written, so the quantized experts were served as unquantized weights. Re-export such checkpoints.
 - Fix DFlash conversion on NoPE targets whose config leaves ``rope_theta`` unset.
+- Fix offline DFlash training failing to reconstruct the target logits when the captured hidden states are stored in a different dtype than the target's weights.
 - Fix Megatron-Core checkpoint saving for quantized grouped MoE experts when tensor and expert parallelism are both enabled.
 - Fix unified HuggingFace export of RADIO-based VLMs retaining post-conversion vision and
   projector names instead of restoring the hub layout; deployment loaders could skip those weights.
diff --git a/modelopt/torch/speculative/config.py b/modelopt/torch/speculative/config.py
index e5db62d2cde..c3f42a5ad14 100644
--- a/modelopt/torch/speculative/config.py
+++ b/modelopt/torch/speculative/config.py
@@ -382,8 +382,8 @@ class DFlashConfig(ModeloptBaseConfig):
         ge=0.0,
         allow_inf_nan=False,
         description=(
-            "LiLiCorr only: absolute weight of the optional calibration term, a cross-entropy "
-            "from the target's distribution renormalized over each slot's k candidates onto "
+            "LiLiCorr only: absolute weight of the optional calibration term, a KL divergence "
+            "from the target's distribution renormalized over each slot's k candidates to "
             "the reranker's. An alternative to w_pen whose gradient does not scale with the "
             "target's logit gap. Requires the target's logits, like w_pen. 0.0 (default) is a "
             "no-op, and it is not part of the three weights' all-or-nothing validation. "
diff --git a/modelopt/torch/speculative/plugins/hf_lilicorr.py b/modelopt/torch/speculative/plugins/hf_lilicorr.py
index 13a171c65c4..b6d8b2f2d8c 100644
--- a/modelopt/torch/speculative/plugins/hf_lilicorr.py
+++ b/modelopt/torch/speculative/plugins/hf_lilicorr.py
@@ -59,8 +59,8 @@
 - ``L_pen``: the head's own probability mass on the competing candidates, each
   weighted by the target model's logit gap to the ground truth — so a candidate the
   target finds plausible is penalized lightly and a confident wrong one hard.
-- ``L_cal`` (optional, off by default): cross-entropy from the target's distribution
-  renormalized over the ``k`` candidates onto the head's. An alternative to ``L_pen``
+- ``L_cal`` (optional, off by default): KL divergence from the target's distribution
+  renormalized over the ``k`` candidates to the head's. An alternative to ``L_pen``
   whose gradient does not scale with the target's logit gap.
 
 The weights are absolute, with no outer multiplier, so
@@ -576,10 +576,12 @@ def _calibration_loss(
         supervised,
         denominator,
     ):
-        """Cross-entropy ``H(p|K, q)`` from the target's candidate distribution to the head's.
+        """``KL(p|K || q)`` from the target's candidate distribution to the head's.
 
         ``p|K`` is the target renormalized over the ``k`` candidates, so the gathered
-        candidate logits suffice and the full-vocabulary normalizer is not needed.
+        candidate logits suffice and the full-vocabulary normalizer is not needed. The
+        gradient is the cross-entropy's, and the logged value floors at 0 rather than at
+        ``H(p|K)``, which varies with the data and the temperature.
         """
         candidate_target_logits = self._candidate_target_logits(
             candidate_ids=candidate_ids,
@@ -588,10 +590,11 @@ def _calibration_loss(
             requested_by="dflash_lilicorr_w_cal",
         )
         target_probs = F.softmax(candidate_target_logits, dim=-1)
+        target_log_probs = F.log_softmax(candidate_target_logits, dim=-1)
         potentials = torch.stack(node_potentials, dim=1)
         head_log_probs = F.log_softmax(potentials.float(), dim=-1)
 
-        per_slot = -(target_probs * head_log_probs).sum(dim=-1)
+        per_slot = (target_probs * (target_log_probs - head_log_probs)).sum(dim=-1)
         return (per_slot * supervised).sum() / denominator
 
     @torch.no_grad()
diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index 739b3157dca..c6c4edbdef7 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -149,7 +149,8 @@ def from_offline_dict(
             if out_hiddens is None:
                 raise KeyError("base_model_hidden_states")
             out_hiddens = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm)
-            logits = base_model_lm_head(out_hiddens)
+            # A producer can store the hidden states in a wider dtype than the target's weights.
+            logits = base_model_lm_head(out_hiddens.to(base_model_lm_head.weight.dtype))
         return cls(
             target_hidden=d["aux_hidden_states"],
             logits=logits,
diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
index bf2fc8c256d..515a483a3fe 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
@@ -336,7 +336,7 @@ def test_calibration_enters_the_loss_with_its_weight(self, w_cal):
         )
 
     def test_calibration_loss_is_minimal_when_the_head_matches_the_target(self):
-        """Zero gradient when the reranker's candidate distribution already matches the target's."""
+        """Zero loss and gradient when the reranker's candidate distribution matches the target's."""
         model = _converted(dflash_lilicorr_w_cal=0.5)
         torch.manual_seed(0)
         num_slots = BLOCK_SIZE - 1
@@ -351,20 +351,71 @@ def test_calibration_loss_is_minimal_when_the_head_matches_the_target(self):
         }
         at_target = model._candidate_target_logits(**inputs, requested_by="test")
 
-        def grad_at(potentials):
+        def loss_and_grad_at(potentials):
             potentials = potentials.clone().requires_grad_(True)
             supervised = torch.ones(2, num_slots)
-            model._calibration_loss(
+            loss = model._calibration_loss(
                 node_potentials=list(potentials.unbind(1)),
                 supervised=supervised,
                 denominator=supervised.sum(),
                 **inputs,
-            ).backward()
-            return potentials.grad
+            )
+            loss.backward()
+            return loss.item(), potentials.grad
 
         # A per-slot shift leaves the softmax, and so the optimum, unchanged.
-        assert grad_at(at_target + 3.0).abs().max() < 1e-6
-        assert grad_at(torch.zeros_like(at_target)).abs().max() > 1e-3
+        loss, grad = loss_and_grad_at(at_target + 3.0)
+        assert loss == pytest.approx(0.0, abs=1e-6)
+        assert grad.abs().max() < 1e-6
+        loss, grad = loss_and_grad_at(torch.zeros_like(at_target))
+        assert loss > 1e-3
+        assert grad.abs().max() > 1e-3
+
+    def test_offline_forward_matches_online(self):
+        """Target logits rebuilt from a captured hidden state give the online loss.
+
+        The hidden state is stored in float64, wider than the target's weights.
+        """
+        models = []
+        for offline in (False, True):
+            model = get_tiny_llama(num_hidden_layers=4)
+            model.config.num_orig_hidden_layers = model.config.num_hidden_layers
+            config = _get_lilicorr_config(
+                topk=VOCAB_SIZE, dflash_lilicorr_w_cal=0.5, dflash_offline=offline
+            )
+            mtsp.convert(model, [("dflash", config)])
+            models.append(model)
+        online, offline = models
+        assert not offline.load_state_dict(online.state_dict(), strict=False).missing_keys
+
+        batch = _make_batch(VOCAB_SIZE)
+        online.eval()
+        with torch.no_grad():
+            hidden_states = online(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                output_hidden_states=True,
+            ).hidden_states
+        base_model_outputs = {
+            "aux_hidden_states": torch.cat(
+                [hidden_states[lid + 1] for lid in online.target_layer_ids], dim=-1
+            ),
+            "base_model_hidden_states": hidden_states[-1].double(),
+        }
+
+        online.train()
+        offline.train()
+        torch.manual_seed(1)
+        online_out = online(**batch)
+        torch.manual_seed(1)
+        offline_out = offline(**batch, base_model_outputs=base_model_outputs)
+
+        assert offline_out.loss.item() == pytest.approx(online_out.loss.item(), abs=1e-5)
+        for term in ("lilicorr_penalty", "lilicorr_calibration"):
+            assert online_out.lilicorr_metrics[term] > 0.0
+            assert offline_out.lilicorr_metrics[term] == pytest.approx(
+                online_out.lilicorr_metrics[term], abs=1e-5
+            )
 
     @pytest.mark.parametrize(
         ("kwargs", "degenerate"),

From b1944e41cca530e9ec531529bc17c9550dde31d7 Mon Sep 17 00:00:00 2001
From: mrusanovsky <mrusanovsky@nvidia.com>
Date: Thu, 1 Oct 2026 06:15:52 -0700
Subject: [PATCH 4/4] review: cast the captured hidden state before the final
 norm

Offline reconstruction now casts to the target's dtype before re-applying the final norm, so a pre-norm capture is normalized in the same dtype online training uses. The offline parity test also covers a pre-norm capture.

Signed-off-by: mrusanovsky <mrusanovsky@nvidia.com>
---
 .../torch/speculative/plugins/modeling_dflash.py  |  6 ++++--
 .../torch/speculative/plugins/test_hf_lilicorr.py | 15 ++++++++++++---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py
index c6c4edbdef7..cc1a9b5277a 100644
--- a/modelopt/torch/speculative/plugins/modeling_dflash.py
+++ b/modelopt/torch/speculative/plugins/modeling_dflash.py
@@ -148,9 +148,11 @@ def from_offline_dict(
             out_hiddens = d.get("base_model_hidden_states")
             if out_hiddens is None:
                 raise KeyError("base_model_hidden_states")
-            out_hiddens = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm)
             # A producer can store the hidden states in a wider dtype than the target's weights.
-            logits = base_model_lm_head(out_hiddens.to(base_model_lm_head.weight.dtype))
+            # Cast before the final norm too, which online training runs in the target's dtype.
+            out_hiddens = out_hiddens.to(base_model_lm_head.weight.dtype)
+            out_hiddens = _maybe_apply_base_final_norm(out_hiddens, d, base_model_norm)
+            logits = base_model_lm_head(out_hiddens)
         return cls(
             target_hidden=d["aux_hidden_states"],
             logits=logits,
diff --git a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
index 515a483a3fe..d1402987550 100644
--- a/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
+++ b/tests/unit/torch/speculative/plugins/test_hf_lilicorr.py
@@ -371,10 +371,12 @@ def loss_and_grad_at(potentials):
         assert loss > 1e-3
         assert grad.abs().max() > 1e-3
 
-    def test_offline_forward_matches_online(self):
+    @pytest.mark.parametrize("prenorm", [False, True], ids=["post_norm", "pre_norm"])
+    def test_offline_forward_matches_online(self, prenorm):
         """Target logits rebuilt from a captured hidden state give the online loss.
 
-        The hidden state is stored in float64, wider than the target's weights.
+        The hidden state is stored in float64, wider than the target's weights. A pre-norm
+        capture also runs the target's final norm on it.
         """
         models = []
         for offline in (False, True):
@@ -389,6 +391,10 @@ def test_offline_forward_matches_online(self):
         assert not offline.load_state_dict(online.state_dict(), strict=False).missing_keys
 
         batch = _make_batch(VOCAB_SIZE)
+        norm_inputs = []
+        hook = online._base_model_norm.register_forward_pre_hook(
+            lambda _module, args: norm_inputs.append(args[0])
+        )
         online.eval()
         with torch.no_grad():
             hidden_states = online(
@@ -396,11 +402,14 @@ def test_offline_forward_matches_online(self):
                 attention_mask=batch["attention_mask"],
                 output_hidden_states=True,
             ).hidden_states
+        hook.remove()
+        final_hidden = norm_inputs[-1] if prenorm else hidden_states[-1]
         base_model_outputs = {
             "aux_hidden_states": torch.cat(
                 [hidden_states[lid + 1] for lid in online.target_layer_ids], dim=-1
             ),
-            "base_model_hidden_states": hidden_states[-1].double(),
+            "base_model_hidden_states": final_hidden.double(),
+            "base_hidden_prenorm": prenorm,
         }
 
         online.train()