feat(intel-xpu): enable distributed Wan tensor-parallel inference

configs/platforms/intel_xpu/dist_infer/wan22_ti2v_t2v_tp.json

@@ -0,0 +1,27 @@
+{
+    "infer_steps": 50,
+    "text_len": 512,
+    "target_video_length": 120,
+    "target_height": 704,
+    "target_width": 1280,
+    "cpu_offload": false,
+    "offload_granularity": "block",
+    "t5_cpu_offload": true,
+    "vae_cpu_offload": true,
+    "num_channels_latents": 48,
+    "vae_stride": [4, 16, 16],
+    "self_attn_1_type": "torch_sdpa",
+    "cross_attn_1_type": "torch_sdpa",
+    "cross_attn_2_type": "torch_sdpa",
+    "sample_guide_scale": 5.0,
+    "sample_shift": 5.0,
+    "enable_cfg": true,
+    "fps": 24,
+    "rope_type": "torch_real_rope",
+    "parallel": {
+        "seq_p_size": 1,
+        "tensor_p_size": 2,
+        "tp_load_mode": "local",
+        "cfg_p_size": 1
+    }
+}

lightx2v/common/ops/mm/mm_weight.py

@@ -2522,6 +2522,14 @@ def __init__(
         )
         self._row_split_bias = None
 
+    def _extract_row_split_bias(self, clone=False):
+        if self.split_dim != "row":
+            return
+        bias = getattr(self._mm, "bias", None)
+        if bias is not None:
+            self._row_split_bias = bias.clone() if clone else bias
+            self._mm.bias = None
+
     def load(self, weight_dict):
         """Load weights using internal MMWeight's load method.
 
@@ -2534,8 +2542,34 @@ def load(self, weight_dict):
         """
         self._mm.load(weight_dict)
         if self.split_dim == "row" and self.bias_name is not None and self.bias_name in weight_dict:
-            self._row_split_bias = self._mm.bias.clone()
-            self._mm.bias = None
+            # Preserve the original resident-weight behavior. Buffer-only
+            # modules do not have a materialized bias yet and are handled when
+            # load_state_dict fills the buffer.
+            self._extract_row_split_bias(clone=True)
+
+    def state_dict(self, destination=None):
+        return self._mm.state_dict(destination)
+
+    def load_state_dict(self, destination, block_index, adapter_block_index=None):
+        result = self._mm.load_state_dict(destination, block_index, adapter_block_index)
+        self._extract_row_split_bias()
+        return result
+
+    def load_state_dict_from_disk(self, block_index, adapter_block_index=None):
+        result = self._mm.load_state_dict_from_disk(block_index, adapter_block_index)
+        self._extract_row_split_bias()
+        return result
+
+    def to_cuda(self, non_blocking=False):
+        result = self._mm.to_cuda(non_blocking)
+        self._extract_row_split_bias()
+        return result
+
+    def to_cpu(self, non_blocking=False):
+        if self._row_split_bias is not None:
+            self._mm.bias = self._row_split_bias
+            self._row_split_bias = None
+        return self._mm.to_cpu(non_blocking)
 
     def apply(self, input_tensor):
         """Apply matrix multiplication with tensor parallel support."""

lightx2v/models/networks/wan/model.py

@@ -104,10 +104,42 @@ def _split_weight_for_tp(self, key, weight, tp_size):
             raise ValueError(f"Cannot split {key} shape {tuple(weight.shape)} across tensor parallel size {tp_size} on dimension {split_dim}")
         return list(torch.chunk(weight, tp_size, dim=split_dim))
 
