Repository navigation
[OMNIML-6180] Let the T5 and Mixtral hf_ptq tests share the GPU - #2926
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. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
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 4 functions across 1 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely identifies the main change: enabling the T5 and Mixtral hf_ptq tests to share the GPU. 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. Security Anti-Patterns The PR changes only tests/examples/hf_ptq/test_llm_ptq.py. The added helper setscalib_batch_sizeand callscommand.run; it adds notorch.load(weights_only=False), `numpy.load(allow_pickle=Tru…Full details: Security Anti-Patterns
-
Explanation
The PR changes only
tests/examples/hf_ptq/test_llm_ptq.py. The added helper setscalib_batch_sizeand callscommand.run; it adds notorch.load(weights_only=False),numpy.load(allow_pickle=True), hardcodedtrust_remote_code=True,eval,exec, or# nosec. It also changes no dependency files.
-
✨ 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.
|
/ok to test ec4646c |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2926 +/- ##
=======================================
Coverage 79.66% 79.66%
=======================================
Files 659 659
Lines 74400 74400
=======================================
+ Hits 59272 59273 +1
+ Misses 15128 15127 -1
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:
|
Under pytest-xdist, test_ptq_t5 and test_ptq_mixtral took the exclusive GPU lock, so the lane waited for them one at a time: 25 s and 114 s of calls in a PR run (547 s of pytest). They are small (t5-small, a 2B two-expert TinyLlama MoE), and the only reason to run alone was the calibration batch size: without one, hf_ptq runs a forward pass sized to fill all free GPU memory. The TinyLlama tests already set the batch size to the calibration size and share the GPU, so give these two the same treatment through one helper. Test ids are unchanged. The Qwen3-VL tests keep the exclusive lock: the TRT-LLM multimodal deploy takes 60% of the free GPU memory for its KV cache. Testing: collection gives the same 21 node ids as main; --setup-plan shows shared_gpu for t5, mixtral and llama and exclusive_gpu for qwen_vl; the repo's lock tests pass; with fake commands that sleep for the call times of a recent PR run (scaled by 0.1) under -n 4 and the real lock, the exclusive tests never overlap another test and the last test ends 92 s earlier at real scale (486 s -> 394 s). Not run on a GPU: the real lane. Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
ec4646c to
456b2ed
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
/claude review |
Claude review summaryScope: full review. The trigger comment was a bare FindingsCRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 0 What I verified
Risk assessmentLow. Test-only, no production code path touched, mechanically equivalent to logic already in CI for 17 sibling cases, and the behavior change (fixed calibration batch instead of a memory probe for two tests) is measured in two independent CI run pairs in the description. The one residual unknown — peak GPU memory for Mixtral running next to three TinyLlama workers — is explicitly called out as unmeasured by the author, with a concrete follow-up path, and two full GPU lane runs completed without OOM or exceeding the 300 s per-test cap. That is a reasonable place to leave it for a CI-runtime change; I see nothing to block on. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
LGTM: this is a small, correct test-only change. It moves test_ptq_t5 and test_ptq_mixtral onto the shared GPU lock and pulls the existing TinyLlama batch-size logic into _run_sharing_the_gpu.
Needs action:
- Sign off on the existing-test edit in
tests/examples/hf_ptq/test_llm_ptq.py. T5 and Mixtral now use a fixed calibration batch instead of the memory probe, in both PR CI and the nightly. That's justified by the scheduling change, and the PR's CI runs show no OOM.
No action needed:
- The helper is a verbatim extraction of the previous
test_ptq_llamalogic, so the llama cases behave exactly as before. The explicitcalib_batch_size=4for the AutoQuantize case is still respected. - The Qwen3-VL tests keep
exclusive_gpu, which fits the TRT-LLM KV cache reservation described in the PR. - The
_example_tests_runner.ymlcomment is now slightly out of date. The author has deliberately left that for the next workflow change. - I read the tests but didn't run them.
|
What does this PR do?
Type of change: CI (test runtime, tests only)
PR runs of the
hf_ptqexample lane use 4 pytest-xdist workers on the runner's one GPU (#2771). Small-model tests share the GPU, and tests that need it alone take an exclusive lock (tests/examples/hf_ptq/conftest.py). While an exclusive test runs, the other three workers can only run tests that need no GPU. Four tests were exclusive:test_ptq_t5,test_ptq_mixtraland the twotest_qwen_vlcases. Their calls took 25 s, 114 s, 114 s and 82 s, 335 s in total, one after the other, in a lane whose pytest time is 547 s (job 852 s; run 39662683232, job 119443637479, on a 1-GPU PR runner, 238 passed).T5 (t5-small) and Mixtral (
TinyLlama-1.1B-Chat-v1.0-x2-MoE, a 2B two-expert MoE) do not need the GPU alone. The reason they did is the calibration batch size: without one,hf_ptq.pyruns a forward pass sized to fill all free GPU memory. The TinyLlama tests already pass a batch size equal to the calibration size and share the GPU. This PR gives T5 and Mixtral the same treatment through one helper (_run_sharing_the_gpu), used by all three tests. Test ids are unchanged. The two Qwen3-VL tests keep the exclusive lock: the TRT-LLM multimodal deploy reserves 60% of the free GPU memory for its KV cache (60 GiB on a 96 GiB GPU, measured in #2771).Measured in CI in two pairs of runs of the
hf_ptqjob (table below): pytest -77 s (-14%) on equally fast runners, and -24 s (-4%) after the rebase, where this PR's runner was about 10% slower than the baseline's (roughly -70 s scaled by that). The simulation had predicted -92 s.Time improvement (measured in CI)
Seconds, from the
--durationslines and the result-line timestamps of thehf_ptqjob on a 1-GPU runner (COVERAGE_CORE=sysmon), in two pairs of runs.hf_ptqjob of [OMNIML-6180] Run the small example lanes of a job together #2924's run (run 39737194943; main's unchanged test file, started within 10 s of the other); after = this PR's run (run 39737199847). Both have 241 tests: the three extra ones (test_convert_ggml_to_int8.py) came with main, take no GPU lock and cost both runs alike.test_ptq_t5call (exclusive, now shared)test_ptq_mixtralcall (exclusive, now shared)test_qwen_vl[fp8]call (exclusive, unchanged)test_qwen_vl[nvfp4]call (exclusive, unchanged)test_ptq_llamacalls, summed (they now run next to T5 and Mixtral)setupdurations pytest lists for the LLM and VLM tests, summed (waiting for the GPU lock)Run tests;Example 1: hf_ptqin #2924's run)The job also includes the container start and the dependency install, which this PR does not change (237 and 26 s before, 251 and 28 s after in pair 1; 243 and 17 s, 255 and 29 s in pair 2).
In pair 1 the unchanged exclusive tests (Qwen3-VL) took the same time in both runs within 2%, so the runners were equally fast. In pair 2 the runner of the "after" run was about 10% slower (the same two tests took 9% and 11% longer), which hides part of the gain: scaled by that slowdown, the pytest saving is roughly 70 s instead of 24 s. In both pairs the mechanism shows the same way: the GPU is held alone for about 140 s less, and the workers wait 293 s and 515 s less in total for the lock. T5 and Mixtral take longer (+2 to +19 s each) because they now run next to the TinyLlama workers, but they no longer block them. Mixtral's call stays under the 300 s per-test cap in both runs, and no test failed or ran out of memory. The two Qwen3-VL tests still hold the GPU alone for about 200 s, which bounds what is left to save.
Usage
Nothing user-facing.
Testing
test_llm_ptq.pyandtest_vlm_ptq.pyas main.--setup-planshowsshared_gpufor t5, mixtral and llama andexclusive_gpufor qwen_vl.tests/examples/hf_ptq/test_gpu_lock.py(the lock itself) passes: 3 tests.-n 4: a plugin (not part of this PR) fakes the GPU facts and replacesrun_hf_ptq_commandwith a sleep of the measured call time of each test (from the 1-GPU run above, scaled by 0.1). On main the exclusive tests never overlap another test, and the first-to-last-test span is 486 s at real scale. On this branch they still never overlap, t5 and mixtral run alongside the llama cases, and the span is 394 s (-92 s). Every fake run of t5, mixtral and llama receivedcalib_batch_size=16(4 for the AutoQuantize case, as before), and qwen_vl received none.What this PR's own GPU run confirmed, and what is left:
trtllm (hf_ptq)runs on the 1-GPU runner (before and after the rebase): T5 and Mixtral stay within the GPU's memory next to the TinyLlama workers (no failure or out-of-memory) and Mixtral's call takes 122 s and 133 s, under the 300 s per-test cap. Peak GPU memory was not logged; to measure it, a cluster agent (lyris, hecate or ptyche) can runpytest tests/examples/hf_ptq/test_llm_ptq.py -n 4at this SHA in the TRT-LLMrelease:1.3.0rc20container on one GPU while loggingnvidia-smi --query-gpu=memory.used, and compare with main.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 456b2ed, no findings)Additional Information
Left out on purpose, each needs a real GPU run or the owners' call:
quickstart_multimodal.pyfrom the container (/app/tensorrt_llm/examples/llm-api/), called fromhuggingface_example.shwith only--model_dirand--modality image, so it uses the script's default KV fraction. Capping it needs a flag whose name forrelease:1.3.0rc20I could not check here (the image pins rc20 because newer releases break the script), plus a proof that a smaller KV cache does not OOM or slow the deploy while three other workers run.KV_CACHE_FREE_GPU_MEMORY_FRACTIONinparser.shis a different knob (the lm_eval path, default 0.8).modelopt.deploy.llm.LLMsetstp = torch.cuda.device_count(), so on 2 GPUs every deploy test runs with tensor parallel 2 and several would overlap on both GPUs. That needs a 2-GPU run to trust._example_tests_runner.yml("small-model tests share it and the rest run alone on it") is now slightly stale (only the Qwen3-VL tests run alone). Editing the workflow would make every example lane run on this PR, so I left it for the next workflow change.Related: #2771 (the GPU lock and
-n 4).Summary by CodeRabbit