h-guo18 · Oct 6, 2026 · Sep 29, 2026 · Sep 29, 2026 · Sep 30, 2026 · Oct 1, 2026
@@ -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 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)*
 
@@ -106,6 +108,8 @@ 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 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.

@@ -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 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. "
             "Ignored unless dflash_architecture_config.projector_type == 'lilicorr'."
         ),
     )

@@ -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:

@@ -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): 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
 ``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,37 @@ 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,
+    ):
+        """``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. 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,
+            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)
+        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)
 
-        # 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 * (target_log_probs - head_log_probs)).sum(dim=-1)
         return (per_slot * supervised).sum() / denominator
 
     @torch.no_grad()

@@ -148,6 +148,9 @@ def from_offline_dict(
             out_hiddens = d.get("base_model_hidden_states")
             if out_hiddens is None:
                 raise KeyError("base_model_hidden_states")
+            # A producer can store the hidden states in a wider dtype than the target's weights.
+            # 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(

@@ -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.