[Bugfix] Adjust Z-Image Tensor Parallelism Diff Threshold (#1808)

Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

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