[AMD][CI] Fix AMD CI (Partial) (#6704)

Signed-off-by: root <root@tw020.mia.tensorwave.lan>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: root <root@tw020.mia.tensorwave.lan>

tests/diffusion/distributed/test_tensor_parallel.py

@@ -9,6 +9,7 @@ Uses ``Tongyi-MAI/Z-Image-Turbo`` as the parity model.
 
 from __future__ import annotations
 
+import gc
 import os
 import time
 from pathlib import Path
@@ -75,8 +76,13 @@ def _run_generate(
     if num_requests < 2:
         raise ValueError("num_requests must be >= 2 (1 warmup + >=1 timed)")
 
+    gc.collect()
     current_omni_platform.empty_cache()
     device_index = current_omni_platform.current_device()
+    with current_omni_platform.device(device_index):
+        free_bytes, total_bytes = current_omni_platform.mem_get_info()
+    initial_used_mb = (total_bytes - free_bytes) / (1024**2)
+
     monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
     monitor.start()
     try:
@@ -113,7 +119,10 @@ def _run_generate(
                 pass
 
             median_time_s = float(np.median(per_request_times_s))
-            peak_memory_mb = monitor.peak_used_mb
+            # DeviceMemoryMonitor samples absolute device usage. Normalize by
+            # this run's starting usage so compiler and allocator caches retained
+            # by the preceding TP=1 run are not charged to the TP=2 run.
+            peak_memory_mb = max(0.0, monitor.peak_used_mb - initial_used_mb)
             return _extract_single_image([last_output]), median_time_s, peak_memory_mb
     finally:
         monitor.stop()