[1/4] fix(speculative): four training fixes behind the Muse-Spark-Open DFlash2 recipe

CHANGELOG.rst

@@ -59,6 +59,7 @@ Changelog
 - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning.
 
 - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing.
+- Streaming speculative-decoding training's vLLM connector (``modelopt.torch.speculative.plugins.rdma_hidden_states_connector``) now requires vLLM 0.22 or later. It is tested with vLLM 0.29 and later.
 
 **Deprecations**
 
@@ -105,6 +106,10 @@ Changelog
 - Fix ``export_hf_checkpoint`` of an FSDP2 model quantized with W4A8 AWQ (``W4A8_AWQ_BETA_CFG``), which failed with an ``AttributeError`` on every rank.
 - AWQ calibration (``awq_lite``, ``awq_full``) of an FSDP2 model sharded over more than one rank is now rejected with ``NotImplementedError`` by ``mtq.quantize``, ``mtq.calibrate`` and ``mtq.auto_quantize`` (before its search starts, its default formats included), and by ``examples/hf_ptq --use_fsdp2`` before loading the model: each rank would fold a different smoothing scale into its weight shard. Calibrate in a single process, or use a non-AWQ algorithm.
 - ``mtq.compress`` now rejects W4A8 weights with ``NotImplementedError``: a compressed W4A8 model could neither run nor be reloaded.
+- Fix streaming speculative-decoding training copying the whole training corpus into every rank and every DataLoader worker, which ran trainer nodes out of host memory on corpora of a few million samples. The corpus now stays memory-mapped.
+- Fix streaming speculative-decoding training capturing another request's hidden states when the vLLM serve ran requests concurrently, and failing on vLLM 0.29 and later. Retrain drafters trained in streaming mode against a serve that ran requests concurrently.
+- Fix DFlash-family drafters (DFlash, DFlash2, DSpark, Domino, LiLiCorr) trained from scratch with ``examples/speculative_decoding/main.py`` having their RoPE frequencies rounded to BF16, so they trained under different rotary phases than vLLM serves them with.
+- Fix DFlash2's candidate-selector loss weighing every block position equally and dividing by its own count of positions, unlike the backbone loss and the SGLang/SpecForge DFlash2 reference. Logged DFlash2 losses change.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^

modelopt/torch/speculative/plugins/hf_dflash.py

@@ -562,7 +562,8 @@ def modify(self, config):
         # lists. After the cast, because `Module.to` casts float buffers and building first
         # would round the RoPE frequencies.
         base_device = self._base_device()
-        if base_device.type != "meta":
+        self._draft_placed = base_device.type != "meta"
+        if self._draft_placed:
             self.dflash_module.to(device=base_device, dtype=self._base_model.dtype)
             self.dflash_module._maybe_init_rotary_emb(device=base_device)
 
@@ -592,10 +593,12 @@ def restore_draft_precision(self):
         placement entirely. A no-op on a freshly converted model.
         """
         base_device = self._base_device()
-        if base_device.type == "meta":
+        # Placing an already-placed draft again would round its fp32 rotary buffer to bf16.
+        if self._draft_placed or base_device.type == "meta":
             return
         self.dflash_module.to(device=base_device, dtype=self._base_model.dtype)
         self.dflash_module._maybe_init_rotary_emb(device=base_device)
+        self._draft_placed = True
 
     # Draft-module entries that legitimately come from the base model rather than the
     # exported draft checkpoint, so their absence (or presence) is not an error.

modelopt/torch/speculative/plugins/hf_dflash2.py

@@ -29,17 +29,19 @@
 - Take the backbone's top-k candidates per block position.
 - Score each candidate against its *teacher-forced* predecessor token, so the
   positions train in parallel exactly as the backbone does.
-- When the gold token is missing from the top-k, substitute it into the last
-  candidate slot. Without this the selector sees no positive class on the hard
-  positions and never learns those edges.
+- Supervise only the strict top-k, the candidate set serving builds. A gold token
+  outside it is a backbone recall miss, not a selector example, and carries no
+  selector gradient.
+- Weight each position as the backbone loss does (decay or D-PACE) and divide by the
+  backbone's weight sum, misses included.
 """
 
 import torch
 import torch.nn.functional as F
 from transformers import PreTrainedModel
 
 from ..dflash.conversion import DFlash2DMRegistry
