[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",
+    "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/flux2/infer/transformer_infer.py

@@ -13,6 +13,14 @@ def __init__(self, config):
         self.infer_conditional = True
         self.clean_cuda_cache = self.config.get("clean_cuda_cache", False)
 
+        self.tp_group = None
+        self.tp_rank = 0
+        self.tp_size = 1
+        if self.config.get("tensor_parallel", False):
+            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)
+
         self.inner_dim = config.get("num_attention_heads", 24) * config.get("attention_head_dim", 64)
 
         if self.config.get("seq_parallel", False):
@@ -58,7 +66,7 @@ def infer_double_stream_block(
         image_rotary_emb,
         img_attn_hook=None,
     ):
-        heads = self.config["num_attention_heads"]
+        heads = self.config["num_attention_heads"] // self.tp_size
         head_dim = self.config["attention_head_dim"]
 
         (shift_msa, scale_msa, gate_msa), (shift_mlp, scale_mlp, gate_mlp) = self._split_double_modulation(temb_mod_img)
@@ -169,7 +177,7 @@ def infer_single_stream_block(
         image_rotary_emb,
         num_txt_tokens=0,
     ):
-        heads = self.config["num_attention_heads"]
+        heads = self.config["num_attention_heads"] // self.tp_size
         head_dim = self.config["attention_head_dim"]
 
         if encoder_hidden_states is not None:

lightx2v/models/networks/flux2/model.py

@@ -12,6 +12,7 @@
 from lightx2v.models.networks.flux2.weights.pre_weights import Flux2DevPreWeights, Flux2PreWeights
 from lightx2v.models.networks.flux2.weights.transformer_weights import Flux2TransformerWeights
 from lightx2v.utils.custom_compiler import compiled_method
+from lightx2v_platform.base import global_var
 
 
 class _Flux2TransformerModelBase(BaseTransformerModel):
@@ -28,6 +29,204 @@ def __init__(self, config, model_path, device):
         self._init_weights()
         self._init_infer()
 
