diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3d0112d84d6..00770449df8 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``, 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)* 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..2112cec3d3f 100644 --- a/modelopt_recipes/general/speculative_decoding/dflash2.yaml +++ b/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 diff --git a/tests/examples/speculative_decoding/test_dataloader_in_order.py b/tests/examples/speculative_decoding/test_dataloader_in_order.py new file mode 100644 index 00000000000..39e2480fec6 --- /dev/null +++ b/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 diff --git a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 5f2e8adc045..01ed2713068 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/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) diff --git a/tests/unit/examples/test_speculative_decoding_main.py b/tests/unit/examples/test_speculative_decoding_main.py index f0fb2c6113d..591cadc50ff 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/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)