fix tests

Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>

Author: h-guo18

.github/workflows/example_tests.yml

@@ -63,7 +63,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [llm_distill, llm_qat, llm_sparsity, speculative_decoding]
+        example: [llm_distill, llm_qat, llm_sparsity]
     uses: ./.github/workflows/_example_tests_runner.yml
     secrets: inherit
     with:
@@ -77,7 +77,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        example: [llm_distill, llm_qat, llm_sparsity, speculative_decoding]
+        example: [llm_distill, llm_qat, llm_sparsity]
     uses: ./.github/workflows/_example_tests_runner.yml
     secrets: inherit
     with:
@@ -86,6 +86,28 @@ jobs:
       pip_install_extras: "[hf,dev-test]"
       runner: linux-amd64-gpu-h100-latest-2
 
+  ##### Speculative Decoding Example Tests (requires 25.08 image) #####
+  speculative-decoding-pr:
+    needs: [check-file-changes, wait-checks]
+    if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true'
+    uses: ./.github/workflows/_example_tests_runner.yml
+    secrets: inherit
+    with:
+      docker_image: "nvcr.io/nvidia/pytorch:25.08-py3"
+      example: speculative_decoding
+      pip_install_extras: "[hf,dev-test]"
+      runner: linux-amd64-gpu-l4-latest-1
+
+  speculative-decoding-non-pr:
+    if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
+    uses: ./.github/workflows/_example_tests_runner.yml
+    secrets: inherit
+    with:
+      docker_image: "nvcr.io/nvidia/pytorch:25.08-py3"
+      example: speculative_decoding
+      pip_install_extras: "[hf,dev-test]"
+      runner: linux-amd64-gpu-h100-latest-2
+
   ##### TensorRT-LLM Example Tests #####
   trtllm-pr:
     needs: [check-file-changes, wait-checks]
@@ -150,14 +172,15 @@ jobs:
   example-pr-required-check:
     # Run even if example tests are skipped
     if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }}
-    needs: [check-file-changes, torch-pr, trtllm-pr, onnx-pr]
+    needs: [check-file-changes, torch-pr, speculative-decoding-pr, trtllm-pr, onnx-pr]
     runs-on: ubuntu-latest
     steps:
       - name: Required GPU tests did not succeed
         if: |
           needs.check-file-changes.result != 'success' ||
           (needs.check-file-changes.outputs.any_changed == 'true' && (
             needs.torch-pr.result != 'success' ||
+            needs.speculative-decoding-pr.result != 'success' ||
             needs.trtllm-pr.result != 'success' ||
             needs.onnx-pr.result != 'success'
           ))

examples/speculative_decoding/README.md

@@ -30,7 +30,7 @@ This example focuses on training with Hugging Face. To train with Megatron‑LM,
 
 ### Docker
 
