[shengliangxu]

What does this PR do?

Type of change: CI (test runtime, tests only)

Five cuts to the example lanes, each with its reasoning:

• test_lm_eval_hf evaluates one MMLU subtask. It ran --tasks mmlu --limit 0.1: the 57-subtask group, about 5.6k loglikelihood requests, on a tiny Qwen3 whose max_position_embeddings is 32, so its runtime is dominated by loading one HF dataset config per subtask (over 331 CI runs: min 198.9 s, median 212.5 s). It runs mmlu_abstract_algebra (100 test documents, 5-shot) through the same lm_eval_hf.py path and the same accuracy gate (exactly one task, lower bound 0.1 unchanged; the label distribution of that split is 22/26/31/21, so even a model that always predicts one letter scores at least 0.21).
• The llm_qat QAT and QAD tests quantize once per module. test_qwen3_qat_nvfp4 (3 backends), test_qwen3_qad_nvfp4 (2) and test_qwen3_qlora_nvfp4 ran 6 quantize.py and 12 cold tokenized-dataset builds per module (104-121 s each, qlora 169 s). A module-scoped fixture runs quantize.py once (guarded by assert_unmodified_tree, like the other shared tiny-model fixtures) and all quantize and train calls share one dataset cache; the qlora test keeps its own PTQ because it passes --compress True.
• The 1-GPU EAGLE3 compile smoke test skips the Triton GEMM autotune (TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=ATEN, the mechanism test_diffusers.py already uses). test_llama_eagle3[1-False] takes 175.5 s against 40.6 s for [1-True] (the same flow without compile), about 135 s of it torch.compile benchmarking Triton GEMM templates for a draft model that is one tiny decoder layer. The mode stays max-autotune, and the shortcut applies only when the run has one GPU: [2-False] and any run on more GPUs (the nightly) keep the full autotune.
• The duplicate tiny-llama offline EAGLE3 training test is dropped (test_offline_eagle3_training[tinyllama], 34.9 s per job, almost all launcher start-up). test_offline_eagle_training already runs that mode on the same model with hidden states from compute_hidden_states_hf.py, and the FakeBaseModel case keeps the synthetic path. The one thing the dropped case covered that stage 2 did not, the collator truncating dumps, moves there (training_seq_len=16, so 28-token dumps are truncated), and the collator is also unit tested on CPU.
• megatron_bridge tests that use no GPU skip the Megatron state reset. The conftest resets Megatron's global state around every test (destroy_model_parallel, gc.collect(), empty_cache), about 1.3 s per test; 77 tests in the directory are pure Python (the mlflow wiring and checkpoint-path tests) and leave nothing to reset (the 82 tests that took under 3 s each add up to 118 s in the series' runs before this PR and to 0.3 s after it). The reset now runs only for tests that request num_gpus; the environment restore stays for every test. The job runs at 68-75 minutes against its 75-minute nightly limit.

Time improvement (measured in CI)

Example lanes, pytest time of the changed tests and of the lane; before = #2846 (plus #2848 where its runner was as fast), after = #2851, #2854 and #2855 (median). #2849 and #2850 ran on runners 1.3-1.5x slower for every test and are left out of the before column.

