[diffusion] CI: guard the allocated vram peak with reporting the reserved one (#38172)

This commit is contained in:
Mick
2026-09-06 17:22:16 +08:00
committed by GitHub
parent 8ef646a5c6
commit a9944aec01
8 changed files with 328 additions and 29 deletions
@@ -153,12 +153,17 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
self.pipeline: ComposedPipelineBase = None
self.init_device_and_model()
load_snapshot = None if current_platform.is_cpu() else capture_memory_snapshot()
self._load_peak_reserved_mb = (
0.0
if current_platform.is_cpu()
else capture_memory_snapshot().peak_reserved_mb
load_snapshot.peak_reserved_mb if load_snapshot is not None else 0.0
)
# Allocated peaks are the guarded figure: reserved peaks also track the
# allocator's pool history and move a few percent for identical work.
self._load_peak_allocated_mb = (
load_snapshot.peak_allocated_mb if load_snapshot is not None else 0.0
)
self._runtime_peak_reserved_mb = 0.0
self._runtime_peak_allocated_mb = 0.0
self.sp_group = get_sp_group()
self.sp_cpu_group = self.sp_group.cpu_group
self.tp_group = get_tp_group()
@@ -731,12 +736,15 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
if current_platform.is_cpu():
return
peak_reserved_mb = capture_memory_snapshot().peak_reserved_mb
snapshot = capture_memory_snapshot()
self._runtime_peak_reserved_mb = max(
self._runtime_peak_reserved_mb, peak_reserved_mb
self._runtime_peak_reserved_mb, snapshot.peak_reserved_mb
)
self._runtime_peak_allocated_mb = max(
self._runtime_peak_allocated_mb, snapshot.peak_allocated_mb
)
if self.is_output_rank:
output_batch.peak_memory_mb = peak_reserved_mb
output_batch.peak_memory_mb = snapshot.peak_reserved_mb
def _record_replica_peak_memory(self, output_metrics: list[Any]) -> None:
"""Record replica-wide loading and runtime allocator peaks."""
@@ -744,7 +752,12 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
return
peaks = torch.tensor(
[self._load_peak_reserved_mb, self._runtime_peak_reserved_mb],
[
self._load_peak_reserved_mb,
self._runtime_peak_reserved_mb,
self._load_peak_allocated_mb,
self._runtime_peak_allocated_mb,
],
dtype=torch.float64,
device=current_platform.get_device(self.local_rank),
)
@@ -753,13 +766,28 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
return
snapshot = capture_memory_snapshot()
load_peak_mb, runtime_peak_mb = peaks.tolist()
(
load_peak_mb,
runtime_peak_mb,
load_peak_allocated_mb,
runtime_peak_allocated_mb,
) = peaks.tolist()
for metrics in output_metrics:
metrics.record_memory_snapshot(
"load_peak", replace(snapshot, peak_reserved_mb=load_peak_mb)
"load_peak",
replace(
snapshot,
peak_reserved_mb=load_peak_mb,
peak_allocated_mb=load_peak_allocated_mb,
),
)
metrics.record_memory_snapshot(
"runtime_peak", replace(snapshot, peak_reserved_mb=runtime_peak_mb)
"runtime_peak",
replace(
snapshot,
peak_reserved_mb=runtime_peak_mb,
peak_allocated_mb=runtime_peak_allocated_mb,
),
)
def _forward_group(self, batch: list[Req]) -> OutputBatch:
@@ -83,8 +83,8 @@ def _generate_diffusion_markdown_report(results: list) -> str:
# Main performance table
markdown = header
markdown += "| Test Suite | Test Name | Modality | E2E (ms) | Avg Denoise (ms) | Median Denoise (ms) | Load Peak VRAM (MiB) | Runtime Peak VRAM (MiB) |\n"
markdown += "| ---------- | --------- | -------- | -------- | ---------------- | ------------------- | -------------------- | ----------------------- |\n"
markdown += "| Test Suite | Test Name | Modality | E2E (ms) | Avg Denoise (ms) | Median Denoise (ms) | Load Peak VRAM (MiB) | Runtime Peak VRAM (MiB) | Load Peak Alloc (MiB) | Runtime Peak Alloc (MiB) |\n"
markdown += "| ---------- | --------- | -------- | -------- | ---------------- | ------------------- | -------------------- | ----------------------- | --------------------- | ------------------------ |\n"
for entry in sorted(results, key=lambda x: (x["class_name"], x["test_name"])):
modality = entry.get("modality", "image")
@@ -93,7 +93,9 @@ def _generate_diffusion_markdown_report(results: list) -> str:
f"{entry['e2e_ms']:.2f} | {entry['avg_denoise_ms']:.2f} | "
f"{entry['median_denoise_ms']:.2f} | "
f"{entry.get('load_peak_vram_mb', 0):.0f} | "
f"{entry.get('runtime_peak_vram_mb', 0):.0f} |\n"
f"{entry.get('runtime_peak_vram_mb', 0):.0f} | "
f"{entry.get('load_peak_allocated_mb', 0):.0f} | "
f"{entry.get('runtime_peak_allocated_mb', 0):.0f} |\n"
)
# Video-specific metrics table (if any video tests)
@@ -133,7 +135,7 @@ def pytest_sessionfinish(session):
# Print to stdout (existing behavior)
print("\n\n" + "=" * 35 + " Performance Summary " + "=" * 35)
print(
f"{'Test Suite':<30} | {'Test Name':<20} | {'E2E (ms)':>12} | {'Avg Denoise (ms)':>18} | {'Median Denoise (ms)':>20} | {'Load Peak (MiB)':>15} | {'Runtime Peak (MiB)':>18}"
f"{'Test Suite':<30} | {'Test Name':<20} | {'E2E (ms)':>12} | {'Avg Denoise (ms)':>18} | {'Median Denoise (ms)':>20} | {'Load Peak (MiB)':>15} | {'Runtime Peak (MiB)':>18} | {'Load Alloc (MiB)':>16} | {'Runtime Alloc (MiB)':>19}"
)
print(
"-" * 30
@@ -149,6 +151,10 @@ def pytest_sessionfinish(session):
+ "-" * 15
+ "-+-"
+ "-" * 18
+ "-+-"
+ "-" * 16
+ "-+-"
+ "-" * 19
)
for entry in sorted_results:
@@ -156,7 +162,9 @@ def pytest_sessionfinish(session):
f"{entry['class_name']:<30} | {entry['test_name']:<20} | {entry['e2e_ms']:>12.2f} | "
f"{entry['avg_denoise_ms']:>18.2f} | {entry['median_denoise_ms']:>20.2f} | "
f"{entry.get('load_peak_vram_mb', 0):>15.0f} | "
f"{entry.get('runtime_peak_vram_mb', 0):>18.0f}"
f"{entry.get('runtime_peak_vram_mb', 0):>18.0f} | "
f"{entry.get('load_peak_allocated_mb', 0):>16.0f} | "
f"{entry.get('runtime_peak_allocated_mb', 0):>19.0f}"
)
print("=" * 130)
@@ -62,6 +62,8 @@
"expected_median_denoise_ms": 345.95,
"load_peak_vram_mb": 1126.0,
"runtime_peak_vram_mb": 13180.0,
"load_peak_allocated_mb": 1019.0,
"runtime_peak_allocated_mb": 8912.0,
"estimated_full_test_time_s": 94.0
},
"wan2_1_t2v_1.3b": {
@@ -86,6 +88,8 @@
"expected_median_denoise_ms": 427.94,
"load_peak_vram_mb": 4934.0,
"runtime_peak_vram_mb": 18246.0,
"load_peak_allocated_mb": 903.0,
"runtime_peak_allocated_mb": 6934.0,
"estimated_full_test_time_s": 160.9
},
"turbo_wan2_1_t2v_1.3b": {
@@ -109,6 +113,8 @@
"expected_median_denoise_ms": 160.64,
"load_peak_vram_mb": 928.0,
"runtime_peak_vram_mb": 14216.0,
"load_peak_allocated_mb": 903.0,
"runtime_peak_allocated_mb": 6925.0,
"load_peak_host_anon_mb": 6759.33,
"runtime_peak_host_anon_mb": 6759.33,
"estimated_full_test_time_s": 200.3
@@ -135,6 +141,8 @@
"expected_median_denoise_ms": 265.1,
"load_peak_vram_mb": 13120.0,
"runtime_peak_vram_mb": 17776.0,
"load_peak_allocated_mb": 12852.0,
"runtime_peak_allocated_mb": 15133.0,
"estimated_full_test_time_s": 329.8
},
"flux_image_t2i_layerwise_cpu_offload_5090": {
@@ -157,6 +165,8 @@
"expected_median_denoise_ms": 779.5,
"load_peak_vram_mb": 4916.0,
"runtime_peak_vram_mb": 19146.0,
"load_peak_allocated_mb": 4880.0,
"runtime_peak_allocated_mb": 14216.0,
"load_peak_host_anon_mb": 10559.66,
"runtime_peak_host_anon_mb": 10559.66,
"estimated_full_test_time_s": 172.8
@@ -209,6 +219,8 @@
"expected_median_denoise_ms": 4863.02,
"load_peak_vram_mb": 5052.0,
"runtime_peak_vram_mb": 12278.0,
"load_peak_allocated_mb": 4995.0,
"runtime_peak_allocated_mb": 11483.0,
"load_peak_host_anon_mb": 32768.0,
"runtime_peak_host_anon_mb": 32768.0,
"estimated_full_test_time_s": 283.0
@@ -105,6 +105,8 @@
"expected_median_denoise_ms": 249.01,
"load_peak_vram_mb": 41840.0,
"runtime_peak_vram_mb": 46938.0,
"load_peak_allocated_mb": 41798.0,
"runtime_peak_allocated_mb": 44855.0,
"estimated_full_test_time_s": 133.1
},
"qwen_image_t2i_2_gpus": {
@@ -173,6 +175,8 @@
"expected_median_denoise_ms": 190.02,
"load_peak_vram_mb": 41840.0,
"runtime_peak_vram_mb": 46938.0,
"load_peak_allocated_mb": 41798.0,
"runtime_peak_allocated_mb": 44855.0,
"estimated_full_test_time_s": 133.2
},
"qwen_image_t2i_2_gpus_extra_high": {
@@ -312,6 +316,8 @@
"expected_median_denoise_ms": 130.22,
"load_peak_vram_mb": 23594.0,
"runtime_peak_vram_mb": 28540.0,
"load_peak_allocated_mb": 23572.0,
"runtime_peak_allocated_mb": 25667.0,
"estimated_full_test_time_s": 127.4
},
"flux_2_image_t2i": {
@@ -381,6 +387,8 @@
"expected_median_denoise_ms": 445.86,
"load_peak_vram_mb": 64344.0,
"runtime_peak_vram_mb": 70560.0,
"load_peak_allocated_mb": 63963.0,
"runtime_peak_allocated_mb": 65404.0,
"estimated_full_test_time_s": 145.2
},
"flux_2_klein_image_t2i": {
@@ -404,6 +412,8 @@
"expected_median_denoise_ms": 24.14,
"load_peak_vram_mb": 8504.0,
"runtime_peak_vram_mb": 13732.0,
"load_peak_allocated_mb": 8498.0,
"runtime_peak_allocated_mb": 10793.0,
"estimated_full_test_time_s": 120.5
},
"flux_2_klein_base_image_t2i": {
@@ -473,6 +483,8 @@
"expected_median_denoise_ms": 117.28,
"load_peak_vram_mb": 8504.0,
"runtime_peak_vram_mb": 13754.0,
"load_peak_allocated_mb": 8498.0,
"runtime_peak_allocated_mb": 10800.0,
"estimated_full_test_time_s": 124.4
},
"flux_2_ti2i": {
@@ -542,6 +554,8 @@
"expected_median_denoise_ms": 889.93,
"load_peak_vram_mb": 64344.0,
"runtime_peak_vram_mb": 70564.0,
"load_peak_allocated_mb": 63963.0,
"runtime_peak_allocated_mb": 65388.0,
"estimated_full_test_time_s": 168.9
},
"flux_2_ti2i_multi_image_cache_dit": {
@@ -611,6 +625,8 @@
"expected_median_denoise_ms": 53.69,
"load_peak_vram_mb": 64344.0,
"runtime_peak_vram_mb": 71904.0,
"load_peak_allocated_mb": 63963.0,
"runtime_peak_allocated_mb": 67003.0,
"estimated_full_test_time_s": 148.6
},
"flux_image_t2i_2_gpus": {
@@ -679,6 +695,8 @@
"expected_median_denoise_ms": 81.01,
"load_peak_vram_mb": 23284.0,
"runtime_peak_vram_mb": 28148.0,
"load_peak_allocated_mb": 23262.0,
"runtime_peak_allocated_mb": 25541.0,
"estimated_full_test_time_s": 125.5
},
"zimage_image_t2i": {
@@ -706,6 +724,8 @@
"expected_median_denoise_ms": 90.0,
"load_peak_vram_mb": 13120.0,
"runtime_peak_vram_mb": 17776.0,
"load_peak_allocated_mb": 12852.0,
"runtime_peak_allocated_mb": 15133.0,
"estimated_full_test_time_s": 116.3
},
"zimage_image_t2i_fp8": {
@@ -733,6 +753,8 @@
"expected_median_denoise_ms": 108.73,
"load_peak_vram_mb": 11344.0,
"runtime_peak_vram_mb": 15868.0,
"load_peak_allocated_mb": 10941.0,
"runtime_peak_allocated_mb": 13254.0,
"estimated_full_test_time_s": 123.7
},
"zimage_image_t2i_multi_lora": {
@@ -760,6 +782,8 @@
"expected_median_denoise_ms": 89.93,
"load_peak_vram_mb": 14196.0,
"runtime_peak_vram_mb": 17612.0,
"load_peak_allocated_mb": 13059.0,
"runtime_peak_allocated_mb": 15307.0,
"estimated_full_test_time_s": 162.1
},
"zimage_image_t2i_2_gpus": {
@@ -787,6 +811,8 @@
"expected_median_denoise_ms": 54.3,
"load_peak_vram_mb": 13120.0,
"runtime_peak_vram_mb": 17602.0,
"load_peak_allocated_mb": 12852.0,
"runtime_peak_allocated_mb": 15131.0,
"estimated_full_test_time_s": 75.0
},
"qwen_image_edit_ti2i": {
@@ -856,6 +882,8 @@
"expected_median_denoise_ms": 636.16,
"load_peak_vram_mb": 43244.0,
"runtime_peak_vram_mb": 48250.0,
"load_peak_allocated_mb": 41926.0,
"runtime_peak_allocated_mb": 44966.0,
"estimated_full_test_time_s": 153.6
},
"joyai_image_edit_ti2i": {
@@ -915,6 +943,8 @@
"expected_median_denoise_ms": 673.29,
"load_peak_vram_mb": 35324.0,
"runtime_peak_vram_mb": 45124.0,
"load_peak_allocated_mb": 33291.0,
"runtime_peak_allocated_mb": 40644.0,
"estimated_full_test_time_s": 117.6
},
"qwen_image_t2i_cache_dit_enabled": {
@@ -983,6 +1013,8 @@
"expected_median_denoise_ms": 53.48,
"load_peak_vram_mb": 41840.0,
"runtime_peak_vram_mb": 47034.0,
"load_peak_allocated_mb": 41798.0,
"runtime_peak_allocated_mb": 44951.0,
"estimated_full_test_time_s": 124.9
},
"wan2_1_t2v_1.3b_teacache_enabled": {
@@ -1051,6 +1083,8 @@
"expected_median_denoise_ms": 98.28,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 18078.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 9730.0,
"estimated_full_test_time_s": 126.0
},
"wan2_1_t2v_1.3b": {
@@ -1119,6 +1153,8 @@
"expected_median_denoise_ms": 144.57,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 18078.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 9673.0,
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_cfg_parallel": {
@@ -1187,6 +1223,8 @@
"expected_median_denoise_ms": 74.05,
"load_peak_vram_mb": 5422.0,
"runtime_peak_vram_mb": 11972.0,
"load_peak_allocated_mb": 5274.0,
"runtime_peak_allocated_mb": 9251.0,
"estimated_full_test_time_s": 127.6
},
"turbo_wan2_1_t2v_1.3b": {
@@ -1210,6 +1248,8 @@
"expected_median_denoise_ms": 72.78,
"load_peak_vram_mb": 26256.0,
"runtime_peak_vram_mb": 34340.0,
"load_peak_allocated_mb": 24985.0,
"runtime_peak_allocated_mb": 31308.0,
"estimated_full_test_time_s": 124.7
},
"ltx_2_two_stage_t2v": {
@@ -1279,6 +1319,8 @@
"expected_median_denoise_ms": 109.14,
"load_peak_vram_mb": 43920.0,
"runtime_peak_vram_mb": 57528.0,
"load_peak_allocated_mb": 43411.0,
"runtime_peak_allocated_mb": 54736.0,
"estimated_full_test_time_s": 345.4
},
"wan2_2_ti2v_5b": {
@@ -1347,6 +1389,8 @@
"expected_median_denoise_ms": 256.21,
"load_peak_vram_mb": 17610.0,
"runtime_peak_vram_mb": 39856.0,
"load_peak_allocated_mb": 11629.0,
"runtime_peak_allocated_mb": 26066.0,
"estimated_full_test_time_s": 141.7
},
"qwen_image_edit_2509_ti2i": {
@@ -1406,6 +1450,8 @@
"expected_median_denoise_ms": 945.14,
"load_peak_vram_mb": 43244.0,
"runtime_peak_vram_mb": 47874.0,
"load_peak_allocated_mb": 41926.0,
"runtime_peak_allocated_mb": 44960.0,
"estimated_full_test_time_s": 160.2
},
"qwen_image_layered_i2i": {
@@ -1472,6 +1518,8 @@
"expected_median_denoise_ms": 664.26,
"load_peak_vram_mb": 58204.0,
"runtime_peak_vram_mb": 61846.0,
"load_peak_allocated_mb": 55943.0,
"runtime_peak_allocated_mb": 60609.0,
"estimated_full_test_time_s": 161.5
},
"fastwan2_2_ti2v_5b": {
@@ -1493,6 +1541,8 @@
"expected_median_denoise_ms": 112.67,
"load_peak_vram_mb": 17610.0,
"runtime_peak_vram_mb": 39054.0,
"load_peak_allocated_mb": 11629.0,
"runtime_peak_allocated_mb": 26051.0,
"estimated_full_test_time_s": 125.2
},
"fast_hunyuan_video": {
@@ -1518,6 +1568,8 @@
"expected_median_denoise_ms": 1010.9,
"load_peak_vram_mb": 27320.0,
"runtime_peak_vram_mb": 35926.0,
"load_peak_allocated_mb": 26241.0,
"runtime_peak_allocated_mb": 31922.0,
"estimated_full_test_time_s": 77.0
},
"wan2_2_i2v_a14b_2gpu": {
@@ -1577,6 +1629,8 @@
"expected_median_denoise_ms": 1714.58,
"load_peak_vram_mb": 6440.0,
"runtime_peak_vram_mb": 26322.0,
"load_peak_allocated_mb": 6177.0,
"runtime_peak_allocated_mb": 19978.0,
"estimated_full_test_time_s": 264.1
},
"wan2_1_i2v_14b_480P_2gpu": {
@@ -1647,6 +1701,8 @@
"expected_median_denoise_ms": 656.68,
"load_peak_vram_mb": 35350.0,
"runtime_peak_vram_mb": 42674.0,
"load_peak_allocated_mb": 34033.0,
"runtime_peak_allocated_mb": 38569.0,
"estimated_full_test_time_s": 243.0
},
"wan2_1_i2v_14b_720P_2gpu": {
@@ -1717,6 +1773,8 @@
"expected_median_denoise_ms": 1933.7,
"load_peak_vram_mb": 35350.0,
"runtime_peak_vram_mb": 51230.0,
"load_peak_allocated_mb": 34033.0,
"runtime_peak_allocated_mb": 44644.0,
"estimated_full_test_time_s": 248.9
},
"wan2_2_t2v_a14b_2gpu": {
@@ -1775,6 +1833,8 @@
"expected_median_denoise_ms": 2020.71,
"load_peak_vram_mb": 6440.0,
"runtime_peak_vram_mb": 21524.0,
"load_peak_allocated_mb": 6175.0,
"runtime_peak_allocated_mb": 15472.0,
"estimated_full_test_time_s": 204.3
},
"wan2_1_t2v_14b_2gpu": {
@@ -1843,6 +1903,8 @@
"expected_median_denoise_ms": 467.69,
"load_peak_vram_mb": 31246.0,
"runtime_peak_vram_mb": 37788.0,
"load_peak_allocated_mb": 29869.0,
"runtime_peak_allocated_mb": 33843.0,
"estimated_full_test_time_s": 173.2
},
"wan2_2_t2v_a14b_lora_2gpu": {
@@ -1901,6 +1963,8 @@
"expected_median_denoise_ms": 1416.85,
"load_peak_vram_mb": 6784.0,
"runtime_peak_vram_mb": 21450.0,
"load_peak_allocated_mb": 6499.0,
"runtime_peak_allocated_mb": 14775.0,
"estimated_full_test_time_s": 181.8
},
"wan2_1_t2v_1_3b_lora_1gpu": {
@@ -1969,6 +2033,8 @@
"expected_median_denoise_ms": 147.68,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 18350.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 9999.0,
"estimated_full_test_time_s": 129.6
},
"wan2_1_i2v_14b_lora_2gpu": {
@@ -2039,6 +2105,8 @@
"expected_median_denoise_ms": 1922.21,
"load_peak_vram_mb": 36266.0,
"runtime_peak_vram_mb": 51132.0,
"load_peak_allocated_mb": 34407.0,
"runtime_peak_allocated_mb": 44992.0,
"estimated_full_test_time_s": 248.2
},
"flux_2_image_t2i_2_gpus": {
@@ -2108,6 +2176,8 @@
"expected_median_denoise_ms": 246.92,
"load_peak_vram_mb": 33110.0,
"runtime_peak_vram_mb": 38262.0,
"load_peak_allocated_mb": 32851.0,
"runtime_peak_allocated_mb": 34824.0,
"estimated_full_test_time_s": 135.3
},
"qwen_image_edit_2511_ti2i": {
@@ -2167,6 +2237,8 @@
"expected_median_denoise_ms": 555.79,
"load_peak_vram_mb": 43244.0,
"runtime_peak_vram_mb": 47874.0,
"load_peak_allocated_mb": 41926.0,
"runtime_peak_allocated_mb": 44950.0,
"estimated_full_test_time_s": 143.7
},
"fsdp-inference": {
@@ -2194,6 +2266,8 @@
"expected_median_denoise_ms": 81.75,
"load_peak_vram_mb": 7176.0,
"runtime_peak_vram_mb": 12634.0,
"load_peak_allocated_mb": 7033.0,
"runtime_peak_allocated_mb": 9311.0,
"estimated_full_test_time_s": 122.7
},
"hunyuan3d_shape_gen": {
@@ -2263,6 +2337,8 @@
"expected_median_denoise_ms": 30.2,
"load_peak_vram_mb": 876.0,
"runtime_peak_vram_mb": 6094.0,
"load_peak_allocated_mb": 737.0,
"runtime_peak_allocated_mb": 5120.0,
"estimated_full_test_time_s": 420.1
},
"wan2_1_t2v_1.3b_frame_interp_2x": {
@@ -2331,6 +2407,8 @@
"expected_median_denoise_ms": 147.3,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 18080.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 9673.0,
"estimated_full_test_time_s": 129.3
},
"flux_2_image_t2i_upscaling_4x": {
@@ -2400,6 +2478,8 @@
"expected_median_denoise_ms": 438.64,
"load_peak_vram_mb": 64344.0,
"runtime_peak_vram_mb": 70620.0,
"load_peak_allocated_mb": 63963.0,
"runtime_peak_allocated_mb": 65854.0,
"estimated_full_test_time_s": 145.1
},
"wan2_1_t2v_1.3b_upscaling_4x": {
@@ -2468,6 +2548,8 @@
"expected_median_denoise_ms": 144.74,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 24380.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 10508.0,
"estimated_full_test_time_s": 129.3
},
"wan2_1_t2v_1.3b_frame_interp_2x_upscaling_4x": {
@@ -2536,6 +2618,8 @@
"expected_median_denoise_ms": 144.67,
"load_peak_vram_mb": 7814.0,
"runtime_peak_vram_mb": 24388.0,
"load_peak_allocated_mb": 3641.0,
"runtime_peak_allocated_mb": 10531.0,
"estimated_full_test_time_s": 129.4
},
"ltx_2.3_one_stage_ti2v": {
@@ -2587,6 +2671,8 @@
"expected_median_denoise_ms": 494.9,
"load_peak_vram_mb": 46278.0,
"runtime_peak_vram_mb": 49390.0,
"load_peak_allocated_mb": 45460.0,
"runtime_peak_allocated_mb": 47725.0,
"estimated_full_test_time_s": 144.2
},
"ltx_2.3_two_stage_t2v_2gpus": {
@@ -2646,6 +2732,8 @@
"expected_median_denoise_ms": 145.66,
"load_peak_vram_mb": 47234.0,
"runtime_peak_vram_mb": 59544.0,
"load_peak_allocated_mb": 46411.0,
"runtime_peak_allocated_mb": 56951.0,
"estimated_full_test_time_s": 160.0
},
"ltx_2_3_two_stage_ti2v_2gpus": {
@@ -2705,6 +2793,8 @@
"expected_median_denoise_ms": 242.79,
"load_peak_vram_mb": 47234.0,
"runtime_peak_vram_mb": 60286.0,
"load_peak_allocated_mb": 46411.0,
"runtime_peak_allocated_mb": 56951.0,
"estimated_full_test_time_s": 170.0
},
"longlive2_t2v": {
@@ -2723,6 +2813,8 @@
"expected_median_denoise_ms": 56.74,
"load_peak_vram_mb": 34022.0,
"runtime_peak_vram_mb": 61190.0,
"load_peak_allocated_mb": 33866.0,
"runtime_peak_allocated_mb": 57075.0,
"estimated_full_test_time_s": 153.1
},
"longlive2_i2v": {
@@ -2741,6 +2833,8 @@
"expected_median_denoise_ms": 151.9,
"load_peak_vram_mb": 34022.0,
"runtime_peak_vram_mb": 62510.0,
"load_peak_allocated_mb": 33866.0,
"runtime_peak_allocated_mb": 57274.0,
"estimated_full_test_time_s": 149.4
},
"lingbot_video_moe_t2v": {
@@ -2759,6 +2853,8 @@
"expected_median_denoise_ms": 0.0,
"load_peak_vram_mb": 42814.0,
"runtime_peak_vram_mb": 73256.0,
"load_peak_allocated_mb": 41599.0,
"runtime_peak_allocated_mb": 71172.0,
"estimated_full_test_time_s": 126.0
},
"ltx_2_3_hq_pipeline": {
@@ -2805,6 +2901,8 @@
"expected_median_denoise_ms": 747.62,
"load_peak_vram_mb": 56048.0,
"runtime_peak_vram_mb": 59804.0,
"load_peak_allocated_mb": 55239.0,
"runtime_peak_allocated_mb": 57444.0,
"estimated_full_test_time_s": 363.2
},
"qwen_image_t2i_cache_dit_scm_config_diffusers_1gpu": {
@@ -2946,6 +3044,8 @@
"expected_median_denoise_ms": 2332.99,
"load_peak_vram_mb": 45428.0,
"runtime_peak_vram_mb": 63312.0,
"load_peak_allocated_mb": 15139.0,
"runtime_peak_allocated_mb": 22887.0,
"estimated_full_test_time_s": 208.0
},
"mova_360p_tp2": {
@@ -465,6 +465,10 @@ class DiffusionServerBase:
summary,
expected_load_peak_vram_mb,
expected_runtime_peak_vram_mb,
expected_load_peak_allocated_mb=scenario.load_peak_allocated_mb,
expected_runtime_peak_allocated_mb=(
scenario.runtime_peak_allocated_mb
),
)
validator.validate_peak_host_anon(
summary,
@@ -555,6 +559,8 @@ class DiffusionServerBase:
summary,
scenario.load_peak_vram_mb,
scenario.runtime_peak_vram_mb,
expected_load_peak_allocated_mb=scenario.load_peak_allocated_mb,
expected_runtime_peak_allocated_mb=scenario.runtime_peak_allocated_mb,
)
except AssertionError as e:
logger.error(f"Peak VRAM validation failed for {case.id}:\n{e}")
@@ -573,6 +579,8 @@ class DiffusionServerBase:
"median_denoise_ms": summary.median_denoise_ms,
"load_peak_vram_mb": summary.load_peak_vram_mb,
"runtime_peak_vram_mb": summary.runtime_peak_vram_mb,
"load_peak_allocated_mb": summary.load_peak_allocated_mb,
"runtime_peak_allocated_mb": summary.runtime_peak_allocated_mb,
"stage_metrics": summary.stage_metrics,
"sampled_steps": summary.sampled_steps,
}
@@ -606,7 +614,9 @@ class DiffusionServerBase:
f"avg_denoise={summary.avg_denoise_ms:.2f}ms, "
f"median_denoise={summary.median_denoise_ms:.2f}ms, "
f"load_peak_vram={summary.load_peak_vram_mb:.0f}MiB, "
f"runtime_peak_vram={summary.runtime_peak_vram_mb:.0f}MiB"
f"runtime_peak_vram={summary.runtime_peak_vram_mb:.0f}MiB, "
f"load_peak_alloc={summary.load_peak_allocated_mb:.0f}MiB, "
f"runtime_peak_alloc={summary.runtime_peak_allocated_mb:.0f}MiB"
),
]
if scenario is not None:
@@ -670,6 +680,10 @@ class DiffusionServerBase:
{
"load_peak_vram_mb": round(summary.load_peak_vram_mb, 2),
"runtime_peak_vram_mb": round(summary.runtime_peak_vram_mb, 2),
"load_peak_allocated_mb": round(summary.load_peak_allocated_mb, 2),
"runtime_peak_allocated_mb": round(
summary.runtime_peak_allocated_mb, 2
),
"load_peak_host_anon_mb": round(summary.load_peak_host_anon_mb, 2),
"runtime_peak_host_anon_mb": round(
summary.runtime_peak_host_anon_mb, 2
@@ -590,22 +590,66 @@ class PerformanceValidator:
summary: PerformanceSummary,
expected_load_peak_vram_mb: float,
expected_runtime_peak_vram_mb: float,
expected_load_peak_allocated_mb: float | None = None,
expected_runtime_peak_allocated_mb: float | None = None,
) -> None:
assert summary.load_peak_vram_mb > 0, "Load peak VRAM metric missing"
assert summary.runtime_peak_vram_mb > 0, "Runtime peak VRAM metric missing"
self._assert_le(
self._assert_peak_vram(
"Load Peak VRAM",
summary.load_peak_vram_mb,
expected_load_peak_vram_mb,
self.tolerances.load_peak_vram,
min_abs_tolerance=128.0,
unit=" MiB",
reserved=summary.load_peak_vram_mb,
allocated=summary.load_peak_allocated_mb,
expected_reserved=expected_load_peak_vram_mb,
expected_allocated=expected_load_peak_allocated_mb,
tolerance=self.tolerances.load_peak_vram,
)
self._assert_le(
self._assert_peak_vram(
"Runtime Peak VRAM",
summary.runtime_peak_vram_mb,
expected_runtime_peak_vram_mb,
self.tolerances.runtime_peak_vram,
reserved=summary.runtime_peak_vram_mb,
allocated=summary.runtime_peak_allocated_mb,
expected_reserved=expected_runtime_peak_vram_mb,
expected_allocated=expected_runtime_peak_allocated_mb,
tolerance=self.tolerances.runtime_peak_vram,
)
def _assert_peak_vram(
self,
name: str,
*,
reserved: float,
allocated: float,
expected_reserved: float,
expected_allocated: float | None,
tolerance: float,
) -> None:
"""Enforce the allocated peak when the baseline has one, else reserved.
Reserved peaks include the caching allocator's pool, which follows the
allocation history of everything run before the request (warmup shapes,
load-time leftovers) and moves a few percent for identical work. The
allocated peak is what the model and its activations actually use.
"""
if expected_allocated is not None and allocated > 0:
self._assert_le(
f"{name} (allocated)",
allocated,
expected_allocated,
tolerance,
min_abs_tolerance=128.0,
unit=" MiB",
)
logger.info(
"%s reserved %.0f MiB (baseline %.0f MiB, reported only)",
name,
reserved,
expected_reserved,
)
return
self._assert_le(
name,
reserved,
expected_reserved,
tolerance,
min_abs_tolerance=128.0,
unit=" MiB",
)
@@ -117,6 +117,10 @@ class ScenarioConfig:
estimated_full_test_time_s: float | None = None
load_peak_vram_mb: float | None = None
runtime_peak_vram_mb: float | None = None
# Allocated peaks; when present they are the enforced VRAM figure and the
# reserved peaks above are reported only (reserved tracks pool history).
load_peak_allocated_mb: float | None = None
runtime_peak_allocated_mb: float | None = None
# Anonymous-host budget caps; None skips the check (older baselines).
load_peak_host_anon_mb: float | None = None
runtime_peak_host_anon_mb: float | None = None
@@ -136,6 +140,8 @@ class ScenarioConfig:
estimated_full_test_time_s=optional_float("estimated_full_test_time_s"),
load_peak_vram_mb=optional_float("load_peak_vram_mb"),
runtime_peak_vram_mb=optional_float("runtime_peak_vram_mb"),
load_peak_allocated_mb=optional_float("load_peak_allocated_mb"),
runtime_peak_allocated_mb=optional_float("runtime_peak_allocated_mb"),
load_peak_host_anon_mb=optional_float("load_peak_host_anon_mb"),
runtime_peak_host_anon_mb=optional_float("runtime_peak_host_anon_mb"),
)
@@ -445,6 +451,8 @@ class PerformanceSummary:
all_denoise_steps: dict[int, float]
load_peak_vram_mb: float = 0.0
runtime_peak_vram_mb: float = 0.0
load_peak_allocated_mb: float = 0.0
runtime_peak_allocated_mb: float = 0.0
load_peak_host_anon_mb: float = 0.0
runtime_peak_host_anon_mb: float = 0.0
frames_per_second: float | None = None
@@ -482,6 +490,14 @@ class PerformanceSummary:
runtime_peak_vram_mb = float(
record.memory_snapshots.get("runtime_peak", {}).get("peak_reserved_mb", 0.0)
)
load_peak_allocated_mb = float(
record.memory_snapshots.get("load_peak", {}).get("peak_allocated_mb", 0.0)
)
runtime_peak_allocated_mb = float(
record.memory_snapshots.get("runtime_peak", {}).get(
"peak_allocated_mb", 0.0
)
)
load_peak_host_anon_mb = float(
record.memory_snapshots.get("load_peak", {}).get("peak_host_anon_mb", 0.0)
)
@@ -501,6 +517,8 @@ class PerformanceSummary:
all_denoise_steps=per_step,
load_peak_vram_mb=load_peak_vram_mb,
runtime_peak_vram_mb=runtime_peak_vram_mb,
load_peak_allocated_mb=load_peak_allocated_mb,
runtime_peak_allocated_mb=runtime_peak_allocated_mb,
load_peak_host_anon_mb=load_peak_host_anon_mb,
runtime_peak_host_anon_mb=runtime_peak_host_anon_mb,
)
@@ -40,8 +40,11 @@ def test_performance_summary_separates_load_and_runtime_peaks():
_perf_record(
{
"after_forward": {"peak_reserved_mb": 1024.0},
"load_peak": {"peak_reserved_mb": 4096.0},
"runtime_peak": {"peak_reserved_mb": 3072.0},
"load_peak": {"peak_reserved_mb": 4096.0, "peak_allocated_mb": 3900.0},
"runtime_peak": {
"peak_reserved_mb": 3072.0,
"peak_allocated_mb": 2900.0,
},
}
),
step_fractions=(),
@@ -49,6 +52,8 @@ def test_performance_summary_separates_load_and_runtime_peaks():
assert summary.load_peak_vram_mb == 4096.0
assert summary.runtime_peak_vram_mb == 3072.0
assert summary.load_peak_allocated_mb == 3900.0
assert summary.runtime_peak_allocated_mb == 2900.0
def test_worker_records_replica_load_and_runtime_peaks():
@@ -57,11 +62,13 @@ def test_worker_records_replica_load_and_runtime_peaks():
worker.is_output_rank = True
worker._load_peak_reserved_mb = 4096.0
worker._runtime_peak_reserved_mb = 0.0
worker._load_peak_allocated_mb = 3000.0
worker._runtime_peak_allocated_mb = 0.0
output = OutputBatch()
metrics = RequestMetrics("request")
replica_group = Mock()
replica_group.all_reduce.return_value = torch.tensor(
[5120.0, 3584.0], dtype=torch.float64
[5120.0, 3584.0, 3500.0, 2560.0], dtype=torch.float64
)
snapshots = [
MemorySnapshot(0.0, 0.0, 2048.0, 3072.0),
@@ -85,8 +92,11 @@ def test_worker_records_replica_load_and_runtime_peaks():
assert output.peak_memory_mb == 3072.0
assert worker._load_peak_reserved_mb == 4096.0
assert worker._runtime_peak_reserved_mb == 3072.0
assert worker._runtime_peak_allocated_mb == 2048.0
assert metrics.memory_snapshots["load_peak"].peak_reserved_mb == 5120.0
assert metrics.memory_snapshots["runtime_peak"].peak_reserved_mb == 3584.0
assert metrics.memory_snapshots["load_peak"].peak_allocated_mb == 3500.0
assert metrics.memory_snapshots["runtime_peak"].peak_allocated_mb == 2560.0
def test_baseline_config_loads_per_scenario_peak_vram(tmp_path):
@@ -115,6 +125,7 @@ def test_baseline_config_loads_per_scenario_peak_vram(tmp_path):
"expected_median_denoise_ms": 1.0,
"load_peak_vram_mb": 1234.5,
"runtime_peak_vram_mb": 2345.6,
"runtime_peak_allocated_mb": 2000.5,
}
},
}
@@ -127,6 +138,8 @@ def test_baseline_config_loads_per_scenario_peak_vram(tmp_path):
scenario = config.scenarios["case"]
assert scenario.load_peak_vram_mb == 1234.5
assert scenario.runtime_peak_vram_mb == 2345.6
assert scenario.load_peak_allocated_mb is None
assert scenario.runtime_peak_allocated_mb == 2000.5
assert config.tolerances.load_peak_vram == 0.01
assert config.tolerances.runtime_peak_vram == 0.02
@@ -175,6 +188,68 @@ def test_peak_vram_validation_uses_independent_tolerances():
validator.validate_peak_vram(runtime_regression, 10_000.0, 10_000.0)
def test_peak_vram_validation_enforces_allocated_when_baselined():
validator = PerformanceValidator(
scenario=ScenarioConfig({}, {}, 0.0, 0.0, 0.0),
tolerances=ToleranceConfig(
0.0,
0.0,
0.0,
0.0,
0.0,
load_peak_vram=0.01,
runtime_peak_vram=0.02,
),
step_fractions=(),
)
# reserved drifts with allocator pool history: +5% is not a regression
# once the allocated peak is baselined and unchanged
reserved_drift = PerformanceSummary(
0.0,
0.0,
0.0,
{},
[],
{},
{},
load_peak_vram_mb=10_000.0,
runtime_peak_vram_mb=10_500.0,
load_peak_allocated_mb=8_000.0,
runtime_peak_allocated_mb=8_000.0,
)
allocated_regression = PerformanceSummary(
0.0,
0.0,
0.0,
{},
[],
{},
{},
load_peak_vram_mb=10_000.0,
runtime_peak_vram_mb=10_000.0,
load_peak_allocated_mb=8_000.0,
runtime_peak_allocated_mb=8_300.0,
)
with patch.object(current_platform, "is_hip", return_value=False):
validator.validate_peak_vram(
reserved_drift,
10_000.0,
10_000.0,
expected_runtime_peak_allocated_mb=8_000.0,
)
with pytest.raises(AssertionError, match=r"Runtime Peak VRAM \(allocated\)"):
validator.validate_peak_vram(
allocated_regression,
10_000.0,
10_000.0,
expected_runtime_peak_allocated_mb=8_000.0,
)
# without an allocated baseline the reserved figure is still enforced
with pytest.raises(AssertionError, match="Runtime Peak VRAM"):
validator.validate_peak_vram(reserved_drift, 10_000.0, 10_000.0)
@pytest.mark.parametrize(
("load_peak_vram_mb", "runtime_peak_vram_mb", "message"),
[