Repository navigation
[OMNIML-6180] Run only the fsdp2 llm_qat training cases on PRs - #2925
Conversation
|
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. |
|
/ok to test 7b7f0c8 |
Codecov Report✅ All modified and coverable lines are covered by tests. 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
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:
|
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>
7b7f0c8 to
992db89
Compare
|
/claude review |
Claude review summaryScope: full review. The trigger comment was a bare Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1 What I verified rather than assumed
Note on CodeRabbit's walkthroughIts summary states QAD "removes DDP from its backend cases." That is incorrect — The one suggestionNon-blocking, left inline: the PR-detection snippet is now duplicated verbatim across two example-test files, and both already import from On the open question for code ownersNot a code finding, but worth an explicit opinion since the description asks: keeping 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. |
There was a problem hiding this comment.
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_qatcoverage 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( |
There was a problem hiding this comment.
Can we use pytest.mark.manual marker directly?
There was a problem hiding this comment.
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
left a comment
There was a problem hiding this comment.
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
|
What does this PR do?
Type of change: CI (test runtime, tests only)
On pull requests the
llm_qatexample lane runs one training backend per test family: thefsdp2cases oftest_qwen3_qat_nvfp4andtest_qwen3_qad_nvfp4. Thedeepspeedcase of both and theddpcase oftest_qwen3_qat_nvfp4keep 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 anaccelerate launchoftrain.pyand takes 99-109 s on a normal-speed runner, 146-157 s on a slow one.fsdp2is 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 intransformers_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 onfsdp2or QLoRA cases (three at the 300 s per-test cap on slow runners, one real error intest_qwen3_qad_nvfp4[fsdp2]); none is ondeepspeedorddp.test_qwen3_qlora_nvfp4still trains onddpon PRs.The skip is a
pytest.mark.skipifon the three parameters, using the sameGITHUB_REFcheck astests/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 themanualmarker 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:
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
fsdp2QAT (100 s),fsdp2QAD (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);fsdp2QAT 152 s,fsdp2QAD 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; theRun testsstep 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 unchangedllm_qatjob 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,fsdp2QAT 128 s against 83 s).Usage
Nothing user-facing. On a
pull-request/Nref (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
run_example_commandreplaced 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]andtest_qwen3_qad_nvfp4[fsdp2]run;test_qwen3_qat_nvfp4[deepspeed],[ddp]andtest_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 unsetGITHUB_REF: all five run (5 passed, 1 skipped, the same existing skip).fsdp2cases, 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
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.ddpQAT. The DDP wrapper itself still runs on PRs through the QLoRA test.fsdp2QAT and QAD, QLoRA, the pre-tokenize test and the unit-style tests of the lane are unchanged.Open question for the code owners: keep
fsdp2for both families (this PR), or spread the backends so every one stays on PRs once (fsdp2for QAT,deepspeedfor QAD), which keeps DeepSpeed's distillation path on PRs but movesfsdp2distillation 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.).CONTRIBUTING.md: no new dependency/claude reviewapproved 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