Skip to content

[OMNIML-6180] Run only the fsdp2 llm_qat training cases on PRs - #2925

Merged
shengliangxu merged 1 commit into
mainfrom
shengliangx/ci-qat-one-backend
Oct 22, 2026
Merged

shengliangxu merged 1 commit into
mainfrom
shengliangx/ci-qat-one-backend

Conversation

@shengliangxu

@shengliangxu shengliangxu commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

What does this PR do?

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

On pull requests the llm_qat example lane runs one training backend per test family: the fsdp2 cases of test_qwen3_qat_nvfp4 and test_qwen3_qad_nvfp4. The deepspeed case of both and the ddp case of test_qwen3_qat_nvfp4 keep their ids and run in the scheduled CI (nightly and on demand) as before.

Why: PR runners have one GPU (linux-amd64-gpu-rtxpro6000-latest-1; the nightly has two), so a backend case on a PR exercises the backend wrapper and ModelOpt's patches but not sharding. Each of the three cases is an accelerate launch of train.py and takes 99-109 s on a normal-speed runner, 146-157 s on a slow one. fsdp2 is the one that stays because it is the documented default and has the most ModelOpt-specific code (_patch_fsdp2_post_backward, _patch_accelerate_for_fsdp2_fix, the sharded save in transformers_trainer.py, and the FSDP2-only path of distillation, which rejects FSDP1). In the CI history since 10-06 all four failures of this lane are on fsdp2 or QLoRA cases (three at the 300 s per-test cap on slow runners, one real error in test_qwen3_qad_nvfp4[fsdp2]); none is on deepspeed or ddp. test_qwen3_qlora_nvfp4 still trains on ddp on PRs.

The skip is a pytest.mark.skipif on the three parameters, using the same GITHUB_REF check as tests/examples/torch_onnx/test_torch_quant_to_onnx.py (refs/heads/pull-request/). The example workflow does not pass --run-manual (only the unit workflow does), so the manual marker is not an option here. No workflow file changes.

Time improvement (measured in CI, per run of the lane on 1 GPU, COVERAGE_CORE=sysmon)

Per-test times from the result-line gaps of the lane's job logs on 2026-10-21:

run runner pytest QAT deepspeed QAT ddp QAD deepspeed skipped on PRs after
e5220c3 (20:36 UTC) normal 680 s 99 s 100 s 109 s 308 s about 372 s
7607678 (19:07 UTC) normal 687 s 105 s 101 s 109 s 314 s about 373 s
d1dad23 (20:57 UTC) slow 984 s 151 s 147 s 157 s 454 s about 530 s

So about 310 s of pytest (45%) per PR run on a normal runner and about 450 s on a slow one; the lane's job goes from about 900 to about 600 s (normal) and from about 1,320 to about 900 s (slow). The remaining cases on a PR are fsdp2 QAT (100 s), fsdp2 QAD (106 s), QLoRA (145 s) and the pre-tokenize test (18 s). Nightly and local runs are unchanged. Measured in this PR's own run (run 39686596776, 1-GPU runner): pytest 464 s, 13 passed and 4 skipped (the three nightly-only cases with the reason "Only fsdp2 runs on PRs; runs in the scheduled CI", plus the existing LoRA skip); fsdp2 QAT 152 s, fsdp2 QAD 103 s, QLoRA 178 s, pre-tokenize 28 s. This runner was slower than the normal-speed runners of the table (the four unchanged cases took 461 s together against 369 s on a normal runner, 1.25x), so at this speed the three skipped cases would have cost about 385 s (308 s at normal speed): the saving is about 45% of the lane's pytest time, as projected. The job as a whole took 1,260 s only because this run's container initialization took 655 s; the Run tests step was 492 s. After the rebase onto main f983619, this PR's run (run 39737195847) took 594 s for the job and 342 s for pytest (13 passed, 4 skipped), against 1,144 s and 761 s (16 passed, 1 skipped) for the unchanged llm_qat job of #2924's run on the same main (run 39737194943, started within 14 s of it): -550 s for the job and -419 s (-55%) for pytest. The three skipped cases took 121 + 105 + 92 = 318 s in the unchanged job; the rest of that gap is runner speed (the unchanged QLoRA test took 135 s against 108 s, fsdp2 QAT 128 s against 83 s).

Usage

Nothing user-facing. On a pull-request/N ref (what copy-pr-bot runs for a PR) the three cases are reported as skipped with "Only fsdp2 runs on PRs; runs in the scheduled CI".

