From af9784c03f9163e41f656af710e8021f781034b2 Mon Sep 17 00:00:00 2001 From: Didan Deng <33117903+wtomin@users.noreply.github.com> Date: Wed, 11 Mar 2026 20:30:54 +0800 Subject: [PATCH] [Bugfix] Adjust Z-Image Tensor Parallelism Diff Threshold (#1808) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../test_zimage_parallelism.py | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index 47de955aa..4ee774bfd 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -54,12 +54,13 @@ def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor: def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: - """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1].""" + """Return (mean_abs_diff, p99_abs_diff) over RGB pixels in [0, 1].""" ta = _pil_to_float_rgb_tensor(a) tb = _pil_to_float_rgb_tensor(b) assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" abs_diff = torch.abs(ta - tb) - return abs_diff.mean().item(), abs_diff.max().item() + p99_abs_diff = torch.quantile(abs_diff.flatten(), 0.99).item() + return abs_diff.mean().item(), p99_abs_diff def _extract_single_image(outputs) -> Image.Image: @@ -196,18 +197,18 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): assert tp1_img.width == width and tp1_img.height == height assert tp2_img.width == width and tp2_img.height == height - mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img) + mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) mean_threshold = 3e-2 - max_threshold = 5e-1 + p99_threshold = 2.5e-1 print( "Z-Image TP image diff stats (TP=1 vs TP=2): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " f"tp1_img={tp1_path}, tp2_img={tp2_path}" ) - assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" ) print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") @@ -262,16 +263,16 @@ def test_zimage_vae_patch_parallel_tp2(tmp_path: Path): baseline_img.save(baseline_path) pp2_img.save(pp2_path) - mean_abs_diff, max_abs_diff = _diff_metrics(baseline_img, pp2_img) + mean_abs_diff, p99_abs_diff = _diff_metrics(baseline_img, pp2_img) mean_threshold = 5e-3 - max_threshold = 1e-1 + p99_threshold = 1e-1 print( "Z-Image VAE patch parallel image diff stats (TP=2, pp=1 vs pp=2): " - f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " - f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e}; " f"pp1_img={baseline_path}, pp2_img={pp2_path}" ) - assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( - f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " - f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" + assert mean_abs_diff <= mean_threshold and p99_abs_diff <= p99_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, p99_abs_diff={p99_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, p99<={p99_threshold:.6e})" )