h-guo18 · Oct 30, 2026 · Oct 30, 2026
@@ -31,6 +31,7 @@ Changelog
 - With ``training.dataloader_num_workers`` > 0, ``examples/speculative_decoding/main.py`` now yields each training batch as soon as a worker finishes it, so one slow sample, such as a slow hidden-state fetch in streaming training, no longer stalls every DDP rank. Sample order within the prefetch window is then no longer deterministic; set ``training.dataloader_in_order=true`` to restore it (context-parallel runs always keep it).
 - The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``, and turns ``tf32`` off, which compilation needs on PyTorch 2.9 and 2.10. To train eagerly as before, start the run with ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``; a run resumed from a checkpoint keeps the checkpoint's DFlash settings.
 - Add ``dflash_fused_loss`` to the DFlash config. DFlash, DFlash2 and LiLiCorr training then computes its loss in row chunks fused with the frozen ``lm_head`` projection instead of materializing the draft's full-vocabulary logits, which saves time and memory on large vocabularies.
+- Add DFlash2 training recipes for ``meta-models/Muse-Spark-Open-816B-A42B`` under ``modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/``, with streaming launcher examples, that reproduce the released ``nvidia/Muse-Spark-Open-816B-A42B-DFlash2`` drafter. Streaming and offline training against a Muse-Spark target now apply its embedding norm and logit soft cap, and DFlash exports record the target's ``output_multiplier`` and ``final_logit_softcapping`` for vLLM.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

@@ -453,6 +453,16 @@ def _export_config(self):
             config["dflash_config"]["attention_sink_bias"] = True
             config["attention_sink_bias"] = True
 
+        # The target's logit transform, which vLLM's DFlash2 reads only from dflash_config
+        # (default: none). Read off the head the draft logits were trained through.
+        head = model._base_model_lm_head
+        for key, attr in (
+            ("output_multiplier", "logits_multiplier"),
+            ("final_logit_softcapping", "logits_soft_cap"),
+        ):
+            if (value := getattr(head, attr, None)) is not None:
+                config["dflash_config"][key] = float(value)
+
         # Inject the export-time YaRN rope_scaling from the dflash_export_rope_scaling
         # config field (empty dict disables). Mirrors eagle's eagle_export_rope_scaling.
         export_rope_scaling = getattr(self.model, "dflash_export_rope_scaling", None)

@@ -96,6 +96,7 @@
     _BASE_MODEL_PATHS,
     _EMBED_TOKENS_PATHS,
     _LM_HEAD_PATHS,
+    _SoftCappedLMHead,
     find_final_norm_path,
 )
 from .modeling_final_norm import _apply_base_final_norm
@@ -610,9 +611,10 @@ def _check_fused_loss(self):
                 )
         head = self._base_model_lm_head
         # The fused pass reads lm_head.weight directly: whatever a bias, a quantizer or a
-        # hook adds in lm_head's forward would be silently skipped.
+        # hook adds in lm_head's forward would be silently skipped. The fake base's soft cap
+        # is the one transform it reproduces (see _fused_head_transform).
         plain_forward = (
-            type(head).forward is torch.nn.Linear.forward
+            type(head).forward in (torch.nn.Linear.forward, _SoftCappedLMHead.forward)
             and "forward" not in vars(head)
             and not (head._forward_hooks or head._forward_pre_hooks)
         )
@@ -1122,6 +1124,14 @@ def _fused_head_extras(self):
         """``(statistics with gradient, top_k)`` a variant's own terms read from the fused pass."""
         return set(), 0
 
+    def _fused_head_transform(self):
+        """The logit transform the base ``lm_head`` applies after its projection, if any."""
+        head = self._base_model_lm_head
+        return {
+            "logit_scale": getattr(head, "logits_multiplier", None),
+            "soft_cap": getattr(head, "logits_soft_cap", None),
+        }
+
     def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_targets):
         """Every statistic of the draft's logits the loss reads, from one fused pass."""
         grads, top_k = self._fused_head_extras()
