[2/4] perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe

CHANGELOG.rst

@@ -28,6 +28,8 @@ Changelog
 *Speculative Decoding*
 
 - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint.
+- With ``training.dataloader_num_workers`` > 0, ``examples/speculative_decoding/main.py`` now yields each training batch as soon as a worker finishes it, so one slow sample, such as a slow hidden-state fetch in streaming training, no longer stalls every DDP rank. Sample order within the prefetch window is then no longer deterministic; set ``training.dataloader_in_order=true`` to restore it (context-parallel runs always keep it).
+- The ``general/speculative_decoding/dflash2`` recipe now enables ``dflash_use_torch_compile`` and ``dflash_use_flex_attention``, and turns ``tf32`` off, which compilation needs on PyTorch 2.9 and 2.10. To train eagerly as before, start the run with ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``; a run resumed from a checkpoint keeps the checkpoint's DFlash settings.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

examples/speculative_decoding/eagle_utils.py

@@ -235,6 +235,14 @@ def __init__(
         super().__init__(*args, **kwargs)
         self.lora_lr_multiplier = lora_lr_multiplier
 
+    def get_train_dataloader(self):
+        """Apply ``dataloader_in_order``, which ``Trainer`` forwards only from transformers 5.15."""
+        dataloader = super().get_train_dataloader()
+        # accelerate's wrapper forwards attribute reads to the torch loader, not writes.
+        loader = getattr(dataloader, "base_dataloader", dataloader)
+        loader.in_order = getattr(self.args, "dataloader_in_order", True)
+        return dataloader
+
     def create_optimizer(self, model=None):
         """Override to give LoRA parameters a higher learning rate.
 

examples/speculative_decoding/main.py

@@ -129,6 +129,9 @@ def init_distributed_env(training_args: transformers.TrainingArguments) -> None:
     """
     if training_args.cp_size < 1:
         raise ValueError(f"cp_size must be >= 1, got {training_args.cp_size}.")
+    if training_args.cp_size > 1:
+        # Every rank of a CP group loads the same batches, so all must take them in sampler order.
+        training_args.dataloader_in_order = True
     world_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count() or 1))
     if training_args.dp_shard_size is None:
         training_args.dp_shard_size = world_size // training_args.cp_size
@@ -549,7 +552,8 @@ def train():
     # Leave training_args.ignore_data_skip at its default (False). The dataset is
     # map-style, so HF Trainer's resume skips consumed indices at the batch-sampler
     # level (accelerate.skip_first_batches) without re-fetching them, landing at the
-    # exact data position. Setting it True would restart the data order from the top.
+    # saved data position (to within the batches in flight when dataloader_in_order is
+    # off). Setting it True would restart the data order from the top.
 
     # Tell the draft model the CP degree so it skips the dense eagle mask under CP.
     model.eagle_cp_size = training_args.cp_size

modelopt/torch/speculative/plugins/hf_training_args.py

@@ -135,3 +135,6 @@ class TrainingArguments(BaseModel):
     answer_only_loss: bool = False
     cp_size: int = 1
     dp_shard_size: int | None = None
+    # Off: take batches as workers finish them so one slow fetch does not stall every DDP rank;
+    # CP forces it on. Mirrors HF's field (transformers>=5.15) so older versions accept it.
+    dataloader_in_order: bool = False

modelopt_recipes/general/speculative_decoding/dflash2.yaml

@@ -58,7 +58,9 @@ training:
   max_grad_norm: 1.0
   dataloader_drop_last: true
   bf16: true
-  tf32: true
+  # Off: with dflash_use_torch_compile, TF32 makes Inductor fail at the first step on PyTorch
+  # 2.9 and 2.10.
+  tf32: false
   remove_unused_columns: false
   # Safe default: the selector params are unused when
   # dflash_selector_loss_alpha == 0, which would otherwise trip DDP.
@@ -70,7 +72,10 @@ training:
 dflash:
   dflash_block_size: 16
   dflash_num_anchors: 256
-  dflash_use_torch_compile: false
+  # Both fix the draft's block count at dflash_num_anchors. With very short answers that
+  # padding outweighs their speedup, so set both false there.
+  dflash_use_torch_compile: true
+  dflash_use_flex_attention: true
   dflash_self_logit_distillation: false
   # gamma for exponential loss decay (block_size=16 -> 7).
   dflash_loss_decay_factor: 7.0

tests/examples/speculative_decoding/test_dataloader_in_order.py

