[Bugfix][Test] Assert stable-audio CPU offload savings on max_memory_allocated (#6826)

Signed-off-by: Inesh Reddy <ineshreddy249@gmail.com>
Signed-off-by: Inesh Reddy Chappidi <ineshreddy249@gmail.com>

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"
+        )