[h-guo18]

Stack (merge in this order; each PR's base is the branch of the one before it). [7/7] 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: draft RoPE frequencies stay in fp32
2. [2/7] fix(speculative): make streaming hidden-state capture work on vLLM 0.29 #2522: streaming hidden-state capture on vLLM >= 0.29
3. [3/7] fix(speculative): keep the streaming corpus memory-mapped instead of copying it per rank #3142: streaming corpus stays memory-mapped <- this PR
4. [4/7] fix(speculative): weigh DFlash2's selector loss like the backbone loss #3143: DFlash2 selector loss weighted like the backbone loss
5. [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
6. [3/4] perf(speculative): fuse the DFlash-family loss with the lm_head projection #3010: fused DFlash-family loss
7. [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: Bug fix

StreamingDataset.__init__ turned its entries into a list. The streaming trainer (examples/speculative_decoding/eagle_utils.py) passes an Arrow-backed datasets.Dataset, so list() decoded every row of the corpus into Python objects in every rank. Every forked DataLoader worker then copy-on-wrote the pages it touched.

An indexable source is now kept as is: its rows stay memory-mapped and are shared by the forked workers. Only an iterator is still materialized. The empty-corpus check moved after the assignment, since a generator is always truthy.

Why it matters. The copy grows with the corpus, the rank count and the worker count:

• On a 3.1M-row Muse-Spark corpus the list measured 80 GB per rank. With 4 ranks and 8 DataLoader workers each, it OOM-killed a trainer node with 950 GB of host memory.
• On the 3.48M-row corpus of the released Muse-Spark-Open DFlash2 drafter, a trainer node stopped responding 20 steps into training. With this fix the same run trains normally (see [7/7]).

Usage

No API or config change.

Testing

• New unit test test_an_indexable_corpus_is_kept_and_an_iterator_materialized in tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py. It fails on the previous code and passes with this change; the file's 19 tests pass (CPU, transformers 5.5.4).
• End to end, see [7/7]: Muse-Spark-Open DFlash2 streaming on GB300, 2 serve nodes and 1 trainer node with 4 ranks and 8 workers each, on the full 3.48M-row corpus.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅
• 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

_{Comment @coderabbitai help to get the list of available commands.}

[h-guo18]

Folded into #2617 together with #2617's own fix and the other two small fixes, rebased onto main after #2757. It is now [1/4] of the Muse-Spark-Open DFlash2 recipe stack. The commits are the same, re-signed, with the selector one ported onto #2757's helpers; see #2617 for the combined description and testing.