Add tensor parallel and torch compile (#262)

* feat(tp): add tensor parallel and compile support

* refactor model weights loading

* test(tp): add Qwen Image and Wan I2V cases

---------

Co-authored-by: zhuguoxuan.zgx <zhuguoxuan.zgx@alibaba-inc.com>

Author: nianniann

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

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,37 +65,37 @@ 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"
-
-    cfg_degree = 2 if config.use_cfg_parallel else 1
+    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  # 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}")
+    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}")
+
+    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)
+    )
 
-    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"
+    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}) = {parallel_degree}"
         )
-        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})"
-    )
+
+    if config.tp_degree > 1 and config.use_fsdp:
+        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"

diffsynth_engine/layers/tensor_parallel/__init__.py

@@ -0,0 +1,11 @@
+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.norm import TensorParallelRMSNorm
+
+__all__ = [
+    "ColumnParallelLinear",
+    "ColumnParallelGELU",
+    "RowParallelLinear",
+    "TensorParallelRMSNorm",
+    "TPFeedForward",
+]

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

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

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