Stabilize GPU memory teardown between ROCm CI tests (#49242)

Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com>
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Aarushi Jain
2026-07-25 04:21:36 +00:00
committed by GitHub
co-authored by Andreas Karatzas Andreas Karatzas OpenAI Codex
parent aaaeda98dc
commit 0ba2aa35a8
3 changed files with 126 additions and 93 deletions
+26 -9
View File
@@ -1728,33 +1728,50 @@ def disable_deepgemm_ue8m0(monkeypatch):
is_deep_gemm_e8m0_used.cache_clear()
def _should_clean_gpu_memory_between_tests() -> bool:
setting = os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY")
if setting == "1":
return True
if setting == "0":
return False
# ROCm reclaims VRAM lazily; default to waiting between tests on ROCm CI.
return current_platform.is_rocm()
@pytest.fixture(autouse=True)
def clean_gpu_memory_between_tests():
if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1":
if not _should_clean_gpu_memory_between_tests():
yield
return
# Wait for GPU memory to be cleared before starting the test
import gc
from tests.utils import wait_for_gpu_memory_to_clear
from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle
num_gpus = torch.accelerator.device_count()
if num_gpus > 0:
def _wait_for_settled_gpu_memory() -> None:
if num_gpus <= 0:
return
try:
wait_for_gpu_memory_to_clear(
devices=list(range(num_gpus)),
threshold_ratio=0.1,
)
if current_platform.is_rocm():
wait_for_rocm_memory_to_settle()
else:
wait_for_gpu_memory_to_clear(
devices=list(range(num_gpus)),
threshold_ratio=0.1,
)
except ValueError as e:
logger.info("Failed to clean GPU memory: %s", e)
_wait_for_settled_gpu_memory()
yield
# Clean up GPU memory after the test
if torch.cuda.is_available():
torch.accelerator.empty_cache()
gc.collect()
_wait_for_settled_gpu_memory()
@pytest.fixture
+87 -81
View File
@@ -7,6 +7,7 @@ from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
import vllm
from tests.conftest import VllmRunner
from vllm.assets.image import ImageAsset
from vllm.lora.request import LoRARequest
from vllm.platforms import current_platform
@@ -56,24 +57,29 @@ class Qwen2VLTester:
def __init__(self, config: TestConfig):
self.config = config
self.llm = self._initialize_llm()
def _initialize_llm(self) -> vllm.LLM:
"""Initialize the LLM with given configuration"""
return vllm.LLM(
model=self.config.model_path,
max_num_seqs=self.config.max_num_seqs,
self._runner = VllmRunner(
model_name=config.model_path,
max_num_seqs=config.max_num_seqs,
enable_lora=True,
max_loras=self.config.max_loras,
max_lora_rank=self.config.max_lora_rank,
enable_tower_connector_lora=self.config.enable_tower_connector_lora,
trust_remote_code=True,
gpu_memory_utilization=self.config.gpu_memory_utilization,
mm_processor_kwargs=self.config.mm_processor_kwargs,
mm_processor_cache_gb=self.config.mm_processor_cache_gb,
max_model_len=self.config.max_model_len,
max_loras=config.max_loras,
max_lora_rank=config.max_lora_rank,
enable_tower_connector_lora=config.enable_tower_connector_lora,
gpu_memory_utilization=config.gpu_memory_utilization,
mm_processor_kwargs=config.mm_processor_kwargs,
mm_processor_cache_gb=config.mm_processor_cache_gb,
max_model_len=config.max_model_len,
)
@property
def llm(self) -> vllm.LLM:
return self._runner.get_llm()
def __enter__(self) -> "Qwen2VLTester":
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
self._runner.__exit__(exc_type, exc_value, traceback)
def run_test(
self,
images: list[ImageAsset],
@@ -183,29 +189,29 @@ QWEN3VL_MODEL_PATH = "Qwen/Qwen3-VL-4B-Instruct"
def test_qwen2vl_lora(qwen2vl_lora_files):
"""Test Qwen 2.0 VL model with LoRA"""
config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
tester = Qwen2VLTester(config)
# Test with different LoRA IDs
for lora_id in [1, 2]:
tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
with Qwen2VLTester(config) as tester:
# Test with different LoRA IDs
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id
)
def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
"""Test Qwen 2.0 VL model with LoRA through beam search."""
config = TestConfig(model_path=QWEN2VL_MODEL_PATH, lora_path=qwen2vl_lora_files)
tester = Qwen2VLTester(config)
# Test with different LoRA IDs
for lora_id in [1, 2]:
# NOTE currently, we only test cherry blossom since stop sign
# output is slightly different for v1; - the root cause is likely
# independent of the intent of this test, which is to ensure beam
# search passes through lora through correctly.
tester.run_beam_search_test(
[ImageAsset("cherry_blossom")],
expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS,
lora_id=lora_id,
)
with Qwen2VLTester(config) as tester:
# Test with different LoRA IDs
for lora_id in [1, 2]:
# NOTE currently, we only test cherry blossom since stop sign
# output is slightly different for v1; - the root cause is likely
# independent of the intent of this test, which is to ensure beam
# search passes through lora through correctly.
tester.run_beam_search_test(
[ImageAsset("cherry_blossom")],
expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS,
lora_id=lora_id,
)
@pytest.mark.skipif(
@@ -214,11 +220,12 @@ def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):
def test_qwen25vl_lora(qwen25vl_lora_files):
"""Test Qwen 2.5 VL model with LoRA"""
config = TestConfig(model_path=QWEN25VL_MODEL_PATH, lora_path=qwen25vl_lora_files)
tester = Qwen2VLTester(config)
# Test with different LoRA IDs
for lora_id in [1, 2]:
tester.run_test(TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id)
with Qwen2VLTester(config) as tester:
# Test with different LoRA IDs
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES, expected_outputs=EXPECTED_OUTPUTS, lora_id=lora_id
)
@pytest.mark.skipif(
@@ -234,13 +241,13 @@ def test_qwen25vl_vision_lora(qwen25vl_vision_lora_files):
mm_processor_cache_gb=0,
enable_tower_connector_lora=True,
)
tester = Qwen2VLTester(config)
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS,
lora_id=lora_id,
)
with Qwen2VLTester(config) as tester:
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS,
lora_id=lora_id,
)
def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files):
@@ -253,13 +260,13 @@ def test_qwen3vl_vision_lora(qwen3vl_vision_lora_files):
mm_processor_cache_gb=0,
enable_tower_connector_lora=True,
)
tester = Qwen2VLTester(config)
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS,
lora_id=lora_id,
)
with Qwen2VLTester(config) as tester:
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS,
lora_id=lora_id,
)
def test_qwen2vl_multiple_lora_types(
@@ -287,34 +294,33 @@ def test_qwen2vl_multiple_lora_types(
mm_processor_cache_gb=0,
enable_tower_connector_lora=True,
)
tester = Qwen2VLTester(config)
with Qwen2VLTester(config) as tester:
# Test 1: Language-only LoRA adapter
tester.config.lora_path = qwen2vl_language_lora_files
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_LANGUAGE,
lora_id=lora_id,
lora_name="language_only",
)
# Test 1: Language-only LoRA adapter
tester.config.lora_path = qwen2vl_language_lora_files
for lora_id in [1, 2]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_LANGUAGE,
lora_id=lora_id,
lora_name="language_only",
)
# Test 2: Vision tower + connector LoRA adapter
tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files
for lora_id in [3, 4]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_VISION,
lora_id=lora_id,
lora_name="vision_tower_connector",
)
# Test 2: Vision tower + connector LoRA adapter
tester.config.lora_path = qwen2vl_vision_tower_connector_lora_files
for lora_id in [3, 4]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_VISION,
lora_id=lora_id,
lora_name="vision_tower_connector",
)
# Test 3: Vision tower only LoRA adapter (no connector)
tester.config.lora_path = qwen2vl_vision_tower_lora_files
for lora_id in [5, 6]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR,
lora_id=lora_id,
lora_name="vision_tower",
)
# Test 3: Vision tower only LoRA adapter (no connector)
tester.config.lora_path = qwen2vl_vision_tower_lora_files
for lora_id in [5, 6]:
tester.run_test(
TEST_IMAGES,
expected_outputs=EXPECTED_OUTPUTS_VISION_NO_CONNECTOR,
lora_id=lora_id,
lora_name="vision_tower",
)
@@ -98,7 +98,7 @@ GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
# Trap the SIGINT signal (triggered by Ctrl+C)
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
trap 'kill $(jobs -pr) 2>/dev/null || true' SIGINT SIGTERM EXIT
# Waits for vLLM to start.
wait_for_server() {
@@ -110,10 +110,20 @@ wait_for_server() {
}
# Function to clean up previous instances
wait_for_gpu_memory_release() {
if [[ "$SMI_BIN" == *"rocm"* ]]; then
PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()"
fi
}
cleanup_instances() {
echo "Cleaning up any running vLLM instances..."
pkill -f "vllm serve" || true
pkill -f "toy_proxy_server.py" || true
pkill -TERM -f "vllm serve" || true
sleep 3
pkill -9 -f "vllm serve" || true
sleep 2
wait_for_gpu_memory_release
}
get_num_gpus() {
@@ -132,6 +142,7 @@ get_num_gpus() {
# Function to run tests for a specific model
run_tests_for_model() {
local model_name=$1
cleanup_instances
echo "================================"
echo "Testing model: $model_name"
echo "================================"
@@ -303,7 +314,6 @@ run_tests_for_model() {
# Clean up before running next model
cleanup_instances
sleep 3
}
# Run tests for each model