[Bug][EAGLE] Crash when using mix hidden states

[benchislett]

Before submitting an issue, please make sure it hasn't been already addressed by searching through the existing and past issues.

Describe the bug

• When enabling mix hidden states mode in EAGLE3 training, I get a crash about some autograd error:

[rank1]: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [CUDABFloat16Type [2, 4096, 2880]], which is output 0 of IndexPutBackward0, is at version 3; expected version 2 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

Steps/Code to reproduce bug

Developed on a WIP branch for GPT-OSS EAGLE3 training. A few local changes which may be related. Creating the issue for tracking in case I cannot solve it.

Expected behavior

Who can help?

• ?

System information

• Container used (if applicable): ?
• OS (e.g., Ubuntu 22.04, CentOS 7, Windows 10): ?
• CPU architecture (x86_64, aarch64): ?
• GPU name (e.g. H100, A100, L40S): ?
• GPU memory size: ?
• Number of GPUs: ?
• Library versions (if applicable):
  • Python: ?
  • ModelOpt version or commit hash: ?
  • CUDA: ?
  • PyTorch: ?
  • Transformers: ?
  • TensorRT-LLM: ?
  • ONNXRuntime: ?
  • TensorRT: ?
• Any other details that may help: ?

[ChenhanYu]

Triage Analysis

Field	Value
Classification	Bug
Severity	High
Complexity	Moderate
Auto-fixable	No

Summary

When enabling eagle_mix_hidden_states mode in EAGLE3 training, a PyTorch autograd error occurs due to in-place modification of tensors needed for gradient computation. The error indicates that an IndexPutBackward0 operation has modified a tensor version, breaking the gradient computation chain.

Root Cause / Approach

The eagle_mix_hidden_states feature likely performs in-place operations on hidden state tensors (such as index assignments or modifications) during forward pass. When gradients are computed in backward pass, PyTorch detects that a tensor used in the computation graph has been modified by an IndexPutBackward0 operation at a different version than expected, violating the constraint that tensors in the computation graph must remain immutable during gradient computation.

Suggested Fix

1. Identify where hidden states are being mixed/modified in-place when eagle_mix_hidden_states=True. 2. Replace in-place operations (e.g., tensor[...] = value) with out-of-place alternatives (e.g., tensor.clone() or torch.cat() operations). 3. Ensure all tensor modifications during the forward pass create new tensors rather than modifying existing ones that are part of the autograd graph. 4. Add unit tests specifically for eagle_mix_hidden_states=True to prevent regression.

Relevant Files

• modelopt/torch/speculative/eagle/eagle_model.py
• modelopt/torch/speculative/config.py
• examples/speculative_decoding/main.py
• examples/speculative_decoding/eagle_utils.py
• modelopt/torch/speculative/eagle/conversion.py

---

Auto-triaged by pensieve /magic-triage