From 8ac354f166cc424392c9a9a61981e8a2a05dca3d Mon Sep 17 00:00:00 2001
From: "yuanbai.zy" <yuanbai.zy@alibaba-inc.com>
Date: Tue, 28 Jul 2026 17:54:35 +0800
Subject: [PATCH 1/3] feat(tp): add tensor parallel and compile support

---
 diffsynth_engine/args.py                      |  14 +-
 diffsynth_engine/configs/base.py              |  70 +++++----
 .../layers/tensor_parallel/__init__.py        |  13 ++
 .../layers/tensor_parallel/feed_forward.py    |  58 +++++++
 .../layers/tensor_parallel/linear.py          | 139 +++++++++++++++++
 .../layers/tensor_parallel/load_plan.py       |  29 ++++
 .../layers/tensor_parallel/norm.py            |  57 +++++++
 diffsynth_engine/models/base.py               |  31 +++-
 .../qwen_image/transformer_qwenimage.py       |  39 +++--
 .../models/wan/transformer_wan.py             |  44 ++++--
 .../models/wan/transformer_wan_animate.py     |  21 ++-
 .../models/wan/transformer_wan_vace.py        |   8 +-
 diffsynth_engine/pipelines/base.py            |  58 ++++++-
 diffsynth_engine/utils/load_utils.py          | 145 ++++++++++++------
 14 files changed, 612 insertions(+), 114 deletions(-)
 create mode 100644 diffsynth_engine/layers/tensor_parallel/__init__.py
 create mode 100644 diffsynth_engine/layers/tensor_parallel/feed_forward.py
 create mode 100644 diffsynth_engine/layers/tensor_parallel/linear.py
 create mode 100644 diffsynth_engine/layers/tensor_parallel/load_plan.py
 create mode 100644 diffsynth_engine/layers/tensor_parallel/norm.py

diff --git a/diffsynth_engine/args.py b/diffsynth_engine/args.py
index d73a030a..d4452553 100644
--- a/diffsynth_engine/args.py
+++ b/diffsynth_engine/args.py
@@ -110,14 +110,21 @@ def parse_cli_args() -> Dict[str, Any]:
         help="Sparge attention topk parameter (default: 0.5)",
     )
 
