diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 205fad8fd..873f355c9 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -16,7 +16,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform AUDIO_MODEL: dict[str, dict[str, int | None]] = { - "stabilityai/stable-audio-open-1.0": {"cuda": 100, "rocm": None}, + "stabilityai/stable-audio-open-1.0": {"cuda": 150, "rocm": None}, } IMAGE_MODELS: dict[str, dict[str, int | None]] = { @@ -35,6 +35,11 @@ MODEL_MARKS = { _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} +# Models whose thresholds are calibrated against max_memory_allocated. All +# other models keep asserting on the polled device-wide figure until their +# thresholds have been re-measured on that metric on both CUDA and ROCm. +_ALLOCATED_THRESHOLD_MODELS = {"stabilityai/stable-audio-open-1.0"} + AUDIO_MODEL_PARAMS: dict[str, dict[str, Any]] = { "runner_params": {}, @@ -50,7 +55,7 @@ IMAGE_MODEL_PARAMS: dict[str, dict[str, Any]] = { } -def inference(model_name: str, offload: bool = True) -> tuple[float, Any]: +def inference(model_name: str, offload: bool = True) -> tuple[float, float, Any]: gc.collect() current_omni_platform.empty_cache() device_index = current_omni_platform.current_device() @@ -81,12 +86,17 @@ def inference(model_name: str, offload: bool = True) -> tuple[float, Any]: ), ) peak = monitor.peak_used_mb + # Exact in-process peak of live tensors. Unlike the polled device-wide + # figure above, this excludes caching-allocator slack and CUDA context, + # which differ between the offload and no-offload paths (the offload hook + # calls empty_cache() on every swap) and drift between runs. + peak_allocated = current_omni_platform.max_memory_allocated(device=device_index) / (1024**2) monitor.stop() gc.collect() current_omni_platform.empty_cache() - return peak, output + return peak, peak_allocated, output def check_audio_determinism(audio1, audio2, atol=1e-2): @@ -112,9 +122,9 @@ def check_audio_determinism(audio1, audio2, atol=1e-2): ) def test_cpu_offload_diffusion_model(model_name: str): try: - offload_peak_memory, output_offload = inference(model_name, offload=True) + offload_peak_memory, offload_peak_allocated, output_offload = inference(model_name, offload=True) cleanup_dist_env_and_memory() - no_offload_peak_memory, output_no_offload = inference(model_name, offload=False) + no_offload_peak_memory, no_offload_peak_allocated, output_no_offload = inference(model_name, offload=False) except ValueError as exc: # omni_snapshot_download wraps GatedRepoError in a ValueError; skip instead of failing. if "Access to model" in str(exc) and "is restricted" in str(exc): @@ -125,16 +135,14 @@ def test_cpu_offload_diffusion_model(model_name: str): pytest.fail(f"Inference failed: {exc}") except Exception: pytest.fail("Inference failed") - print(f"Offload peak memory: {offload_peak_memory} MB") - print(f"No offload peak memory: {no_offload_peak_memory} MB") + print(f"Offload peak memory: {offload_peak_memory} MB (allocated: {offload_peak_allocated:.1f} MB)") + print(f"No offload peak memory: {no_offload_peak_memory} MB (allocated: {no_offload_peak_allocated:.1f} MB)") if model_name == "stabilityai/stable-audio-open-1.0": audio_offload = output_offload[0].multimodal_output.get("audio") audio_no_offload = output_no_offload[0].multimodal_output.get("audio") check_audio_determinism(audio_offload, audio_no_offload, atol=1e-2) - # Set platform-specific VRAM saving thresholds to account - # for varying runtime memory overhead and fragmentation between CUDA and ROCm. is_rocm = torch.version.hip is not None platform = "rocm" if is_rocm else "cuda" threshold = MODELS[model_name][platform] @@ -142,7 +150,22 @@ def test_cpu_offload_diffusion_model(model_name: str): pytest.skip(f"Threshold not defined for {platform} on {model_name}") assert threshold is not None - assert offload_peak_memory + threshold < no_offload_peak_memory, ( - f"Offload peak memory {offload_peak_memory} MB should be less than " - f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB" - ) + if model_name in _ALLOCATED_THRESHOLD_MODELS: + # Threshold is a lower bound on the peak *allocated* saving, i.e. the + # weights the offloader keeps off-GPU at the moment of peak usage + # (stable-audio: the fp16 T5 encoder, ~209 MB). max_memory_allocated + # is used instead of the polled device-wide figure because the latter + # includes caching-allocator slack that differs between the two paths + # and drifts by tens of MB between runs. + assert offload_peak_allocated + threshold < no_offload_peak_allocated, ( + f"Offload peak allocated memory {offload_peak_allocated:.1f} MB should be less than " + f"no offload peak allocated memory {no_offload_peak_allocated:.1f} MB by {threshold} MB" + ) + else: + # Threshold is calibrated against the polled device-wide figure and + # accounts for runtime memory overhead and fragmentation differences + # between CUDA and ROCm. + assert offload_peak_memory + threshold < no_offload_peak_memory, ( + f"Offload peak memory {offload_peak_memory} MB should be less than " + f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB" + )