From 1e81510ef870ff233da5d7eda73c3db75fd170a4 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Thu, 6 Aug 2026 07:21:28 -0700 Subject: [PATCH] Fix EAGLE-3 context-parallel training and re-enable its tests cp_size>1 has been broken since accelerate 1.13 and the tests never ran: the guard compared Version("2.10.0a0") to Version("2.10.0"), which is False on every NGC alpha torch build. - main.py: rebuild the FSDP2 plugin accelerate requires for cp_size>1; the --fsdp launcher flags that supplied it were dropped. Pass the CP degree to the draft model. - modeling_eagle.py: apply the first input norm inside layers[0]'s own forward, where FSDP2 has unsharded its weights, and only stash the input embeds on the path whose pre-hook consumes them. - hf_eagle.py: skip the dense eagle mask under CP and warn that padded positions are therefore unmasked; stop `or 0` replacing a 0.0 loss tensor with an int that detaches the graph. - eagle_utils.py: key TTT-mask injection off the backward call's grad_out kwarg, since newer torch omits attn_bias on the forward call. - utils.py: CUDNN-only SDPA under CP; MATH decomposes the mask add and breaks on DTensors. Drop the speculative_decoding 26.01 container override, added when the lane ran 25.06 and spec-dec needed something newer. Later bumps moved the default past it, so it had become a ceiling holding spec-dec on a 6-month-old image. Verified in nvcr.io/nvidia/pytorch:26.07-py3 on 2 GPUs: all 4 test_llama_eagle3 cases pass, and the full speculative_decoding suite is 16 passed / 2 skipped (the 2 being pre-existing --run-manual tests). Co-Authored-By: Claude Opus 5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .github/workflows/example_tests.yml | 9 +++---- CHANGELOG.rst | 6 +++-- examples/speculative_decoding/eagle_utils.py | 5 ++-- examples/speculative_decoding/main.py | 11 ++++++-- .../torch/speculative/plugins/hf_eagle.py | 25 ++++++++++++++++--- .../speculative/plugins/modeling_eagle.py | 7 ++++-- modelopt/torch/speculative/utils.py | 9 +++++-- .../speculative_decoding/test_eagle.py | 3 --- 8 files changed, 52 insertions(+), 23 deletions(-) diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 7ecec4f6699..f7e06262e1f 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -33,23 +33,20 @@ jobs: pyproject.toml tests/examples/** - ##### PyTorch Example Tests (speculative_decoding requires 26.01 image) ##### + ##### PyTorch Example Tests ##### torch: needs: [pr-gate] if: needs.pr-gate.outputs.any_changed == 'true' strategy: fail-fast: false matrix: - example: [gpt-oss, llm_distill, llm_qat, llm_sparsity, specdec_bench] - include: - - example: speculative_decoding - docker_image: "26.01" + example: [gpt-oss, llm_distill, llm_qat, llm_sparsity, specdec_bench, speculative_decoding] uses: ./.github/workflows/_example_tests_runner.yml permissions: contents: read secrets: inherit with: - docker_image: "nvcr.io/nvidia/pytorch:${{ matrix.docker_image || '26.06' }}-py3" + docker_image: "nvcr.io/nvidia/pytorch:26.06-py3" example: ${{ matrix.example }} timeout_minutes: 30 pip_install_extras: "[hf,dev-test]" diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5d01388d58f..d7dc3668a85 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,7 +1,7 @@ Changelog ========= -0.47 (2026-xx-xx) +0.47 (2026-09-xx) ^^^^^^^^^^^^^^^^^ **New Features** @@ -21,7 +21,9 @@ Changelog **Bug Fixes** -0.46 (2026-08-xx) +- Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. + +0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ **New Features** diff --git a/examples/speculative_decoding/eagle_utils.py b/examples/speculative_decoding/eagle_utils.py index 68c6db45235..2ba2635fbca 100644 --- a/examples/speculative_decoding/eagle_utils.py +++ b/examples/speculative_decoding/eagle_utils.py @@ -518,8 +518,9 @@ def patched_op(*args, **kwargs): raise RuntimeError( f"Failed to capture loop variables in patched _templated_ring_attention: {e}" ) from e - # Set attn mask to permuted TTT mask - if "attn_bias" in kwargs: + # Set attn mask to permuted TTT mask. Newer torch omits the attn_bias kwarg on the + # forward call, so key off grad_out instead to tell forward from backward. + if patch_enbabled and "grad_out" not in kwargs: kwargs["attn_bias"] = _get_sharded_ttt_msk( i, rank, size, query.shape[2], ttt_step, query.dtype ) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index 1750a3b1095..ce9460637ac 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -186,8 +186,12 @@ def train(): raise ValueError(f"data.mode={recipe.data.mode!r} requires data.data_path.") if training_args.cp_size > 1: patch_ring_attention_for_ttt() - # Specific patch to accelerate 1.12.0. Removable after move to 1.13.0 - training_args.parallelism_config.sp_backend = None + # accelerate requires an fsdp_plugin when cp_size > 1; the --fsdp launcher flags that + # used to provide one were dropped from launch_train.sh. + if not training_args.fsdp_plugin_args: + training_args.fsdp = "full_shard" + training_args.fsdp_config = {"fsdp_version": 2} + training_args.fsdp_plugin_args = training_args._process_fsdp_args() if is_master(): pprint(recipe) @@ -308,6 +312,9 @@ def train(): # 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. + # Tell the draft model the CP degree so it skips the dense eagle mask under CP. + model.eagle_cp_size = training_args.cp_size + trainer = EagleTrainerWithAccLog( model=model, processing_class=tokenizer, diff --git a/modelopt/torch/speculative/plugins/hf_eagle.py b/modelopt/torch/speculative/plugins/hf_eagle.py index 207961dab7d..67adab385d2 100644 --- a/modelopt/torch/speculative/plugins/hf_eagle.py +++ b/modelopt/torch/speculative/plugins/hf_eagle.py @@ -26,7 +26,7 @@ from transformers.models.llama.modeling_llama import LlamaDecoderLayer from transformers.utils import ModelOutput -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import print_rank_0, warn_rank_0 from ...export.plugins.hf_spec_export import EagleExporter, SpeculativeDecodingExporter from ..eagle.conversion import EagleDMRegistry @@ -67,6 +67,9 @@ def default_eagle_aux_layer_ids(num_layers: int) -> list[int]: class HFEagleModel(EagleModel): """Eagle Model Class for huggingface models.""" + # Context-parallel degree, set by the training script when it launches with cp_size > 1. + eagle_cp_size: int = 1 + @property def _base_model(self): return self.get_submodule(self.base_model_path) @@ -192,7 +195,7 @@ def get_exporter(self) -> SpeculativeDecodingExporter: def _enable_cp_ttt(self): if self.training and not self.eagle_mix_hidden_states: - return enable_cp_ttt_patch() + return enable_cp_ttt_patch(self.eagle_cp_size) return contextlib.nullcontext() def _set_default_aux_hidden_state_layers(self): @@ -762,12 +765,21 @@ def forward( # ====Run eagle forward with extra training-time-test steps==== num_ttt_steps = self.eagle_ttt_steps if self.training else 1 for ttt_step in range(num_ttt_steps): - # TODO: (hg) during cp training, this mask is not used. Maybe turn it off then. eagle_attention_mask = ( eagle_attn_mask_0 if self.eagle_mix_hidden_states or ttt_step == 0 else self._get_ttt_attention_mask(b, seq_length, ttt_step) ) + # Under CP the dense mask is unused and fatal (plain tensor vs DTensor scores); + # causal masking comes from is_causal and TTT masking from the ring-attention patch. + if self.eagle_cp_size > 1: + warn_rank_0( + "Context-parallel EAGLE training does not mask padded positions: the dense " + "mask cannot be applied to the sharded (DTensor) sequence, so the draft model " + "attends to any pad tokens. Pack or truncate samples to a fixed length under " + "cp_size > 1." + ) + eagle_attention_mask = None with self._enable_cp_ttt(), self._nvtx_range("eagle_forward"): _, eagle_output_hiddens, eagle_logits, eagle_cache = self._eagle_forward( eagle_input_hiddens, @@ -820,7 +832,12 @@ def forward( loss = None assert not self.training, "At least one loss must be computed for training." else: - loss = (base_outputs.loss or 0) + (eagle_loss or 0) + # Test for None, not truthiness: a 0.0 loss tensor is falsy, and `or 0` would + # replace it with an int and detach the graph. + loss = None + for term in (base_outputs.loss, eagle_loss): + if term is not None: + loss = term if loss is None else loss + term return ModelOutput( loss=loss, diff --git a/modelopt/torch/speculative/plugins/modeling_eagle.py b/modelopt/torch/speculative/plugins/modeling_eagle.py index d6bf525edd2..2fed24ac35f 100644 --- a/modelopt/torch/speculative/plugins/modeling_eagle.py +++ b/modelopt/torch/speculative/plugins/modeling_eagle.py @@ -118,7 +118,7 @@ def _eagle3_attention_forward_pre_hook(self, module, args, kwargs): if self._input_embeds is None: raise ValueError("self._input_embeds is None") - input_embeds = self._input_embeds + input_embeds = self.layers[0].input_layernorm(self._input_embeds) self._input_embeds = None kwargs["hidden_states"] = torch.cat( (input_embeds, self.layers[0].hidden_norm(kwargs["hidden_states"])), dim=-1 @@ -160,7 +160,10 @@ def forward( # In EAGLE-3, we save input embeddings to attribute, and use it in first decoder layer by hook function # Also, we normalize input embeddings and hidden states before concatenating them. # The default input norm in first layer attn will be disabled. - self._input_embeds = self.layers[0].input_layernorm(inputs_embeds) + # Stash raw embeds and normalize in the attention pre-hook instead: FSDP2 only + # unshards layers[0] weights inside its own forward. Nothing consumes them without + # the hook, so don't hold the reference past this forward. + self._input_embeds = inputs_embeds if self.config.use_aux_hidden_state else None if self.config.eagle_decoder_type == "llama": # rotary_emb must be pre-initialized by the caller (see HFEagleModel); diff --git a/modelopt/torch/speculative/utils.py b/modelopt/torch/speculative/utils.py index 8c5418bb6b8..b0ab85b6390 100644 --- a/modelopt/torch/speculative/utils.py +++ b/modelopt/torch/speculative/utils.py @@ -552,12 +552,17 @@ def ttt_msk_func(b, h, q_idx, kv_idx): @contextlib.contextmanager -def enable_cp_ttt_patch(): +def enable_cp_ttt_patch(cp_size: int = 1): """Context manager to enable CP TTT patch.""" import modelopt.torch.speculative.plugins.hf_eagle modelopt.torch.speculative.plugins.hf_eagle.ENABLE_CP_TTT_PATCH = True - with sdpa_kernel([SDPBackend.CUDNN_ATTENTION, SDPBackend.MATH]): + # Under CP, restrict to cudnn: MATH decomposes SDPA and breaks on DTensors. Elsewhere keep + # MATH, since cudnn is unavailable on CPU and for some head dims. + backends = [SDPBackend.CUDNN_ATTENTION] + if cp_size == 1: + backends.append(SDPBackend.MATH) + with sdpa_kernel(backends): try: yield finally: diff --git a/tests/examples/speculative_decoding/test_eagle.py b/tests/examples/speculative_decoding/test_eagle.py index 00f0adc33c3..0decc6b03d0 100644 --- a/tests/examples/speculative_decoding/test_eagle.py +++ b/tests/examples/speculative_decoding/test_eagle.py @@ -20,7 +20,6 @@ import safetensors.torch import torch from _test_utils.examples.run_command import MODELOPT_ROOT, run_example_command -from packaging.version import Version from transformers import AutoConfig from modelopt.torch.export.plugins.hf_spec_export import LLAMA_EAGLE_SINGLE_LAYER @@ -127,8 +126,6 @@ def test_llama_eagle3(tiny_llama_path, """Test Eagle3 training with a tiny llama model, using different cp_size values.""" if cp_size == 2 and num_gpus < 2: pytest.skip(f"cp_size=2 requires at least 2 GPUs, but only {num_gpus} found.") - if cp_size == 2 and not Version(torch.__version__) >= Version("2.10.0"): - pytest.skip("cp_size=2 requires torch 2.10.0") output_dir = str(eagle_output_dir / f"eagle-tinyllama-cp{cp_size}-mix{mix_hidden_states}") overrides = [