[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 <- this PR
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

---

What does this PR do?

Type of change: new feature

Two speculative-decoding training speedups.

1. Out-of-order training batches. New training.dataloader_in_order, default false. With DataLoader workers, a rank takes whichever batch a worker finishes first instead of waiting for its slowest outstanding sample (head-of-line blocking). Per-sample fetch latency varies, and DDP lockstep turns one rank's stall into a stall on every rank.

• HF Trainer forwards PyTorch's in_order only from transformers 5.15. accelerate's DataLoaderShard forwards attribute reads, not writes, to the torch loader it wraps, so prepared.in_order = False silently does nothing. EagleTrainerWithAccLog.get_train_dataloader therefore sets it on the underlying loader. The field shares HF's name, so on transformers >= 5.15 it overrides HF's own default.
• Behavior change: per-rank sample order within the prefetch window is no longer deterministic; training.dataloader_in_order=true restores it. Nothing changes with dataloader_num_workers=0, the online-mode default.
• Why off by default for every recipe, not only streaming ones. Recipes add workers because per-sample loading is slow, and slow loading is seldom uniform: a streaming fetch over RDMA, or an offline read from a shared filesystem. Under DDP lockstep one late sample then stalls every rank. The gain below was measured on streaming. The cost is bounded: every epoch still visits each sample once, and only the order inside the prefetch window (num_workers x prefetch_factor batches per rank) changes. HF's own default from transformers 5.15 is true, and this field deliberately overrides it.
• Context parallelism always keeps in-order delivery. accelerate gives every rank of a CP group the same data shard, and each rank loads it on its own. Out of order, those ranks would split different sequences at the same step, so init_distributed_env sets the flag to true when cp_size > 1.
• Resume: accelerate's skip_first_batches rebuilds the loader and carries in_order over (accelerate >= 1.3; ModelOpt requires >= 1.6). The resume position is exact only to within the batches in flight at the save, at most num_workers x prefetch_factor per rank.

2. Compile + FlexAttention in the DFlash2 recipe: dflash_use_torch_compile and dflash_use_flex_attention go false -> true in general/speculative_decoding/dflash2.yaml, the configuration the multi-node numbers below were measured with, and tf32 goes true -> false (see the first caveat). dflash, domino, dspark, lilicorr* and the launcher examples are unchanged; they are possible follow-ups.

One GPU (RTX PRO 6000, PyTorch 2.10, tf32=false), the shipped online example's shape (Qwen3-8B, seq 4096, 512 anchors, SWA 2048, bs 1), 40 steps:

draft attention	eager	compile
SDPA, dense mask (current recipe)	1612 ms/step	1540 ms/step (1.05x)
FlexAttention	1088 ms/step (1.48x)	985 ms/step (1.64x)

Measured on DFlash2 streaming training. A large MoE target is served by vLLM replicas (TP4/EP4), and the trainer pulls hidden states over RDMA. 8 trainer nodes x 4 GB300, DDP, global batch 64, seq 4096, 8 workers per rank, same data_seed throughout. These numbers come from a fork of this repo, not from main:

• in_order=False, 8 serve + 8 trainer nodes: 0.655 -> 0.523 s/step (+25%). Serve replicas went from 26% idle to 3%, and the worst per-rank batch wait from 5.9 s to ~2 s.
• Compile gains nothing when serve-bound (8+8: 0.676 vs 0.696 s/step).
  • Once the trainer is the bottleneck it gains ~13%: 0.445 -> 0.386 s/step at 12+8 and 0.446 -> 0.389 at 16+8.
  • Trainer in-step time went from 0.416 to ~0.35 s.
  • No recompiles: the step rate stayed flat over 2.5K steps.
  • Those runs had dflash_use_flex_attention on, as this recipe now does, with DDP and bf16.
• Both changes together at 12 serve + 8 train ran at 0.379 s/step, vs the original 0.655. Part of that gap comes from the extra serve nodes.
• Correctness vs the baseline over ~2.5K steps: mean |loss diff| <= 0.016, and mean acceptance diff within +-0.003, i.e. noise. Exact per-step equality is not expected, because samples are reordered within the prefetch window.

Compile caveats, found while testing this PR:

• PyTorch 2.9 / 2.10 wheels: with tf32 on, Inductor raises on the first step ("mix of the legacy and new APIs to set the TF32 status"). transformers enables TF32 through torch.backends.fp32_precision, and Inductor 2.9/2.10 then reads the legacy allow_tf32 flag, which raises. Inductor reads it whenever the compiled graph has an fp32 CUDA matmul, and the draft's rotary embedding is one. PyTorch 2.8 and >= 2.11 are fine. PyTorch 2.10 is ModelOpt's floor, so the recipe now turns tf32 off. DFlash2 trains in bf16, so TF32 only reached fp32 matmuls such as that one.
• With compile (or flex) on, the draft's block count is pinned to dflash_num_anchors, while eager sizes it to the batch. With short answers the padding costs more than compile saves: in a toy run with ~35-token answers at seq 512, steps were 3.6x slower.
• There is one recompile, at step 2: _get_attn_fn caches the attention function on the module inside the compiled region.

Usage

# Restore deterministic per-rank sample order
python main.py --config <recipe> ... training.dataloader_in_order=true
# Train DFlash2 eagerly with SDPA and TF32, as before (when starting a run: a resumed run keeps the checkpoint's DFlash settings)
python main.py --config modelopt_recipes/general/speculative_decoding/dflash2.yaml ... dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false training.tf32=true

Testing

• tests/examples/speculative_decoding/test_dataloader_in_order.py::test_dataloader_in_order_reaches_the_torch_loader[False/True]
  • Runs the real HF Trainer and accelerate with 2 workers; index 0 is held back until another batch has been yielded.
  • Asserts that each sample arrives exactly once, out of sampler order with the flag off and in order with it on.
  • Fails on transformers 5.5 without the override (5.18 forwards the flag itself).
  • It starts DataLoader workers, so it runs in the speculative_decoding example lane rather than tests/unit.
• test_context_parallelism_keeps_batches_in_order[1/2]: init_distributed_env forces in-order delivery for cp_size=2 only. The cp_size=2 case fails without the guard.
• tests/gpu/torch/speculative/plugins/test_hf_dflash.py::TestDFlashFlexAttentionGPU::test_compiled_flex_matches_eager_flex: the recipe's compile + flex pair runs forward and backward on a tiny model with the recipe's DFlash2 draft (conv projector, candidate selector) and objective, with the loss within bf16 tolerance of eager flex and finite grads. It passes on PyTorch 2.10, the floor, and on 2.13 (about 15 s, mostly compilation).
• Related unit tests pass on transformers 5.5.4 and 5.18.0: entrypoint, recipe loader, offline data, eagle config, external-draft sparse.
• GPU smoke on 1x RTX PRO 6000: DFlash2 recipe, 40 steps, same seed, compile off vs on.
  • Setup: Qwen3-8B target at the shipped online example's shape (Daring-Anteater, seq 4096, 512 anchors, SWA 2048, bs 1).
  • torch 2.11 with tf32 on, the recipe's setting before this PR: 1573 -> 1499 ms/step (1.05x), peak memory 57.0 -> 53.8 GiB. Losses for steps 1-4 match within 8e-5 relative; after that bf16 noise grows as training goes on.
  • torch 2.10 with tf32 off, as the recipe now ships: also 1.05x. Adding flex: compile is 1.11x over eager-with-flex, 1.64x over the current recipe.
• Resume with 4 workers and out-of-order delivery: trained 10 steps, saved a checkpoint, resumed to 20.
• Not run: the final recipe end to end on PyTorch 2.11 or later; the table above is the recipe as shipped (compile + flex, tf32 off) on 2.10. Not run in this tree: multi-node DDP (the numbers above come from the fork), FSDP, and acceptance-length eval.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ❌ Per-rank sample order is no longer deterministic within the prefetch window by default (training.dataloader_in_order=true restores it), and the DFlash2 recipe now compiles, uses FlexAttention and turns TF32 off.
• 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?: ❌ The Claude review's findings are addressed or answered in de31718 and 4f3724e; re-review pending.

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

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review: 3

• main
• release/.*
• feature/.*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

Run configuration

• Source mode: Auto
• Repository: NVIDIA/Model-Optimizer
• Branch: haoguo/spec-train-inorder-compile
• File: .coderabbit.yaml
• Commit: 1ef4963
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 6469145c-725b-4d00-90ee-6208c7897b31

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• Trigger review

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-3001/
Built to branch gh-pages at 2026-10-25 05:33 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.80%. Comparing base (fb3ccd0) to head (edd263c).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3001   +/-   ##
=======================================
  Coverage   79.79%   79.80%           
=======================================
  Files         661      661           
  Lines       74935    74936    +1     
=======================================
+ Hits        59798    59799    +1     
  Misses      15137    15137           

Flag	Coverage Δ
unit	61.12% <100.00%> (-0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[h-guo18]

/claude review

[ChenhanYu]

Pensieve comment review using chenhany's profile (revision 12).

Comment: DFlash2 now enables a compile/TF32 combination documented to crash on supported PyTorch 2.10.

Needs action:

• Fix modelopt_recipes/general/speculative_decoding/dflash2.yaml:73 with compatible defaults or a version guard; extend GPU coverage to execute the shipped compile+FlexAttention configuration through forward/backward (inline).
• Move the worker-based regression at tests/unit/examples/test_speculative_decoding_main.py:344 to tests/gpu*, following the base AGENTS.md and CONTRIBUTING.md multiprocessing rule (inline).

No action needed:

• Loader unwrapping, argument plumbing and the CP ordering guard look consistent. The ordering regression checks actual delivery and preservation of every sample.
• Existing test edits update scope documentation and add cases without weakening assertions. The diff adds 16 core source lines; no licensing changes were identified.
• Baseline: 687720b5d17c40c9d9fd7d28cc6622d5df99cdf6. Tests were not executed; current-head CI results were unavailable. Duplicate checking covers the supplied bounded history only.

Review baseline: https://github.com/NVIDIA/Model-Optimizer/blob/687720b5d17c40c9d9fd7d28cc6622d5df99cdf6/.agents/review-guidelines/code-review.md

[blocking] modelopt_recipes/general/speculative_decoding/dflash2.yaml:73 — Enabling compilation while retaining training.tf32: true at line 61 introduces a default-path failure on PyTorch 2.10, explicitly documented in the added CHANGELOG.rst:32 and the PR's reported testing. PyTorch 2.10 remains supported (pyproject.toml:40 specifies torch>=2.10); main.py forwards these settings unchanged, and DFlashModule.compile_body unconditionally uses Inductor. Documenting an override does not preserve the previously working recipe defaults. Set a compatible TF32 default, keep compilation opt-in on affected versions, or add a compatibility guard. Extend existing GPU coverage to execute the resulting shipped compile+FlexAttention configuration through forward/backward on the supported Torch floor.

[warning] tests/unit/examples/test_speculative_decoding_main.py:356 — With dataloader_num_workers=2, this new parametrized unit test starts fresh worker processes for each case and does not use the shared cpu_dist_workers fixture. The base AGENTS.md 'Adding tests' rule and CONTRIBUTING.md:196-200 explicitly place process-starting tests in tests/gpu*, with the shared CPU fixture as the exception. The reported fork timings and non-fork skip do not satisfy that placement rule. Move the worker-driven ordering regression and its helpers to tests/gpu*; the lightweight CP argument-wiring test can remain here.

Review run: 0b09cb37cfc94ea1a6a76cb2ad251341

[claude]

The out-of-order dataloader work is clean and well tested, but the DFlash2 recipe as shipped now fails on the first step on torch 2.10 — the minimum version pyproject.toml allows.

Needs action:

• Set tf32: false in modelopt_recipes/general/speculative_decoding/dflash2.yaml, or raise an actionable error from the compile path: per your own PR body, tf32: true + dflash_use_torch_compile: true makes Inductor raise on torch 2.9/2.10, and your numbers show tf32 buys nothing here.
• Add a comment beside the two newly flipped dflash: flags naming the short-answer regression (compile pins blocks to dflash_num_anchors; 3.6x slower in your toy run) and dflash.dflash_use_torch_compile=false. Today that caveat is only in the PR body.
• Default dataloader_in_order to True and set it false in the streaming recipes that benefit, or explain in the PR body why the repo-wide flip is preferred — the gain was measured in one mode, the determinism loss lands on all of them.
• Warn via warn_rank_0 when cp_size > 1 overrides an explicit dataloader_in_order=false in main.py:132.

No action needed:

• Two unit tests cover the flag reaching the torch loader and the CP guard; both fail without the change. 81 added lines, well under the size budget. No licensing surface.
• Baseline .agents/review-guidelines/code-review.md at 687720b5; reviewed head 1ef49638; all 6 changed files read. Stacked on haoguo/dflash-rope-fp32, so the diff is against that base.
• Verification limits: no tests or CI run, and the sandbox here blocked inspecting the installed accelerate/transformers, so the base_dataloader reach-through and skip_first_batches carry-over are assessed from source and your reported runs, not re-verified.

[shengliangxu]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

Changes requested: the DFlash2 default fails on supported PyTorch 2.10, and the documented eager opt-out is ignored on HF-checkpoint resume.

Needs action:

• Fix TF32/Inductor compatibility in dflash2.yaml and add a real GPU forward/backward test for the combined compile+FlexAttention default (inline at line 73).
• Honor compile/FlexAttention overrides on HF-checkpoint resume in main.py, preserving weights and saved ModelOpt configuration; add resume coverage (inline at CHANGELOG.rst:32).
• Move the two-worker loader integration test out of tests/unit into tests/gpu*, following CONTRIBUTING’s process-spawning-test rule (inline at test_speculative_decoding_main.py:356).

[shengliangxu]

Bot comment.

[P2] The dflash.dflash_use_torch_compile=false workaround does not take effect when resuming a consolidated HF checkpoint created with compilation enabled. main.train() detects checkpoints automatically and its HF-format branch restores the saved ModelOpt config through from_pretrained; only the fresh/sharded-checkpoint branch consumes recipe.dflash. Thus the restored model retains dflash_use_torch_compile=True and _compiled_body, despite the CLI override. The Usage snippet's FlexAttention opt-out is likewise ignored on that path. Please apply these execution flags to the restored model without rebuilding its draft weights, synchronize the saved ModelOpt configuration, and add an HF-resume test that switches a compile+flex checkpoint back to eager SDPA.

[h-guo18]

That's right, but it isn't specific to these two flags. On the HF-format resume path, from_pretrained restores the whole DFlash config saved with the checkpoint, and recipe.dflash is never read, so every DFlash setting given at resume is ignored, not only compile and flex. Deciding which settings may override a checkpoint on resume, and keeping the saved config in sync, is a change to resume semantics. I'd rather do that in its own PR than inside this one.

For this PR, 4f3724e narrows the CHANGELOG entry: the override applies when a run starts, and a resumed run keeps the checkpoint's DFlash settings. The PR description's usage line now says the same.

[h-guo18]

Thanks for the Pensieve review. Both items are addressed in de31718:

• TF32 with compile on PyTorch 2.10. The DFlash2 recipe now sets tf32: false, so it runs as shipped on the 2.10 floor. A new GPU test, test_compiled_flex_matches_eager_flex, runs the recipe's compile + flex pair forward and backward and checks the loss against eager flex. It passes on PyTorch 2.10 and 2.13.
• Worker-based test. It starts DataLoader workers, so it moved from tests/unit to tests/examples/speculative_decoding/test_dataloader_in_order.py. That lane runs when the trainer code under examples/speculative_decoding changes. tests/gpu would miss such changes, because its lane does not watch examples/. The lightweight CP test stays in tests/unit.

The Claude review's threads each have a reply. The PRs above this one in the stack are rebased onto it.

[h-guo18]

/claude review

[claude]

Approving: the TF32/Inductor crash, the missing compile+flex GPU coverage and the tests/unit worker test are all fixed, and the 16 added logic lines hold up.

No action needed:

• ✔️ Resolved since the last review: tf32: false with the Inductor reason inline, the short-answer caveat beside the flipped dflash: flags, the default-flip rationale in the PR body, and the worker test out of tests/unit.
• hf_dflash.py:706 confirms compile or flex pins the anchor count, so the new recipe comment is accurate; the new GPU test's dflash_use_torch_compile=True does override the shared helper's eager default, and tests/examples allows 300 s.
• Only tools/launcher/common/eagle3/train_eagle_streaming.sh sets dataloader_num_workers > 0 today, so the out-of-order default changes behavior just for the streaming path it was measured on.
• tests/examples/speculative_decoding already hosts process-spawning tests and its lane watches examples/, which meets CONTRIBUTING's intent of keeping tests/unit fast.
• The warn_rank_0 idea for the CP override is still unimplemented at main.py:133; you replied in-thread and inline threads aren't readable from this run, so I'm not re-raising it.
• Baseline .agents/review-guidelines/code-review.md at 687720b5; head 4f3724ef; all 8 changed files read; the new test's header matches LICENSE_HEADER. Limits: no tests or CI run, and the installed accelerate/transformers were again not inspectable, so the base_dataloader reach-through rests on the new ordering test and your reported runs.

[ChenhanYu]

Automatic approval delegated by chenhany (mode: approve).

Pensieve approve review using chenhany's profile (revision 12).

Approve: the prior compatibility and coverage concerns are resolved; no new material defects were found.

No action needed:

• ✔️ Resolved since 5727331218: the TF32 default, real compile+FlexAttention forward/backward coverage, and worker-test placement. The short-answer caveat and default-ordering rationale are also present.
• The worker regression belongs in the example integration lane: its filter watches examples/speculative_decoding/, unlike the GPU-test filter, and process startup is outside tests/unit. It checks both ordering modes and preservation of every sample.
• The new GPU case exercises actual Inductor compilation with DFlash2 components, compares loss against eager FlexAttention, and checks finite gradients. Existing test assertions are unchanged.
• The fresh-run opt-out documentation now explicitly excludes checkpoint resume; preserving saved DFlash settings predates this PR. The optional CP-warning suggestion has a reasonable documented rationale for omission.
• The diff adds 16 core source lines; the new test header matches LICENSE_HEADER.
• Baseline: 687720b5d17c40c9d9fd7d28cc6622d5df99cdf6. Tests were not executed; current-head CI results were unavailable. Duplicate checking covers the supplied bounded history only.

Review baseline: https://github.com/NVIDIA/Model-Optimizer/blob/687720b5d17c40c9d9fd7d28cc6622d5df99cdf6/.agents/review-guidelines/code-review.md

Review run: c9bfb4fbccad432db9e6b4dcae836cba