Testing

  • On a CPU host with run_example_command replaced by a recorder (scratch plugin, not part of the PR), pytest tests/examples/llm_qat/test_llm_qat.py -k "qat_nvfp4 or qad_nvfp4":
    • GITHUB_REF=refs/heads/pull-request/2900: test_qwen3_qat_nvfp4[fsdp2] and test_qwen3_qad_nvfp4[fsdp2] run; test_qwen3_qat_nvfp4[deepspeed], [ddp] and test_qwen3_qad_nvfp4[deepspeed] are skipped with the reason (2 passed, 4 skipped, the 4th being the existing LoRA skip).
    • GITHUB_REF=refs/heads/main (the nightly) and an unset GITHUB_REF: all five run (5 passed, 1 skipped, the same existing skip).
  • Pre-commit over the range against main passes. No workflow file changed, so no actionlint run.
  • Run on a GPU runner: this PR's Example run shows the two fsdp2 cases, QLoRA and the pre-tokenize test passing and the three nightly-only cases skipped (numbers above). Still to confirm: the next nightly runs all five.

Coverage that PRs lose

  • DeepSpeed ZeRO-3 training for QAT and QAD: make_deepspeed_compatible, the ZeRO-3 state-dict loading patch (_load_params_and_buffers_into_zero3_model), the parameter gathering in distillation (deepspeed.zero.GatheredParameters) and the AWQ/SmoothQuant guard for DeepSpeed. Both DeepSpeed cases still run nightly, on two GPUs, which is the only place sharding is exercised anyway.
  • Full-parameter ddp QAT. The DDP wrapper itself still runs on PRs through the QLoRA test.
  • fsdp2 QAT and QAD, QLoRA, the pre-tokenize test and the unit-style tests of the lane are unchanged.

Open question for the code owners: keep fsdp2 for both families (this PR), or spread the backends so every one stays on PRs once (fsdp2 for QAT, deepspeed for QAD), which keeps DeepSpeed's distillation path on PRs but moves fsdp2 distillation to the nightly. A regression in a skipped backend would surface the next night rather than on the PR.

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 cases are skipped on PRs; ids unchanged)
  • Did you update Changelog?: N/A (CI and tests)
  • Did you get Claude approval on this PR?: ✅ (/claude review approved 992db89; its one suggestion, a shared PR-detection helper, is postponed)

Additional Information

Part of the GPU-CI cost reduction under OMNIML-6177.

Summary by CodeRabbit

  • Tests
    • Some quantization-aware training checks now run only in scheduled CI, including DeepSpeed cases and the NVFP4 DDP case.
    • DDP is no longer included among the QAD backend test cases.
    • QLoRA test coverage remains unchanged.

@copy-pr-bot

copy-pr-bot Bot commented Oct 21, 2026

Copy link
Copy Markdown

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

coderabbitai Bot commented Oct 21, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

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: ce767201-a909-481b-912b-31d7d8252562
📥 Commits

Reviewing files that changed from the base of the PR and between f983619 and 992db89.

📒 Files selected for processing: 1
  • tests/examples/llm_qat/test_llm_qat.py

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

📝 Walkthrough

The LLM QAT test module detects pull request runs and skips designated nightly-only backend cases on those runs. The QLoRA test remains unchanged.

Changes: LLM QAT test selection

Layer / File(s) Summary
Pull request backend filtering
tests/examples/llm_qat/test_llm_qat.py
The test module detects pull request runs and skips marked cases on those runs. NVFP4 QAT marks DeepSpeed and DDP as nightly-only. NVFP4 QAD marks DeepSpeed as nightly-only and removes DDP from its backend cases. The fsdp2 case remains unmarked.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~8 minutes

Change: Other

Suggested reviewers: kevalmorabia97

Merge Risk: ⚪ Minimal · up to 992db

The targeted test selection change is mergeable with normal CI checks.

Pre-merge checks: Passed 5 · Failed 1

Failed Failed checks: 1 warning