lane		before	after	change
llm_eval	test_lm_eval_hf	130 s	19 s	-111 s
llm_eval	lane	332 s	212 s	-120 s (-36%)
llm_qat	six QAT, QAD and QLoRA tests (call time)	563 s	428 s	-135 s (the shared quantize adds 20 s to the first test's setup)
llm_qat	lane	591 s	475 s	-116 s (-20%)
speculative_decoding	test_llama_eagle3[1-False]	131 s	87 s	-44 s
speculative_decoding	dropped test_offline_eagle3_training[tinyllama]	20 s	-	-20 s
speculative_decoding	lane	407 s	341 s	-66 s (-16%)
megatron	the 82 megatron_bridge tests that use no GPU (the state reset)	118 s	0.3 s	-118 s
megatron	lane	583 s	460 s	-123 s (-21%)

• A mmlu_limit=1 change for test_qwen3_eval_fp8 was dropped from this PR. The runs above are of its earlier head, which had it: the test took 175-246 s in all 7 runs, before and after (187 s against 176 s on comparable runners), so the CI history it was based on (mean 723.7 s, 2026-09-25 to 10-09) is not reproduced by the current lane and the change saved nothing measurable. The llm_eval lane figure above is test_lm_eval_hf's.
• That is 425 s of lane time per PR run. The Example workflow ends with its slowest lane, torch_onnx (18.1-18.8 min of job time) until [OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853, so the gain reaches the workflow's wall time only with that PR.

Method: every PR of the series has its own CI run, so "before" is the run of the PR(s) below that lack the change and "after" the runs that carry it. Times come from the job logs (--durations, or the gap between consecutive result lines). Where runner speed matters, the tests that no PR of the series touches scale each run, and the raw and the scaled medians are both shown. One test varies by 15-30% from run to run on its own, so smaller differences are not evidence.

Usage

Nothing user-facing.

Testing

• llm_eval: ran the test's command on CPU with the tiny Qwen3 it builds: old configuration 223.4 s, mmlu/acc 0.2343; new 22.4-24.0 s, mmlu_abstract_algebra/acc 0.27, gate passes.
• llm_qat: with run_example_command faked, quantize.py runs 6 -> 2 times and cold dataset builds go 12 -> 2; real quantize.py and train.py (QAT, then QAD from the checkpoint) ran on CPU with the int4 and fp8 recipes (NVFP4 needs CUDA) and left the checkpoint unmodified.
• EAGLE: stage 1 and 2 of the offline flow ran on CPU with training_seq_len 64 and 16 (the collator pads 28 -> 64 and truncates 28 -> 16); with run_example_command faked, only [1-False] on one GPU gets the autotune variable (that case on 2 GPUs and every other case keep the full autotune) and the only removed test id is test_offline_eagle3_training[tinyllama].
• megatron_bridge: a throwaway pytest module (not in this PR) against the patched conftest counts the resets: plain tests trigger none, tests that request num_gpus trigger one before and one after; it fails against the old conftest.
• Not run: the example lanes themselves (they need GPUs, TRT-LLM and model downloads), so the accuracy with CUDA bf16 numerics, the savings and the autotune variable are checked by this PR's example jobs. Pre-commit passes.
• trust_remote_code=True in test_offline_eagle3_training is not new: it is the line on main (the Kimi and MiniMax repos it loads need remote code), and this PR only renames its variable after dropping the tinyllama case.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: N/A (CI and tests)
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: no new dependency
• Did you write any new necessary tests?: N/A (existing tests made cheaper; one duplicate test dropped, see the description)
• Did you update Changelog?: N/A (CI and tests)
• Did you get Claude approval on this PR?: N/A (optional, not run)

Additional Information

Slice 7 of 11 of the GPU-CI speed-up work, split out of the reference draft #2730 (not for merge; it holds the whole series with its measurements). Slices 1, 2, 5, 8, 9 and 10 (#2845, #2846, #2849, #2852, #2853, #2854) have merged. The remaining slices (#2847, #2848, #2850, #2851, #2855) are independent PRs on main: each applies to main on its own, and merging all of them onto main in different orders gives the same tree. The times above are from the series' CI runs before 2026-10-21, when all slices formed one linear stack; a PR number in a method note names that PR's run from then.

1. [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist #2845 [OMNIML-6178] [1/11] Run the unit tests in parallel with pytest-xdist (merged)
2. [OMNIML-6178] [2/11] Use the sysmon coverage core in CI #2846 [OMNIML-6178] [2/11] Use the sysmon coverage core in CI (merged)
3. [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs #2847 [OMNIML-6179] [3/11] Run manual-marked tests on nightly and on-demand GPU runs
4. [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time #2848 [OMNIML-6181] [4/11] Give gpu_megatron's slow cold starts more time
5. [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models #2849 [OMNIML-6179] [5/11] Make the ONNX tests cheaper: CPU-only ones to the unit lane, small models (merged)
6. [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data #2850 [OMNIML-6179] [6/11] Use local tokenizers and calibration text in GPU tests, shrink puzzletron data
7. [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests #2851 [OMNIML-6180] [7/11] Trim the llm_eval, llm_qat, EAGLE and megatron_bridge example tests (this PR)
8. [OMNIML-6180] [8/11] Trim the DFlash regression test #2852 [OMNIML-6180] [8/11] Trim the DFlash regression test (merged)
9. [OMNIML-6180] [9/11] Trim the torch_onnx example lane #2853 [OMNIML-6180] [9/11] Trim the torch_onnx example lane (merged)
10. [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies #2854 [OMNIML-6181] [10/11] Make the distributed worker pool fail fast when a worker dies (merged)
11. [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules #2855 [OMNIML-6181] [11/11] Share one Megatron worker pool across the gpu_megatron modules

Summary by CodeRabbit

• Tests
  • Focused language-model evaluation coverage on the abstract algebra task while retaining five-shot evaluation.
  • Streamlined quantization-aware training scenarios by reusing cached datasets and checkpoints, reducing repeated setup.
  • Improved Megatron test isolation by applying state cleanup only where needed.
  • Refined EAGLE test configurations for single-GPU runs, remote model coverage, and offline training with shorter sequence lengths.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

• Last reviewed commit: 9bcfae8
• Latest commit, not reviewed: 5b4b6ea

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 5ba1a897-b7a8-40da-918b-02f3a52718e3

📥 Commits

Reviewing files that changed from the base of the PR and between 5529023 and 9bcfae8.

📒 Files selected for processing: 2

• tests/examples/llm_qat/test_llm_qat.py
• tests/examples/speculative_decoding/test_eagle.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 5 remain after this review.

📝 Walkthrough

The pull request updates tests for LLM evaluation, QAT workflows, Megatron state isolation, and EAGLE speculative decoding. It changes fixture reuse, task selection, conditional state resets, environment handling, and offline-training inputs.

Changes: LLM evaluation and workflow fixtures

Layer / File(s)	Summary
Update evaluation and LLM workflow fixtures tests/examples/llm_eval/test_llm_eval.py, tests/examples/llm_qat/test_llm_qat.py	The evaluation test selects mmlu_abstract_algebra without a limit. Module-scoped fixtures provide a shared dataset cache and protected NVFP4 PTQ checkpoint. QAT and QAD use the checkpoint, while QAT, QAD, and QLoRA use the cache.

Megatron test state isolation

Layer / File(s)	Summary
Conditionally reset Megatron state tests/examples/megatron_bridge/conftest.py	The fixture resets Megatron state only for tests that use num_gpus. Environment variables remain restored around every test.

EAGLE speculative decoding tests

Layer / File(s)	Summary
Adjust EAGLE runtime configuration tests/examples/speculative_decoding/test_eagle.py	The single-GPU case without mixed hidden states sets the TorchInductor GEMM autotuning backend to ATEN.
Update EAGLE offline-training inputs tests/examples/speculative_decoding/test_eagle.py, tests/examples/speculative_decoding/test_eagle_offline_ptq.py	The offline-training test removes the tiny-Llama case and uses model_source directly. The offline PTQ test sets training.training_seq_len to 16.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Other

Suggested reviewers: realasma, kevalmorabia97

Merge Risk: ⚪ Minimal · up to 9bcfa

The test-only changes preserve the inspected contracts and are suitable to merge.

Caution

Pre-merge checks failed

Please resolve all errors before merging. Addressing warnings is optional.

• Ignore

Failed checks: 1 error · 1 warning

Check name	Status	Explanation	Resolution
Security Anti-Patterns		CRITICAL: The PR adds AutoConfig.from_pretrained(model_source, trust_remote_code=True) in tests/examples/speculative_decoding/test_eagle.py:261. SECURITY.md prohibits hardcoded `trust_remote_cod…	Add a trust_remote_code parameter with default False to the test or its model-loading path. Pass that parameter to AutoConfig.from_pretrained and use it for the training override. Set it to True only for explicitly trusted model sou…
Docstring Coverage		Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely describes the main change: reducing runtime for the llm_eval, llm_qat, EAGLE, and megatron_bridge example tests.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.

Full details: Security Anti-Patterns

Explanation

CRITICAL: The PR adds AutoConfig.from_pretrained(model_source, trust_remote_code=True) in tests/examples/speculative_decoding/test_eagle.py:261. SECURITY.md prohibits hardcoded trust_remote_code=True because it can execute code from an untrusted model source. The function has no caller-configurable trust_remote_code parameter, and the training override also remains hardcoded as model.trust_remote_code=true at line 274.

Resolution

Add a trust_remote_code parameter with default False to the test or its model-loading path. Pass that parameter to AutoConfig.from_pretrained and use it for the training override. Set it to True only for explicitly trusted model sources, with the opt-in documented. The current diff shows no new dependency or other prohibited security pattern that changes this result.

✨ Finishing Touches: 💡 2

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.10%. Comparing base (6d76b59) to head (5b4b6ea).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2851      +/-   ##
==========================================
- Coverage   70.16%   70.10%   -0.07%     
==========================================
  Files         659      659              
  Lines       74081    74081              
==========================================
- Hits        51978    51931      -47     
- Misses      22103    22150      +47     

Flag	Coverage Δ
examples-gpt-oss	13.51% <ø> (-0.01%)	⬇️
examples-hf_ptq	23.19% <ø> (-0.04%)	⬇️
examples-llm_distill	13.59% <ø> (-0.01%)	⬇️
examples-llm_eval	17.24% <ø> (+<0.01%)	⬆️
examples-llm_qat	17.48% <ø> (-0.01%)	⬇️
examples-llm_sparsity	15.76% <ø> (-0.01%)	⬇️
examples-megatron_bridge	26.43% <ø> (-0.13%)	⬇️
examples-specdec_bench	13.29% <ø> (-0.01%)	⬇️
examples-speculative_decoding	17.47% <ø> (-0.01%)	⬇️
unit	60.59% <ø> (+<0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shengliangxu]

/ok to test 2519292

[shengliangxu]

/ok to test f8fa147

[coderabbitai]

Pre-merge checks failed. Please resolve the failing checks before merging.

[cjluo-nv]

Bot review (gpt-6.1-sol) — DM the bot to share feedback.

LGTM: the test reductions are justified and preserve the relevant smoke paths; modifying and removing existing tests requires human sign-off.

Needs action:

• Confirm human sign-off on the narrowed MMLU smoke and duplicate EAGLE removal; offline real-model training and collator truncation remain covered.

No action needed:

• Shared QAT/QAD PTQ setup uses the existing immutability guard; QLoRA retains separate compressed PTQ.
• Current Megatron example runners request num_gpus; pure-Python tests retain environment restoration.
• Tests were inspected, not executed locally.

[shengliangxu]

The flagged line is not new: AutoConfig.from_pretrained(model_path, trust_remote_code=True) is on main (test_eagle.py:253), next to the model.trust_remote_code=true override. This PR drops the duplicate tinyllama case and only renames model_path to model_source. The remaining cases (Kimi-K2.5, Kimi-K2-Thinking, MiniMax-M2.5) ship their config classes as repo code (auto_map → AutoConfig), so AutoConfig.from_pretrained fails without trust_remote_code=True; a default-False parameter would be True in every case. SECURITY.md's rule targets library code that silently opts users in; this test names the repos itself, and 20 other test lines on main use the same call. The docstring-coverage warning is optional and the PR only edits existing test functions. No code change; please ignore the pre-merge check.