[feat]: add tp by mixin，support flux2-dev，wan2.2

configs/flux2/flux2_dev_tp.json

@@ -0,0 +1,16 @@
+{
+    "model_cls": "flux2_dev",
+    "task": "t2i",
+    "infer_steps": 50,
+    "sample_guide_scale": 4.0,
+    "vae_scale_factor": 16,
+    "feature_caching": "None",
+    "enable_cfg": false,
+    "patch_size": 2,
+    "tokenizer_max_length": 512,
+    "rope_type": "flashinfer_rope",
+    "text_encoder_out_layers": [10, 20, 30],
+    "parallel": {
+        "tensor_p_size": 4
+    }
+}

configs/wan22/wan_moe_t2v_tp.json

@@ -0,0 +1,24 @@
+{
+    "infer_steps": 40,
+    "target_video_length": 81,
+    "text_len": 512,
+    "target_height": 720,
+    "target_width": 1280,
+    "self_attn_1_type": "flash_attn3",
+    "cross_attn_1_type": "flash_attn3",
+    "cross_attn_2_type": "flash_attn3",
+    "sample_guide_scale": [
+        4.0,
+        3.0
+    ],
+    "sample_shift": 12.0,
+    "enable_cfg": true,
+    "cpu_offload": false,
+    "t5_cpu_offload": false,
+    "vae_cpu_offload": false,
+    "boundary": 0.875,
+    "use_tiling_vae": true,
+    "parallel": {
+        "tensor_p_size": 4
+    }
+}

lightx2v/models/networks/base_model.py

