[3/7] fix(speculative): keep the streaming corpus memory-mapped instead of copying it per rank

CHANGELOG.rst

@@ -97,6 +97,7 @@ Changelog
 - Fix the column-major ``DequantizeLinear`` surgery leaving a negative or default ``axis`` attribute unchanged while transposing the weight it refers to, which changed the model's output. Re-run the surgery on affected models.
 - Fix QAT of Hugging Face fused MoE experts with non-reentrant activation checkpointing, as in the ``examples/llm_qat`` FSDP2 config: the experts' ``F.linear`` interception stayed installed after the recompute, and training crashed with ``aten.mm.default got mixed torch.Tensor and DTensor``.
 - Fix quantization-aware training of Diffusers models failing in the backward pass with ``NotImplementedError`` once their attention modules are quantized, with or without MHA quantization.
+- Fix streaming speculative-decoding training copying the whole training corpus into every rank and every DataLoader worker, which ran trainer nodes out of host memory on corpora of a few million samples. The corpus now stays memory-mapped.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^

modelopt/torch/speculative/plugins/hf_streaming_dataset.py

@@ -218,12 +218,16 @@ def __init__(
             config: Tuning knobs (timeout, answer_only_loss, ...); defaults to
                 ``self.config_cls()``. See :class:`StreamingConfig`.
         """
-        if not entries:
-            raise ValueError("entries is empty")
         self.tokenizer = tokenizer
         self.config = config if config is not None else self.config_cls()
-        # Materialize to a plain list so DataLoader worker processes fork it cheaply.
-        self.entries = list(entries)
+        # Keep an indexable source (e.g. an Arrow-backed datasets.Dataset) as is: list()
+        # would copy every row into each rank, and each forked worker would copy it again.
+        if hasattr(entries, "__len__") and hasattr(entries, "__getitem__"):
+            self.entries = entries
+        else:
+            self.entries = list(entries)
+        if len(self.entries) == 0:
+            raise ValueError("entries is empty")
         # Per-process consecutive-failure counter for the circuit breaker. Reset to 0
         # on every successful fetch; tripped only by fetch failures (not unfit entries).
         self._consecutive_fail = 0

tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py

@@ -56,6 +56,16 @@ def test_empty_corpus_raises():
         StreamingDataset([], tokenizer=MagicMock(), config=StreamingConfig())
 
 
+def test_an_indexable_corpus_is_kept_and_an_iterator_materialized():
+    """An Arrow-backed corpus stays memory-mapped instead of being copied into every rank."""
+    datasets = pytest.importorskip("datasets")
+    corpus = datasets.Dataset.from_list(_entries(5))
+    ds = StreamingDataset(corpus, tokenizer=MagicMock(), config=StreamingConfig())
+    assert ds.entries is corpus
+    ds = StreamingDataset(iter(_entries(5)), tokenizer=MagicMock(), config=StreamingConfig())
+    assert ds.entries == _entries(5)
+
+
 def test_len_matches_corpus():
     ds = StreamingDataset(_entries(37), tokenizer=MagicMock(), config=StreamingConfig())
     assert len(ds) == 37