-from .hf_dflash import HFDFlashModel, _block_targets, _block_tokens
+from .hf_dflash import HFDFlashModel, _block_tokens
 from .modeling_dflash2 import DFlash2Module
 
 __all__ = ["HFDFlash2Model"]
@@ -106,15 +108,20 @@ def get_exporter(self):
 
         return DFlash2Exporter(self)
 
-    def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mask):
+    def _selector_loss(
+        self, logits, target_ids, hidden, predecessor_ids, supervised_mask, loss_weights, weight_sum
+    ):
         """Cross-entropy over the selector's candidate set, and its top-1 accuracy.
 
         Args:
             logits: Backbone logits per block position ``[B, N, block_size, V]``.
             target_ids: Gold token ids ``[B, N, block_size]``.
             hidden: Backbone hidden states ``[B, N, block_size, H]``.
             predecessor_ids: Teacher-forced predecessor ids ``[B, N, block_size]``.
-            weight_mask: Per-position loss weights ``[B, N, block_size]``.
+            supervised_mask: The positions the backbone loss supervises, unweighted
+                ``[B, N, block_size]``.
+            loss_weights: The backbone loss's per-position weights ``[B, N, block_size]``.
+            weight_sum: The backbone loss's denominator, the sum of ``loss_weights``.
 
         Returns:
             ``(loss, accuracy, coverage)`` — coverage is the fraction of supervised
@@ -130,7 +137,7 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
         # token the backbone did not propose is a backbone recall failure, not a
         # selector classification example: substituting it in would teach the selector
         # to override the unary ranking on a set it will never be shown. Those
-        # positions carry no selector gradient and leave the denominator instead.
+        # positions carry no selector gradient.
         gold_matches = candidate_ids == target_ids.unsqueeze(-1)
         gold_in_topk = gold_matches.any(dim=-1)
         gold_slot = gold_matches.long().argmax(dim=-1)
@@ -139,26 +146,27 @@ def _selector_loss(self, logits, target_ids, hidden, predecessor_ids, weight_mas
             candidate_ids, unary_logits, hidden, predecessor_ids
         )
 
-        covered = weight_mask * gold_in_topk.to(weight_mask.dtype)
-        flat_weights = covered.reshape(-1)
-        denominator = flat_weights.sum() + 1e-6
+        in_topk = gold_in_topk.to(loss_weights.dtype)
         per_token = F.cross_entropy(
             selector_logits.float().reshape(-1, top_k),
             gold_slot.reshape(-1),
             reduction="none",
         )
-        loss = (per_token * flat_weights).sum() / denominator
+        # The backbone's denominator, misses included: dividing by the covered positions
+        # alone would scale the term up whenever backbone recall drops.
+        loss = (per_token * (loss_weights * in_topk).reshape(-1)).sum() / weight_sum
 
         with torch.no_grad():
+            covered = (supervised_mask * in_topk).reshape(-1)
             chosen = selector_logits.argmax(dim=-1).reshape(-1)
-            accuracy = (
-                (chosen == gold_slot.reshape(-1)).float() * flat_weights
-            ).sum() / denominator
+            accuracy = ((chosen == gold_slot.reshape(-1)).float() * covered).sum() / (
+                covered.sum() + 1e-6
+            )
             # Coverage keeps the full supervised mask as its denominator: it measures
             # how often the selector was given a solvable problem at all.
             # clamp, not an epsilon: a fully covered batch must read exactly 1.0.
-            supervised = weight_mask.reshape(-1).sum()
-            coverage = flat_weights.sum() / supervised.clamp(min=1.0)
+            supervised = supervised_mask.reshape(-1).sum()
+            coverage = covered.sum() / supervised.clamp(min=1.0)
         # Detached tensors, not Python scalars: .item() would force a CPU-GPU sync on
         # every training step. The trainer converts them at the logging boundary.
         return loss, accuracy.detach(), coverage.detach()
@@ -235,15 +243,10 @@ def _compute_loss(
         block_size = self.dflash_block_size
         n_blocks = anchor_positions.shape[1]
 
-        # The backbone's supervision mask, without its position weighting (decay/D-PACE):
-        # that shapes *where* the backbone spends capacity, while the selector should learn
-        # every position's transition equally.
-        targets = _block_targets(
-            input_ids, anchor_positions, block_keep_mask, loss_mask, block_size, self._label_shift
-        )
+        target_ids = _block_tokens(input_ids, anchor_positions, block_size, self._label_shift)
 
         # Teacher-forced predecessor of block position k is the real token at anchor+k-1.
-        # Only offsets 1..block_size-1 are supervised (the mask zeroes slot 0), so the
+        # Only offsets 1..block_size-1 are supervised (the backbone's mask zeroes slot 0), so the
         # first supervised position, k=1, has the anchor's own token as its predecessor.
         # Slot 0's entry here resolves to anchor-1 and never contributes.
         #
@@ -254,12 +257,16 @@ def _compute_loss(
             input_ids, anchor_positions, block_size, self._label_shift - 1
         )
 
+        # The backbone's own mask and position weights: a block position's choice only
+        # matters if every earlier one was accepted, which decay / D-PACE model.
         selector_loss, selector_accuracy, selector_coverage = self._selector_loss(
             logits.reshape(bsz, n_blocks, block_size, -1),
-            targets.token_ids,
+            target_ids,
             draft_hidden.reshape(bsz, n_blocks, block_size, -1),
             predecessor_ids,
-            targets.mask,
+            terms.supervised_mask.view_as(target_ids),
+            terms.weights.view_as(target_ids),
+            terms.weight_sum,
         )
         self._selector_metrics = {
             "selector_accuracy": selector_accuracy,

modelopt/torch/speculative/plugins/hf_streaming_dataset.py

@@ -41,7 +41,8 @@
 import os
 import re
 import time
-from typing import TypedDict
+from collections.abc import Iterable, Sequence
+from typing import TypedDict, cast
 
 import httpx
 import torch
@@ -199,7 +200,7 @@ class StreamingDataset(Dataset):
 
     def __init__(
         self,
-        entries: list[dict],
+        entries: Iterable[dict],
         tokenizer,
         config: StreamingConfig | None = None,
     ):
@@ -213,17 +214,24 @@ def __init__(
         Args:
             entries: Untokenized per-sample dicts from the input jsonl. Schema is
                 subclass-defined (see :meth:`_tokenize_entry`); passed to :meth:`_fetch`.
+                A sized, indexable source (e.g. an Arrow-backed ``datasets.Dataset``) is
+                kept as is, so its rows stay memory-mapped; any other iterable is
+                materialized into a list.
             tokenizer: HF tokenizer; used for client-side tokenization and the
                 server/client loss-mask alignment in :meth:`_fetch`.
             config: Tuning knobs (timeout, answer_only_loss, ...); defaults to
                 ``self.config_cls()``. See :class:`StreamingConfig`.
         """
-        if not entries:
-            raise ValueError("entries is empty")
         self.tokenizer = tokenizer
         self.config = config if config is not None else self.config_cls()
-        # Materialize to a plain list so DataLoader worker processes fork it cheaply.
-        self.entries = list(entries)
+        # Keep an indexable source (e.g. an Arrow-backed datasets.Dataset) as is: list()
+        # would copy every row into each rank, and each forked worker would copy it again.
+        if hasattr(entries, "__len__") and hasattr(entries, "__getitem__"):
+            self.entries = cast("Sequence[dict]", entries)
+        else:
+            self.entries = list(entries)
+        if len(self.entries) == 0:
+            raise ValueError("entries is empty")
         # Per-process consecutive-failure counter for the circuit breaker. Reset to 0
         # on every successful fetch; tripped only by fetch failures (not unfit entries).
         self._consecutive_fail = 0
@@ -405,7 +413,7 @@ class EagleVllmStreamingDataset(StreamingDataset):
 
     def __init__(
         self,
-        entries: list[dict],
+        entries: Iterable[dict],
         tokenizer,
         config: EagleVllmStreamingConfig,
     ):