@@ -91,6 +91,19 @@ def __init__(self, model_path, config, device, model_type=None, lora_path=None,
         self.dit_quantized = self.config.get("dit_quantized", False)
         if self.dit_quantized:
             self._check_dit_quantized()
+        self._init_tensor_parallel()
+
+    def _init_tensor_parallel(self):
+        if self.config.get("tensor_parallel", False):
+            self.use_tp = True
+            self.tp_group = self.config.get("device_mesh").get_group(mesh_dim="tensor_p")
+            self.tp_rank = dist.get_rank(self.tp_group)
+            self.tp_size = dist.get_world_size(self.tp_group)
+        else:
+            self.use_tp = False
+            self.tp_group = None
+            self.tp_rank = 0
+            self.tp_size = 1
 
     def _check_dit_quantized(self):
         """Check if the model is quantized.
@@ -291,13 +304,10 @@ def _should_load_weights(self):
             # Single GPU mode
             return True
         elif dist.is_initialized():
-            if self.config.get("load_from_rank0", False):
-                # Multi-GPU mode, only rank 0 loads
-                if dist.get_rank() == 0:
-                    logger.info(f"Loading weights from {self.model_path}")
-                    return True
-            else:
-                return True
+            if self.use_tp or self.config.get("load_from_rank0", False):
+                # TP or explicit rank0-only loading: only rank 0 reads from disk
+                return dist.get_rank() == 0
+            return True
         return False
 
     def _apply_weights(self, weight_dict=None):
@@ -424,6 +434,10 @@ def _load_ckpt(self, unified_dtype, sensitive_layer):
         Returns:
             dict: Dictionary of weights
         """
+        _tp_dev = self.use_tp
+        if _tp_dev:
+            _saved_device, self.device = self.device, torch.device("cpu")
+
         if self.config.get("dit_original_ckpt", None):
             safetensors_path = self.config["dit_original_ckpt"]
         else:
@@ -456,6 +470,8 @@ def _load_ckpt(self, unified_dtype, sensitive_layer):
             file_weights = self._load_safetensor_to_dict(file_path, unified_dtype, sensitive_layer)
             weight_dict.update(file_weights)
 
+        if _tp_dev:
+            self.device = _saved_device
         return weight_dict
 
     def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
@@ -468,6 +484,10 @@ def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
         Returns:
             dict: Dictionary of weights
         """
+        _tp_dev = self.use_tp
+        if _tp_dev:
+            _saved_device, self.device = self.device, torch.device("cpu")
+
         remove_keys = self.remove_keys if hasattr(self, "remove_keys") else []
 
         if self.config.get("dit_quantized_ckpt", None):
@@ -485,6 +505,8 @@ def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
             else:
                 gguf_path = safetensors_path
             weight_dict = self._load_gguf_ckpt(gguf_path)
+            if _tp_dev:
+                self.device = _saved_device
             return weight_dict
 
         if os.path.isdir(safetensors_path):
@@ -527,6 +549,8 @@ def _load_quant_ckpt(self, unified_dtype, sensitive_layer):
                 for k, v in calib_data["absmax"].items():
                     weight_dict[k.replace(".weight", ".input_absmax")] = v.to(self.device)
 
+        if _tp_dev:
+            self.device = _saved_device
         return weight_dict
 
     def _load_gguf_ckpt(self, gguf_path):

lightx2v/models/networks/wan/infer/transformer_infer.py

@@ -1,4 +1,5 @@
 import torch
+import torch.distributed as dist
 from loguru import logger
 
 from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer
@@ -23,7 +24,12 @@ def __init__(self, config):
         self.blocks_num = config["num_layers"]
         self.phases_num = 3
         self.has_post_adapter = False
-        self.num_heads = config["num_heads"]
+        if config.get("tensor_parallel", False):
+            _tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p")
+            _tp_size = dist.get_world_size(_tp_group)
+        else:
+            _tp_size = 1
+        self.num_heads = config["num_heads"] // _tp_size
         self.head_dim = config["dim"] // config["num_heads"]
         self.window_size = config.get("window_size", (-1, -1))
         self.parallel_attention = None

lightx2v/models/networks/wan/model.py

@@ -28,6 +28,7 @@
 from lightx2v.utils.custom_compiler import compiled_method
 from lightx2v.utils.envs import *
 from lightx2v.utils.utils import *
+from lightx2v_platform.base.global_var import AI_DEVICE
 
 
 class WanModel(BaseTransformerModel):
@@ -52,6 +53,104 @@ def __init__(self, model_path, config, device, model_type="wan2.1", lora_path=No
         self._init_weights()
         self._init_infer()
 
+    # ------------------------------------------------------------------ TP --
+    def _rank_device(self):
+        if dist.is_initialized():
+            return torch.device(f"{AI_DEVICE}:{dist.get_rank()}")
+        return torch.device(AI_DEVICE)
+
+    def _get_split_type(self, key):
+        if not key.endswith(".weight"):
+            return None
+        col_infixes = (
+            ".self_attn.q.",
+            ".self_attn.k.",
+            ".self_attn.v.",
+            ".self_attn.norm_q.",
+            ".self_attn.norm_k.",
+            ".cross_attn.q.",
+            ".cross_attn.k.",
+            ".cross_attn.v.",
+            ".cross_attn.norm_q.",
+            ".cross_attn.norm_k.",
+            ".cross_attn.k_img.",
+            ".cross_attn.v_img.",
+        )
+        row_infixes = (".self_attn.o.", ".cross_attn.o.", ".ffn.2.")
+        if any(s in key for s in col_infixes):
+            return "col"
+        if any(s in key for s in row_infixes):
+            return "row"
+        if ".ffn.0." in key:
+            return "col"
+        return None
+
+    def _split_bias_for_tp(self, bias, split_type, tp_size):
+        if split_type == "col":
+            return list(torch.chunk(bias, tp_size, dim=0))
+        raise ValueError(f"Unsupported bias split_type: {split_type}")
+
+    def _split_weight_for_tp(self, key, weight, tp_size):
+        split_type = self._get_split_type(key)
+        if split_type == "col":
+            return list(torch.chunk(weight, tp_size, dim=0))
+        if split_type == "row":
+            return list(torch.chunk(weight, tp_size, dim=1))
+        raise ValueError(f"Unknown split_type for {key}")
+
+    def _load_weights_from_rank0(self, weight_dict, is_weight_loader):
+        if not self.use_tp:
+            return super()._load_weights_from_rank0(weight_dict, is_weight_loader)
+
+        src_rank = 0
+        target_device = self._rank_device()
+
+        if is_weight_loader:
+            processed, meta, processed_bias = {}, {}, set()
+            for key, tensor in weight_dict.items():
+                split_type = self._get_split_type(key)
+                if key.endswith(".weight") and split_type is not None:
+                    shards = self._split_weight_for_tp(key, tensor, self.tp_size)
+                    for r, shard in enumerate(shards):
+                        processed[f"{key}__tp_{r}"] = shard.contiguous()
+                    meta[key] = {"shape": shards[0].shape, "dtype": shards[0].dtype, "is_tp": True}
+                    bias_key = key.replace(".weight", ".bias")
+                    if bias_key in weight_dict and split_type == "col":
+                        bias_shards = self._split_bias_for_tp(weight_dict[bias_key], split_type, self.tp_size)
+                        for r, shard in enumerate(bias_shards):
+                            processed[f"{bias_key}__tp_{r}"] = shard.contiguous()
+                        meta[bias_key] = {"shape": bias_shards[0].shape, "dtype": bias_shards[0].dtype, "is_tp": True}
+                        processed_bias.add(bias_key)
+                elif key not in processed_bias:
+                    processed[key] = tensor
+                    meta[key] = {"shape": tensor.shape, "dtype": tensor.dtype, "is_tp": False}
+            obj_list = [meta]
+        else:
+            obj_list = [None]
+
+        dist.broadcast_object_list(obj_list, src=src_rank)
+        synced_meta = obj_list[0]
+
+        distributed = {k: torch.empty(m["shape"], dtype=m["dtype"], device=target_device) for k, m in synced_meta.items()}
+
+        for key in sorted(synced_meta.keys()):
+            m = synced_meta[key]
+            if m["is_tp"]:
+                for r in range(self.tp_size):
+                    buf = processed[f"{key}__tp_{r}"].to(target_device) if is_weight_loader else torch.empty(m["shape"], dtype=m["dtype"], device=target_device)
+                    dist.broadcast(buf, src=src_rank, group=self.tp_group)
+                    if r == self.tp_rank:
+                        distributed[key].copy_(buf)
+                    del buf
+            else:
+                if is_weight_loader:
+                    distributed[key].copy_(processed[key].to(target_device))
+                dist.broadcast(distributed[key], src=src_rank, group=self.tp_group)
+
+        return distributed
+
+    # ------------------------------------------------------------------ TP --
+
     def _init_infer_class(self):
         self.pre_infer_class = WanPreInfer
         self.post_infer_class = WanPostInfer