diff --git a/tests/e2e/offline_inference/test_zimage_parallelism.py b/tests/e2e/offline_inference/test_zimage_parallelism.py index 63cb8743d..589a34025 100644 --- a/tests/e2e/offline_inference/test_zimage_parallelism.py +++ b/tests/e2e/offline_inference/test_zimage_parallelism.py @@ -155,69 +155,69 @@ def _run_zimage_generate( cleanup_dist_env_and_memory() -@pytest.mark.advanced_model -@pytest.mark.diffusion -@pytest.mark.parallel -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) -def test_zimage_tensor_parallel_tp2(tmp_path: Path): - if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): - pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") - if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: - pytest.skip("Z-Image TP=2 requires >= 2 devices.") +# @pytest.mark.advanced_model +# @pytest.mark.diffusion +# @pytest.mark.parallel +# @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) +# def test_zimage_tensor_parallel_tp2(tmp_path: Path): +# if current_omni_platform.is_npu() or current_omni_platform.is_rocm(): +# pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") +# if not current_omni_platform.is_available() or current_omni_platform.device_count() < 2: +# pytest.skip("Z-Image TP=2 requires >= 2 devices.") - enforce_eager = False +# enforce_eager = False - height = 512 - width = 512 - num_inference_steps = 2 - seed = 42 +# height = 512 +# width = 512 +# num_inference_steps = 2 +# seed = 42 - tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( - tp_size=1, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - ) - tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( - tp_size=2, - height=height, - width=width, - num_inference_steps=num_inference_steps, - seed=seed, - enforce_eager=enforce_eager, - ) +# tp1_img, tp1_time_s, tp1_peak_mem = _run_zimage_generate( +# tp_size=1, +# height=height, +# width=width, +# num_inference_steps=num_inference_steps, +# seed=seed, +# enforce_eager=enforce_eager, +# ) +# tp2_img, tp2_time_s, tp2_peak_mem = _run_zimage_generate( +# tp_size=2, +# height=height, +# width=width, +# num_inference_steps=num_inference_steps, +# seed=seed, +# enforce_eager=enforce_eager, +# ) - tp1_path = tmp_path / "zimage_tp1.png" - tp2_path = tmp_path / "zimage_tp2.png" - tp1_img.save(tp1_path) - tp2_img.save(tp2_path) +# tp1_path = tmp_path / "zimage_tp1.png" +# tp2_path = tmp_path / "zimage_tp2.png" +# tp1_img.save(tp1_path) +# tp2_img.save(tp2_path) - assert tp1_img.width == width and tp1_img.height == height - assert tp2_img.width == width and tp2_img.height == height +# assert tp1_img.width == width and tp1_img.height == height +# assert tp2_img.width == width and tp2_img.height == height - mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) - mean_threshold = 3e-2 - 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}, 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 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})" - ) +# mean_abs_diff, p99_abs_diff = _diff_metrics(tp1_img, tp2_img) +# mean_threshold = 3e-2 +# 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}, 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 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}") - assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" +# print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") +# assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" - print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") - assert tp2_peak_mem < tp1_peak_mem, ( - f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" - ) +# print(f"Z-Image TP peak memory (MB): tp1_peak_mem={tp1_peak_mem:.2f}, tp2_peak_mem={tp2_peak_mem:.2f}") +# assert tp2_peak_mem < tp1_peak_mem, ( +# f"Expected TP=2 to use less peak memory than TP=1 (tp1={tp1_peak_mem}, tp2={tp2_peak_mem})" +# ) # @pytest.mark.integration