+    def _rank_device(self):
+        ai_device = global_var.AI_DEVICE
+        if ai_device is None:
+            return torch.device("cpu")
+        if dist.is_initialized():
+            return torch.device(f"{ai_device}:{dist.get_rank()}")
+        return torch.device(ai_device)
+
+    def _sync_device(self):
+        ai_device = global_var.AI_DEVICE
+        device_module = getattr(torch, ai_device, None) if ai_device else None
+        if device_module is not None and hasattr(device_module, "synchronize"):
+            device_module.synchronize()
+
+    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)
+        if self.cpu_offload:
+            raise NotImplementedError("Flux2 tensor parallel weight loading does not support cpu_offload yet.")
+
+        global_src_rank = 0
+        target_device = self._rank_device()
+
+        if is_weight_loader:
+            processed_weight_dict = {}
+            meta_dict = {}
+            processed_bias_keys = 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:
+                    split_weights = self._split_weight_for_tp(key, tensor, self.tp_size)
+                    for rank_idx, split_weight in enumerate(split_weights):
+                        rank_key = f"{key}__tp_rank_{rank_idx}"
+                        processed_weight_dict[rank_key] = split_weight.contiguous()
+                    meta_dict[key] = {"shape": split_weights[0].shape, "dtype": split_weights[0].dtype, "is_tp": True}
+
+                    bias_key = key.replace(".weight", ".bias")
+                    if bias_key in weight_dict and split_type in ("col", "ff_fused_col", "single_fused_col"):
+                        bias_splits = self._split_bias_for_tp(bias_key, weight_dict[bias_key], split_type, self.tp_size)
+                        for rank_idx, split_bias in enumerate(bias_splits):
+                            processed_weight_dict[f"{bias_key}__tp_rank_{rank_idx}"] = split_bias.contiguous()
+                        meta_dict[bias_key] = {"shape": bias_splits[0].shape, "dtype": bias_splits[0].dtype, "is_tp": True}
+                        processed_bias_keys.add(bias_key)
+                elif key not in processed_bias_keys:
+                    processed_weight_dict[key] = tensor
+                    meta_dict[key] = {"shape": tensor.shape, "dtype": tensor.dtype, "is_tp": False}
+
+            obj_list = [meta_dict]
+            dist.broadcast_object_list(obj_list, src=global_src_rank)
+            synced_meta_dict = obj_list[0]
+            weight_dict = processed_weight_dict
+        else:
+            obj_list = [None]
+            dist.broadcast_object_list(obj_list, src=global_src_rank)
+            synced_meta_dict = obj_list[0]
+
+        distributed_weight_dict = {key: torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device) for key, meta in synced_meta_dict.items()}
+        dist.barrier()
+
+        for key in sorted(synced_meta_dict.keys()):
+            meta = synced_meta_dict[key]
+            if meta.get("is_tp", False):
+                for rank_idx in range(self.tp_size):
+                    if is_weight_loader:
+                        src_tensor = weight_dict[f"{key}__tp_rank_{rank_idx}"].to(target_device, non_blocking=True)
+                    else:
+                        src_tensor = torch.empty(meta["shape"], dtype=meta["dtype"], device=target_device)
+                    dist.broadcast(src_tensor, src=global_src_rank)
+                    if rank_idx == self.tp_rank:
+                        distributed_weight_dict[key].copy_(src_tensor, non_blocking=True)
+                    del src_tensor
+            else:
+                if is_weight_loader:
+                    distributed_weight_dict[key].copy_(weight_dict[key].to(target_device, non_blocking=True), non_blocking=True)
+                dist.broadcast(distributed_weight_dict[key], src=global_src_rank)
+
+        self._sync_device()
+        return distributed_weight_dict
+
+    def _get_split_type(self, key):
+        if ".norm_" in key:
+            return None
+        if key.endswith(".weight") and "single_transformer_blocks." in key and ".attn.to_qkv_mlp_proj." in key:
+            return "single_fused_col"
+        if key.endswith(".weight") and "single_transformer_blocks." in key and ".attn.to_out." in key:
+            return "single_fused_row"
+        if key.endswith(".weight") and (".ff.linear_in." in key or ".ff_context.linear_in." in key):
+            return "ff_fused_col"
+        col_patterns = (
+            ".attn.to_q.",
+            ".attn.to_k.",
+            ".attn.to_v.",
+            ".attn.add_q_proj.",
+            ".attn.add_k_proj.",
+            ".attn.add_v_proj.",
+        )
+        row_patterns = (
+            ".attn.to_out.0.",
+            ".attn.to_add_out.",
+            ".ff.linear_out.",
+            ".ff_context.linear_out.",
+        )
+        if any(pattern in key for pattern in col_patterns):
+            return "col"
+        if any(pattern in key for pattern in row_patterns):
+            return "row"
+        return None
+
+    def _split_bias_for_tp(self, key, bias, split_type, tp_size):
+        if split_type == "col":
+            assert bias.shape[0] % tp_size == 0, f"bias dimension ({bias.shape[0]}) must be divisible by tp_size ({tp_size}) for {key}"
+            return list(torch.chunk(bias, tp_size, dim=0))
+
+        if split_type == "ff_fused_col":
+            assert bias.shape[0] % 2 == 0, f"invalid fused SwiGLU bias dim for {key}: {bias.shape[0]}"
+            ffn_dim = bias.shape[0] // 2
+            assert ffn_dim % tp_size == 0, f"ffn_dim ({ffn_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            gate, up = torch.split(bias, [ffn_dim, ffn_dim], dim=0)
+            gate_chunks = torch.chunk(gate, tp_size, dim=0)
+            up_chunks = torch.chunk(up, tp_size, dim=0)
+            return [torch.cat([gate_chunks[rank_idx], up_chunks[rank_idx]], dim=0) for rank_idx in range(tp_size)]
+
+        if split_type == "single_fused_col":
+            inner_dim = self.config["num_attention_heads"] * self.config["attention_head_dim"]
+            ffn_dim_twice = bias.shape[0] - 3 * inner_dim
+            assert ffn_dim_twice > 0 and ffn_dim_twice % 2 == 0, f"invalid fused qkv/mlp bias dim for {key}: {bias.shape[0]}"
+            ffn_dim = ffn_dim_twice // 2
+            assert inner_dim % tp_size == 0, f"inner_dim ({inner_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            assert ffn_dim % tp_size == 0, f"ffn_dim ({ffn_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            q, k, v, mlp_1, mlp_2 = torch.split(bias, [inner_dim, inner_dim, inner_dim, ffn_dim, ffn_dim], dim=0)
+            chunks = [torch.chunk(part, tp_size, dim=0) for part in (q, k, v, mlp_1, mlp_2)]
+            return [torch.cat([part_chunks[rank_idx] for part_chunks in chunks], dim=0) for rank_idx in range(tp_size)]
+
+        raise ValueError(f"Unsupported Flux2 TP bias split type {split_type} for {key}")
+
+    def _split_weight_for_tp(self, key, weight, tp_size):
+        split_type = self._get_split_type(key)
+        if split_type is None:
+            return [weight] * tp_size
+
+        if split_type == "col":
+            assert weight.shape[0] % tp_size == 0, f"out_dim ({weight.shape[0]}) must be divisible by tp_size ({tp_size}) for {key}"
+            return list(torch.chunk(weight, tp_size, dim=0))
+
+        if split_type == "row":
+            assert weight.shape[1] % tp_size == 0, f"in_dim ({weight.shape[1]}) must be divisible by tp_size ({tp_size}) for {key}"
+            return list(torch.chunk(weight, tp_size, dim=1))
+
+        if split_type == "ff_fused_col":
+            assert weight.shape[0] % 2 == 0, f"invalid fused SwiGLU out_dim for {key}: {weight.shape[0]}"
+            ffn_dim = weight.shape[0] // 2
+            assert ffn_dim % tp_size == 0, f"ffn_dim ({ffn_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            gate, up = torch.split(weight, [ffn_dim, ffn_dim], dim=0)
+            gate_chunks = torch.chunk(gate, tp_size, dim=0)
+            up_chunks = torch.chunk(up, tp_size, dim=0)
+            return [torch.cat([gate_chunks[rank_idx], up_chunks[rank_idx]], dim=0) for rank_idx in range(tp_size)]
+
+        inner_dim = self.config["num_attention_heads"] * self.config["attention_head_dim"]
+        if split_type == "single_fused_col":
+            ffn_dim_twice = weight.shape[0] - 3 * inner_dim
+            assert ffn_dim_twice > 0 and ffn_dim_twice % 2 == 0, f"invalid fused qkv/mlp out_dim for {key}: {weight.shape[0]}"
+            ffn_dim = ffn_dim_twice // 2
+            assert inner_dim % tp_size == 0, f"inner_dim ({inner_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            assert ffn_dim % tp_size == 0, f"ffn_dim ({ffn_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            q, k, v, mlp_1, mlp_2 = torch.split(weight, [inner_dim, inner_dim, inner_dim, ffn_dim, ffn_dim], dim=0)
+            return [
+                torch.cat(
+                    [
+                        torch.chunk(q, tp_size, dim=0)[rank_idx],
+                        torch.chunk(k, tp_size, dim=0)[rank_idx],
+                        torch.chunk(v, tp_size, dim=0)[rank_idx],
+                        torch.chunk(mlp_1, tp_size, dim=0)[rank_idx],
+                        torch.chunk(mlp_2, tp_size, dim=0)[rank_idx],
+                    ],
+                    dim=0,
+                )
+                for rank_idx in range(tp_size)
+            ]
+
+        if split_type == "single_fused_row":
+            ffn_dim = weight.shape[1] - inner_dim
+            assert ffn_dim > 0, f"invalid fused output in_dim for {key}: {weight.shape[1]}"
+            assert inner_dim % tp_size == 0, f"inner_dim ({inner_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            assert ffn_dim % tp_size == 0, f"ffn_dim ({ffn_dim}) must be divisible by tp_size ({tp_size}) for {key}"
+            attn, mlp = torch.split(weight, [inner_dim, ffn_dim], dim=1)
+            return [
+                torch.cat(
+                    [
+                        torch.chunk(attn, tp_size, dim=1)[rank_idx],
+                        torch.chunk(mlp, tp_size, dim=1)[rank_idx],
+                    ],
+                    dim=1,
+                )
+                for rank_idx in range(tp_size)
+            ]
+
+        raise ValueError(f"Unsupported Flux2 TP split type {split_type} for {key}")
+
     def _init_infer(self):
         self.transformer_infer = self.transformer_infer_class(self.config)
         self.pre_infer = self.pre_infer_class(self.config)