From a35712e7d6767f1c9de8456b47ef9bead111d859 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:20:38 +0000 Subject: [PATCH] perf(speculative): deliver training batches out of order; compile + FlexAttention in the DFlash2 recipe Add training.dataloader_in_order (default false) to the speculative-decoding training arguments. With dataloader workers, a rank then takes whichever batch a worker finishes first instead of waiting for its slowest outstanding sample, and DDP no longer turns one rank's slow fetch into a stall on every rank. HF Trainer forwards the flag itself only from transformers 5.15, and accelerate's DataLoaderShard forwards attribute reads, not writes, to the torch loader it wraps. EagleTrainerWithAccLog.get_train_dataloader therefore sets it on that loader. The field shares HF's name, so on 5.15+ it overrides HF's own default. Context parallelism keeps in-order delivery: accelerate gives every rank of a CP group the same data shard, and each loads it on its own, so out of order they would split different sequences at the same step. Measured on DFlash2 streaming training, 8 trainer nodes x 4 GPUs, DDP, global batch 64, seq 4096, 8 workers per rank, same data_seed: - in_order=False at 8 serve + 8 trainer nodes: 0.655 -> 0.523 s/step; serve replicas 26% -> 3% idle; worst per-rank batch wait 5.9 s -> ~2 s. - Over ~2.5K steps: mean |loss diff| <= 0.016, mean acceptance diff within +-0.003. Also turn on dflash_use_torch_compile and dflash_use_flex_attention in the DFlash2 recipe, the configuration the numbers below were measured with. With the trainer as the bottleneck (12 or 16 serve + 8 trainer nodes) compile gave 0.445 -> 0.386 and 0.446 -> 0.389 s/step with no recompiles over 2.5K steps; when serve-bound it gave nothing. On one GPU at the shipped online example's shape (PyTorch 2.10, tf32 off) the recipe goes from 1612 to 985 ms/step (1.64x). On PyTorch 2.9/2.10 the recipe's tf32: true makes Inductor's TF32 check raise. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 2 + examples/speculative_decoding/eagle_utils.py | 8 +++ examples/speculative_decoding/main.py | 6 +- .../speculative/plugins/hf_training_args.py | 3 + .../general/speculative_decoding/dflash2.yaml | 3 +- .../test_speculative_decoding_main.py | 62 ++++++++++++++++++- 6 files changed, 81 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3d0112d84d6..4f71f5c4c2b 100755 --- a/CHANGELOG.rst +++ b/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)* diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index ea0d001045d..f1527a0e533 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/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. diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 142d058ebb5..9df976d36e5 100644 --- a/examples/speculative_decoding/main.py +++ b/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 diff --git a/modelopt/torch/speculative/plugins/hf_training_args.py b/modelopt/torch/speculative/plugins/hf_training_args.py index fed44d06228..a77a6e970a5 100644 --- a/modelopt/torch/speculative/plugins/hf_training_args.py +++ b/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 diff --git a/modelopt_recipes/general/speculative_decoding/dflash2.yaml b/modelopt_recipes/general/speculative_decoding/dflash2.yaml index 1f02e44b1aa..33681ed5f1a 100644 --- a/modelopt_recipes/general/speculative_decoding/dflash2.yaml +++ b/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 diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index f0fb2c6113d..8749653af09 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/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)