Fix EAGLE-3 context-parallel training and re-enable its tests - #2086
Merged
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe workflows update PyTorch images, consolidate TensorRT-LLM jobs, and extend PR-gate file detection. EAGLE context-parallel setup, masking, normalization, loss aggregation, and test conditions are updated. ChangesCI and example test alignment
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant TrainingScript
participant HFEagleModel
participant TTTPatch
participant LossAggregation
TrainingScript->>TrainingScript: configure FSDP2 for context parallelism
TrainingScript->>HFEagleModel: set eagle_cp_size
HFEagleModel->>TTTPatch: pass context-parallel size
TTTPatch->>HFEagleModel: apply forward-path attention behavior
HFEagleModel->>LossAggregation: provide base and EAGLE losses
LossAggregation->>LossAggregation: aggregate non-None losses
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.github/workflows/example_tests.yml:
- Line 49: Update the docker_image value in the workflow configuration to use
the fixed PyTorch image tag 26.07 directly, removing the undefined
matrix.docker_image fallback expression. Keep the existing
nvcr.io/nvidia/pytorch image and -py3 suffix unchanged.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2fcc8176-1976-4ecd-89a2-3caa1909bc9a
📒 Files selected for processing (5)
.github/workflows/example_tests.yml.github/workflows/gpu_tests.yml.github/workflows/regression_tests.ymltests/examples/gpt-oss/test_gpt_oss_qat.pytests/examples/speculative_decoding/test_eagle.py
💤 Files with no reviewable changes (1)
- tests/examples/gpt-oss/test_gpt_oss_qat.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2086 +/- ##
==========================================
- Coverage 78.73% 78.08% -0.65%
==========================================
Files 522 522
Lines 60157 60167 +10
==========================================
- Hits 47363 46980 -383
- Misses 12794 13187 +393
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kevalmorabia97
changed the title
Collaborator
Author
|
/claude review |
kevalmorabia97
changed the title
Contributor
|
Contributor
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/speculative_decoding/main.py`:
- Around line 189-194: Declare the FSDP-related members used by the fallback
block on HfTrainingArguments, including fsdp_plugin_args, fsdp, and fsdp_config,
with appropriate defaults and types. Keep the existing assignments and
_process_fsdp_args() flow unchanged.
- Around line 189-194: Update the FSDP initialization block for cp_size > 1 to
avoid relying on fsdp_plugin_args or _process_fsdp_args(), which are unavailable
in Transformers 4.57.x. Use the supported version-compatible mechanism to create
the FSDP plugin, and merge required defaults into any existing fsdp_config
without overwriting supplied settings in Transformers 5.0.
In `@modelopt/torch/speculative/plugins/hf_eagle.py`:
- Around line 770-773: Update the CP branch in the Eagle attention setup to
preserve padding constraints from LanguageDataCollator’s zero-valued
attention_mask entries without passing the incompatible dense mask to DTensor
scores. Add a CP-compatible padding bias and ensure padded positions cannot be
attended to, or explicitly reject padded CP batches with coverage for that
behavior; keep causal and TTT masking intact.
In `@modelopt/torch/speculative/plugins/modeling_eagle.py`:
- Around line 163-165: The EagleModule.forward path currently stores
inputs_embeds even when _eagle3_attention_forward_pre_hook is not registered,
retaining unnecessary tensors and autograd graphs. Guard the _input_embeds
assignment with config.use_aux_hidden_state (or the established hook-enabled
condition), and clear _input_embeds on the non-hook path so no stale reference
remains after forward.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8520879e-8307-47c7-80bb-68e29fa16a5b
📒 Files selected for processing (9)
.github/workflows/example_tests.yml.github/workflows/gpu_tests.ymlCHANGELOG.rstexamples/speculative_decoding/eagle_utils.pyexamples/speculative_decoding/main.pymodelopt/torch/speculative/plugins/hf_eagle.pymodelopt/torch/speculative/plugins/modeling_eagle.pymodelopt/torch/speculative/utils.pytests/examples/speculative_decoding/test_eagle.py
💤 Files with no reviewable changes (1)
- tests/examples/speculative_decoding/test_eagle.py
Collaborator
Author
|
/claude review |
kevalmorabia97
force-pushed
the
kmorabia/ci-example-lanes
branch
from
August 6, 2026 13:53
f426949 to
270ad75
Compare
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 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
force-pushed
the
kmorabia/ci-example-lanes
branch
from
August 6, 2026 14:23
2ecd662 to
1e81510
Compare
kevalmorabia97
changed the title
h-guo18
approved these changes
Aug 6, 2026
kevalmorabia97
added a commit
that referenced
this pull request
Aug 7, 2026
…ate the CI gate (#2090) ### What does this PR do? Type of change: CI/CD improvement Follow-up to #2086, which carried the speculative-decoding fix; this PR is the CI half. **Lanes now run only when their own files change.** A one-line edit to any example started all 12 example lanes, and any `modelopt/**` change started every GPU suite. - Each lane gates itself inside `_example_tests_runner.yml` / `_gpu_tests_runner.yml`, deriving its watch list from the example or suite name it was already given (`examples/<name>/**`, `tests/examples/<name>/**`, `tests/<suite>/**`). **Adding a new example stays a one-line matrix entry** — no central mapping to update. - The five cross-example dependencies are declared as `watch_extra` next to the example that needs them: `hf_ptq` → `llm_eval` (`huggingface_example.sh` runs lm_eval from `../llm_eval`), `torch_trt` → `onnx_ptq`, `speculative_decoding` → `hf_ptq` + `dataset`, `llm_qat` / `gpt-oss` → `dataset`. - `gpu_tests.yml` is split into caller + runner to match. This shape is forced, not stylistic: job-level `if:` cannot read `matrix`, and `container:` images are pulled before any step runs, so gating inside the job would still pull 10–20 GB and hold a GPU runner for every skipped suite. - `modelopt/**`, `modelopt_recipes/**`, `pyproject.toml` and `tests/_test_utils/**` still run every lane. The gate now watches the last two, which example tests depend on but it previously ignored. **Docs-only changes no longer start GPU jobs.** The gate ignores `**.md`, `**.rst`, `**.png` and `**.ipynb` by default, so a README edit short-circuits the whole workflow. Nothing executes notebooks (no `nbmake`/`nbval` in the repo), and `.sh`/`.yaml`/`.txt` stay watched since examples run them. **One file holds the gate logic.** `.github/actions/changed-files-gate` is a composite action doing the merge-base + changed-files comparison and, optionally, the `^linux$` wait. `_pr_gate.yml` and `_wait_for_checks.yml` are both deleted: each top-level workflow keeps a 12-line `pr-gate` job that is pure wiring, and the runners use the action as steps so a lane shows `gate` + `run-test` rather than three checks. Calling a reusable workflow always materializes all of its jobs, including skipped ones, which is what made the per-lane check list noisy. `unit_tests.yml` also drops its DCO wait: DCO can be marked passing manually, so blocking the matrix on it only delayed feedback. The `^linux$` wait remains, which is the gate that actually protects GPU runners. **Also, from the original lane consolidation:** - **One TensorRT-LLM lane.** `trtllm-pr` and `trtllm-non-pr` merge into a single `trtllm` job gated like the others, so `llm_eval` now runs on PRs, where it was nightly-only. - **`gpt-oss` moves to the TensorRT-LLM image.** Its deploy step needs `tensorrt_llm`, which `pytorch` doesn't have, so `deploy_gpt_oss_trtllm` silently skipped in CI. Its `importorskip` is dropped now that the lane guarantees the dependency. - **Containers bumped where no reason was documented:** pytorch `26.06`/`26.01` → `26.07` (torch example lane, regression). Left pinned with their existing in-file reasons: pytorch `26.05` for the gpu lane (`EXPLICIT_BATCH` removed in TensorRT 11), tensorrt `26.05` for the onnx lane (`torch-tensorrt` needs `libnvinfer.so.10`), vllm `v0.20.0` (legacy FusedMoE coverage). TensorRT-LLM stays on `1.3.0rc20`: rc21–rc23 ship a `quickstart_multimodal.py` importing `MultimodalConfig` before `tensorrt_llm.llmapi` exported it (fixed upstream in NVIDIA/TensorRT-LLM#17112, one day after rc23 was cut), which fails the `hf_ptq` VLM deploy smoke test. Also clarifies the changelog line in the PR template to spell out when an entry is expected. **Two silent-failure fixes found in review:** - Every gate used `any_changed`, which is ACMR and excludes deletions, so a delete-only PR (removing an example, a test, or library code) ran nothing. Now `any_modified` (ACMRD), fixed in `example_tests.yml`, `_pr_gate.yml` and `unit_tests.yml`. - `git merge-base` was piped into `tee`, so a failure returned `tee`'s exit status and emitted an empty base, quietly changing which lanes run instead of failing. **Nightly secret scanning is fixed.** `code_quality.yml` excludes trufflehog's `lob` detector: its `(live|test)_[a-zA-Z0-9_]{35}` pattern matches any pytest function whose name is exactly 35 characters after `test_` — 60 of them in this repo, e.g. `test_all_zero_activation_yields_no_scale` — and reports them as *verified* secrets. This only ever failed nightly because the action scans all history on `schedule` and only the PR diff on `pull_request`. ### Testing Workflow YAML validated locally and the selection logic checked by hand across scenarios (`examples/diffusers/**` → onnx lane only; `examples/dataset/**` → `llm_qat`, `speculative_decoding`, `gpt-oss`; `examples/llm_eval/**` → `hf_ptq` + `llm_eval`; `modelopt/**` and nightly → everything). Gating behavior itself can only be exercised by a real PR run — the failure mode to watch for is a lane skipping when it should have run. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A — CI configuration - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ — not yet run <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **CI Improvements** - Improved change detection for model recipes, test utilities, workflow actions, and documentation-only updates. - Streamlined pull request checks and status monitoring. - Consolidated TensorRT-LLM example validation and refined conditional test execution. - Updated test environments to newer PyTorch releases. - Refined GPU, regression, and unit test triggers. - Deployment tests no longer automatically skip when TensorRT-LLM is unavailable. - Updated secret scanning configuration. - **Documentation** - Updated pull request checklist guidance to include deprecations and critical bug fixes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
kevalmorabia97
added a commit
that referenced
this pull request
Aug 7, 2026
test_llama_eagle3[2-False] timed out at the 300s examples default in the first nightly after #2086 re-enabled the cp_size=2 cases, which the old version guard had skipped since February. Measured on the 2-GPU runner: [2-False] 304s, [1-False] 173s, [2-True] 51s, [1-True] 39s. Only [2-False] exceeds the default, and only by 4s: it pays context parallelism plus torch.compile, which mix_hidden_states=False enables. The marker is scoped to that one case so the 300s cap keeps protecting the rest. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
kevalmorabia97
added a commit
that referenced
this pull request
Aug 7, 2026
### What does this PR do? Type of change: Bug fix (CI) `test_llama_eagle3[2-False]` timed out at the 300s examples default in the [first nightly](https://github.com/NVIDIA/Model-Optimizer/actions/runs/31137639327) after #2086 re-enabled the `cp_size=2` cases. Those cases had never run in CI before — the old guard compared `Version("2.10.0a0")` against `Version("2.10.0")`, which is False on every NGC alpha build, so they had been silently skipped since February. Measured on the 2-GPU runner: | case | duration | |---|---| | `[2-False]` | **304s** — over the 300s default by 4s | | `[1-False]` | 173s | | `[2-True]` | 51s | | `[1-True]` | 39s | Only `[2-False]` exceeds the cap, and only barely: it pays context parallelism *plus* `torch.compile`, which `mix_hidden_states=False` enables. It is a margin problem, not a hang. The marker is scoped to that single param so the 300s cap keeps protecting every other example test. Note 360s is roughly 18% headroom over a single sample on a shared runner. If it ever trips again the signal will be a duration just above 360s, which distinguishes normal variance from a real regression. ### Testing Measured by temporarily running the spec-dec lane ungated on the 2-GPU runner in this PR (**those overrides have been reverted**; the diff is now one file). That run: **16 passed, 2 skipped** in 1016s, with `[2-False]` at 304.27s. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A — adjusts an existing test's timeout - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A - Did you get Claude approval on this PR?: ❌ — not yet run <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Tests** * Reformatted test parameter definitions for improved readability. * Preserved all existing test combinations and timeout settings. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix
EAGLE-3 context-parallel training (
--cp_size > 1) is fixed, and its tests run again. CP has been broken sinceaccelerate1.13, and the tests never caught it: the guard comparedVersion("2.10.0a0")againstVersion("2.10.0"), which is False on every NGC alpha torch build, sotest_llama_eagle3[cp_size=2]has never actually run in CI.Five fixes:
main.py— rebuild the FSDP2 plugin accelerate requires forcp_size > 1. The--fsdp full_shard --fsdp_configlauncher flags that used to supply it were dropped fromlaunch_train.sh, so CP could not start at all. Also pass the CP degree to the draft model.modeling_eagle.py— apply the draft model's first input norm insidelayers[0]'s own forward, where FSDP2 has actually unsharded its weights, and only stash the input embeds on the path whose pre-hook consumes them.hf_eagle.py— skip the dense eagle attention mask under CP (causal masking comes fromis_causal, TTT masking from the ring-attention patch), and warn that padded positions are therefore unmasked. Also stop(eagle_loss or 0)replacing a0.0loss tensor with a plainint, which detached the graph.eagle_utils.py— key TTT-mask injection off the backward call'sgrad_outkwarg, since newer torch omitsattn_biason the forward call, silently disabling TTT masking.utils.py— CUDNN-only SDPA under CP; theMATHbackend decomposes SDPA and breaks on DTensors. Scoped tocp_size > 1, since this context manager wraps every training forward and CPU has no cudnn backend.Drops the
speculative_decoding26.01 container override. It was added when the lane ran 25.06 and spec-dec needed something newer — a floor. Later bumps moved the default past it, so it had silently become a ceiling holding spec-dec on a 6-month-old image.Testing
Ran
tests/examples/speculative_decodinginnvcr.io/nvidia/pytorch:26.07-py3on 2 GPUs, reproducing the CI install steps (pip uninstall -y nvidia-modelopt,pip install -e ".[hf,dev-test]", example requirements): 16 passed, 2 skipped — the 2 skipped being pre-existing--run-manualtests. All fourtest_llama_eagle3cases pass, including bothcp_size=2ones.Before your PR is "Ready for review"
CONTRIBUTING.md: N/Acp_size=2tests are re-enabled