-Please use the PyTorch docker image (e.g., `nvcr.io/nvidia/pytorch:25.06-py3`) or visit our [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more information.
+Please use the PyTorch docker image (e.g., `nvcr.io/nvidia/pytorch:25.08-py3`) or visit our [installation docs](https://nvidia.github.io/Model-Optimizer/getting_started/2_installation.html) for more information.
 
 Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install dataset and example-specific dependencies.
 
@@ -56,7 +56,7 @@ See [other-datasets](#other-datasets) section for other dataset options and inst
 ## Getting Started: Simplified Workflow
 
 ```bash
-bash train_eagle3_and_export.sh --base_model meta-llama/Llama-3.2-1B-Instruct --num_gpu 4
+bash train_eagle3_and_export.sh --base_model meta-llama/Llama-3.2-1B-Instruct
 ```
 
 This one-line command runs a minimal example workflow of training and exporting an EAGLE draft model in Modelopt. Specifically, it
@@ -74,12 +74,11 @@ For small base models that fit in GPU memory, we can collocate them with draft m
 ./launch_train.sh --model $BASE_MODEL \
             --output_dir $OUTPUT_DIR \
             --data input_conversations/daring-anteater.jsonl  \
-            --num_gpu $NUM_GPU \
             --num_epochs $NUM_EPOCH \
             --eagle_config eagle_config.json
 ```
 
-This command will launch `main.py` with `accelerate`. See [section: interact with modelopt.torch.speculative](#interact-with-modelopttorchspeculative) for more details.
+FSDP2 is used by default. To enable context parallelism for long-context training, specify `--cp_size n`.
 The saved modelopt checkpoint is similar in architecture to HF models. It can be further optimized through **ModelOpt**, e.g., PTQ and QAT.
 
 ## Training Draft Model with Offline Base Model
@@ -118,7 +117,6 @@ Once we finish dumping hidden states, launch offline training with an extra `--o
 ./launch_train.sh --model $BASE_MODEL \
             --output_dir $OUTPUT_DIR \
             --data $DATA \
-            --num_gpu $NUM_GPU \
             --num_epochs $NUM_EPOCH \
             --eagle_config eagle_config.json \
             --offline-data $HIDDEN_STATES_DIR

examples/speculative_decoding/eagle_utils.py

@@ -28,14 +28,15 @@
 import torch
 import transformers
 from datasets import load_dataset
+from packaging.version import Version
 from PIL import Image
 from scripts.ar_validate import validate_ar
 from torch.distributed.tensor.experimental._attention import _SDPAMerger
 from torch.utils.data import Dataset
 from transformers import AutoProcessor, Trainer, TrainerCallback
 from transformers.trainer_pt_utils import LabelSmoother
 
-import modelopt.torch.speculative.plugins.transformers
+import modelopt
 from modelopt.torch.speculative.utils import get_ttt_msk_func
 from modelopt.torch.utils import print_rank_0
 from modelopt.torch.utils.distributed import is_master
@@ -577,27 +578,6 @@ def on_step_end(self, args, state, control, **kwargs):
         return control
 
 
-def _compute_ttt_attention_mask(batch_size, seq_length, ttt_step, dtype) -> torch.Tensor:
-    """Return TTT attention_mask tensor of type BlockMask or Tensor depends on eagle attn impl."""
-
-    msk_func = get_ttt_msk_func(seq_length, ttt_step)
-
-    dtypemin = torch.finfo(dtype).min
-    q_len = seq_length
-    kv_len = seq_length * (1 + ttt_step)
-    # Return tensor mask for non-flex attention
-    tensor_mask = msk_func(
-        None,
-        None,
-        torch.arange(q_len).view(1, 1, q_len, 1),
-        torch.arange(kv_len).view(1, 1, 1, kv_len),
-    ).to(torch.cuda.current_device())
-    tensor_mask = torch.full_like(
-        tensor_mask, 0, dtype=dtype, device=torch.cuda.current_device()
-    ).masked_fill(~tensor_mask, dtypemin)
-    return tensor_mask
-
-
 def get_patched_templated_ring_attn(orig_templated_attn: Callable):
     """
     Return patched version of
@@ -659,14 +639,14 @@ def patched_templated_attn(*args, **kwargs):
         patch_enbabled = modelopt.torch.speculative.plugins.transformers.ENABLE_CP_TTT_PATCH
 
         if patch_enbabled and original_op != torch.ops.aten._scaled_dot_product_cudnn_attention:
-            raise ValueError(f"CP TTT only supports cuddn attention now. Got: {original_op}")
+            raise ValueError(f"CP TTT only supports cudnn attention now. Got: {original_op}")
 
         # Unset is_causal to use custom attn mask
         if patch_enbabled:
             kwargs["is_causal"] = False
 
         def patched_op(*args, **kwargs):
-            # Inpect the parent frame to get current shard info
+            # Inspect the parent frame to get current shard info
             # This is sensitive to torch _templated_ring_attention impl
             try:
                 frame: FrameType = inspect.currentframe()
@@ -678,7 +658,9 @@ def patched_op(*args, **kwargs):
                 i = f_back.f_locals["i"]
                 ttt_step = (key.shape[2] // query.shape[2]) - 1
             except Exception as e:
-                print(f"Failed to capture loop variables in patched _templated_ring_attention: {e}")
+                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:
                 kwargs["attn_bias"] = _get_sharded_ttt_msk(
@@ -696,8 +678,18 @@ def patch_ring_attention_for_ttt():
     """Patch torch ring attention to support context parallelism for TTT."""
     # Torch Ring Attention only supports no mask or causal mask. We apply the following patches to enable TTT mask.
 
+    if not (
+        Version(torch.__version__) > Version("2.7.1")
+        and Version(torch.__version__) < Version("2.9.0")
+    ):
+        raise RuntimeError(
+            f"Context parallel TTT only supported for PyTorch 2.8.0 now. "
+            f"Got {torch.__version__}. "
+            f"Please use nvcr.io/nvidia/pytorch:25.08-py3 or torch 2.8.0 or cp_size=1."
+        )
+
     # 1. Disable load balance, which is designed for causal mask.
-    # This affect how buffers are sharded. So need to be done permenantly before accelerate/hf trainer init.
+    # This affect how buffers are sharded. So need to be done permanently before accelerate/hf trainer init.
     torch.distributed.tensor.experimental._attention._cp_options.enable_load_balance = False
 
     # 2. Patch templated ring attention for TTT mask.
@@ -717,9 +709,7 @@ def patch_ring_attention_for_ttt():
     # 3. Patch merger to skip the blank shard to avoid difference in output.
     original_sdpa_merger_step = _SDPAMerger.step
 
-    def patched_sdpa_merger_step(
-        self, out: torch.Tensor, lse: torch.Tensor, partial: bool
-    ) -> torch.Tensor:
+    def patched_sdpa_merger_step(self, out: torch.Tensor, lse: torch.Tensor, partial: bool):
         if lse.sum() <= 0:
             return
         return original_sdpa_merger_step(self, out, lse, partial)

examples/speculative_decoding/launch_train.sh

@@ -170,6 +170,15 @@ else
   VLM_ARGS=""
 fi
 
+if [[ "$GPU_COUNT" -gt 1 ]]; then
+  #Use FSDP2 when multi GPU available
+  FSDP_ARGS="--fsdp 'full_shard' --fsdp_config fsdp_config.json"
+else
+  #Otherwise, single GPU training
+  FSDP_ARGS=""
+fi
+
+
 # Disable tokenizers parallelism to avoid warning
 export TOKENIZERS_PARALLELISM=False
 CMD="accelerate launch --mixed_precision bf16 main.py \
@@ -201,8 +210,7 @@ CMD="accelerate launch --mixed_precision bf16 main.py \
     $VLM_ARGS \
     $OFFLINE_TRAINING_ARGS \
     $SPECULATIVE_ARGS \
-    --fsdp 'full_shard' \
-    --fsdp_config fsdp_config.json \
+    $FSDP_ARGS \
     --cp_size $CP_SIZE \
     --dp_shard_size $DP_SHARD_SIZE \
 "

examples/speculative_decoding/requirements.txt

@@ -1,4 +1,2 @@
 accelerate==1.12.0
-torch==2.8.0
 transformers==5.0.0rc1
-wandb

examples/speculative_decoding/train_eagle3_and_export.sh

@@ -17,23 +17,18 @@
 
 set -eo pipefail
 
-# Set default values for BASE_MODEL, NUM_GPU, and DATA
+# Set default values for BASE_MODEL and DATA
 BASE_MODEL=meta-llama/Llama-3.2-1B-Instruct
-NUM_GPU=1
 DATA=input_conversations/daring-anteater.jsonl
 
-# Parse input arguments --base_model, --num_gpu, and --data
+# Parse input arguments --base_model and --data
 while [[ $# -gt 0 ]]; do
   key="$1"
   case $key in
     --base_model)
       BASE_MODEL="$2"
       shift; shift
       ;;
-    --num_gpu)
-      NUM_GPU="$2"
-      shift; shift
-      ;;
     --data)
       DATA="$2"
       shift; shift
@@ -49,15 +44,6 @@ while [[ $# -gt 0 ]]; do
   esac
 done
 
-
-if [[ "$NUM_GPU" == 1 ]]; then
-  export CUDA_VISIBLE_DEVICES=0
-else
-  # Export as 0,1,...,N-1 for NUM_GPU GPUs
-  devs="$(seq -s, 0 $((NUM_GPU-1)))"
-  export CUDA_VISIBLE_DEVICES="$devs"
-fi
-
 if [[ "$OFFLINE_DATA_PATH" != "" ]]; then
   OFFLINE_DATA_ARGS="--offline-data $OFFLINE_DATA_PATH"
 else
@@ -73,7 +59,6 @@ mkdir -p "$(dirname "$OUTPUT_DIR")"
             --output_dir $OUTPUT_DIR \
             $OFFLINE_DATA_ARGS \
             --data $DATA \
-            --num_gpu $NUM_GPU \
             --num_epochs 2 \
             --eagle_config eagle_config.json
 

modelopt/torch/speculative/plugins/transformers.py

@@ -34,6 +34,8 @@
 from typing import Any
 
 import torch
+import transformers
+from packaging.version import Version
 from torch import nn
 from torch.nn import CrossEntropyLoss
 from torch.nn.attention.flex_attention import BlockMask, create_block_mask
@@ -61,11 +63,22 @@
     temporary_set_config_value,
 )
 
+__all__ = ["HFARValidation", "HFEagleModel", "HFMedusaModel"]
+
 IGNORE_TOKEN_ID = LabelSmoother.ignore_index
 ENABLE_CP_TTT_PATCH = False
+# module variable to cache attention mask for cp ttt
 CACHED_SHARD_TTT_MASKS = {}
 
 
+def _get_empty_cache(config):
+    """Return an empty cache. Handle different versions of transformers for unit tests."""
+    if Version(transformers.__version__) >= Version("4.54"):
+        return DynamicCache(config=config)
+    else:
+        return DynamicCache()
+
+
 @MedusaDMRegistry.register({PreTrainedModel: "hf.PreTrainedModel"})
 class HFMedusaModel(MedusaModel):
     """Medusa Model Class for huggingface models."""
@@ -852,9 +865,7 @@ def forward(
                 base_model_logits = base_outputs["base_model_logits"]
             else:
                 base_model_logits = self.lm_head(base_model_hidden_states)
-            base_model_loss = None
-            past_key_values = DynamicCache()  # Dummy cache
-
+            base_model_loss, past_key_values = None, None
         else:
             base_model_hidden_states, base_model_logits, base_model_loss, past_key_values = (
                 self._base_model_forward(
@@ -869,9 +880,9 @@ def forward(
             )
 
         if not isinstance(past_key_values, Cache):
-            past_key_values = DynamicCache(config=self._base_llm_config)
+            past_key_values = _get_empty_cache(self._base_llm_config)
         if not isinstance(eagle_cache, Cache):
-            eagle_cache = DynamicCache(config=self.eagle_module.config)
+            eagle_cache = _get_empty_cache(self.eagle_module.config)
 
         # ====Run eagle forward====
         eagle_loss = None
@@ -908,7 +919,7 @@ def forward(
                 if ttt_step == 0
                 else self._get_ttt_attention_mask(b, seq_length, ttt_step)
             )
-            with enable_cp_ttt_patch():
+            with enable_cp_ttt_patch() if self.training else contextlib.nullcontext():
                 _, eagle_input_hidden_states, eagle_logits, eagle_cache = self._eagle_forward(
                     eagle_input_hidden_states,
                     inputs_embeds,

modelopt/torch/speculative/utils.py

@@ -30,8 +30,6 @@
 from torch.nn.attention import SDPBackend, sdpa_kernel
 from transformers.cache_utils import DynamicCache
 
-import modelopt.torch.speculative.plugins.transformers
-
 KIMI_K2_REPO_ID = "moonshotai/Kimi-K2-Thinking"
 KIMI_K2_PACKAGE_NAME = "kimi_k2_temp"
 
@@ -462,6 +460,8 @@ def ttt_msk_func(b, h, q_idx, kv_idx):
 @contextlib.contextmanager
 def enable_cp_ttt_patch():
     """Context manager to enable CP TTT patch."""
+    import modelopt.torch.speculative.plugins.transformers
+
     modelopt.torch.speculative.plugins.transformers.ENABLE_CP_TTT_PATCH = True
     with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
         try: