From 1ef49638ae57a1fb9955edfc6749d29409a9d7ec 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 1/3] 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) From de317180b58fb6a7d429580ef6abed882735174c Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:29:56 +0000 Subject: [PATCH 2/3] fix(speculative): turn tf32 off in the compiled DFlash2 recipe; move the worker test to examples Review follow-ups. - The recipe turns tf32 off. With dflash_use_torch_compile on, TF32 makes Inductor fail at the first step on PyTorch 2.9 and 2.10, and 2.10 is the lowest version ModelOpt supports. DFlash2 trains in bf16, so TF32 only reached fp32 matmuls such as the draft's rotary embedding. - The recipe now says, next to the compile and flex flags, that both fix the draft's block count at dflash_num_anchors, which costs more than it saves on very short answers. - A GPU test runs the recipe's compile + flex pair forward and backward on a tiny model and checks the loss against eager flex. It passes on PyTorch 2.10 and 2.13. - The dataloader ordering test starts DataLoader workers, so it moves from tests/unit to tests/examples/speculative_decoding, the lane that runs when the trainer code under examples/ changes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 2 +- .../general/speculative_decoding/dflash2.yaml | 6 +- .../test_dataloader_in_order.py | 91 +++++++++++++++++++ .../speculative/plugins/test_hf_dflash.py | 18 ++++ .../test_speculative_decoding_main.py | 53 +---------- 5 files changed, 116 insertions(+), 54 deletions(-) create mode 100644 tests/examples/speculative_decoding/test_dataloader_in_order.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f71f5c4c2b..e5606c930c7 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -29,7 +29,7 @@ Changelog - 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. +- 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, set ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``. *Megatron Framework (M-LM / M-Bridge)* diff --git a/modelopt_recipes/general/speculative_decoding/dflash2.yaml b/modelopt_recipes/general/speculative_decoding/dflash2.yaml index 33681ed5f1a..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,6 +72,8 @@ training: dflash: dflash_block_size: 16 dflash_num_anchors: 256 + # 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 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..7bf834af559 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -395,3 +395,21 @@ 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.""" + eager = self._model(dflash_use_flex_attention=True) + compiled = self._model(dflash_use_flex_attention=True, dflash_use_torch_compile=True) + 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 8749653af09..591cadc50ff 100644 --- a/tests/unit/examples/test_speculative_decoding_main.py +++ b/tests/unit/examples/test_speculative_decoding_main.py @@ -13,14 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""``setup_external_draft`` and trainer wiring in the speculative-decoding entrypoint. +"""``setup_external_draft`` 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 @@ -321,56 +320,6 @@ def test_sparse_path_skips_the_base_head_attach(main_mod, stub_tokenizers): ) -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.""" From 4f3724efac52c931639d70cf9f569c2fa3852d9f Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Sun, 1 Nov 2026 09:51:53 +0000 Subject: [PATCH 3/3] test(speculative): test the compiled DFlash2 recipe with its own draft; scope the eager override Review follow-ups. - test_compiled_flex_matches_eager_flex now builds the recipe's DFlash2 draft (conv projector and candidate selector) and objective (lambda LK loss, no self-distillation), sized for the test model, instead of a plain DFlash draft. It passes on PyTorch 2.10 and 2.13. - The CHANGELOG entry says the compile/flex override applies when a run starts: a run resumed from a consolidated checkpoint restores the DFlash settings saved with it, as it already did for every DFlash setting. Co-Authored-By: Claude Opus 5.5 Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- CHANGELOG.rst | 2 +- .../speculative/plugins/test_hf_dflash.py | 19 +++++++++++++++++-- 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index e5606c930c7..00770449df8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -29,7 +29,7 @@ Changelog - 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, set ``dflash.dflash_use_torch_compile=false dflash.dflash_use_flex_attention=false``. +- 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/tests/gpu/torch/speculative/plugins/test_hf_dflash.py b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py index 7bf834af559..01ed2713068 100644 --- a/tests/gpu/torch/speculative/plugins/test_hf_dflash.py +++ b/tests/gpu/torch/speculative/plugins/test_hf_dflash.py @@ -398,8 +398,23 @@ def test_backward_produces_finite_grads(self): def test_compiled_flex_matches_eager_flex(self): """The DFlash2 recipe's compile + flex pair trains: eager flex's loss, finite grads.""" - eager = self._model(dflash_use_flex_attention=True) - compiled = self._model(dflash_use_flex_attention=True, dflash_use_torch_compile=True) + # 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()