[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.
• 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). With older accelerate, the first partial epoch after a resume runs in order. 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. 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: the recipe's tf32: true makes Inductor raise 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, and so is the launcher's default TRT-LLM container (an NGC 2.11 pre-release). Workaround: training.tf32=false, or dflash.dflash_use_torch_compile=false.
• 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, as before
python main.py --config modelopt_recipes/general/speculative_decoding/dflash2.yaml ... dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false

Testing

• tests/unit/examples/test_speculative_decoding_main.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).
  • Takes 0.07 s / 1.05 s on 4 cores, and is skipped where workers don't start by fork (Windows, Python 3.14).
• 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.
• 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, recipe as shipped (tf32 on): 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, tf32=false: 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 (compile + flex, tf32 on) end to end on PyTorch 2.11; the table above is 2.10 with tf32=false. 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 and uses FlexAttention.
• 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?: ❌ Draft

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: a35712e
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: ae918452-8d35-43f0-8d52-db051496b8e8

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.