From 456b2edce5909e7c4f0966ca281f14bc2b97575e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 21 Oct 2026 22:13:20 +0000 Subject: [PATCH] [OMNIML-6180] Let the T5 and Mixtral hf_ptq tests share the GPU 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 --- tests/examples/hf_ptq/test_llm_ptq.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/examples/hf_ptq/test_llm_ptq.py b/tests/examples/hf_ptq/test_llm_ptq.py index 4de5d95d622..a820158fa9a 100644 --- a/tests/examples/hf_ptq/test_llm_ptq.py +++ b/tests/examples/hf_ptq/test_llm_ptq.py @@ -19,11 +19,18 @@ from _test_utils.examples.models import MIXTRAL_PATH, T5_PATH, TINY_LLAMA_PATH +def _run_sharing_the_gpu(command: PTQCommand, model_path: str): + # Sharing the GPU, so set a batch size: hf_ptq sizes one to fill the free GPU memory. + if command.calib_batch_size is None: + command = dataclasses.replace(command, calib_batch_size=command.calib) + command.run(model_path) + + # Encoder-decoder representative; BART shares the same export path. @pytest.mark.parametrize("command", [PTQCommand(quant="fp8", min_sm=89)], ids=PTQCommand.param_str) -@pytest.mark.usefixtures("exclusive_gpu") +@pytest.mark.usefixtures("shared_gpu") def test_ptq_t5(command): - command.run(T5_PATH) + _run_sharing_the_gpu(command, T5_PATH) @pytest.mark.parametrize( @@ -33,9 +40,9 @@ def test_ptq_t5(command): ], ids=PTQCommand.param_str, ) -@pytest.mark.usefixtures("exclusive_gpu") +@pytest.mark.usefixtures("shared_gpu") def test_ptq_mixtral(command): - command.run(MIXTRAL_PATH) + _run_sharing_the_gpu(command, MIXTRAL_PATH) # Per-format export on tiny models (int8_weight_only, nvfp4_awq_lite, ...) is covered by @@ -77,7 +84,4 @@ def test_ptq_mixtral(command): ) @pytest.mark.usefixtures("shared_gpu") def test_ptq_llama(command): - # Sharing the GPU, so set a batch size: hf_ptq sizes one to fill the free GPU memory. - if command.calib_batch_size is None: - command = dataclasses.replace(command, calib_batch_size=command.calib) - command.run(TINY_LLAMA_PATH) + _run_sharing_the_gpu(command, TINY_LLAMA_PATH)