From f18cb9a945b2627a307978975e47eaf195581af6 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:23:56 +0000 Subject: [PATCH] fix(speculative): keep the streaming corpus memory-mapped instead of copying it per rank StreamingDataset turned its entries into a list. The trainer passes an Arrow-backed datasets.Dataset, so list() decoded every row into Python objects in every rank, and every forked DataLoader worker then copy-on-wrote the pages it touched. An indexable source is now kept as is; only an iterator is still materialized. On a 3.1M-row Muse-Spark corpus the list measured 80 GB per rank, and 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. Kept as Arrow, the rows are memory-mapped and shared by the forked workers. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 1 + .../speculative/plugins/hf_streaming_dataset.py | 12 ++++++++---- .../speculative/plugins/test_hf_streaming_dataset.py | 10 ++++++++++ 3 files changed, 19 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index a38f163db99..fda53e5ef53 100755 --- a/CHANGELOG.rst +++ b/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) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py b/modelopt/torch/speculative/plugins/hf_streaming_dataset.py index 5e02097c96a..f226899baaf 100644 --- a/modelopt/torch/speculative/plugins/hf_streaming_dataset.py +++ b/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 diff --git a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py b/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py index 00b5a1f6ab5..bdd4517384b 100644 --- a/tests/unit/torch/speculative/plugins/test_hf_streaming_dataset.py +++ b/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