@@ -0,0 +1,91 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""``training.dataloader_in_order`` reaches the torch loader the speculative-decoding trainer iterates.
+
+It starts DataLoader worker processes, so it lives here rather than in ``tests/unit``.
+"""
+
+import importlib.util
+import multiprocessing
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+_EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "speculative_decoding"
+
+
+@pytest.fixture(scope="module")
+def main_mod():
+    """Import the entrypoint by path; its sibling modules need the example dir on sys.path."""
+    sys.path.insert(0, str(_EXAMPLE_DIR))
+    try:
+        spec = importlib.util.spec_from_file_location("sd_main", _EXAMPLE_DIR / "main.py")
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+    finally:
+        sys.path.pop(0)
+    return module
+
+
+class _FirstIndexWaits(torch.utils.data.Dataset):
+    """Index 0 waits on ``release``, so only a loader free to reorder yields the others first."""
+
+    def __init__(self, release, timeout):
+        self.release, self.timeout = release, timeout
+
+    def __len__(self):
+        return 8
+
+    def __getitem__(self, idx):
+        if idx == 0:
+            self.release.wait(self.timeout)
+        return idx
+
+
+def _collate_ids(samples):
+    return {"idx": torch.tensor(samples)}
+
+
+@pytest.mark.parametrize("in_order", [False, True])
+def test_dataloader_in_order_reaches_the_torch_loader(main_mod, tmp_path, in_order):
+    """``Trainer`` forwards the flag only from transformers 5.15; accelerate's wrapper drops writes."""
+    release = multiprocessing.Event()
+    # In order, nothing can overtake index 0 to release it, so it waits out a short timeout.
+    dataset = _FirstIndexWaits(release, timeout=1 if in_order else 10)
+    trainer = main_mod.EagleTrainerWithAccLog(
+        model=torch.nn.Linear(1, 1),
+        args=main_mod.HfTrainingArguments(
+            output_dir=str(tmp_path),
+            per_device_train_batch_size=1,
+            dataloader_num_workers=2,
+            dataloader_in_order=in_order,
+            train_sampling_strategy="sequential",
+            remove_unused_columns=False,
+            report_to="none",
+            use_cpu=True,
+        ),
+        train_dataset=dataset,
+        data_collator=_collate_ids,
+    )
+    order = []
+    for batch in trainer.get_train_dataloader():
+        order += batch["idx"].tolist()
+        release.set()
+    assert sorted(order) == list(range(8))
+    assert (order == sorted(order)) == in_order, order

tests/gpu/torch/speculative/plugins/test_hf_dflash.py

@@ -395,3 +395,36 @@ def test_backward_produces_finite_grads(self):
         ]
         assert grads, "no draft gradients were produced"
         assert all(torch.isfinite(g).all() for g in grads)
+
+    def test_compiled_flex_matches_eager_flex(self):
+        """The DFlash2 recipe's compile + flex pair trains: eager flex's loss, finite grads."""
+        # The recipe's DFlash2 draft and objective, sized for this model.
+        dflash2 = {
+            "dflash_self_logit_distillation": False,
+            "dflash_lk_loss_type": "lambda",
+            "dflash_architecture_config": {
+                "num_hidden_layers": 2,
+                "projector_type": "dflash2",
+                "conv_kernel_size": 2,
+                "conv_group_size": 16,
+                "selector_rank": 8,
+                "selector_top_k": 5,
+            },
+        }
+        eager = self._model(dflash_use_flex_attention=True, **dflash2)
+        compiled = self._model(
+            dflash_use_flex_attention=True, dflash_use_torch_compile=True, **dflash2
+        )
+        compiled.dflash_module.load_state_dict(eager.dflash_module.state_dict())
+        input_ids, attention_mask = self._inputs()
+
+        torch.manual_seed(1234)
+        loss_eager = eager(input_ids=input_ids, attention_mask=attention_mask).loss
+        torch.manual_seed(1234)
+        loss = compiled(input_ids=input_ids, attention_mask=attention_mask).loss
+        torch.testing.assert_close(loss, loss_eager, rtol=2e-2, atol=2e-2)
+
+        loss.backward()
+        grads = [p.grad for p in compiled.dflash_module.parameters() if p.grad is not None]
+        assert grads, "no draft gradients were produced"
+        assert all(torch.isfinite(g).all() for g in grads)

tests/unit/examples/test_speculative_decoding_main.py

@@ -318,3 +318,12 @@ def test_sparse_path_skips_the_base_head_attach(main_mod, stub_tokenizers):
     assert getattr(draft, "_base_lm_head", None) is None, (
         "the sparse path attached the base head it does not use"
     )
+
+
+@pytest.mark.parametrize("cp_size", [1, 2])
+def test_context_parallelism_keeps_batches_in_order(main_mod, monkeypatch, cp_size):
+    """Ranks of a CP group load the same batches; taken out of order they would split different ones."""
+    monkeypatch.setenv("WORLD_SIZE", "2")
+    args = types.SimpleNamespace(cp_size=cp_size, dp_shard_size=None, dataloader_in_order=False)
+    main_mod.init_distributed_env(args)
+    assert args.dataloader_in_order == (cp_size > 1)