+    def _should_load_weights(self):
+        if self.use_tp and self._use_local_tp_load():
+            return True
+        return super()._should_load_weights()
+
+    def _use_local_tp_load(self):
+        mode = self.config.get("parallel", {}).get("tp_load_mode", "broadcast")
+        if mode not in ("broadcast", "local"):
+            raise ValueError(f"Unsupported Wan TP load mode: {mode!r}; expected 'broadcast' or 'local'.")
+        return mode == "local"
+
+    def _shard_weights_locally(self, weight_dict):
+        local_weights = {}
+        processed_bias = set()
+        storage_device = self.device
+        for key, tensor in weight_dict.items():
+            split_type = self._get_split_type(key)
+            if key.endswith(".weight") and split_type is not None:
+                local_weights[key] = self._split_weight_for_tp(key, tensor, self.tp_size)[self.tp_rank].contiguous().to(storage_device)
+                bias_key = key.replace(".weight", ".bias")
+                if bias_key in weight_dict and split_type == "col":
+                    local_weights[bias_key] = self._split_bias_for_tp(weight_dict[bias_key], split_type, self.tp_size)[self.tp_rank].contiguous().to(storage_device)
+                    processed_bias.add(bias_key)
+            elif key not in processed_bias:
+                local_weights[key] = tensor.to(storage_device)
+        return local_weights
+
     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._use_local_tp_load():
+            if not is_weight_loader or weight_dict is None:
+                raise RuntimeError("Wan TP local loading requires every rank to load the shared checkpoint")
+            return self._shard_weights_locally(weight_dict)
+
         src_rank = 0
         target_device = self._rank_device()
 

lightx2v_platform/base/intel_xpu.py

@@ -7,8 +7,13 @@
 - Distributed training with Intel oneCCL backend
 """
 
+import os
+import platform
+
 import torch
+import torch.distributed as dist
 from loguru import logger
+from packaging.version import Version
 
 from lightx2v_platform.registry_factory import PLATFORM_DEVICE_REGISTER
 
@@ -48,15 +53,22 @@ def get_device() -> str:
         """Get the device type string. Returns 'xpu' for Intel XPU."""
         return "xpu"
 
-    # @staticmethod
-    # def init_parallel_env():
-    #     """
-    #     Initialize distributed parallel environment for Intel XPU.
-
-    #     Uses Intel oneCCL backend for distributed training.
-    #     """
-    #     dist.init_process_group(backend="ccl")
-    #     torch.xpu.set_device(dist.get_rank())
+    @staticmethod
+    def init_parallel_env():
+        """Initialize a single-node distributed environment for Intel XPU."""
+        if (
+            platform.system() != "Linux"
+            and Version(torch.__version__) < Version("2.10.0+xpu")
+        ):
+            raise RuntimeError(
+                "Intel XPU distributed initialization on non-Linux systems requires "
+                "PyTorch >= 2.10.0+xpu. "
+                f"Found PyTorch {torch.__version__}."
+            )
+
+        local_rank = int(os.environ["LOCAL_RANK"])
+        torch.xpu.set_device(local_rank)
+        dist.init_process_group(backend="xccl")
 
 
 # Register alias "xpu" for backward compatibility

scripts/platforms/intel_xpu/dist_infer/run_wan22_ti2v_t2v_tp.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
+REPO_ROOT=$(cd -- "${SCRIPT_DIR}/../../../.." && pwd)
+
+lightx2v_path=${LIGHTX2V_PATH:-${REPO_ROOT}}
+model_path=${MODEL_PATH:-/llm/models/Wan2.2-TI2V-5B}
+config_json=${CONFIG_JSON:-${lightx2v_path}/configs/platforms/intel_xpu/dist_infer/wan22_ti2v_t2v_tp.json}
+output_path=${OUTPUT_PATH:-${lightx2v_path}/save_results/output_wan22_ti2v_t2v_tp.mp4}
+
+export LIGHTX2V_XPU_DEVICE_MAP=${LIGHTX2V_XPU_DEVICE_MAP:-0,1}
+export PYTHONFAULTHANDLER=${PYTHONFAULTHANDLER:-1}
+export PYTHONPATH=${PYTHONPATH:-}
+
+source "${lightx2v_path}/scripts/base/base.sh"
+mkdir -p "$(dirname -- "${output_path}")"
+
+torchrun --standalone --nproc_per_node=2 -m lightx2v.infer \
+  --model_cls wan2.2 \
+  --task t2v \
+  --model_path "${model_path}" \
+  --config_json "${config_json}" \
+  --prompt "${PROMPT:-Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage}" \
+  --negative_prompt "${NEGATIVE_PROMPT:-色调艳丽，过曝，静态，细节模糊不清，字幕，低质量，JPEG压缩残留，畸形，多余的手指，杂乱的背景}" \
+  --seed "${SEED:-42}" \
+  --save_result_path "${output_path}"