From e7213003cbf64d3f35b97d711eb595aa9e47039c Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Tue, 3 Mar 2026 14:57:34 -0600 Subject: [PATCH] [ROCm][CI] Fix TP size issue for `test_gpt_oss` (#35887) Signed-off-by: Micah Williamson --- tests/models/quantization/test_gpt_oss.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/models/quantization/test_gpt_oss.py b/tests/models/quantization/test_gpt_oss.py index 6fab653d009..7599a5a5ee4 100644 --- a/tests/models/quantization/test_gpt_oss.py +++ b/tests/models/quantization/test_gpt_oss.py @@ -21,6 +21,8 @@ import lm_eval import pytest from packaging import version +from vllm.utils.torch_utils import cuda_device_count_stateless + MODEL_ACCURACIES = { # Full quantization: attention linears and MoE linears "amd/gpt-oss-20b-WFP8-AFP8-KVFP8": 0.89, @@ -83,6 +85,9 @@ class EvaluationConfig: def test_gpt_oss_attention_quantization( model_name: str, tp_size: int, expected_accuracy: float ): + if tp_size > cuda_device_count_stateless(): + pytest.skip("Not enough GPUs to run this test case") + model_args = EvaluationConfig(model_name).get_model_args(tp_size) extra_run_kwargs = {