@@ -1145,6 +1155,7 @@ def _fused_head_stats(self, draft_hidden, base_outputs, label_positions, flat_ta
             grads=grads,
             top_k=top_k,
             **teacher,
+            **self._fused_head_transform(),
         )
 
     def forward(

@@ -17,6 +17,8 @@
 
 import json
 import os
+import types
+import warnings
 
 import torch
 import torch.nn as nn
@@ -96,6 +98,92 @@ def find_final_norm_path(model) -> str | None:
 # Single-file safetensors names to try, in order.  Mistral native checkpoints
 # use ``consolidated.safetensors`` instead of the HF-standard ``model.safetensors``.
 _SAFETENSORS_SINGLE_FILENAMES = ["model.safetensors", "consolidated.safetensors"]
+# Muse-Spark applies a weightless RMSNorm to its token embeddings and a multiplier and tanh
+# soft cap to its logits. Both live in the model's code, not its config, hence the table.
+_MUSE_SPARK_MODEL_TYPES = frozenset({"muse_spark", "muse_spark_text"})
+
+
+class _NormedEmbedding(nn.Embedding):
+    """``nn.Embedding`` followed by an RMSNorm with no gain, computed in fp32.
+
+    The draft embeds its block tokens through ``embed_tokens``, so the norm has to live in
+    the module every caller resolves rather than at the call sites.
+    """
+
+    def __init__(self, *args, eps: float, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.eps = eps
+
+    def forward(self, input_ids):
+        h = super().forward(input_ids)
+        h32 = h.float()
+        return (h32 * torch.rsqrt(h32.pow(2).mean(-1, keepdim=True) + self.eps)).to(h.dtype)
+
+
+class _SoftCappedLMHead(nn.Linear):
+    """``lm_head`` followed by ``cap * tanh(multiplier * logits / cap)``, computed in fp32.
+
+    Draft and teacher logits both come through this head, so the draft trains against the
+    distribution the target actually samples from.
+    """
+
+    def __init__(self, *args, logits_multiplier=None, logits_soft_cap=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logits_multiplier = logits_multiplier
+        self.logits_soft_cap = logits_soft_cap
+
+    def forward(self, x):
+        logits = super().forward(x)
+        out = logits.float()
+        if self.logits_multiplier is not None:
+            out = out * self.logits_multiplier
+        if self.logits_soft_cap is not None:
+            out = self.logits_soft_cap * torch.tanh(out / self.logits_soft_cap)
+        return out.to(logits.dtype)
+
+
+def _select_base_transforms(base_cfg) -> dict:
+    """FakeBaseConfig fields for what the base does outside ``embed_tokens`` and ``lm_head``."""
+    if getattr(base_cfg, "model_type", None) not in _MUSE_SPARK_MODEL_TYPES:
+        return {}
+    return {
+        "embed_norm": True,
+        "logits_multiplier": getattr(base_cfg, "output_multiplier", None),
+        "logits_soft_cap": getattr(base_cfg, "final_logit_softcapping", None),
+    }
+
+
+def _load_base_config(source: str, trust_remote_code: bool):
+    """The checkpoint's config, or its raw ``config.json`` if transformers lacks the model_type.
+
+    FakeBaseModel reads only scalars, so the JSON is enough for them. Anything the model does
+    outside embed_tokens and lm_head is not in the JSON, which is why _select_base_transforms
+    and the final-norm table are keyed by model_type.
+    """
+    try:
+        return transformers.AutoConfig.from_pretrained(source, trust_remote_code=trust_remote_code)
+    except ValueError as exc:
+        if "does not recognize this architecture" not in str(exc):
+            raise
+    if os.path.isdir(source):
+        path = os.path.join(source, "config.json")
+    else:
+        path = hf_hub_download(repo_id=source, filename="config.json")
+    with open(path) as f:
+        raw = json.load(f)
+    warnings.warn(
+        f"transformers does not know model_type {raw.get('model_type')!r}; reading {path} "
+        "directly for the dimensions FakeBaseModel needs.",
+        stacklevel=2,
+    )
+    return types.SimpleNamespace(
+        **{
+            key: types.SimpleNamespace(**value)
+            if key in _VLM_CONFIG_ATTRS and isinstance(value, dict)
+            else value
+            for key, value in raw.items()
+        }
+    )
 
 
 class FakeBaseConfig(PretrainedConfig):
@@ -118,6 +206,9 @@ def __init__(
         rms_norm_eps=1e-6,
         rope_theta=None,
         final_norm_type=None,
+        embed_norm=False,
+        logits_multiplier=None,
+        logits_soft_cap=None,
         **kwargs,
     ):
         """Initialize FakeBaseConfig with minimal model configuration parameters."""
@@ -127,6 +218,11 @@ def __init__(
         # (model whose final-norm type we don't know). See _FINAL_NORM_CLASSES /
         # _FINAL_NORM_TYPE_BY_MODEL_TYPE. Persisted so a reloaded config rebuilds the same norm.
         self.final_norm_type = final_norm_type
+        # The base's embedding and logit transforms (see _select_base_transforms), persisted
+        # so a reloaded config rebuilds the same modules.
+        self.embed_norm = embed_norm
+        self.logits_multiplier = logits_multiplier
+        self.logits_soft_cap = logits_soft_cap
         self.num_hidden_layers = num_hidden_layers
         # Mirror the original base layer count. The non-fake offline path loads with
         # num_hidden_layers=0 and stashes the real count here (see utils.load_vlm_or_llm);
@@ -184,10 +280,27 @@ def __init__(self, config: FakeBaseConfig, **kwargs):
         self.model = nn.Module()
         self.model.layers = nn.ModuleList()
         self.model.dtype = config.dtype
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, dtype=config.dtype)
-        self.lm_head = nn.Linear(
-            config.hidden_size, config.vocab_size, bias=False, dtype=config.dtype
-        )
+        if config.embed_norm:
+            self.embed_tokens = _NormedEmbedding(
+                config.vocab_size, config.hidden_size, dtype=config.dtype, eps=config.rms_norm_eps
+            )
+        else:
+            self.embed_tokens = nn.Embedding(
+                config.vocab_size, config.hidden_size, dtype=config.dtype
+            )
+        if config.logits_multiplier is not None or config.logits_soft_cap is not None:
+            self.lm_head = _SoftCappedLMHead(
+                config.hidden_size,
+                config.vocab_size,
+                bias=False,
+                dtype=config.dtype,
+                logits_multiplier=config.logits_multiplier,
+                logits_soft_cap=config.logits_soft_cap,
+            )
+        else:
+            self.lm_head = nn.Linear(
+                config.hidden_size, config.vocab_size, bias=False, dtype=config.dtype
+            )
         # Final pre-lm_head norm, applied before lm_head when reconstructing base logits for
         # self-logit-distillation (vLLM-captured final hidden states are un-normed). Built ONLY
         # when the base model's final-norm type is known (config.final_norm_type set); otherwise
@@ -212,9 +325,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
         """
         from modelopt.torch.export.plugins.hf_spec_export import _get_rope_theta
 
-        orig_config = transformers.AutoConfig.from_pretrained(
-            source, trust_remote_code=trust_remote_code
-        )
+        orig_config = _load_base_config(source, trust_remote_code)
         # For vlms, detect language model config based on _VLM_CONFIG_ATTRS
         base_cfg = next(
             (
@@ -242,6 +353,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
             final_norm_type=_select_final_norm_type(
                 getattr(base_cfg, "model_type", None), base_cfg
             ),
+            **_select_base_transforms(base_cfg),
         )
         model = cls(config)
         # Load lm_head, embed_tokens, and (for known models) the final norm into the model.

@@ -104,6 +104,9 @@ def __init__(self, hidden_size, eps=1e-6, dtype=torch.bfloat16):
     # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1.
     "minimax_m3_vl_text": "gemma_rmsnorm",
     "gpt_oss": "gpt_oss_rmsnorm",
+    # Muse-Spark's final norm uses its weight as stored: no (1 + weight) as in Gemma.
+    "muse_spark": "rmsnorm",
+    "muse_spark_text": "rmsnorm",
 }
 
 

@@ -0,0 +1,64 @@
+# Muse-Spark-Open-816B-A42B DFlash2
+
+These two recipes train the DFlash2 drafter released as
+[nvidia/Muse-Spark-Open-816B-A42B-DFlash2](https://huggingface.co/nvidia/Muse-Spark-Open-816B-A42B-DFlash2)
+for [meta-models/Muse-Spark-Open-816B-A42B](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B).
+
+| Stage | Recipe | Launcher example | Steps | Sequence | Learning rate |
+| --- | --- | --- | --- | --- | --- |
+| 1, from scratch | `dflash2.yaml` | `hf_streaming_dflash2.yaml` | 390,000 of 434,744 (8 epochs) | 4,096 | 1e-3, 4% warmup, linear |
+| 2, long context | `dflash2_long_context.yaml` | `hf_streaming_dflash2_long_context.yaml` | 25,000 of 59,502 (6 epochs) | 12,288 | 3e-5, 1,500 warmup steps, linear |
+
+The launcher examples are under `tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/`.
+Stage 2 starts from the export of stage 1's checkpoint-390000, and the release is stage 2's
+checkpoint-25000. Neither is the last checkpoint of its run, and the recipes keep only the last
+three, so stop each run once its checkpoint is written, or copy the checkpoint out. Both stages
+stream hidden states from the NVFP4 target,
+[meta-models/Muse-Spark-Open-816B-A42B-NVFP4](https://huggingface.co/meta-models/Muse-Spark-Open-816B-A42B-NVFP4),
+served by vLLM.
+
+## Training data
+
+The drafter learns from the target's own outputs. The prompts come from
+[nvidia/Nemotron-Post-Training-Dataset-v2](https://huggingface.co/datasets/nvidia/Nemotron-Post-Training-Dataset-v2);
+its responses are not used.
+
+- The target answered each prompt at temperature 1, with up to 8,192 output tokens, at one of
+  its five reasoning-effort levels: xhigh for 78% of the samples, and minimal, low, medium and
+  high for about 200,000 each.
+- Each response is split into its reasoning and its answer. Responses that reached the token
+  limit without an answer, and a few malformed ones, were dropped, leaving 3.48 million samples.
+- Each row is a `conversations` list: a system turn that sets the reasoning effort the way
+  the model's chat template writes it, the user turn, and an assistant turn whose `content` is
+  the answer and whose `reasoning` is the reasoning.
+
+Stage 2 uses the samples whose conversation, rendered with the chat template, is longer than
+4,096 tokens: 634,723 samples, 18% of the corpus.
+
+Synthesizing the corpus took about 525 node-hours on 4-GPU GB300 nodes, one TP4 replica of the
+target per node.
+
+## Compute
+
+Both stages ran on 22 nodes of 4 GB300 GPUs: 14 vLLM serve replicas of the target and 8
+trainer nodes.
+
+| Stage | Run time | GPU-hours |
+| --- | --- | --- |
+| 1 (390,000 steps) | 37.7 h | ~3,320 |
+| 2 (25,000 steps) | 6.0 h | ~530 |
+
+## Results
+
+Acceptance length with vLLM, draft length 7, reasoning effort xhigh, 4,096 output tokens,
+concurrency 32, on one GB300 node (TP4, EP4). At temperature 1 the draft samples with
+`draft_sample_method: probabilistic`.
+
+| Benchmark | Greedy | Temperature 1 |
+| --- | --- | --- |
+| MT-Bench | 4.01 | 3.66 |
+| SPEED-Bench, qualitative | 4.08 | 3.68 |
+| SPEED-Bench, throughput 32K (first 100 per category) | 3.59 | 3.26 |
+
+Against stage 1's checkpoint-390000, stage 2 raised SPEED-Bench acceptance by 0.5-1.9% and
+moved MT-Bench by -0.9% (greedy) and +0.4% (temperature 1).