[Feat] Enable DiT tensor parallel for Diffusion Pipeline(Z-Image) (#735)

Signed-off-by: dongbo910220 <1275604947@qq.com>
Signed-off-by: dongbo910220 <32610838+dongbo910220@users.noreply.github.com>
Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>

.buildkite/pipeline.yml

@@ -77,10 +77,11 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Parallelism Test"
-    timeout_in_minutes: 15
+    timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:

docs/user_guide/acceleration/parallelism_acceleration.md

@@ -12,27 +12,56 @@ The following parallelism methods are currently supported in vLLM-Omni:
 
 3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step.
 
+4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT.
+
 The following table shows which models are currently supported by parallelism method:
 
 ### ImageGen
 
-| Model | Model Identifier | Ulysses-SP |  Ring-SP |CFG-Parallel |
-|-------|------------------|-----------|-------------|-------------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------|--------------------------|
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
+| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+
+!!! note "Why Z-Image is TP=2 only"
+    Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
+    For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**.
 
 ### VideoGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP |
-|-------|------------------|-----------|---------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ |
+| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
+|-------|------------------|------------|---------|--------------------------|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+
+### Tensor Parallelism
+
+Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`.
+
+#### Offline Inference
+
+```python
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+omni = Omni(
+    model="Tongyi-MAI/Z-Image-Turbo",
+    parallel_config=DiffusionParallelConfig(tensor_parallel_size=2),
+)
+
+outputs = omni.generate(
+    prompt="a cat reading a book",
+    num_inference_steps=9,
+    width=512,
+    height=512,
+)
+```
 
 ### Sequence Parallelism
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -88,6 +88,17 @@ def parse_args() -> argparse.Namespace:
         choices=[1, 2],
         help="Number of GPUs used for classifier free guidance parallel size.",
     )
+    parser.add_argument(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
+    parser.add_argument(
+        "--tensor_parallel_size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     return parser.parse_args()
 
 
@@ -131,7 +142,10 @@ def main():
 
     # assert args.ring_degree == 1, "Ring attention is not supported yet"
     parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
     )
 
     omni = Omni(
@@ -141,6 +155,7 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
 
     # Time profiling for generation
@@ -150,7 +165,8 @@ def main():
     print(f"  Inference steps: {args.num_inference_steps}")
     print(f"  Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+        f"  Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
+        f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
     )
     print(f"  Image size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")

tests/diffusion/models/z_image/test_zimage_tp_constraints.py

@@ -0,0 +1,44 @@
+import pytest
+
+from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints
+
+
+def test_validate_zimage_tp_constraints_tp2_ok():
+    ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints(
+        dim=3840,
+        n_heads=30,
+        n_kv_heads=30,
+        in_channels=16,
+        all_patch_size=(2,),
+        all_f_patch_size=(1,),
+        tensor_parallel_size=2,
+    )
+    assert ffn_hidden_dim == 10240
+    assert final_out_dims == [64]
+    assert supported_tp == [1, 2]
+
+
+def test_validate_zimage_tp_constraints_tp4_fails_on_heads():
+    with pytest.raises(ValueError, match=r"n_heads % tensor_parallel_size"):
+        validate_zimage_tp_constraints(
+            dim=3840,
+            n_heads=30,
+            n_kv_heads=30,
+            in_channels=16,
+            all_patch_size=(2,),
+            all_f_patch_size=(1,),
+            tensor_parallel_size=4,
+        )
+
+
+def test_validate_zimage_tp_constraints_tp3_fails_on_ffn_hidden_dim():
+    with pytest.raises(ValueError, match=r"ffn_hidden_dim % tensor_parallel_size"):
+        validate_zimage_tp_constraints(
+            dim=3840,
+            n_heads=30,
+            n_kv_heads=30,
+            in_channels=16,
+            all_patch_size=(2,),
+            all_f_patch_size=(1,),
+            tensor_parallel_size=3,
+        )

tests/e2e/offline_inference/test_zimage_tensor_parallel.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
+
+# ruff: noqa: E402
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from vllm_omni import Omni
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.outputs import OmniRequestOutput
+from vllm_omni.utils.platform_utils import is_npu, is_rocm
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+PROMPT = "a photo of a cat sitting on a laptop keyboard"
+
+
+def _get_zimage_model() -> str:
+    # Allow overriding the model for local/offline environments.
+    # Can be either a HuggingFace repo id or a local path.
+    return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo")
+
+
+def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor:
+    """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3]."""
+    arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0
+    return torch.from_numpy(arr)
+
+
+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]."""
+    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()
+
+
+def _extract_single_image(outputs) -> Image.Image:
+    first_output = outputs[0]
+    assert first_output.final_output_type == "image"
+    if not hasattr(first_output, "request_output") or not first_output.request_output:
+        raise ValueError("No request_output found in OmniRequestOutput")
+
+    req_out = first_output.request_output[0]
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
+
+    images = req_out.images
+    if images is None or len(images) != 1:
+        raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}")
+    return images[0]
+
+
+def _run_zimage_generate(
+    *, tp_size: int, height: int, width: int, num_inference_steps: int, seed: int
+) -> tuple[Image.Image, float]:
+    m = Omni(
+        model=_get_zimage_model(),
+        parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size),
+    )
+    try:
+        # NOTE: Omni closes itself when a generate() call is exhausted.
+        # To avoid measuring teardown time (process shutdown, memory cleanup),
+        # we measure the latency to produce *subsequent* outputs within a single
+        # generator run.
+        #
+        # This also serves as a warmup: the first output may include extra
+        # compilation/caching overhead, while later outputs are closer to
+        # steady-state inference.
+        num_requests = 4  # 1 warmup + 3 timed
+        gen = m.generate(
+            [PROMPT] * num_requests,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,
+            seed=seed,
+            num_outputs_per_prompt=1,
+            py_generator=True,
+        )
+
+        warmup_output = next(gen)
+        t_prev = time.perf_counter()
+        per_request_times_s: list[float] = []
+        last_output = warmup_output
+        for _ in range(num_requests - 1):
+            last_output = next(gen)
+            t_now = time.perf_counter()
+            per_request_times_s.append(t_now - t_prev)
+            t_prev = t_now
+
+        # Ensure the generator is fully consumed so it can clean up.
+        for _ in gen:
+            pass
+
+        median_time_s = float(np.median(per_request_times_s))
+        return _extract_single_image([last_output]), median_time_s
+    finally:
+        m.close()
+        cleanup_dist_env_and_memory()
+
+
+@pytest.mark.integration
+def test_zimage_tensor_parallel_tp2(tmp_path: Path):
+    if is_npu() or is_rocm():
+        pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.")
+    if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
+        pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.")
+
+    height = 256
+    width = 256
+    num_inference_steps = 2
+    seed = 42
+
+    tp1_img, tp1_time_s = _run_zimage_generate(
+        tp_size=1,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+    )
+    tp2_img, tp2_time_s = _run_zimage_generate(
+        tp_size=2,
+        height=height,
+        width=width,
+        num_inference_steps=num_inference_steps,
+        seed=seed,
+    )
+
+    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
+
+    mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img)
+    mean_threshold = 3e-2
+    max_threshold = 3.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"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})"
+    )
+
+    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})"

vllm_omni/diffusion/models/z_image/z_image_transformer.py

@@ -21,6 +21,7 @@ from collections.abc import Iterable
 import torch
 import torch.nn as nn
 from torch.nn.utils.rnn import pad_sequence
+from vllm.logger import init_logger
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -31,11 +32,107 @@ from vllm.model_executor.layers.linear import (
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 
 ADALN_EMBED_DIM = 256
 SEQ_MULTI_OF = 32
 
+logger = init_logger(__name__)
+
+
+def _positive_divisors(n: int) -> set[int]:
+    if n <= 0:
+        return set()
+    divs: set[int] = set()
+    for d in range(1, int(math.isqrt(n)) + 1):
+        if n % d == 0:
+            divs.add(d)
+            divs.add(n // d)
+    return divs
+
+
+def _get_tensor_parallel_size_from_context() -> int:
+    if not is_forward_context_available():
+        return 1
+    try:
+        od_config = get_forward_context().omni_diffusion_config
+        if od_config is None:
+            return 1
+        return int(od_config.parallel_config.tensor_parallel_size)
+    except Exception:
+        return 1
+
+
+def validate_zimage_tp_constraints(
+    *,
+    dim: int,
+    n_heads: int,
+    n_kv_heads: int,
+    in_channels: int,
+    all_patch_size: tuple[int, ...],
+    all_f_patch_size: tuple[int, ...],
+    tensor_parallel_size: int,
+) -> tuple[int, list[int], list[int]]:
+    """Validate Z-Image TP constraints without requiring a distributed context.
+
+    Returns:
+        (ffn_hidden_dim, final_out_dims, supported_tp_candidates)
+    """
+    tp_size = int(tensor_parallel_size)
+    if tp_size <= 0:
+        raise ValueError(f"tensor_parallel_size must be > 0, got {tp_size}")
+    if dim % n_heads != 0:
+        raise ValueError(f"dim must be divisible by n_heads, got dim={dim}, n_heads={n_heads}")
+    if dim % tp_size != 0:
+        supported = sorted(_positive_divisors(dim))
+        raise ValueError(
+            f"Z-Image requires dim % tensor_parallel_size == 0, but got dim={dim}, tp={tp_size}. "
+            f"Supported tp candidates by dim: {supported}"
+        )
+    if n_heads % tp_size != 0:
+        supported = sorted(_positive_divisors(n_heads))
+        raise ValueError(
+            f"Z-Image requires n_heads % tensor_parallel_size == 0, but got n_heads={n_heads}, tp={tp_size}. "
+            f"Supported tp candidates by n_heads: {supported}"
+        )
+    if n_kv_heads % tp_size != 0:
+        supported = sorted(_positive_divisors(n_kv_heads))
+        raise ValueError(
+            f"Z-Image requires n_kv_heads % tensor_parallel_size == 0, but got n_kv_heads={n_kv_heads}, "
+            f"tp={tp_size}. Supported tp candidates by n_kv_heads: {supported}"
+        )
+
+    ffn_hidden_dim = int(dim / 3 * 8)
+    if ffn_hidden_dim % tp_size != 0:
+        supported = sorted(_positive_divisors(ffn_hidden_dim))
+        raise ValueError(
+            "Z-Image requires ffn_hidden_dim % tensor_parallel_size == 0 (for TP-sharded MLP), but got "
+            f"ffn_hidden_dim={ffn_hidden_dim}, tp={tp_size}. Supported tp candidates by ffn_hidden_dim: {supported}"
+        )
+
+    final_out_dims = [
+        int(patch_size) * int(patch_size) * int(f_patch_size) * int(in_channels)
+        for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size)
+    ]
+    bad_final_out_dims = [d for d in final_out_dims if d % tp_size != 0]
+    if bad_final_out_dims:
+        supported = sorted(_positive_divisors(math.gcd(*final_out_dims)))
+        raise ValueError(
+            "Z-Image requires final projection out_features divisible by tensor_parallel_size, but got "
+            f"final_out_dims={final_out_dims}, tp={tp_size}. "
+            f"Supported tp candidates by final_out_dims gcd: {supported}"
+        )
+
+    supported_tp_candidates = sorted(
+        _positive_divisors(n_heads)
+        & _positive_divisors(n_kv_heads)
+        & _positive_divisors(dim)
+        & _positive_divisors(ffn_hidden_dim)
+        & _positive_divisors(math.gcd(*final_out_dims))
+    )
+    return ffn_hidden_dim, final_out_dims, supported_tp_candidates
+
 
 class TimestepEmbedder(nn.Module):
     def __init__(self, out_size, mid_size=None, frequency_embedding_size=256):
@@ -90,8 +187,8 @@ class ZImageAttention(nn.Module):
     ) -> None:
         super().__init__()
         self.dim = dim
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
+        self.total_num_heads = num_heads
+        self.total_num_kv_heads = num_kv_heads
         self.head_dim = dim // num_heads
         self.qk_norm = qk_norm
 
@@ -99,7 +196,7 @@ class ZImageAttention(nn.Module):
             hidden_size=dim,
             head_size=self.head_dim,
             total_num_heads=num_heads,
-            disable_tp=True,
+            total_num_kv_heads=num_kv_heads,
             bias=False,
         )
 
@@ -107,13 +204,27 @@ class ZImageAttention(nn.Module):
         self.norm_q = RMSNorm(self.head_dim, eps=eps)
         self.norm_k = RMSNorm(self.head_dim, eps=eps)
 
-        self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)])
+        # NOTE: QKV is column-parallel on heads, so attention output is sharded
+        # on the last dim (dim / tp). Use row-parallel output projection to
+        # all-reduce back to full dim.
+        self.to_out = nn.ModuleList(
+            [
+                RowParallelLinear(
+                    dim,
+                    dim,
+                    bias=False,
+                    input_is_parallel=True,
+                    return_bias=False,
+                )
+            ]
+        )
 
         self.attn = Attention(
-            num_heads=num_heads,
+            num_heads=self.to_qkv.num_heads,
             head_size=self.head_dim,
             softmax_scale=1.0 / (self.head_dim**0.5),
             causal=False,
+            num_kv_heads=self.to_qkv.num_kv_heads,
         )
         self.rope = RotaryEmbedding(is_neox_style=False)
 
@@ -125,11 +236,13 @@ class ZImageAttention(nn.Module):
         sin: torch.Tensor,
     ):
         qkv, _ = self.to_qkv(hidden_states)
-        query, key, value = qkv.chunk(3, dim=-1)
+        q_size = self.to_qkv.num_heads * self.head_dim
+        kv_size = self.to_qkv.num_kv_heads * self.head_dim
+        query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1)
 
-        query = query.unflatten(-1, (self.num_heads, -1))
-        key = key.unflatten(-1, (self.num_heads, -1))
-        value = value.unflatten(-1, (self.num_heads, -1))
+        query = query.unflatten(-1, (self.to_qkv.num_heads, -1))
+        key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
+        value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1))
 
         query = self.norm_q(query)
         key = self.norm_k(key)
@@ -170,7 +283,6 @@ class FeedForward(nn.Module):
             dim,
             [hidden_dim] * 2,
             bias=False,
-            disable_tp=True,
             return_bias=False,
         )
         self.act = SiluAndMul()
@@ -178,7 +290,7 @@ class FeedForward(nn.Module):
             hidden_dim,
             dim,
             bias=False,
-            disable_tp=True,
+            input_is_parallel=True,
             return_bias=False,
         )
 
@@ -379,6 +491,28 @@ class ZImageTransformer2DModel(nn.Module):
 
         assert len(all_patch_size) == len(all_f_patch_size)
 
+        tp_size = _get_tensor_parallel_size_from_context()
+        ffn_hidden_dim, final_out_dims, supported_tp_candidates = validate_zimage_tp_constraints(
+            dim=dim,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            in_channels=self.out_channels,
+            all_patch_size=tuple(all_patch_size),
+            all_f_patch_size=tuple(all_f_patch_size),
+            tensor_parallel_size=tp_size,
+        )
+
+        logger.info_once(
+            "Z-Image init: dim=%d n_heads=%d n_kv_heads=%d ffn_hidden_dim=%d final_out_dims=%s tp=%d (supported_tp=%s)",
+            dim,
+            n_heads,
+            n_kv_heads,
+            ffn_hidden_dim,
+            tuple(final_out_dims),
+            tp_size,
+            tuple(supported_tp_candidates),
+        )
+
         all_x_embedder = {}
         all_final_layer = {}
         for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):