Check name Status Explanation Resolution
Docstring Coverage Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 1 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed Passed checks: 5 passed
Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly and concisely describes the main change: limiting llm_qat pull request runs to FSDP2 training cases.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns Passed PASS. The pull request changes only tests/examples/llm_qat/test_llm_qat.py. The additions read GITHUB_REF and add pytest skip markers. They do not add torch.load(..., weights_only=False), `numpy…
Full details: Security Anti-Patterns

Explanation

PASS. The pull request changes only tests/examples/llm_qat/test_llm_qat.py. The additions read GITHUB_REF and add pytest skip markers. They do not add torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, external-input eval()/exec(), or # nosec. No pyproject.toml or requirements.txt files changed.

✨ Finishing Touches: 💡 1
📝 Generate docstrings: 💡
  • 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.

@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/ok to test 7b7f0c8

@codecov

codecov Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.66%. Comparing base (f983619) to head (992db89).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #2925      +/-   ##
==========================================
- Coverage   79.66%   79.66%   -0.01%     
==========================================
  Files         659      659              
  Lines       74400    74400              
==========================================
- Hits        59272    59270       -2     
- Misses      15128    15130       +2     
Flag Coverage Δ
examples-llm_distill 13.56% <ø> (-0.06%) ⬇️
examples-llm_qat 17.41% <ø> (-0.11%) ⬇️
examples-llm_sparsity 16.80% <ø> (-0.07%) ⬇️
examples-specdec_bench 13.27% <ø> (-0.06%) ⬇️
examples-speculative_decoding 17.42% <ø> (-0.14%) ⬇️
unit 60.84% <ø> (ø)

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.

PR runners have one GPU, so the deepspeed and ddp cases of test_qwen3_qat_nvfp4 and the deepspeed
case of test_qwen3_qad_nvfp4 exercise the backend wrapper and ModelOpt's patches but not sharding,
and each takes 100-109 s on a normal runner (146-157 s on a slow one). On pull requests only the
fsdp2 cases run: fsdp2 is the documented default and has the most ModelOpt-specific code (the
post-backward patch, the accelerate fix, the sharded save, distillation that requires FSDP2). The
other cases keep their ids and run in the scheduled CI, selected with the same GITHUB_REF check the
torch_onnx test uses. QLoRA still trains on ddp on PRs.

Testing: on a CPU host with run_example_command stubbed, the PR environment runs the two fsdp2
cases and skips the other three with the reason; GITHUB_REF=refs/heads/main and an unset
GITHUB_REF run all five. Not run on a GPU here.

Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
@shengliangxu
shengliangxu force-pushed the shengliangx/ci-qat-one-backend branch from 7b7f0c8 to 992db89 Compare October 22, 2026 06:56
@shengliangxu
shengliangxu marked this pull request as ready for review October 22, 2026 06:57
@shengliangxu
shengliangxu requested a review from a team as a code owner October 22, 2026 06:57
@shengliangxu

Copy link
Copy Markdown
Collaborator Author

/claude review

Comment thread tests/examples/llm_qat/test_llm_qat.py
@claude

claude Bot commented Oct 22, 2026

Copy link
Copy Markdown
Contributor

Claude review summary

Scope: full review. The trigger comment was a bare /claude review with no scoping instructions. 1 changed file (tests/examples/llm_qat/test_llm_qat.py, +12/-3), reviewed in full along with the surrounding helpers (_run_train, the nvfp4_ptq_dir/dataset_cache_dir fixtures, the QLoRA test) and the CI plumbing it depends on (.github/workflows/example_tests.yml, _example_tests_runner.yml, tests/examples/torch_onnx/test_torch_quant_to_onnx.py). No modelopt/ or examples/ source changed, so the algorithm/mode-state/export categories do not apply.

Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

What I verified rather than assumed

  • PR detection is correct and consistent. _PR_CI matches the existing precedent at test_torch_quant_to_onnx.py:43 character for character, and the same refs/heads/pull-request/ prefix is what the workflow itself gates on in 15+ places (example_tests.yml:5 triggers push only on pull-request/[0-9]+). The nightly schedule and workflow_dispatch runs carry refs/heads/main, so all five cases still run there — the stated intent holds.
  • GITHUB_REF actually reaches pytest. _example_tests_runner.yml:42 runs the job under container:, so the default GITHUB_* env vars are present in the test step. The PR's own run (13 passed / 4 skipped, the three nightly-only cases plus the pre-existing LoRA skip) confirms this empirically rather than by reasoning — good evidence to include.
  • Test ids are preserved. pytest.param("deepspeed", marks=...) keeps the id deepspeed, so nothing selecting these cases by id (-k, nightly reports, flake tracking) breaks. This matters because the PR explicitly promises unchanged ids.
  • The comment's coverage claim is accurate. "QLoRA still trains on ddp on PRs" checks out — test_qwen3_qlora_nvfp4 passes backend=\"ddp\" (test_llm_qat.py:285), so the DDP wrapper path survives on PRs even with the QAT ddp case skipped. GRADIENT_CHECKPOINTING_BACKENDS therefore still has a live PR consumer and is not dead on PRs.
  • No shared-fixture fallout. nvfp4_ptq_dir and dataset_cache_dir are module-scoped and still requested by the surviving fsdp2 cases, so skipping params does not leave an unbuilt PTQ checkpoint or change what the remaining cases see.
  • No reuse opportunity missed in the workflow. The workflow's NON_PR (example_tests.yml:92) is a step-local env in the pr-gate job and never reaches the test container, so deriving the flag from GITHUB_REF in the test file is the right call, not a duplication of an available signal.

Note on CodeRabbit's walkthrough

Its summary states QAD "removes DDP from its backend cases." That is incorrect — test_qwen3_qad_nvfp4 was already parametrized over fsdp2/deepspeed only; no ddp case existed to remove. The PR description has it right. Flagging so the inaccurate note doesn't get carried into the merge commit or release summary.

The one suggestion

Non-blocking, left inline: the PR-detection snippet is now duplicated verbatim across two example-test files, and both already import from tests/_test_utils/examples/run_command.py, which is a natural shared home as more lanes get trimmed under OMNIML-6177. Fine to defer to whoever adds the third copy.

On the open question for code owners

Not a code finding, but worth an explicit opinion since the description asks: keeping fsdp2 for both families (this PR) looks like the better of the two options. The alternative (fsdp2 QAT + deepspeed QAD) buys DeepSpeed's GatheredParameters distillation path on PRs but gives up the FSDP2-only distillation path — and FSDP2 is both the documented default and the side with the ModelOpt-specific code (_patch_fsdp2_post_backward, _patch_accelerate_for_fsdp2_fix, the sharded save). The CI history cited (all four failures since 10-06 on fsdp2/QLoRA, none on deepspeed or ddp) points the same way. The residual risk is real but bounded: a regression confined to DeepSpeed or full-parameter DDP now surfaces the next night instead of on the PR, and since the torch lane only runs when {examples,tests/examples}/llm_qat/** and friends change, a modelopt/ change that breaks make_deepspeed_compatible was already not gated by this lane.

Risk: low. Test-selection only, no source or workflow changes, behavior on nightly and local runs unchanged, and the timing claims are backed by measured CI runs rather than estimates.

@claude claude Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude review passed — no blocking issues found. LGTM

@cjluo-nv cjluo-nv left a comment •

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

LGTM: the PR-only skips are correctly scoped and justified, but the existing-test coverage tradeoff needs human sign-off.

Needs action:

  • Confirm the llm_qat coverage tradeoff: retain FSDP2 for both PR test families and defer DeepSpeed QAT/QAD and full-parameter DDP QAT to scheduled CI.

No action needed:

  • The ref check matches the example workflow and existing ONNX tests; main and local runs retain all five backend cases.
  • Parameter IDs, test bodies, and assertions are preserved. QLoRA continues exercising DDP on PRs.
  • Tests were inspected, not executed.

# sharding. PRs run fsdp2 (the default, with the most ModelOpt-specific code); the other backends run
# in the scheduled CI. QLoRA still trains on ddp on PRs.
_PR_CI = os.environ.get("GITHUB_REF", "").startswith("refs/heads/pull-request/")
_NIGHTLY_ONLY = pytest.mark.skipif(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we use pytest.mark.manual marker directly?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not directly, for now. The example workflow never passes --run-manual (only gpu_tests.yml and unit_tests.yml do, on non-PR runs), so a bare pytest.mark.manual would skip these three cases on the nightly too, and nothing would run them. Making it work means the example runner passes --run-manual off-PR, which would also start the two manual cases in tests/examples/speculative_decoding/test_eagle.py ("skip redundant test, too slow") on every nightly, and local runs would skip the three cases unless --run-manual is given. I can do that in a follow-up (the runner passes --run-manual off-PR, the examples use the marker, and we decide what to do with the two eagle cases). Here the GITHUB_REF check matches the existing precedent in test_torch_quant_to_onnx.py (#2853).

@kevalmorabia97 kevalmorabia97 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Would be nice to do a follow-up to make it generic like unit and gpu tests.
Perhaps we do pytest.mark.nightly for nightly ones and pytest.mark.manual for the speculative ones which we may want to skip in nightly CI too

@shengliangxu
shengliangxu merged commit 5945aa8 into main Oct 22, 2026
49 checks passed
@shengliangxu
shengliangxu deleted the shengliangx/ci-qat-one-backend branch October 22, 2026 19:55
@github-actions

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-10-22 19:55 UTC

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants