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:
co-authored by
Andreas Karatzas
Andreas Karatzas
OpenAI Codex
parent
aaaeda98dc
commit
0ba2aa35a8
+26
-9
@@ -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
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user