[h-guo18]

Stack (merge in this order; each PR's base is the branch of the one before it). [4/4] reproduces nvidia/Muse-Spark-Open-816B-A42B-DFlash2 and needs all of them.

1. [1/4] fix(speculative): four training fixes behind the Muse-Spark-Open DFlash2 recipe #2617: four training fixes: draft RoPE in fp32, streaming capture on vLLM >= 0.29, a memory-mapped streaming corpus, DFlash2's selector loss weighted like the backbone's
2. [2/4] perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe #3001: out-of-order training batches; compile + FlexAttention in the DFlash2 recipe
3. [3/4] perf(speculative): fuse the DFlash-family loss with the lm_head projection #3010: fused DFlash-family loss
4. [4/4] feat(speculative): Muse-Spark-Open DFlash2 recipes that reproduce the released drafter #3144: Muse-Spark-Open DFlash2 recipes <- this PR

---

What does this PR do?

Type of change: new example, plus the target support it needs

This PR adds the recipe that trained nvidia/Muse-Spark-Open-816B-A42B-DFlash2, so the release can be reproduced from main once this stack lands.

• modelopt_recipes/models/meta-models/Muse-Spark-Open-816B-A42B/speculative_decoding/
  • dflash2.yaml is stage 1: trained from scratch for 8 epochs at 4K tokens. The release continues from its step 390,000.
  • dflash2_long_context.yaml is stage 2: it continues the stage-1 export at 12K tokens on the corpus's long samples. The release is its step 25,000.
  • README.md covers the training data, the compute and the results.
• tools/launcher/examples/meta-models/Muse-Spark-Open-816B-A42B/ has one streaming launcher example per stage: 14 vLLM serve replicas of the NVFP4 target and 8 trainer nodes, all 4-GPU GB300.

Training against this target needed three changes on main:

1. Fake base. Muse-Spark applies a weightless RMSNorm to its token embeddings and cap * tanh(multiplier * logits / cap) to its logits (multiplier 0.177, cap 20). The fake base now builds both into embed_tokens and lm_head, so the draft's block embeddings, the draft's logits and the teacher's logits all match the target's. Two more pieces:
   • transformers has no muse_spark config class, so the fake base reads the raw config.json when AutoConfig does not know the model type. It needs only scalars from it.
   • Muse-Spark's final norm is registered as a plain RMSNorm, so the teacher's final hidden state is normed before the head.
2. Export. DFlash exports write the target's output_multiplier and final_logit_softcapping into dflash_config, which is the only place vLLM's DFlash2 reads them. Without them the served draft's logits are 5.7x too sharp and uncapped. On an earlier Muse-Spark-Open drafter, adding the two keys raised MT-Bench acceptance length by 9.3% at temperature 1 and by 1.4% greedy.
3. Fused loss. dflash_fused_loss ([3/4] perf(speculative): fuse the DFlash-family loss with the lm_head projection #3010) accepts the soft-capped head and passes its transform to fused_head_stats, which already supported one.

The streaming launcher also gets SERVE_ENABLE_EP for expert-parallel serve replicas. It is its own knob because multi-token SERVE_EXTRA_ARGS values do not survive nemo_run's environment export, which is also why SERVE_BLOCK_SIZE is separate.

Results

The released drafter, which is step 25,000 of stage 2. Acceptance length with vLLM, draft length 7, reasoning effort xhigh, 4,096 output tokens, concurrency 32, on one GB300 node (TP4, EP4):

Benchmark	Greedy	Temperature 1
MT-Bench	4.01	3.66
SPEED-Bench, qualitative	4.08	3.68
SPEED-Bench, throughput 32K	3.59	3.26

Compute: stage 1 took 37.7 h and stage 2 took 6.0 h on 88 GB300 GPUs, about 3,850 GPU-hours in all. Synthesizing the training corpus took about 2,100 more.

Usage

On the cluster login node, with the usual launcher environment (SLURM_HOST, SLURM_ACCOUNT, SLURM_PARTITION, SLURM_HF_LOCAL, SLURM_JOB_DIR):

cd tools/launcher
uv run launch.py --yaml examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2.yaml --yes
# Stop it once checkpoint-390000 is written and export that checkpoint, then:
uv run launch.py --yaml examples/meta-models/Muse-Spark-Open-816B-A42B/hf_streaming_dflash2_long_context.yaml --yes
# Stop it once checkpoint-25000 is written: that checkpoint, exported, is the release.

Fill in the placeholders first: the corpus paths, the stage-1 export path and a vLLM image that serves Muse-Spark-Open. Neither checkpoint is the last of its run, and both recipes keep only the last three checkpoints, so a run left to finish deletes it.

Testing

• Unit tests on CPU, with transformers 5.5.4 and 5.18.0:
  • New: test_muse_spark_base_reproduces_its_embedding_and_logit_transforms (a fake base built from a raw Muse-Spark config.json), test_export_config_carries_the_target_logit_transform, and a soft-capped-head case in the fused-loss test. Each fails without its change.
  • The surrounding speculative-decoding unit tests pass (719), as do test_examples_resolve (77) and the recipe tests (474).
• After the stack was folded to four PRs and rebased onto main after refactor(speculative): one home each for the DFlash family's shared rules #2757: 881 related unit tests pass at this PR (Python 3.12, transformers 5.18.0). test_hf_spec_rope_export.py's exporter stub now carries an lm_head. The export reads the target's logit transform off it, and without it 6 of that file's 16 tests failed, before the rebase too.
• GPU, one GPU, torch 2.10: tests/gpu/torch/speculative/plugins/test_hf_dflash.py has 17 passed and 1 skipped, including the new dflash2-kd-soft-capped-head case.
• The recipe checker and the launcher yaml checker pass on the new files.
• The corpus renders identically under the checkpoint's own chat template and under the copy the release was trained with, which changes only the default reasoning effort. I checked 360 rows; every row states its effort in its system turn.
• The configuration these recipes resolve to, as main.py prints it at startup in the end-to-end runs below (minus their step, checkpoint and data-order overrides), matches what the two release runs printed. The remaining differences are spellings of the same setting: the token loss is ce here and unset there, both meaning the distillation cross-entropy; the fused loss and out-of-order batch delivery were environment variables there; and the chat template is the one above. Stage 2 also saves every 2,500 steps instead of 1,250.
• On a scaled-down Muse-Spark model, this stack and the tree that trained the release build bit-identical initial draft weights. For the same weights, batch and anchors, they compute the same loss in fp32 (7.037086), and their gradients agree to within numerical noise.
• End to end on GB300 with 2 serve nodes, 1 trainer node and the full corpus, against the tree that trained the release, with the same data order:
  • Stage 2, starting from the release's stage-1 checkpoint: mean loss over steps 100-400 is 1.835, against 1.830 in the release run.
  • Stage 1, from scratch, 1,500 steps on a 1,500-step schedule, two runs on each tree: this stack matches the release tree within run-to-run variation. Mean loss over steps 1000-1500 is 4.348 and 4.272 on the release tree, and 4.265 and 4.214 on this stack. Rerunning a tree with the same data order and initialization moved this mean by up to 1.8%: the global batch here is 8, so a run's trajectory is noisy. A run of this stack that holds the draft in fp32 under autocast, as the release tree does, gives 4.256.
  • The exported dflash_config carries output_multiplier 0.1768 and final_logit_softcapping 20.0.
  • Without [1/4]'s corpus fix, the trainer node of this run stopped responding 20 steps in.
• Not run: the full 22-node schedule from main, about 3,850 GPU-hours. The acceptance lengths above are those of the released drafter, which was trained on the release tree.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ (a fake base for any other model is built as before)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

🤖 Generated with Claude Code

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts