[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``. With PyTorch 2.9 or 2.10 its ``tf32: true`` makes Inductor fail on the first step; set ``training.tf32=false`` or ``dflash.dflash_use_torch_compile=false`` there.
 
 *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

@@ -70,7 +70,8 @@ training:
 dflash:
   dflash_block_size: 16
   dflash_num_anchors: 256
-  dflash_use_torch_compile: false
+  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/unit/examples/test_speculative_decoding_main.py

@@ -13,13 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""``setup_external_draft`` wiring in the speculative-decoding entrypoint.
+"""``setup_external_draft`` and trainer wiring in the speculative-decoding entrypoint.
 
 Lives under ``tests/unit`` rather than ``tests/examples/`` so CI runs it on every
 change: the example lanes cover a fixed allowlist.
 """
 
 import importlib.util
+import multiprocessing
 import sys
 import types
 from pathlib import Path
@@ -318,3 +319,62 @@ 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"
     )
+
+
+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."""
+    if multiprocessing.get_start_method() != "fork":
+        pytest.skip("workers started without fork re-import torch, too slow for a unit test")
+    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
+
+
+@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)