diff --git a/configs/flux2/flux2_dev_tp.json b/configs/flux2/flux2_dev_tp.json new file mode 100644 index 000000000..0ed626cfa --- /dev/null +++ b/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 + } +} diff --git a/configs/wan22/wan_moe_t2v_tp.json b/configs/wan22/wan_moe_t2v_tp.json new file mode 100755 index 000000000..88393761d --- /dev/null +++ b/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 + } +} diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 4145a3ec1..44d98124c 100644 --- a/lightx2v/models/networks/base_model.py +++ b/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): diff --git a/lightx2v/models/networks/wan/infer/transformer_infer.py b/lightx2v/models/networks/wan/infer/transformer_infer.py index 27dbce1c4..c6abe71d8 100755 --- a/lightx2v/models/networks/wan/infer/transformer_infer.py +++ b/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 diff --git a/lightx2v/models/networks/wan/model.py b/lightx2v/models/networks/wan/model.py index b80b3a7da..911da8258 100755 --- a/lightx2v/models/networks/wan/model.py +++ b/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 diff --git a/lightx2v/models/networks/wan/weights/transformer_weights.py b/lightx2v/models/networks/wan/weights/transformer_weights.py index 563fd74f6..d718b3d91 100755 --- a/lightx2v/models/networks/wan/weights/transformer_weights.py +++ b/lightx2v/models/networks/wan/weights/transformer_weights.py @@ -1,4 +1,5 @@ import torch +import torch.distributed as dist from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.models.networks.wan.infer.utils import WanCausalRope # noqa: F401 @@ -12,6 +13,39 @@ ) +def _mm_weight(config, weight_name, bias_name, split_dim=None, create_cuda_buffer=False, create_cpu_buffer=False, lazy_load=False, lazy_load_file=None, lora_prefix="", lora_path=""): + mm_type = config.get("dit_quant_scheme", "Default") + if config.get("do_mm_calib", False): + mm_type = "Calib" + if config.get("tensor_parallel", False) and split_dim is not None: + tp_group = config["device_mesh"].get_group(mesh_dim="tensor_p") + return MM_WEIGHT_REGISTER["TensorParallel"]( + weight_name=weight_name, + bias_name=bias_name, + mm_type=mm_type, + tp_group=tp_group, + tp_rank=dist.get_rank(tp_group), + tp_size=dist.get_world_size(tp_group), + split_dim=split_dim, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=lazy_load, + lazy_load_file=lazy_load_file, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + return MM_WEIGHT_REGISTER[mm_type]( + weight_name, + bias_name, + create_cuda_buffer, + create_cpu_buffer, + lazy_load, + lazy_load_file, + lora_prefix=lora_prefix, + lora_path=lora_path, + ) + + class WanTransformerWeights(WeightModule): def __init__(self, config, lazy_load_path=None, lora_path=None): super().__init__() @@ -285,55 +319,63 @@ def __init__( LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) + p = f"{block_prefix}.{self.block_index}" self.add_module( "self_attn_q", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.self_attn.q.weight", - f"{block_prefix}.{self.block_index}.self_attn.q.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{p}.self_attn.q.weight", + f"{p}.self_attn.q.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) - self.add_module( "self_attn_k", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.self_attn.k.weight", - f"{block_prefix}.{self.block_index}.self_attn.k.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{p}.self_attn.k.weight", + f"{p}.self_attn.k.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "self_attn_v", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.self_attn.v.weight", - f"{block_prefix}.{self.block_index}.self_attn.v.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{p}.self_attn.v.weight", + f"{p}.self_attn.v.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "self_attn_o", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.self_attn.o.weight", - f"{block_prefix}.{self.block_index}.self_attn.o.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{p}.self_attn.o.weight", + f"{p}.self_attn.o.bias", + split_dim="row", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), @@ -520,54 +562,63 @@ def __init__( lora_path=lora_path, ), ) + cp = f"{block_prefix}.{self.block_index}" self.add_module( "cross_attn_q", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.q.weight", - f"{block_prefix}.{self.block_index}.cross_attn.q.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.q.weight", + f"{cp}.cross_attn.q.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "cross_attn_k", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.k.weight", - f"{block_prefix}.{self.block_index}.cross_attn.k.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.k.weight", + f"{cp}.cross_attn.k.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "cross_attn_v", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.v.weight", - f"{block_prefix}.{self.block_index}.cross_attn.v.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.v.weight", + f"{cp}.cross_attn.v.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "cross_attn_o", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.o.weight", - f"{block_prefix}.{self.block_index}.cross_attn.o.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.o.weight", + f"{cp}.cross_attn.o.bias", + split_dim="row", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), @@ -601,26 +652,30 @@ def __init__( if self.config["task"] in ["i2v", "flf2v", "animate", "s2v", "rs2v"] and self.config.get("use_image_encoder", True) and self.config["model_cls"] != "wan2.1_sf_mtxg2": self.add_module( "cross_attn_k_img", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.k_img.weight", - f"{block_prefix}.{self.block_index}.cross_attn.k_img.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.k_img.weight", + f"{cp}.cross_attn.k_img.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "cross_attn_v_img", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.cross_attn.v_img.weight", - f"{block_prefix}.{self.block_index}.cross_attn.v_img.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{cp}.cross_attn.v_img.weight", + f"{cp}.cross_attn.v_img.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), @@ -668,28 +723,33 @@ def __init__( LN_WEIGHT_REGISTER[config.get("layer_norm_type", "torch")](), ) + fp = f"{block_prefix}.{self.block_index}" self.add_module( "ffn_0", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.ffn.0.weight", - f"{block_prefix}.{self.block_index}.ffn.0.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{fp}.ffn.0.weight", + f"{fp}.ffn.0.bias", + split_dim="col", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), ) self.add_module( "ffn_2", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{block_prefix}.{self.block_index}.ffn.2.weight", - f"{block_prefix}.{self.block_index}.ffn.2.bias", - create_cuda_buffer, - create_cpu_buffer, - self.lazy_load, - self.lazy_load_file, + _mm_weight( + config, + f"{fp}.ffn.2.weight", + f"{fp}.ffn.2.bias", + split_dim="row", + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + lazy_load=self.lazy_load, + lazy_load_file=self.lazy_load_file, lora_prefix=block_prefix, lora_path=lora_path, ), diff --git a/lightx2v/models/runners/default_runner.py b/lightx2v/models/runners/default_runner.py index 0ae53d6a1..e2ff9a03c 100755 --- a/lightx2v/models/runners/default_runner.py +++ b/lightx2v/models/runners/default_runner.py @@ -11,6 +11,7 @@ from requests.exceptions import RequestException from lightx2v.models.runners.base_runner import BaseRunner +from lightx2v.models.runners.tp_runner_mixin import TPRunnerMixin from lightx2v.server.metrics import monitor_cli from lightx2v.utils.envs import * from lightx2v.utils.generate_task_id import generate_task_id @@ -83,7 +84,7 @@ def resize_image(img, resize_mode="adaptive", resolution="480p", bucket_shape=No return cropped_img, target_h, target_w -class DefaultRunner(BaseRunner): +class DefaultRunner(TPRunnerMixin, BaseRunner): def __init__(self, config): super().__init__(config) self.has_prompt_enhancer = False diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py index ce22b7e38..032adf932 100644 --- a/lightx2v/models/runners/flux2/flux2_runner.py +++ b/lightx2v/models/runners/flux2/flux2_runner.py @@ -59,79 +59,6 @@ def load_model(self): self.vae = self.load_vae() self.model = self.load_transformer() - def _use_tp_rank0_io(self): - return self.config.get("tensor_parallel", False) and dist.is_initialized() - - def _is_rank0(self): - return not dist.is_initialized() or dist.get_rank() == 0 - - def _rank_device(self): - if dist.is_initialized(): - return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") - return torch.device(AI_DEVICE) - - def _tensor_meta_tree(self, obj): - if torch.is_tensor(obj): - return {"__tensor__": True, "shape": tuple(obj.shape), "dtype": obj.dtype} - if isinstance(obj, dict): - return {"__dict__": [(key, self._tensor_meta_tree(value)) for key, value in obj.items()]} - if isinstance(obj, list): - return {"__list__": [self._tensor_meta_tree(value) for value in obj]} - if isinstance(obj, tuple): - return {"__tuple__": [self._tensor_meta_tree(value) for value in obj]} - return {"__object__": obj} - - def _materialize_meta_tree(self, meta, device): - if meta.get("__tensor__", False): - return torch.empty(meta["shape"], dtype=meta["dtype"], device=device) - if "__dict__" in meta: - return {key: self._materialize_meta_tree(value, device) for key, value in meta["__dict__"]} - if "__list__" in meta: - return [self._materialize_meta_tree(value, device) for value in meta["__list__"]] - if "__tuple__" in meta: - return tuple(self._materialize_meta_tree(value, device) for value in meta["__tuple__"]) - return meta["__object__"] - - def _broadcast_tensor_tree(self, obj, src_obj, src=0): - if torch.is_tensor(obj): - if self._is_rank0(): - obj = src_obj.to(self._rank_device(), non_blocking=True) - dist.broadcast(obj, src=src) - return obj - if isinstance(obj, dict): - return {key: self._broadcast_tensor_tree(value, src_obj[key] if self._is_rank0() else None, src=src) for key, value in obj.items()} - if isinstance(obj, list): - return [self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)] - if isinstance(obj, tuple): - return tuple(self._broadcast_tensor_tree(value, src_obj[idx] if self._is_rank0() else None, src=src) for idx, value in enumerate(obj)) - return obj - - def _broadcast_rank0_payload(self, payload): - if not self._use_tp_rank0_io(): - return payload - - meta_list = [self._tensor_meta_tree(payload) if self._is_rank0() else None] - dist.broadcast_object_list(meta_list, src=0) - payload_tree = payload if self._is_rank0() else self._materialize_meta_tree(meta_list[0], self._rank_device()) - payload_tree = self._broadcast_tensor_tree(payload_tree, payload if self._is_rank0() else None, src=0) - dist.barrier() - return payload_tree - - def _load_rank0_text_encoder(self): - if not self._is_rank0(): - return - if self.text_encoders is None or self.config.get("lazy_load", False) or self.config.get("unload_modules", False): - self.text_encoders = self.load_text_encoder() - - def _unload_rank0_text_encoder(self): - if not self._is_rank0(): - return - if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): - del self.text_encoders - self.text_encoders = None - torch_device_module.empty_cache() - gc.collect() - def _load_rank0_vae(self): if not self._is_rank0(): return diff --git a/lightx2v/models/runners/tp_runner_mixin.py b/lightx2v/models/runners/tp_runner_mixin.py new file mode 100644 index 000000000..c883c0e7b --- /dev/null +++ b/lightx2v/models/runners/tp_runner_mixin.py @@ -0,0 +1,137 @@ +""" +Mixin that gives any Runner tensor-parallel rank-0 I/O support: + + - rank predicates (_use_tp_rank0_io, _is_rank0, _rank_device) + - broadcast any nested dict/list/tuple-of-tensors from rank 0 to all TP ranks + - rank-0-only text-encoder load / unload helpers +""" + +import gc + +import torch +import torch.distributed as dist + +from lightx2v_platform.base.global_var import AI_DEVICE + +torch_device_module = getattr(torch, AI_DEVICE) + +# --------------------------------------------------------------------------- +# Module-level helpers for the broadcast primitive. +# These are pure functions (no Runner state needed), so they live at module +# level rather than as instance methods. +# --------------------------------------------------------------------------- + +_TENSOR_TAG = "__T__" +_OBJECT_TAG = "__O__" + + +def _pack(obj): + """Recursively replace tensors with (shape, dtype) metadata stubs.""" + if torch.is_tensor(obj): + return {_TENSOR_TAG: (tuple(obj.shape), obj.dtype)} + if isinstance(obj, dict): + return {k: _pack(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + packed = [_pack(v) for v in obj] + return packed if isinstance(obj, list) else tuple(packed) + return {_OBJECT_TAG: obj} + + +def _alloc(meta, device): + """Allocate empty tensors to mirror the structure returned by _pack.""" + if isinstance(meta, dict): + if _TENSOR_TAG in meta: + shape, dtype = meta[_TENSOR_TAG] + return torch.empty(shape, dtype=dtype, device=device) + if _OBJECT_TAG in meta: + return meta[_OBJECT_TAG] + return {k: _alloc(v, device) for k, v in meta.items()} + if isinstance(meta, list): + return [_alloc(v, device) for v in meta] + if isinstance(meta, tuple): + return tuple(_alloc(v, device) for v in meta) + return meta + + +def _broadcast(skeleton, source, src_rank, rank0, device): + """Fill skeleton with tensors broadcast from rank src_rank. + + skeleton: pre-allocated structure (from _alloc) on every non-rank0 rank; + on rank0 it is the original source. + source: original payload on rank0, None on other ranks. + """ + if torch.is_tensor(skeleton): + if rank0: + skeleton = source.to(device, non_blocking=True) + dist.broadcast(skeleton, src=src_rank) + return skeleton + if isinstance(skeleton, dict): + return {k: _broadcast(skeleton[k], source[k] if rank0 else None, src_rank, rank0, device) for k in skeleton} + if isinstance(skeleton, list): + return [_broadcast(skeleton[i], source[i] if rank0 else None, src_rank, rank0, device) for i in range(len(skeleton))] + if isinstance(skeleton, tuple): + return tuple(_broadcast(skeleton[i], source[i] if rank0 else None, src_rank, rank0, device) for i in range(len(skeleton))) + return skeleton + + +# --------------------------------------------------------------------------- +# Mixin +# --------------------------------------------------------------------------- + + +class TPRunnerMixin: + """Rank-0 I/O helpers for tensor-parallel runners. + + Mix in before DefaultRunner / BaseRunner so that any runner subclass + automatically gets TP support without touching its own load_model or + run_input_encoder logic. + """ + + # --- predicates --------------------------------------------------------- + + def _use_tp_rank0_io(self): + return self.config.get("tensor_parallel", False) and dist.is_initialized() + + def _is_rank0(self): + return not dist.is_initialized() or dist.get_rank() == 0 + + def _rank_device(self): + if dist.is_initialized(): + return torch.device(f"{AI_DEVICE}:{dist.get_rank()}") + return torch.device(AI_DEVICE) + + # --- broadcast ---------------------------------------------------------- + + def _broadcast_rank0_payload(self, payload): + """Broadcast a nested dict/list/tuple-of-tensors from rank 0 to all ranks. + + Uses broadcast_object_list for the metadata (fast, zero-copy for small + dicts) and dist.broadcast for each tensor buffer (direct NCCL transfer). + """ + if not self._use_tp_rank0_io(): + return payload + + device = self._rank_device() + meta = [_pack(payload) if self._is_rank0() else None] + dist.broadcast_object_list(meta, src=0) + skeleton = payload if self._is_rank0() else _alloc(meta[0], device) + result = _broadcast(skeleton, payload if self._is_rank0() else None, 0, self._is_rank0(), device) + dist.barrier() + return result + + # --- encoder lifecycle -------------------------------------------------- + + def _load_rank0_text_encoder(self): + if not self._is_rank0(): + return + if self.text_encoders is None or self.config.get("lazy_load", False) or self.config.get("unload_modules", False): + self.text_encoders = self.load_text_encoder() + + def _unload_rank0_text_encoder(self): + if not self._is_rank0(): + return + if self.text_encoders is not None and (self._use_tp_rank0_io() or self.config.get("lazy_load", False) or self.config.get("unload_modules", False)): + del self.text_encoders + self.text_encoders = None + torch_device_module.empty_cache() + gc.collect() diff --git a/lightx2v/models/runners/wan/wan_runner.py b/lightx2v/models/runners/wan/wan_runner.py index e36e6befd..30551cbea 100755 --- a/lightx2v/models/runners/wan/wan_runner.py +++ b/lightx2v/models/runners/wan/wan_runner.py @@ -285,6 +285,8 @@ def load_text_encoder(self): return text_encoders def get_vae_parallel(self): + if self.config.get("tensor_parallel", False): + return False if isinstance(self.config.get("parallel", False), bool): return self.config.get("parallel", False) if isinstance(self.config.get("parallel", False), dict): @@ -846,6 +848,30 @@ def get_latent_shape_with_target_hw(self): ] return latent_shape + @ProfilingContext4DebugL2("Run Encoders") + def _run_input_encoder_local_t2v(self): + if self._use_tp_rank0_io(): + payload = None + if self._is_rank0(): + self.input_info.latent_shape = self.get_latent_shape_with_target_hw() + self._load_rank0_text_encoder() + text_encoder_output = self.run_text_encoder(self.input_info) + self._unload_rank0_text_encoder() + payload = {"text_encoder_output": text_encoder_output, "latent_shape": self.input_info.latent_shape} + payload = self._broadcast_rank0_payload(payload) + self.input_info.latent_shape = payload["latent_shape"] + return {"text_encoder_output": payload["text_encoder_output"], "image_encoder_output": None} + return super()._run_input_encoder_local_t2v() + + def run_vae_decoder(self, latents): + if self._use_tp_rank0_io(): + images = None + if self._is_rank0(): + images = super().run_vae_decoder(latents) + dist.barrier() + return images + return super().run_vae_decoder(latents) + class MultiModelStruct: def __init__(self, model_list, config, boundary=0.875, num_train_timesteps=1000): diff --git a/scripts/flux2/infer_flux2_dev_dist.sh b/scripts/flux2/infer_flux2_dev_dist.sh new file mode 100644 index 000000000..751c1ce19 --- /dev/null +++ b/scripts/flux2/infer_flux2_dev_dist.sh @@ -0,0 +1,14 @@ +#!/bin/bash +lightx2v_path=/data/nvme1/wushuo/LightX2V +model_path="/data/nvme1/models/black-forest-labs/FLUX.2-dev/" +export CUDA_VISIBLE_DEVICES=0,1,2,3 + +source ${lightx2v_path}/scripts/base/base.sh + +torchrun --nproc_per_node=4 -m lightx2v.infer \ + --model_cls flux2_dev \ + --task t2i \ + --model_path $model_path \ + --prompt "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom." \ + --save_result_path "${lightx2v_path}/save_results/flux2_dev.png" \ + --config_json "${lightx2v_path}/configs/flux2/flux2_dev_tp.json" diff --git a/scripts/wan22/run_wan22_moe_t2v_tp.sh b/scripts/wan22/run_wan22_moe_t2v_tp.sh new file mode 100755 index 000000000..bd4d9bb02 --- /dev/null +++ b/scripts/wan22/run_wan22_moe_t2v_tp.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/data/nvme1/wushuo/LightX2V +model_path=/data/nvme1/models/Wan2.2-T2V-A14B/ + +export CUDA_VISIBLE_DEVICES=0,1,2,3 + +# set environment variables +source ${lightx2v_path}/scripts/base/base.sh + +torchrun --nproc_per_node=4 -m lightx2v.infer \ +--model_cls wan2.2_moe \ +--task t2v \ +--model_path $model_path \ +--config_json ${lightx2v_path}/configs/wan22/wan_moe_t2v_tp.json \ +--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \ +--negative_prompt "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \ +--save_result_path ${lightx2v_path}/save_results/output_lightx2v_wan22_moe_t2v.mp4