Fix EAGLE-3 context-parallel training and re-enable its tests

.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]"

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**

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
                 )

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,

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,

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);

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:

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 = [