+    # Optimization configuration group
+    optimization_group = parser.add_argument_group("Optimization Configuration")
+    optimization_group.add_argument(
+        "--use-torch-compile",
+        action="store_true",
+        help="Compile repeated transformer blocks with torch.compile",
+    )
+
     # Parallelism configuration group
     parallel_group = parser.add_argument_group("Parallelism Configuration")
     parallel_group.add_argument(
         "--parallelism",
         type=int,
         default=1,
-        choices=[1, 2, 4, 8],
-        help="Parallelism degree (default: 1, choices: 1, 2, 4, 8)",
+        help="Total number of inference workers (default: 1)",
     )
     parallel_group.add_argument(
         "--use-cfg-parallel",
@@ -175,6 +182,9 @@ def parse_cli_args() -> Dict[str, Any]:
     args_dict["attn_type"] = attn_type
     args_dict["attn_params"] = _parse_attention_params(attn_type, args.sparge_topk)
 
+    # Optimization configuration
+    args_dict["use_torch_compile"] = args.use_torch_compile
+
     # Parallelism configuration
     args_dict["parallelism"] = args.parallelism
     args_dict["use_cfg_parallel"] = args.use_cfg_parallel
diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py
index b44c472d..09d59c2c 100644
--- a/diffsynth_engine/configs/base.py
+++ b/diffsynth_engine/configs/base.py
@@ -40,6 +40,9 @@ class PipelineConfig:
     attn_type: AttentionType | str = AttentionType.SDPA
     attn_params: Optional[AttentionParams] = None
 
+    # optimization
+    use_torch_compile: bool = False
+
     # parallelism
     parallelism: int = 1
     use_cfg_parallel: bool = False
@@ -62,44 +65,59 @@ def __post_init__(self):
 
 
 def init_parallel_config(config: PipelineConfig):
-    assert config.parallelism in (1, 2, 4, 8), "parallelism must be 1, 2, 4 or 8"
+    if config.parallelism <= 0:
+        raise ValueError(f"parallelism must be a positive integer, got {config.parallelism}")
 
     cfg_degree = 2 if config.use_cfg_parallel else 1
 
-    if config.tp_degree is not None:
-        assert config.sp_ulysses_degree is None and config.sp_ring_degree is None, (
-            "not allowed to enable sequence parallel and tensor parallel together; "
-            "either set sp_ulysses_degree=None, sp_ring_degree=None or set tp_degree=None during pipeline initialization"
-        )
-        assert config.use_fsdp is False, (
-            "not allowed to enable fully sharded data parallel and tensor parallel together; "
-            "either set use_fsdp=False or set tp_degree=None during pipeline initialization"
+    if config.tp_degree is not None and config.tp_degree <= 0:
+        raise ValueError(f"tp_degree must be None or a positive integer, got {config.tp_degree}")
+    if config.sp_ulysses_degree is not None and config.sp_ulysses_degree <= 0:
+        raise ValueError(f"sp_ulysses_degree must be None or a positive integer, got {config.sp_ulysses_degree}")
+    if config.sp_ring_degree is not None and config.sp_ring_degree <= 0:
+        raise ValueError(f"sp_ring_degree must be None or a positive integer, got {config.sp_ring_degree}")
+    if config.use_cfg_parallel and config.parallelism < 2:
+        raise ValueError(
+            f"use_cfg_parallel=True requires parallelism >= 2, got parallelism={config.parallelism}"
         )
-        config.sp_ulysses_degree = 1
-        config.sp_ring_degree = 1
-    elif config.sp_ulysses_degree is None and config.sp_ring_degree is None:
-        # use ulysses if not specified
-        config.sp_ulysses_degree = config.parallelism // cfg_degree
-        config.sp_ring_degree = 1
-        config.tp_degree = 1
-    elif config.sp_ulysses_degree is not None and config.sp_ring_degree is not None:
-        config.tp_degree = 1
-    else:
-        raise ValueError("sp_ulysses_degree and sp_ring_degree must be specified together")
-
-    assert config.parallelism == cfg_degree * config.tp_degree * config.sp_ulysses_degree * config.sp_ring_degree, (
-        f"parallelism ({config.parallelism}) must be equal to cfg_degree ({cfg_degree}) * "
-        f"tp_degree ({config.tp_degree}) * "
-        f"sp_ulysses_degree ({config.sp_ulysses_degree}) * "
-        f"sp_ring_degree ({config.sp_ring_degree})"
+
+    config.tp_degree = config.tp_degree or 1
+    config.sp_ring_degree = config.sp_ring_degree or 1
+    config.sp_ulysses_degree = config.sp_ulysses_degree or (
+        config.parallelism // (cfg_degree * config.tp_degree * config.sp_ring_degree)
     )
 
+    product = cfg_degree * config.tp_degree * config.sp_ulysses_degree * config.sp_ring_degree
+    if product != config.parallelism:
+        raise ValueError(
+            f"parallelism ({config.parallelism}) must equal cfg_degree({cfg_degree}) * "
+            f"tp_degree({config.tp_degree}) * sp_ulysses_degree({config.sp_ulysses_degree}) * "
+            f"sp_ring_degree({config.sp_ring_degree}) = {product}"
+        )
+
+    if config.tp_degree > 1 and config.use_fsdp:
+        raise ValueError("TP and FSDP cannot be enabled together; set use_fsdp=False or tp_degree=None.")
+
     if config.use_vae_parallel:
         assert config.parallelism > 1, "use_vae_parallel requires parallelism > 1"
         if not config.vae_tiled:
             config.vae_tiled = True
             logger.warning("setting vae_tiled to True since use_vae_parallel is enabled")
 
+    if config.tp_degree > 1:
+        logger.info(
+            f"Tensor parallel enabled with tp_degree={config.tp_degree}. "
+            "The model attention heads and sharded MLP dimensions must be divisible by tp_degree."
+        )
+    if config.tp_degree > 1 and config.sp_ulysses_degree > 1:
+        logger.info(
+            f"TP+SP enabled with tp_degree={config.tp_degree} and "
+            f"sp_ulysses_degree={config.sp_ulysses_degree}."
+        )
+
+    if config.use_torch_compile and config.use_fsdp:
+        logger.warning("torch.compile + FSDP may produce graph breaks")
+
 
 def validate_attn_config(config: PipelineConfig):
     attn_backend = get_attn_backend(config.attn_type)
diff --git a/diffsynth_engine/layers/tensor_parallel/__init__.py b/diffsynth_engine/layers/tensor_parallel/__init__.py
new file mode 100644
index 00000000..c211fee6
--- /dev/null
+++ b/diffsynth_engine/layers/tensor_parallel/__init__.py
@@ -0,0 +1,13 @@
+from diffsynth_engine.layers.tensor_parallel.feed_forward import ColumnParallelGELU, TPFeedForward
+from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear
+from diffsynth_engine.layers.tensor_parallel.load_plan import derive_tp_model_weight_load_plan
+from diffsynth_engine.layers.tensor_parallel.norm import TensorParallelRMSNorm
+
+__all__ = [
+    "ColumnParallelLinear",
+    "ColumnParallelGELU",
+    "RowParallelLinear",
+    "TensorParallelRMSNorm",
+    "TPFeedForward",
+    "derive_tp_model_weight_load_plan",
+]
diff --git a/diffsynth_engine/layers/tensor_parallel/feed_forward.py b/diffsynth_engine/layers/tensor_parallel/feed_forward.py
new file mode 100644
index 00000000..15afec72
--- /dev/null
+++ b/diffsynth_engine/layers/tensor_parallel/feed_forward.py
@@ -0,0 +1,58 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear, get_tp_size
+
+
+class ColumnParallelGELU(nn.Module):
+    """Column-parallel linear projection followed by GELU."""
+
+    def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):
+        super().__init__()
+        self.proj = ColumnParallelLinear(dim_in, dim_out, bias=bias, gather_output=False)
+        self.approximate = approximate
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.proj(hidden_states)
+        return F.gelu(hidden_states, approximate=self.approximate)
+
+
+class TPFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int | None = None,
+        mult: float = 4,
+        inner_dim: int | None = None,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu-approximate",
+    ):
+        super().__init__()
+        if activation_fn not in ("gelu", "gelu-approximate"):
+            raise ValueError(f"Unsupported activation_fn={activation_fn!r}; supported: ['gelu', 'gelu-approximate']")
+        approximate = "tanh" if activation_fn == "gelu-approximate" else "none"
+
+        inner_dim = inner_dim if inner_dim is not None else int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+        tp_size = get_tp_size()
+        if inner_dim % tp_size != 0:
+            raise ValueError(f"inner_dim ({inner_dim}) must be divisible by tp_size ({tp_size})")
+
+        self.net = nn.ModuleList(
+            [
+                ColumnParallelGELU(
+                    dim,
+                    inner_dim,
+                    approximate=approximate,
+                    bias=True,
+                ),
+                nn.Dropout(dropout),
+                RowParallelLinear(inner_dim, dim_out, bias=True, input_is_parallel=True),
+            ]
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
diff --git a/diffsynth_engine/layers/tensor_parallel/linear.py b/diffsynth_engine/layers/tensor_parallel/linear.py
new file mode 100644
index 00000000..d02f6ee0
--- /dev/null
+++ b/diffsynth_engine/layers/tensor_parallel/linear.py
@@ -0,0 +1,139 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffsynth_engine.distributed.parallel_state import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+    is_tp_group_initialized,
+)
+
+
+def get_tp_size() -> int:
+    return get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+
+
+def get_tp_rank() -> int:
+    return get_tensor_model_parallel_rank() if is_tp_group_initialized() else 0
+
+
+@torch.compiler.disable
+def tp_all_reduce(output: torch.Tensor) -> torch.Tensor:
+    return get_tp_group().all_reduce(output)
+
+
+@torch.compiler.disable
+def tp_all_gather(output: torch.Tensor, dim: int) -> torch.Tensor:
+    return get_tp_group().all_gather(output, dim=dim)
+
+
+class ColumnParallelLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        gather_output: bool = False,
+        dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+    ):
+        super().__init__()
+        tp_size = get_tp_size()
+        if out_features % tp_size != 0:
+            raise ValueError(
+                f"ColumnParallelLinear: out_features ({out_features}) must be divisible by tp_size ({tp_size})"
+            )
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.gather_output = gather_output
+        self.out_features_per_partition = out_features // tp_size
+        self.tp_size = tp_size
+        self.tp_rank = get_tp_rank()
+
+        factory_kwargs = {"dtype": dtype, "device": device}
+        self.weight = nn.Parameter(torch.empty(self.out_features_per_partition, in_features, **factory_kwargs))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(self.out_features_per_partition, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            bound = 1 / math.sqrt(self.in_features) if self.in_features > 0 else 0
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        output = F.linear(hidden_states, self.weight, self.bias)
+        if self.gather_output and self.tp_size > 1:
+            output = tp_all_gather(output, dim=-1)
+        return output
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"out_per_partition={self.out_features_per_partition}, "
+            f"bias={self.bias is not None}, gather_output={self.gather_output}"
+        )
+
+
+class RowParallelLinear(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        input_is_parallel: bool = True,
+        dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+    ):
+        super().__init__()
+        tp_size = get_tp_size()
+        if in_features % tp_size != 0:
+            raise ValueError(f"RowParallelLinear: in_features ({in_features}) must be divisible by tp_size ({tp_size})")
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.input_is_parallel = input_is_parallel
+        self.in_features_per_partition = in_features // tp_size
+        self.tp_size = tp_size
+        self.tp_rank = get_tp_rank()
+
+        factory_kwargs = {"dtype": dtype, "device": device}
+        self.weight = nn.Parameter(torch.empty(out_features, self.in_features_per_partition, **factory_kwargs))
+        if bias:
+            self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs))
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        bound = 1 / math.sqrt(self.in_features) if self.in_features > 0 else 0
+        nn.init.uniform_(self.weight, -bound, bound)
+        if self.bias is not None:
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.tp_size == 1:
+            return F.linear(hidden_states, self.weight, self.bias)
+
+        if not self.input_is_parallel:
+            hidden_states = hidden_states.chunk(self.tp_size, dim=-1)[self.tp_rank].contiguous()
+
+        output = F.linear(hidden_states, self.weight, None)
+        output = tp_all_reduce(output)
+        if self.bias is not None:
+            output = output + self.bias
+        return output
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"in_per_partition={self.in_features_per_partition}, "
+            f"bias={self.bias is not None}, input_is_parallel={self.input_is_parallel}"
+        )
diff --git a/diffsynth_engine/layers/tensor_parallel/load_plan.py b/diffsynth_engine/layers/tensor_parallel/load_plan.py
new file mode 100644
index 00000000..c265f41d
--- /dev/null
+++ b/diffsynth_engine/layers/tensor_parallel/load_plan.py
@@ -0,0 +1,29 @@
+import torch.nn as nn
+
+from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear
+from diffsynth_engine.layers.tensor_parallel.norm import TensorParallelRMSNorm
+from diffsynth_engine.utils.load_utils import TensorSelection, TensorSelectionPlan, TensorSlice
+
+
+def _slice_selection(dim: int, start: int, end: int) -> TensorSelection:
+    return TensorSelection(slices=(TensorSlice(dim=dim, start=start, end=end),))
+
+
+def derive_tp_model_weight_load_plan(model: nn.Module) -> TensorSelectionPlan:
+    plan: TensorSelectionPlan = {}
+    for name, module in model.named_modules():
+        if isinstance(module, ColumnParallelLinear):
+            start = module.tp_rank * module.out_features_per_partition
+            end = start + module.out_features_per_partition
+            plan[f"{name}.weight"] = _slice_selection(dim=0, start=start, end=end)
+            if module.bias is not None:
+                plan[f"{name}.bias"] = _slice_selection(dim=0, start=start, end=end)
+        elif isinstance(module, RowParallelLinear):
+            start = module.tp_rank * module.in_features_per_partition
+            end = start + module.in_features_per_partition
+            plan[f"{name}.weight"] = _slice_selection(dim=1, start=start, end=end)
+        elif isinstance(module, TensorParallelRMSNorm) and module.weight is not None:
+            start = module.tp_rank * module.hidden_size_per_partition
+            end = start + module.hidden_size_per_partition
+            plan[f"{name}.weight"] = _slice_selection(dim=0, start=start, end=end)
+    return plan
diff --git a/diffsynth_engine/layers/tensor_parallel/norm.py b/diffsynth_engine/layers/tensor_parallel/norm.py
new file mode 100644
index 00000000..10ed965e
--- /dev/null
+++ b/diffsynth_engine/layers/tensor_parallel/norm.py
@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffsynth_engine.layers.tensor_parallel.linear import get_tp_rank, get_tp_size, tp_all_reduce
+
+
+class TensorParallelRMSNorm(nn.Module):
+    """RMSNorm over a hidden dimension sharded across the tensor-parallel group."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float | None = None,
+        elementwise_affine: bool = True,
+        dtype: torch.dtype | None = None,
+        device: torch.device | str | None = None,
+    ):
+        super().__init__()
+        tp_size = get_tp_size()
+        if hidden_size % tp_size != 0:
+            raise ValueError(
+                f"TensorParallelRMSNorm: hidden_size ({hidden_size}) must be divisible by tp_size ({tp_size})"
+            )
+
+        self.hidden_size = hidden_size
+        self.hidden_size_per_partition = hidden_size // tp_size
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        self.tp_size = tp_size
+        self.tp_rank = get_tp_rank()
+
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(self.hidden_size_per_partition, dtype=dtype, device=device))
+        else:
+            self.register_parameter("weight", None)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.tp_size == 1:
+            return F.rms_norm(hidden_states, (self.hidden_size,), self.weight, self.eps)
+        if hidden_states.shape[-1] != self.hidden_size_per_partition:
+            raise ValueError(f"Expected last dimension {self.hidden_size_per_partition}, got {hidden_states.shape[-1]}")
+
+        variance = hidden_states.float().pow(2).sum(dim=-1, keepdim=True)
+        variance = tp_all_reduce(variance) / self.hidden_size
+        eps = self.eps if self.eps is not None else torch.finfo(hidden_states.dtype).eps
+        output = hidden_states * torch.rsqrt(variance + eps).to(hidden_states.dtype)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return (
+            f"hidden_size={self.hidden_size}, "
+            f"hidden_size_per_partition={self.hidden_size_per_partition}, eps={self.eps}, "
+            f"elementwise_affine={self.elementwise_affine}"
+        )
diff --git a/diffsynth_engine/models/base.py b/diffsynth_engine/models/base.py
index 20768ee1..5b8537dc 100644
--- a/diffsynth_engine/models/base.py
+++ b/diffsynth_engine/models/base.py
@@ -5,6 +5,11 @@
 from accelerate import init_empty_weights
 from diffusers.configuration_utils import ConfigMixin
 
+from diffsynth_engine.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+    is_tp_group_initialized,
+)
+from diffsynth_engine.layers.tensor_parallel import derive_tp_model_weight_load_plan
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import CONFIG_NAME
 from diffsynth_engine.utils.load_utils import load_model_weights
@@ -40,16 +45,38 @@ def from_pretrained(
         with init_empty_weights():
             model = cls.from_config(config_dict)
 
+        tensor_selection_plan = None
+        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+        if tp_size > 1:
+            tensor_selection_plan = derive_tp_model_weight_load_plan(model)
+            if not tensor_selection_plan:
+                raise RuntimeError(
+                    f"Tensor parallel world size is {tp_size}, but {cls.__name__} has no tensor-parallel submodules; "
+                    "the model is not TP-aware."
+                )
+
         # avoid precision loss
         if dtype is not None and dtype != torch.float32 and cls._keep_in_fp32_modules:
-            state_dict = load_model_weights(model_path, subfolder, device, dtype=None)
+            state_dict = load_model_weights(
+                model_path,
+                subfolder,
+                device,
+                dtype=None,
+                tensor_selection_plan=tensor_selection_plan,
+            )
             for k, v in state_dict.items():
                 if any(m in k.split(".") for m in cls._keep_in_fp32_modules):
                     state_dict[k] = v.to(dtype=torch.float32)
                 else:
                     state_dict[k] = v.to(dtype=dtype)
         else:
-            state_dict = load_model_weights(model_path, subfolder, device, dtype)
+            state_dict = load_model_weights(
+                model_path,
+                subfolder,
+                device,
+                dtype,
+                tensor_selection_plan=tensor_selection_plan,
+            )
 
         # drop unexpected keys
         if cls._keys_to_ignore_on_load_unexpected:
diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
index 4d92c05d..36b6a4be 100644
--- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
+++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
@@ -22,14 +22,18 @@
 import torch
 import torch.nn as nn
 from diffusers.configuration_utils import register_to_config
-from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
 
+from diffsynth_engine.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+    is_tp_group_initialized,
+)
 from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard
 from diffsynth_engine.forward_context import get_forward_context
 from diffsynth_engine.layers.attention import USPAttention
+from diffsynth_engine.layers.tensor_parallel import ColumnParallelLinear, RowParallelLinear, TPFeedForward
 from diffsynth_engine.models.base import DiffusionModel
 from diffsynth_engine.utils import logging
 
@@ -415,21 +419,30 @@ def __init__(
         self.dim = dim
         self.num_attention_heads = num_attention_heads
         self.attention_head_dim = attention_head_dim
-        self.heads = num_attention_heads
+        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+        if num_attention_heads % tp_size != 0:
+            raise ValueError(
+                f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})"
+            )
+        self.heads = num_attention_heads // tp_size
 
         # Image stream projections
-        self.to_q = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
-        self.to_k = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
-        self.to_v = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
+        self.to_q = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.to_k = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.to_v = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
         self.to_out = nn.ModuleList([])
-        self.to_out.append(nn.Linear(num_attention_heads * attention_head_dim, dim, bias=True))
+        self.to_out.append(
+            RowParallelLinear(num_attention_heads * attention_head_dim, dim, bias=True, input_is_parallel=True)
+        )
         self.to_out.append(nn.Dropout(dropout))
 
         # Text stream projections
-        self.add_q_proj = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
-        self.add_k_proj = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
-        self.add_v_proj = nn.Linear(dim, num_attention_heads * attention_head_dim, bias=True)
-        self.to_add_out = nn.Linear(num_attention_heads * attention_head_dim, dim, bias=True)
+        self.add_q_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.add_k_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.add_v_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.to_add_out = RowParallelLinear(
+            num_attention_heads * attention_head_dim, dim, bias=True, input_is_parallel=True
+        )
 
         # QK normalization
         if qk_norm == "rms_norm":
@@ -446,7 +459,7 @@ def __init__(
         # USPAttention for joint attention computation
         forward_context = get_forward_context()
         self.usp_attn = USPAttention(
-            num_heads=num_attention_heads,
+            num_heads=self.heads,
             head_size=attention_head_dim,
             attn_type=forward_context.attn_type,
         )
@@ -553,7 +566,7 @@ def __init__(
             eps=eps,
         )
         self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        self.img_mlp = TPFeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
 
         # Text processing modules
         self.txt_mod = nn.Sequential(
@@ -563,7 +576,7 @@ def __init__(
         self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         # Text doesn't need separate attention - it's handled by img_attn joint computation
         self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+        self.txt_mlp = TPFeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
 
         self.zero_cond_t = zero_cond_t
 
diff --git a/diffsynth_engine/models/wan/transformer_wan.py b/diffsynth_engine/models/wan/transformer_wan.py
index 526097f9..4095b8bb 100644
--- a/diffsynth_engine/models/wan/transformer_wan.py
+++ b/diffsynth_engine/models/wan/transformer_wan.py
@@ -24,9 +24,19 @@
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
 
+from diffsynth_engine.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+    is_tp_group_initialized,
+)
 from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard
 from diffsynth_engine.forward_context import get_forward_context
 from diffsynth_engine.layers.attention import LocalAttention, USPAttention
+from diffsynth_engine.layers.tensor_parallel import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    TensorParallelRMSNorm,
+    TPFeedForward,
+)
 from diffsynth_engine.models.base import DiffusionModel
 from diffsynth_engine.utils import logging
 
@@ -76,34 +86,42 @@ def __init__(
     ):
         super().__init__()
 
+        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+        if heads % tp_size != 0:
+            raise ValueError(f"heads ({heads}) must be divisible by tp_size ({tp_size})")
+
         self.inner_dim = dim_head * heads
-        self.heads = heads
+        self.heads = heads // tp_size
         self.added_kv_proj_dim = added_kv_proj_dim
         self.cross_attention_dim_head = cross_attention_dim_head
         self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads
 
-        self.to_q = nn.Linear(dim, self.inner_dim, bias=True)
-        self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True)
-        self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True)
+        self.to_q = ColumnParallelLinear(dim, self.inner_dim, bias=True, gather_output=False)
+        self.to_k = ColumnParallelLinear(dim, self.kv_inner_dim, bias=True, gather_output=False)
+        self.to_v = ColumnParallelLinear(dim, self.kv_inner_dim, bias=True, gather_output=False)
         self.to_out = nn.ModuleList(
             [
-                nn.Linear(self.inner_dim, dim, bias=True),
+                RowParallelLinear(self.inner_dim, dim, bias=True, input_is_parallel=True),
                 nn.Dropout(dropout),
             ]
         )
-        self.norm_q = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True)
-        self.norm_k = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True)
+        self.norm_q = TensorParallelRMSNorm(self.inner_dim, eps=eps, elementwise_affine=True)
+        self.norm_k = TensorParallelRMSNorm(self.kv_inner_dim, eps=eps, elementwise_affine=True)
 
         self.add_k_proj = self.add_v_proj = None
         if added_kv_proj_dim is not None:
-            self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True)
-            self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True)
-            self.norm_added_k = nn.RMSNorm(dim_head * heads, eps=eps)
+            self.add_k_proj = ColumnParallelLinear(
+                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False
+            )
+            self.add_v_proj = ColumnParallelLinear(
+                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False
+            )
+            self.norm_added_k = TensorParallelRMSNorm(self.inner_dim, eps=eps)
 
         forward_context = get_forward_context()
         attn_cls = USPAttention if cross_attention_dim_head is None else LocalAttention
         self.attn = attn_cls(
-            num_heads=heads,
+            num_heads=self.heads,
             head_size=dim_head,
             attn_type=forward_context.attn_type,
         )
@@ -345,7 +363,7 @@ def __init__(
         self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 3. Feed-forward
-        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
+        self.ffn = TPFeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
         self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -529,7 +547,7 @@ def forward(
         rotary_emb = self.rope(hidden_states)
 
         hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous()
 
         original_seq_len = hidden_states.shape[1]
 
diff --git a/diffsynth_engine/models/wan/transformer_wan_animate.py b/diffsynth_engine/models/wan/transformer_wan_animate.py
index bb8eb623..c294b2ce 100644
--- a/diffsynth_engine/models/wan/transformer_wan_animate.py
+++ b/diffsynth_engine/models/wan/transformer_wan_animate.py
@@ -23,9 +23,14 @@
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
 
+from diffsynth_engine.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+    is_tp_group_initialized,
+)
 from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard
 from diffsynth_engine.forward_context import get_forward_context
 from diffsynth_engine.layers.attention import LocalAttention
+from diffsynth_engine.layers.tensor_parallel import ColumnParallelLinear, RowParallelLinear
 from diffsynth_engine.models.base import DiffusionModel
 from diffsynth_engine.models.wan.transformer_wan import (
     WanRotaryPosEmbed,
@@ -400,8 +405,12 @@ def __init__(
         cross_attention_dim_head: int | None = None,
     ):
         super().__init__()
+        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+        if heads % tp_size != 0:
+            raise ValueError(f"heads ({heads}) must be divisible by tp_size ({tp_size})")
+
         self.inner_dim = dim_head * heads
-        self.heads = heads
+        self.heads = heads // tp_size
         self.cross_attention_dim_head = cross_attention_dim_head
         self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads
 
@@ -411,10 +420,10 @@ def __init__(
         self.pre_norm_kv = nn.LayerNorm(dim, eps, elementwise_affine=False)
 
         # 2. QKV and Output Projections
-        self.to_q = nn.Linear(dim, self.inner_dim, bias=True)
-        self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True)
-        self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True)
-        self.to_out = nn.Linear(self.inner_dim, dim, bias=True)
+        self.to_q = ColumnParallelLinear(dim, self.inner_dim, bias=True, gather_output=False)
+        self.to_k = ColumnParallelLinear(dim, self.kv_inner_dim, bias=True, gather_output=False)
+        self.to_v = ColumnParallelLinear(dim, self.kv_inner_dim, bias=True, gather_output=False)
+        self.to_out = RowParallelLinear(self.inner_dim, dim, bias=True, input_is_parallel=True)
 
         # 3. QK Norm
         # NOTE: this is applied after the reshape, so only over dim_head rather than dim_head * heads
@@ -424,7 +433,7 @@ def __init__(
         # 4. Attention
         forward_context = get_forward_context()
         self.attn = LocalAttention(
-            num_heads=heads,
+            num_heads=self.heads,
             head_size=dim_head,
             attn_type=forward_context.attn_type,
         )
diff --git a/diffsynth_engine/models/wan/transformer_wan_vace.py b/diffsynth_engine/models/wan/transformer_wan_vace.py
index a44c64e4..8a7ab626 100644
--- a/diffsynth_engine/models/wan/transformer_wan_vace.py
+++ b/diffsynth_engine/models/wan/transformer_wan_vace.py
@@ -19,11 +19,11 @@
 import torch
 import torch.nn as nn
 from diffusers.configuration_utils import register_to_config
-from diffusers.models.attention import FeedForward
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import FP32LayerNorm
 
 from diffsynth_engine.distributed.utils import sequence_parallel_shard, sequence_parallel_unshard
+from diffsynth_engine.layers.tensor_parallel import TPFeedForward
 from diffsynth_engine.models.base import DiffusionModel
 from diffsynth_engine.models.wan.transformer_wan import (
     WanAttention,
@@ -78,7 +78,7 @@ def __init__(
         self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
 
         # 4. Feed-forward
-        self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
+        self.ffn = TPFeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate")
         self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False)
 
         # 5. Output projection
@@ -316,10 +316,10 @@ def forward(
 
         # 2. Patch embedding
         hidden_states = self.patch_embedding(hidden_states)
-        hidden_states = hidden_states.flatten(2).transpose(1, 2)
+        hidden_states = hidden_states.flatten(2).transpose(1, 2).contiguous()
 
         control_hidden_states = self.vace_patch_embedding(control_hidden_states)
-        control_hidden_states = control_hidden_states.flatten(2).transpose(1, 2)
+        control_hidden_states = control_hidden_states.flatten(2).transpose(1, 2).contiguous()
         control_hidden_states_padding = control_hidden_states.new_zeros(
             batch_size, hidden_states.size(1) - control_hidden_states.size(1), control_hidden_states.size(2)
         )
diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py
index 91c2266e..1f511c09 100644
--- a/diffsynth_engine/pipelines/base.py
+++ b/diffsynth_engine/pipelines/base.py
@@ -8,8 +8,16 @@
 from tqdm import tqdm
 
 from diffsynth_engine.configs import PipelineConfig
-from diffsynth_engine.distributed.parallel_state import get_global_rank, is_world_group_initialized
+from diffsynth_engine.distributed.parallel_state import (
+    get_global_rank,
+    get_tensor_model_parallel_world_size,
+    get_ulysses_parallel_world_size,
+    is_sp_group_initialized,
+    is_tp_group_initialized,
+    is_world_group_initialized,
+)
 from diffsynth_engine.forward_context import set_forward_context
+from diffsynth_engine.layers.tensor_parallel import derive_tp_model_weight_load_plan
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.load_utils import fix_state_dict_key, load_model_weights
 
@@ -29,7 +37,26 @@ def __call__(self, *args, **kwargs):
         raise NotImplementedError()
 
     @staticmethod
+    def compile_transformer_blocks(model: nn.Module) -> nn.Module:
+        repeated_blocks = getattr(model, "_repeated_blocks", None)
+        if not repeated_blocks:
+            raise ValueError(f"`_repeated_blocks` is not defined for {type(model).__name__}")
+
+        has_compiled_region = False
+        for submodule in model.modules():
+            if submodule.__class__.__name__ in repeated_blocks:
+                submodule.compile(dynamic=True, fullgraph=False)
+                has_compiled_region = True
+
+        if not has_compiled_region:
+            raise ValueError(
+                f"None of the repeated block classes {repeated_blocks} were found in {type(model).__name__}"
+            )
+        return model
+
+    @classmethod
     def init_transformer(
+        cls,
         model_cls: Type[nn.Module],
         pipeline_config: PipelineConfig,
         empty_weights: bool = False,
@@ -59,6 +86,31 @@ def init_transformer(
                     fully_shard(submodule)
             fully_shard(model)
 
+        tensor_selection_plan = None
+        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+        if tp_size > 1:
+            tensor_selection_plan = derive_tp_model_weight_load_plan(model)
+            if not tensor_selection_plan:
+                raise RuntimeError(
+                    f"Tensor parallel world size is {tp_size}, but {model_cls.__name__} has no tensor-parallel submodules; "
+                    "the model is not TP-aware."
+                )
+
+            num_heads = getattr(model.config, "num_attention_heads", None)
+            if num_heads is not None and num_heads % tp_size != 0:
+                raise ValueError(
+                    f"num_attention_heads ({num_heads}) must be divisible by tp_degree ({tp_size})"
+                )
+            if num_heads is not None and is_sp_group_initialized():
+                ulysses_size = get_ulysses_parallel_world_size()
+                heads_per_tp_partition = num_heads // tp_size
+                if ulysses_size > 1 and heads_per_tp_partition % ulysses_size != 0:
+                    raise ValueError(
+                        f"With tp_degree={tp_size} and sp_ulysses_degree={ulysses_size}, "
+                        f"heads_per_tp_partition ({heads_per_tp_partition}) must be divisible by "
+                        "sp_ulysses_degree."
+                    )
+
         load_device = "cpu" if use_fsdp else pipeline_config.device
         model_dtype = pipeline_config.model_dtype
         keep_in_fp32_modules = getattr(model_cls, "_keep_in_fp32_modules", None)
@@ -70,6 +122,7 @@ def init_transformer(
                 device=load_device,
                 dtype=None,
                 broadcast_from_rank0=not use_fsdp,
+                tensor_selection_plan=tensor_selection_plan,
             )
             for k, v in state_dict.items():
                 if any(m in k.split(".") for m in keep_in_fp32_modules):
@@ -83,6 +136,7 @@ def init_transformer(
                 device=load_device,
                 dtype=model_dtype,
                 broadcast_from_rank0=not use_fsdp,
+                tensor_selection_plan=tensor_selection_plan,
             )
 
         if use_fsdp:
@@ -96,6 +150,8 @@ def init_transformer(
             model.to(device=pipeline_config.device)
 
         del state_dict
+        if pipeline_config.use_torch_compile:
+            model = cls.compile_transformer_blocks(model)
         return model
 
     @staticmethod
diff --git a/diffsynth_engine/utils/load_utils.py b/diffsynth_engine/utils/load_utils.py
index 9c2ab043..9a008e4c 100644
--- a/diffsynth_engine/utils/load_utils.py
+++ b/diffsynth_engine/utils/load_utils.py
@@ -2,7 +2,8 @@
 import os
 import re
 import time
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 
@@ -15,6 +16,22 @@
     SAFETENSORS_WEIGHTS_NAME,
 )
 
+
+@dataclass(frozen=True)
+class TensorSlice:
+    dim: int
+    start: int
+    end: int
+
+
+@dataclass(frozen=True)
+class TensorSelection:
+    slices: Tuple[TensorSlice, ...]
+    contiguous: bool = True
+
+
+TensorSelectionPlan = Dict[str, TensorSelection]
+
 logger = logging.get_logger(__name__)
 
 try:
@@ -48,62 +65,96 @@ def load_safetensors(path: str, device: str = "cpu") -> Dict[str, Any]:
     return state_dict
 
 
+def apply_tensor_selection(
+    state_dict: Dict[str, Any],
+    tensor_selection_plan: TensorSelectionPlan,
+) -> Dict[str, Any]:
+    for key in list(state_dict.keys()):
+        selection = tensor_selection_plan.get(key)
+        if selection is None:
+            continue
+
+        value = state_dict[key]
+        if not isinstance(value, torch.Tensor):
+            continue
+
+        slices = [slice(None)] * value.dim()
+        for tensor_slice in selection.slices:
+            dim = tensor_slice.dim if tensor_slice.dim >= 0 else value.dim() + tensor_slice.dim
+            if dim < 0 or dim >= value.dim():
+                raise ValueError(
+                    f"Cannot slice '{key}' shape={tuple(value.shape)} along dim {tensor_slice.dim}: "
+                    f"tensor has {value.dim()} dimensions."
+                )
+            if tensor_slice.start < 0 or tensor_slice.end > value.shape[dim] or tensor_slice.start > tensor_slice.end:
+                raise ValueError(
+                    f"Invalid slice for '{key}' shape={tuple(value.shape)}: "
+                    f"dim={tensor_slice.dim}, start={tensor_slice.start}, end={tensor_slice.end}."
+                )
+            slices[dim] = slice(tensor_slice.start, tensor_slice.end)
+
+        selected = value[tuple(slices)]
+        state_dict[key] = selected.contiguous() if selection.contiguous else selected
+
+    return state_dict
+
+
+def _list_shard_files(path: str) -> list[str]:
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Model path not found: {path}")
+
+    diffusion_index = os.path.join(path, DIFFUSION_SAFETENSORS_INDEX_NAME)
+    diffusion_weights = os.path.join(path, DIFFUSION_SAFETENSORS_WEIGHTS_NAME)
+    generic_index = os.path.join(path, SAFETENSORS_INDEX_NAME)
+    generic_weights = os.path.join(path, SAFETENSORS_WEIGHTS_NAME)
+
+    if os.path.exists(diffusion_index):
+        index_file = diffusion_index
+    elif os.path.exists(diffusion_weights):
+        return [diffusion_weights]
+    elif os.path.exists(generic_index):
+        index_file = generic_index
+    elif os.path.exists(generic_weights):
+        return [generic_weights]
+    else:
+        raise FileNotFoundError(f"Safetensors index or weights file not found in {path}")
+
+    with open(index_file, "r", encoding="utf-8") as file:
+        index_dict = json.load(file)
+    shard_names = sorted(set(index_dict["weight_map"].values()))
+    if not shard_names:
+        raise ValueError(f"Weight index {index_file} contains an empty weight_map")
+    return [os.path.join(path, name) for name in shard_names]
+
+
 def load_model_weights(
     model_path: str,
     subfolder: Optional[str] = None,
     device: Optional[str] = None,
     dtype: Optional[torch.dtype] = None,
     broadcast_from_rank0: bool = True,
+    tensor_selection_plan: Optional[TensorSelectionPlan] = None,
 ) -> Dict[str, Any]:
     world_group = get_world_group() if is_world_group_initialized() else None
     is_rank_zero = world_group is None or get_global_rank() == 0
 
-    # rank 0 reads all shards
-    state_dict = {}
-    if is_rank_zero:
-        if subfolder is not None:
-            model_path = os.path.join(model_path, subfolder)
-
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"Model path not found: {model_path}")
-
-        _diffusion_index_file = os.path.join(model_path, DIFFUSION_SAFETENSORS_INDEX_NAME)
-        _diffusion_weights_file = os.path.join(model_path, DIFFUSION_SAFETENSORS_WEIGHTS_NAME)
-        _index_file = os.path.join(model_path, SAFETENSORS_INDEX_NAME)
-        _weights_file = os.path.join(model_path, SAFETENSORS_WEIGHTS_NAME)
-
-        index_file, weights_file = None, None
-
-        if os.path.exists(_diffusion_index_file):
-            index_file = _diffusion_index_file
-        elif os.path.exists(_diffusion_weights_file):
-            weights_file = _diffusion_weights_file
-        elif os.path.exists(_index_file):
-            index_file = _index_file
-        elif os.path.exists(_weights_file):
-            weights_file = _weights_file
-        else:
-            raise FileNotFoundError(f"Safetensors index or weights file not found in {model_path}")
-
-        if index_file is not None:
-            with open(index_file, "r", encoding="utf-8") as f:
-                index_dict = json.load(f)
-            weight_map = index_dict["weight_map"]
-            shard_files = sorted(set(weight_map.values()))
-            for shard_file in shard_files:
-                path = os.path.join(model_path, shard_file)
-                shard_dict = load_safetensors(path)
-                for k, v in shard_dict.items():
-                    shard_dict[k] = v.to(device=device, dtype=dtype, non_blocking=True)
-                state_dict.update(shard_dict)
-        else:
-            state_dict = load_safetensors(weights_file)
-            for k, v in state_dict.items():
-                state_dict[k] = v.to(device=device, dtype=dtype, non_blocking=True)
-
-    # rank 0 broadcasts full state dict to all other ranks
-    if broadcast_from_rank0 and world_group is not None:
-        state_dict = world_group.broadcast_tensor_dict(state_dict, src=0)
+    resolved_path = os.path.join(model_path, subfolder) if subfolder is not None else model_path
+    shard_files = _list_shard_files(resolved_path)
+
+    state_dict: Dict[str, Any] = {}
+    for shard_file in shard_files:
+        shard_dict = load_safetensors(shard_file) if is_rank_zero else {}
+
+        if world_group is not None and broadcast_from_rank0:
+            shard_dict = world_group.broadcast_tensor_dict(shard_dict, src=0)
+
+        if tensor_selection_plan is not None:
+            shard_dict = apply_tensor_selection(shard_dict, tensor_selection_plan)
+
+        for key, value in shard_dict.items():
+            if isinstance(value, torch.Tensor):
+                shard_dict[key] = value.to(device=device, dtype=dtype, non_blocking=True)
+        state_dict.update(shard_dict)
 
     return state_dict
 

From 791e816b6237d411eba70c21a14e232b966de4cc Mon Sep 17 00:00:00 2001
From: "zhuguoxuan.zgx" <zhuguoxuan.zgx@alibaba-inc.com>
Date: Thu, 30 Jul 2026 02:26:10 +0800
Subject: [PATCH 2/3] refactor model weights loading

---
 diffsynth_engine/configs/base.py              |  31 +---
 .../layers/tensor_parallel/__init__.py        |   2 -
 .../layers/tensor_parallel/load_plan.py       |  29 ---
 diffsynth_engine/models/base.py               |  17 --
 .../qwen_image/transformer_qwenimage.py       |  16 +-
 .../models/wan/transformer_wan.py             |   8 +-
 diffsynth_engine/pipelines/base.py            |  47 ++---
 diffsynth_engine/utils/load_utils.py          | 169 +++++++++++-------
 8 files changed, 137 insertions(+), 182 deletions(-)
 delete mode 100644 diffsynth_engine/layers/tensor_parallel/load_plan.py

diff --git a/diffsynth_engine/configs/base.py b/diffsynth_engine/configs/base.py
index 09d59c2c..226dff23 100644
--- a/diffsynth_engine/configs/base.py
+++ b/diffsynth_engine/configs/base.py
@@ -68,7 +68,7 @@ def init_parallel_config(config: PipelineConfig):
     if config.parallelism <= 0:
         raise ValueError(f"parallelism must be a positive integer, got {config.parallelism}")
 
-    cfg_degree = 2 if config.use_cfg_parallel else 1
+    cfg_degree = 2 if config.use_cfg_parallel else 1  # TODO: support cfg_degree > 2
 
     if config.tp_degree is not None and config.tp_degree <= 0:
         raise ValueError(f"tp_degree must be None or a positive integer, got {config.tp_degree}")
@@ -76,10 +76,6 @@ def init_parallel_config(config: PipelineConfig):
         raise ValueError(f"sp_ulysses_degree must be None or a positive integer, got {config.sp_ulysses_degree}")
     if config.sp_ring_degree is not None and config.sp_ring_degree <= 0:
         raise ValueError(f"sp_ring_degree must be None or a positive integer, got {config.sp_ring_degree}")
-    if config.use_cfg_parallel and config.parallelism < 2:
-        raise ValueError(
-            f"use_cfg_parallel=True requires parallelism >= 2, got parallelism={config.parallelism}"
-        )
 
     config.tp_degree = config.tp_degree or 1
     config.sp_ring_degree = config.sp_ring_degree or 1
@@ -87,16 +83,19 @@ def init_parallel_config(config: PipelineConfig):
         config.parallelism // (cfg_degree * config.tp_degree * config.sp_ring_degree)
     )
 
-    product = cfg_degree * config.tp_degree * config.sp_ulysses_degree * config.sp_ring_degree
-    if product != config.parallelism:
+    parallel_degree = cfg_degree * config.tp_degree * config.sp_ulysses_degree * config.sp_ring_degree
+    if parallel_degree != config.parallelism:
         raise ValueError(
             f"parallelism ({config.parallelism}) must equal cfg_degree({cfg_degree}) * "
             f"tp_degree({config.tp_degree}) * sp_ulysses_degree({config.sp_ulysses_degree}) * "
-            f"sp_ring_degree({config.sp_ring_degree}) = {product}"
+            f"sp_ring_degree({config.sp_ring_degree}) = {parallel_degree}"
         )
 
     if config.tp_degree > 1 and config.use_fsdp:
-        raise ValueError("TP and FSDP cannot be enabled together; set use_fsdp=False or tp_degree=None.")
+        raise ValueError("TP and FSDP cannot be enabled together; set tp_degree=None or use_fsdp=False .")
+
+    if config.use_torch_compile and config.use_fsdp:
+        logger.warning("torch.compile + FSDP may produce graph breaks")
 
     if config.use_vae_parallel:
         assert config.parallelism > 1, "use_vae_parallel requires parallelism > 1"
@@ -104,20 +103,6 @@ def init_parallel_config(config: PipelineConfig):
             config.vae_tiled = True
             logger.warning("setting vae_tiled to True since use_vae_parallel is enabled")
 
-    if config.tp_degree > 1:
-        logger.info(
-            f"Tensor parallel enabled with tp_degree={config.tp_degree}. "
-            "The model attention heads and sharded MLP dimensions must be divisible by tp_degree."
-        )
-    if config.tp_degree > 1 and config.sp_ulysses_degree > 1:
-        logger.info(
-            f"TP+SP enabled with tp_degree={config.tp_degree} and "
-            f"sp_ulysses_degree={config.sp_ulysses_degree}."
-        )
-
-    if config.use_torch_compile and config.use_fsdp:
-        logger.warning("torch.compile + FSDP may produce graph breaks")
-
 
 def validate_attn_config(config: PipelineConfig):
     attn_backend = get_attn_backend(config.attn_type)
diff --git a/diffsynth_engine/layers/tensor_parallel/__init__.py b/diffsynth_engine/layers/tensor_parallel/__init__.py
index c211fee6..7b9bb2b5 100644
--- a/diffsynth_engine/layers/tensor_parallel/__init__.py
+++ b/diffsynth_engine/layers/tensor_parallel/__init__.py
@@ -1,6 +1,5 @@
 from diffsynth_engine.layers.tensor_parallel.feed_forward import ColumnParallelGELU, TPFeedForward
 from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear
-from diffsynth_engine.layers.tensor_parallel.load_plan import derive_tp_model_weight_load_plan
 from diffsynth_engine.layers.tensor_parallel.norm import TensorParallelRMSNorm
 
 __all__ = [
@@ -9,5 +8,4 @@
     "RowParallelLinear",
     "TensorParallelRMSNorm",
     "TPFeedForward",
-    "derive_tp_model_weight_load_plan",
 ]
diff --git a/diffsynth_engine/layers/tensor_parallel/load_plan.py b/diffsynth_engine/layers/tensor_parallel/load_plan.py
deleted file mode 100644
index c265f41d..00000000
--- a/diffsynth_engine/layers/tensor_parallel/load_plan.py
+++ /dev/null
@@ -1,29 +0,0 @@
-import torch.nn as nn
-
-from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear
-from diffsynth_engine.layers.tensor_parallel.norm import TensorParallelRMSNorm
-from diffsynth_engine.utils.load_utils import TensorSelection, TensorSelectionPlan, TensorSlice
-
-
-def _slice_selection(dim: int, start: int, end: int) -> TensorSelection:
-    return TensorSelection(slices=(TensorSlice(dim=dim, start=start, end=end),))
-
-
-def derive_tp_model_weight_load_plan(model: nn.Module) -> TensorSelectionPlan:
-    plan: TensorSelectionPlan = {}
-    for name, module in model.named_modules():
-        if isinstance(module, ColumnParallelLinear):
-            start = module.tp_rank * module.out_features_per_partition
-            end = start + module.out_features_per_partition
-            plan[f"{name}.weight"] = _slice_selection(dim=0, start=start, end=end)
-            if module.bias is not None:
-                plan[f"{name}.bias"] = _slice_selection(dim=0, start=start, end=end)
-        elif isinstance(module, RowParallelLinear):
-            start = module.tp_rank * module.in_features_per_partition
-            end = start + module.in_features_per_partition
-            plan[f"{name}.weight"] = _slice_selection(dim=1, start=start, end=end)
-        elif isinstance(module, TensorParallelRMSNorm) and module.weight is not None:
-            start = module.tp_rank * module.hidden_size_per_partition
-            end = start + module.hidden_size_per_partition
-            plan[f"{name}.weight"] = _slice_selection(dim=0, start=start, end=end)
-    return plan
diff --git a/diffsynth_engine/models/base.py b/diffsynth_engine/models/base.py
index 5b8537dc..dd9dcd7a 100644
--- a/diffsynth_engine/models/base.py
+++ b/diffsynth_engine/models/base.py
@@ -5,11 +5,6 @@
 from accelerate import init_empty_weights
 from diffusers.configuration_utils import ConfigMixin
 
-from diffsynth_engine.distributed.parallel_state import (
-    get_tensor_model_parallel_world_size,
-    is_tp_group_initialized,
-)
-from diffsynth_engine.layers.tensor_parallel import derive_tp_model_weight_load_plan
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import CONFIG_NAME
 from diffsynth_engine.utils.load_utils import load_model_weights
@@ -45,16 +40,6 @@ def from_pretrained(
         with init_empty_weights():
             model = cls.from_config(config_dict)
 
-        tensor_selection_plan = None
-        tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
-        if tp_size > 1:
-            tensor_selection_plan = derive_tp_model_weight_load_plan(model)
-            if not tensor_selection_plan:
-                raise RuntimeError(
-                    f"Tensor parallel world size is {tp_size}, but {cls.__name__} has no tensor-parallel submodules; "
-                    "the model is not TP-aware."
-                )
-
         # avoid precision loss
         if dtype is not None and dtype != torch.float32 and cls._keep_in_fp32_modules:
             state_dict = load_model_weights(
@@ -62,7 +47,6 @@ def from_pretrained(
                 subfolder,
                 device,
                 dtype=None,
-                tensor_selection_plan=tensor_selection_plan,
             )
             for k, v in state_dict.items():
                 if any(m in k.split(".") for m in cls._keep_in_fp32_modules):
@@ -75,7 +59,6 @@ def from_pretrained(
                 subfolder,
                 device,
                 dtype,
-                tensor_selection_plan=tensor_selection_plan,
             )
 
         # drop unexpected keys
diff --git a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
index 36b6a4be..895c2e31 100644
--- a/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
+++ b/diffsynth_engine/models/qwen_image/transformer_qwenimage.py
@@ -421,9 +421,7 @@ def __init__(
         self.attention_head_dim = attention_head_dim
         tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
         if num_attention_heads % tp_size != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})"
-            )
+            raise ValueError(f"num_attention_heads ({num_attention_heads}) must be divisible by tp_size ({tp_size})")
         self.heads = num_attention_heads // tp_size
 
         # Image stream projections
@@ -437,9 +435,15 @@ def __init__(
         self.to_out.append(nn.Dropout(dropout))
 
         # Text stream projections
-        self.add_q_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
-        self.add_k_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
-        self.add_v_proj = ColumnParallelLinear(dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False)
+        self.add_q_proj = ColumnParallelLinear(
+            dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False
+        )
+        self.add_k_proj = ColumnParallelLinear(
+            dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False
+        )
+        self.add_v_proj = ColumnParallelLinear(
+            dim, num_attention_heads * attention_head_dim, bias=True, gather_output=False
+        )
         self.to_add_out = RowParallelLinear(
             num_attention_heads * attention_head_dim, dim, bias=True, input_is_parallel=True
         )
diff --git a/diffsynth_engine/models/wan/transformer_wan.py b/diffsynth_engine/models/wan/transformer_wan.py
index 4095b8bb..b53aa4f6 100644
--- a/diffsynth_engine/models/wan/transformer_wan.py
+++ b/diffsynth_engine/models/wan/transformer_wan.py
@@ -110,12 +110,8 @@ def __init__(
 
         self.add_k_proj = self.add_v_proj = None
         if added_kv_proj_dim is not None:
-            self.add_k_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False
-            )
-            self.add_v_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False
-            )
+            self.add_k_proj = ColumnParallelLinear(added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False)
+            self.add_v_proj = ColumnParallelLinear(added_kv_proj_dim, self.inner_dim, bias=True, gather_output=False)
             self.norm_added_k = TensorParallelRMSNorm(self.inner_dim, eps=eps)
 
         forward_context = get_forward_context()
diff --git a/diffsynth_engine/pipelines/base.py b/diffsynth_engine/pipelines/base.py
index 1f511c09..73509cbe 100644
--- a/diffsynth_engine/pipelines/base.py
+++ b/diffsynth_engine/pipelines/base.py
@@ -17,9 +17,8 @@
     is_world_group_initialized,
 )
 from diffsynth_engine.forward_context import set_forward_context
-from diffsynth_engine.layers.tensor_parallel import derive_tp_model_weight_load_plan
 from diffsynth_engine.utils import logging
-from diffsynth_engine.utils.load_utils import fix_state_dict_key, load_model_weights
+from diffsynth_engine.utils.load_utils import fix_state_dict_key, load_model_weights, prepare_model_weights
 
 logger = logging.get_logger(__name__)
 
@@ -45,7 +44,7 @@ def compile_transformer_blocks(model: nn.Module) -> nn.Module:
         has_compiled_region = False
         for submodule in model.modules():
             if submodule.__class__.__name__ in repeated_blocks:
-                submodule.compile(dynamic=True, fullgraph=False)
+                submodule.compile()
                 has_compiled_region = True
 
         if not has_compiled_region:
@@ -86,21 +85,11 @@ def init_transformer(
                     fully_shard(submodule)
             fully_shard(model)
 
-        tensor_selection_plan = None
         tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
         if tp_size > 1:
-            tensor_selection_plan = derive_tp_model_weight_load_plan(model)
-            if not tensor_selection_plan:
-                raise RuntimeError(
-                    f"Tensor parallel world size is {tp_size}, but {model_cls.__name__} has no tensor-parallel submodules; "
-                    "the model is not TP-aware."
-                )
-
             num_heads = getattr(model.config, "num_attention_heads", None)
             if num_heads is not None and num_heads % tp_size != 0:
-                raise ValueError(
-                    f"num_attention_heads ({num_heads}) must be divisible by tp_degree ({tp_size})"
-                )
+                raise ValueError(f"num_attention_heads ({num_heads}) must be divisible by tp_degree ({tp_size})")
             if num_heads is not None and is_sp_group_initialized():
                 ulysses_size = get_ulysses_parallel_world_size()
                 heads_per_tp_partition = num_heads // tp_size
@@ -114,30 +103,24 @@ def init_transformer(
         load_device = "cpu" if use_fsdp else pipeline_config.device
         model_dtype = pipeline_config.model_dtype
         keep_in_fp32_modules = getattr(model_cls, "_keep_in_fp32_modules", None)
-        if model_dtype is not None and model_dtype != torch.float32 and keep_in_fp32_modules:
+        keep_in_fp32 = bool(model_dtype is not None and model_dtype != torch.float32 and keep_in_fp32_modules)
+        load_dtype = None if keep_in_fp32 else model_dtype
+        state_dict = prepare_model_weights(
+            model,
+            pipeline_config.model_path,
+            subfolder=subfolder,
+            device=load_device,
+            dtype=load_dtype,
+            broadcast_from_rank0=not use_fsdp,
+        )
+
+        if keep_in_fp32:
             # avoid precision loss: keep modules (time embedder, modulation, norms) in fp32
-            state_dict = load_model_weights(
-                pipeline_config.model_path,
-                subfolder=subfolder,
-                device=load_device,
-                dtype=None,
-                broadcast_from_rank0=not use_fsdp,
-                tensor_selection_plan=tensor_selection_plan,
-            )
             for k, v in state_dict.items():
                 if any(m in k.split(".") for m in keep_in_fp32_modules):
                     state_dict[k] = v.to(dtype=torch.float32)
                 else:
                     state_dict[k] = v.to(dtype=model_dtype)
-        else:
-            state_dict = load_model_weights(
-                pipeline_config.model_path,
-                subfolder=subfolder,
-                device=load_device,
-                dtype=model_dtype,
-                broadcast_from_rank0=not use_fsdp,
-                tensor_selection_plan=tensor_selection_plan,
-            )
 
         if use_fsdp:
             set_model_state_dict(
diff --git a/diffsynth_engine/utils/load_utils.py b/diffsynth_engine/utils/load_utils.py
index 9a008e4c..a0b19ed3 100644
--- a/diffsynth_engine/utils/load_utils.py
+++ b/diffsynth_engine/utils/load_utils.py
@@ -2,12 +2,20 @@
 import os
 import re
 import time
-from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional
 
 import torch
-
-from diffsynth_engine.distributed.parallel_state import get_global_rank, get_world_group, is_world_group_initialized
+import torch.nn as nn
+
+from diffsynth_engine.distributed.parallel_state import (
+    get_global_rank,
+    get_tensor_model_parallel_world_size,
+    get_world_group,
+    is_tp_group_initialized,
+    is_world_group_initialized,
+)
+from diffsynth_engine.layers.tensor_parallel.linear import ColumnParallelLinear, RowParallelLinear
+from diffsynth_engine.layers.tensor_parallel.norm import TensorParallelRMSNorm
 from diffsynth_engine.utils import logging
 from diffsynth_engine.utils.constants import (
     DIFFUSION_SAFETENSORS_INDEX_NAME,
@@ -16,22 +24,6 @@
     SAFETENSORS_WEIGHTS_NAME,
 )
 
-
-@dataclass(frozen=True)
-class TensorSlice:
-    dim: int
-    start: int
-    end: int
-
-
-@dataclass(frozen=True)
-class TensorSelection:
-    slices: Tuple[TensorSlice, ...]
-    contiguous: bool = True
-
-
-TensorSelectionPlan = Dict[str, TensorSelection]
-
 logger = logging.get_logger(__name__)
 
 try:
@@ -65,40 +57,6 @@ def load_safetensors(path: str, device: str = "cpu") -> Dict[str, Any]:
     return state_dict
 
 
-def apply_tensor_selection(
-    state_dict: Dict[str, Any],
-    tensor_selection_plan: TensorSelectionPlan,
-) -> Dict[str, Any]:
-    for key in list(state_dict.keys()):
-        selection = tensor_selection_plan.get(key)
-        if selection is None:
-            continue
-
-        value = state_dict[key]
-        if not isinstance(value, torch.Tensor):
-            continue
-
-        slices = [slice(None)] * value.dim()
-        for tensor_slice in selection.slices:
-            dim = tensor_slice.dim if tensor_slice.dim >= 0 else value.dim() + tensor_slice.dim
-            if dim < 0 or dim >= value.dim():
-                raise ValueError(
-                    f"Cannot slice '{key}' shape={tuple(value.shape)} along dim {tensor_slice.dim}: "
-                    f"tensor has {value.dim()} dimensions."
-                )
-            if tensor_slice.start < 0 or tensor_slice.end > value.shape[dim] or tensor_slice.start > tensor_slice.end:
-                raise ValueError(
-                    f"Invalid slice for '{key}' shape={tuple(value.shape)}: "
-                    f"dim={tensor_slice.dim}, start={tensor_slice.start}, end={tensor_slice.end}."
-                )
-            slices[dim] = slice(tensor_slice.start, tensor_slice.end)
-
-        selected = value[tuple(slices)]
-        state_dict[key] = selected.contiguous() if selection.contiguous else selected
-
-    return state_dict
-
-
 def _list_shard_files(path: str) -> list[str]:
     if not os.path.exists(path):
         raise FileNotFoundError(f"Model path not found: {path}")
@@ -121,10 +79,64 @@ def _list_shard_files(path: str) -> list[str]:
 
     with open(index_file, "r", encoding="utf-8") as file:
         index_dict = json.load(file)
-    shard_names = sorted(set(index_dict["weight_map"].values()))
-    if not shard_names:
+    shard_files = sorted(set(index_dict["weight_map"].values()))
+    if not shard_files:
         raise ValueError(f"Weight index {index_file} contains an empty weight_map")
-    return [os.path.join(path, name) for name in shard_names]
+    return [os.path.join(path, name) for name in shard_files]
+
+
+# tensor parallel
+
+
+def _slice_tensor(
+    tensor: torch.Tensor,
+    name: str,
+    span: tuple[int, int, int],
+) -> torch.Tensor:
+    dim, start, length = span
+    if dim < 0 or dim >= tensor.dim():
+        raise ValueError(f"Cannot shard '{name}' shape={tuple(tensor.shape)} along dim {dim}.")
+    if start < 0 or length < 0 or start + length > tensor.shape[dim]:
+        raise ValueError(
+            f"Invalid shard for '{name}' shape={tuple(tensor.shape)}: dim={dim}, start={start}, length={length}."
+        )
+
+    return tensor.narrow(dim, start, length).contiguous()
+
+
+def _slice_tensor_parallel_weights(
+    model: nn.Module,
+    state_dict: Dict[str, Any],
+) -> Dict[str, Any]:
+    tp_size = get_tensor_model_parallel_world_size() if is_tp_group_initialized() else 1
+    if tp_size <= 1:
+        return state_dict
+
+    spans: Dict[str, tuple[int, int, int]] = {}
+    for name, module in model.named_modules():
+        prefix = f"{name}." if name else ""
+
+        if isinstance(module, ColumnParallelLinear) and module.tp_size > 1:
+            start = module.tp_rank * module.out_features_per_partition
+            length = module.out_features_per_partition
+            spans[f"{prefix}weight"] = (0, start, length)
+            if module.bias is not None:
+                spans[f"{prefix}bias"] = (0, start, length)
+        elif isinstance(module, RowParallelLinear) and module.tp_size > 1:
+            start = module.tp_rank * module.in_features_per_partition
+            length = module.in_features_per_partition
+            spans[f"{prefix}weight"] = (1, start, length)
+        elif isinstance(module, TensorParallelRMSNorm) and module.tp_size > 1 and module.weight is not None:
+            start = module.tp_rank * module.hidden_size_per_partition
+            length = module.hidden_size_per_partition
+            spans[f"{prefix}weight"] = (0, start, length)
+
+    for name, tensor in state_dict.items():
+        span = spans.get(name)
+        if span is not None:
+            state_dict[name] = _slice_tensor(tensor, name, span)
+
+    return state_dict
 
 
 def load_model_weights(
@@ -133,27 +145,50 @@ def load_model_weights(
     device: Optional[str] = None,
     dtype: Optional[torch.dtype] = None,
     broadcast_from_rank0: bool = True,
-    tensor_selection_plan: Optional[TensorSelectionPlan] = None,
 ) -> Dict[str, Any]:
     world_group = get_world_group() if is_world_group_initialized() else None
     is_rank_zero = world_group is None or get_global_rank() == 0
+    model_path = os.path.join(model_path, subfolder) if subfolder is not None else model_path
 
-    resolved_path = os.path.join(model_path, subfolder) if subfolder is not None else model_path
-    shard_files = _list_shard_files(resolved_path)
+    state_dict: Dict[str, Any] = {}
+    for shard_file in _list_shard_files(model_path):
+        shard_dict = load_safetensors(shard_file, device=device or "cpu") if is_rank_zero else {}
+
+        if world_group is not None and broadcast_from_rank0:
+            shard_dict = world_group.broadcast_tensor_dict(shard_dict, src=0)
+
+        for name, tensor in shard_dict.items():
+            if isinstance(tensor, torch.Tensor):
+                shard_dict[name] = tensor.to(dtype=dtype, non_blocking=True)
+        state_dict.update(shard_dict)
+
+    return state_dict
+
+
+def prepare_model_weights(
+    model: nn.Module,
+    model_path: str,
+    subfolder: Optional[str] = None,
+    device: Optional[str] = None,
+    dtype: Optional[torch.dtype] = None,
+    broadcast_from_rank0: bool = True,
+) -> Dict[str, Any]:
+    world_group = get_world_group() if is_world_group_initialized() else None
+    is_rank_zero = world_group is None or get_global_rank() == 0
+    model_path = os.path.join(model_path, subfolder) if subfolder is not None else model_path
 
     state_dict: Dict[str, Any] = {}
-    for shard_file in shard_files:
-        shard_dict = load_safetensors(shard_file) if is_rank_zero else {}
+    for shard_file in _list_shard_files(model_path):
+        shard_dict = load_safetensors(shard_file, device=device or "cpu") if is_rank_zero else {}
 
         if world_group is not None and broadcast_from_rank0:
             shard_dict = world_group.broadcast_tensor_dict(shard_dict, src=0)
 
-        if tensor_selection_plan is not None:
-            shard_dict = apply_tensor_selection(shard_dict, tensor_selection_plan)
+        shard_dict = _slice_tensor_parallel_weights(model, shard_dict)
 
-        for key, value in shard_dict.items():
-            if isinstance(value, torch.Tensor):
-                shard_dict[key] = value.to(device=device, dtype=dtype, non_blocking=True)
+        for name, tensor in shard_dict.items():
+            if isinstance(tensor, torch.Tensor):
+                shard_dict[name] = tensor.to(dtype=dtype, non_blocking=True)
         state_dict.update(shard_dict)
 
     return state_dict

From 6c50f5b74abe0f3f901a715eb592f26d6dbc35f5 Mon Sep 17 00:00:00 2001
From: "zhuguoxuan.zgx" <zhuguoxuan.zgx@alibaba-inc.com>
Date: Thu, 30 Jul 2026 15:42:01 +0800
Subject: [PATCH 3/3] test(tp): add Qwen Image and Wan I2V cases

---
 .../test_qwen_image_edit_plus_2511.py         | 36 ++++++++++++++
 .../test_wan_22_image_to_video.py             | 47 +++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/tests/test_pipelines/test_qwen_image_edit_plus_2511.py b/tests/test_pipelines/test_qwen_image_edit_plus_2511.py
index 83dd4148..48383794 100644
--- a/tests/test_pipelines/test_qwen_image_edit_plus_2511.py
+++ b/tests/test_pipelines/test_qwen_image_edit_plus_2511.py
@@ -59,5 +59,41 @@ def test_multi_image_edit(self):
         self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image_edit_plus_multi_2511.png", threshold=0.97)
 
 
+class TestQwenImageEditPlusPipelineTP(ImageTestCase):
+    @classmethod
+    def setUpClass(cls):
+        model_path = fetch_model("Qwen/Qwen-Image-Edit-2511")
+        config = QwenImagePipelineConfig(
+            model_path=model_path,
+            parallelism=4,
+            tp_degree=4,
+            sp_ulysses_degree=1,
+            sp_ring_degree=1,
+        )
+        cls.engine = DiffSynthEngine.from_pretrained(config)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.engine.shutdown()
+        del cls.engine
+        torch.cuda.empty_cache()
+
+    def test_single_image_edit_tp(self):
+        input_image = self.get_input_image("qwen_image_edit_input.png")
+        prompt = "Replace '通义千问' with '呜哩AI'"
+        negative_prompt = " "
+
+        output = self.engine.generate(
+            image=input_image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            true_cfg_scale=4.0,
+            num_inference_steps=50,
+            generator=torch.Generator(device="cpu").manual_seed(42),
+        )
+        image = output.images[0]
+        self.assertImageEqualAndSaveFailed(image, "qwen_image/qwen_image_edit_plus_single_2511.png", threshold=0.99)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_pipelines/test_wan_22_image_to_video.py b/tests/test_pipelines/test_wan_22_image_to_video.py
index 57ae2584..a300cb60 100644
--- a/tests/test_pipelines/test_wan_22_image_to_video.py
+++ b/tests/test_pipelines/test_wan_22_image_to_video.py
@@ -58,5 +58,52 @@ def test_image_to_video(self):
         self.assertVideoMsSsimEqualAndSaveFailed(output_frames, "wan/wan_22_i2v.mp4", threshold=0.98, fps=16)
 
 
+class TestWan22ImageToVideoPipelineTP(VideoTestCase):
+    @classmethod
+    def setUpClass(cls):
+        model_path = fetch_model("Wan-AI/Wan2.2-I2V-A14B-Diffusers")
+        config = WanPipelineConfig(
+            model_path=model_path,
+            pipeline_class_name="WanImageToVideoPipeline",
+            parallelism=4,
+            tp_degree=4,
+            sp_ulysses_degree=1,
+            sp_ring_degree=1,
+        )
+        cls.engine = DiffSynthEngine.from_pretrained(config)
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.engine.shutdown()
+        del cls.engine
+        torch.cuda.empty_cache()
+
+    def test_image_to_video_tp(self):
+        image = self.get_input_image("wan_22_i2v_input.png")
+        max_area = 480 * 832
+        aspect_ratio = image.height / image.width
+        height = round(np.sqrt(max_area * aspect_ratio)) // 16 * 16
+        width = round(np.sqrt(max_area / aspect_ratio)) // 16 * 16
+        image = image.resize((width, height))
+
+        prompt = "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside."
+        negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+
+        video = self.engine.generate(
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=81,
+            guidance_scale=3.5,
+            num_inference_steps=40,
+            generator=torch.Generator(device="cpu").manual_seed(42),
+        )
+
+        output_frames = video.frames[0]
+        self.assertVideoMsSsimEqualAndSaveFailed(output_frames, "wan/wan_22_i2v.mp4", threshold=0.88, fps=16)
+
+
 if __name__ == "__main__":
     unittest.main()