From 56e9ec5c4ce83d54f24454f13959c1b875f99998 Mon Sep 17 00:00:00 2001 From: wushuo <540295877@qq.com> Date: Thu, 2 Jul 2026 09:01:16 +0000 Subject: [PATCH] [platform]: flux2 & tp --- .../networks/flux2/infer/transformer_infer.py | 12 +- lightx2v/models/networks/flux2/model.py | 239 ++++++++++++++++ .../flux2/weights/transformer_weights.py | 262 +++++------------- lightx2v/models/runners/flux2/flux2_runner.py | 160 ++++++++++- 4 files changed, 466 insertions(+), 207 deletions(-) diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py index b4d120990..869bcae06 100644 --- a/lightx2v/models/networks/flux2/infer/transformer_infer.py +++ b/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: diff --git a/lightx2v/models/networks/flux2/model.py b/lightx2v/models/networks/flux2/model.py index d0c389194..4b5b1d479 100644 --- a/lightx2v/models/networks/flux2/model.py +++ b/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): @@ -24,10 +25,248 @@ def __init__(self, config, model_path, device): super().__init__(model_path, config, device) self.in_channels = self.config.get("transformer_in_channels", self.config.get("in_channels", 64)) self.attention_kwargs = {} + self._init_tensor_parallel() self._init_infer_class() self._init_weights() self._init_infer() + 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 _should_load_weights(self): + if self.config.get("device_mesh") is None: + return True + if dist.is_initialized() and self.use_tp: + return dist.get_rank() == 0 + return super()._should_load_weights() + + def _load_ckpt(self, unified_dtype, sensitive_layer): + if not self.use_tp: + return super()._load_ckpt(unified_dtype, sensitive_layer) + original_device = self.device + self.device = torch.device("cpu") + try: + return super()._load_ckpt(unified_dtype, sensitive_layer) + finally: + self.device = original_device + + def _load_quant_ckpt(self, unified_dtype, sensitive_layer): + if not self.use_tp: + return super()._load_quant_ckpt(unified_dtype, sensitive_layer) + original_device = self.device + self.device = torch.device("cpu") + try: + return super()._load_quant_ckpt(unified_dtype, sensitive_layer) + finally: + self.device = original_device + + 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) diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py index da3ed9141..fccca0c72 100644 --- a/lightx2v/models/networks/flux2/weights/transformer_weights.py +++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py @@ -1,7 +1,49 @@ +import torch.distributed as dist + from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER +def _tp_info(config): + if not config.get("tensor_parallel", False): + return None, 0, 1 + tp_group = config.get("device_mesh").get_group(mesh_dim="tensor_p") + return tp_group, dist.get_rank(tp_group), dist.get_world_size(tp_group) + + +def _mm_weight(config, weight_name, bias_name=None, split_dim=None, create_cuda_buffer=False, create_cpu_buffer=False): + mm_type = config.get("dit_quant_scheme", "Default") + if config.get("tensor_parallel", False) and split_dim is not None: + tp_group, tp_rank, tp_size = _tp_info(config) + return MM_WEIGHT_REGISTER["TensorParallel"]( + weight_name=weight_name, + bias_name=bias_name, + mm_type=mm_type, + tp_group=tp_group, + tp_rank=tp_rank, + tp_size=tp_size, + split_dim=split_dim, + create_cuda_buffer=create_cuda_buffer, + create_cpu_buffer=create_cpu_buffer, + ) + return MM_WEIGHT_REGISTER[mm_type]( + weight_name, + bias_name, + create_cuda_buffer, + create_cpu_buffer, + ) + + +def _rms_weight(config, weight_name, create_cuda_buffer=False, create_cpu_buffer=False): + # Flux2 q/k RMSNorm weights are head_dim-sized, so TP over heads must replicate them. + rms_norm_type = config.get("rms_norm_type", "torch") + return RMS_WEIGHT_REGISTER[rms_norm_type]( + weight_name, + create_cuda_buffer, + create_cpu_buffer, + ) + + class Flux2DoubleBlockWeights(WeightModule): """Weights for a single double-stream transformer block.""" @@ -16,112 +58,20 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe p = f"transformer_blocks.{self.block_idx}" - self.add_module( - "to_q", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_q.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "to_k", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_k.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "to_v", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_v.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "norm_q", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_q.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "norm_k", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_k.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) + self.add_module("to_q", _mm_weight(config, f"{p}.attn.to_q.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("to_k", _mm_weight(config, f"{p}.attn.to_k.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("to_v", _mm_weight(config, f"{p}.attn.to_v.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_q", _rms_weight(config, f"{p}.attn.norm_q.weight", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_k", _rms_weight(config, f"{p}.attn.norm_k.weight", create_cuda_buffer, create_cpu_buffer)) - self.add_module( - "add_q_proj", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.add_q_proj.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "add_k_proj", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.add_k_proj.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "add_v_proj", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.add_v_proj.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "norm_added_q", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_added_q.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "norm_added_k", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_added_k.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) + self.add_module("add_q_proj", _mm_weight(config, f"{p}.attn.add_q_proj.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("add_k_proj", _mm_weight(config, f"{p}.attn.add_k_proj.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("add_v_proj", _mm_weight(config, f"{p}.attn.add_v_proj.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_added_q", _rms_weight(config, f"{p}.attn.norm_added_q.weight", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_added_k", _rms_weight(config, f"{p}.attn.norm_added_k.weight", create_cuda_buffer, create_cpu_buffer)) - self.add_module( - "to_out", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_out.0.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "to_add_out", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_add_out.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) + self.add_module("to_out", _mm_weight(config, f"{p}.attn.to_out.0.weight", None, "row", create_cuda_buffer, create_cpu_buffer)) + self.add_module("to_add_out", _mm_weight(config, f"{p}.attn.to_add_out.weight", None, "row", create_cuda_buffer, create_cpu_buffer)) self.add_module("calculate", ATTN_WEIGHT_REGISTER[self.attn_type]()) @@ -131,43 +81,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe ATTN_WEIGHT_REGISTER[self.config["parallel"].get("seq_p_attn_type", "ulysses")](), ) - self.add_module( - "ff_net_0", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.ff.linear_in.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "ff_net_2", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.ff.linear_out.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - - self.add_module( - "ff_context_net_0", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.ff_context.linear_in.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "ff_context_net_2", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.ff_context.linear_out.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) + self.add_module("ff_net_0", _mm_weight(config, f"{p}.ff.linear_in.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("ff_net_2", _mm_weight(config, f"{p}.ff.linear_out.weight", None, "row", create_cuda_buffer, create_cpu_buffer)) + self.add_module("ff_context_net_0", _mm_weight(config, f"{p}.ff_context.linear_in.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("ff_context_net_2", _mm_weight(config, f"{p}.ff_context.linear_out.weight", None, "row", create_cuda_buffer, create_cpu_buffer)) def to_cuda(self, non_blocking=True): for module in self._modules.values(): @@ -194,42 +111,10 @@ def __init__(self, config, block_idx, create_cuda_buffer=False, create_cpu_buffe p = f"single_transformer_blocks.{self.block_idx}" - self.add_module( - "to_qkv_mlp_proj", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_qkv_mlp_proj.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) - - self.add_module( - "norm_q", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_q.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) - self.add_module( - "norm_k", - RMS_WEIGHT_REGISTER[self.rms_norm_type]( - f"{p}.attn.norm_k.weight", - create_cuda_buffer, - create_cpu_buffer, - ), - ) - - self.add_module( - "to_out", - MM_WEIGHT_REGISTER[self.mm_type]( - f"{p}.attn.to_out.weight", - None, - create_cuda_buffer, - create_cpu_buffer, - ), - ) + self.add_module("to_qkv_mlp_proj", _mm_weight(config, f"{p}.attn.to_qkv_mlp_proj.weight", None, "col", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_q", _rms_weight(config, f"{p}.attn.norm_q.weight", create_cuda_buffer, create_cpu_buffer)) + self.add_module("norm_k", _rms_weight(config, f"{p}.attn.norm_k.weight", create_cuda_buffer, create_cpu_buffer)) + self.add_module("to_out", _mm_weight(config, f"{p}.attn.to_out.weight", None, "row", create_cuda_buffer, create_cpu_buffer)) self.add_module("calculate", ATTN_WEIGHT_REGISTER[self.attn_type]()) @@ -260,8 +145,6 @@ def __init__(self, config): self.num_single_layers = config.get("num_single_layers", 20) self.mm_type = config.get("dit_quant_scheme", "Default") - inner_dim = config.get("num_attention_heads", 24) * config.get("attention_head_dim", 64) - self.double_blocks = WeightModuleList([Flux2DoubleBlockWeights(config, i) for i in range(self.num_layers)]) self.single_blocks = WeightModuleList([Flux2SingleBlockWeights(config, i) for i in range(self.num_single_layers)]) self.register_offload_buffers(config) @@ -269,24 +152,9 @@ def __init__(self, config): self.add_module("double_blocks", self.double_blocks) self.add_module("single_blocks", self.single_blocks) - self.add_module( - "double_stream_modulation_img_linear", - MM_WEIGHT_REGISTER[self.mm_type]( - "double_stream_modulation_img.linear.weight", - ), - ) - self.add_module( - "double_stream_modulation_txt_linear", - MM_WEIGHT_REGISTER[self.mm_type]( - "double_stream_modulation_txt.linear.weight", - ), - ) - self.add_module( - "single_stream_modulation_linear", - MM_WEIGHT_REGISTER[self.mm_type]( - "single_stream_modulation.linear.weight", - ), - ) + self.add_module("double_stream_modulation_img_linear", _mm_weight(config, "double_stream_modulation_img.linear.weight")) + self.add_module("double_stream_modulation_txt_linear", _mm_weight(config, "double_stream_modulation_txt.linear.weight")) + self.add_module("single_stream_modulation_linear", _mm_weight(config, "single_stream_modulation.linear.weight")) def register_offload_buffers(self, config): if config.get("cpu_offload", False) and config.get("offload_granularity", "block") == "block": diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py index 3c9334295..198d3a76e 100644 --- a/lightx2v/models/runners/flux2/flux2_runner.py +++ b/lightx2v/models/runners/flux2/flux2_runner.py @@ -4,6 +4,7 @@ import numpy as np import torch +import torch.distributed as dist from loguru import logger from lightx2v.models.networks.flux2.model import Flux2DevTransformerModel, Flux2KleinTransformerModel @@ -48,10 +49,104 @@ def _get_scheduler_class(self): @ProfilingContext4DebugL2("Load models") def load_model(self): + if self._use_tp_rank0_io(): + self.text_encoders = None + self.vae = None + self.model = self.load_transformer() + return + self.text_encoders = self.load_text_encoder() 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 + if self.vae is None or self.config.get("lazy_load", False) or self.config.get("unload_modules", False): + self.vae = self.load_vae() + + def _unload_rank0_vae(self): + if not self._is_rank0(): + return + if self.vae 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.vae + self.vae = None + torch_device_module.empty_cache() + gc.collect() + def load_vae(self): return Flux2VAE(self.config) @@ -64,6 +159,8 @@ def init_modules(self): assert self.config.get("cpu_offload", False) task = self.config.get("task", "t2i") + if self._use_tp_rank0_io() and task == "i2i": + raise NotImplementedError("Flux2 tensor parallel currently supports t2i only; i2i needs rank0 VAE encode broadcast.") if task == "i2i": self.run_input_encoder = self._run_input_encoder_local_i2i self.run_dit = self._run_dit_local_i2i @@ -73,6 +170,24 @@ def init_modules(self): @ProfilingContext4DebugL2("Run Encoders") def _run_input_encoder_local_t2i(self): + if self._use_tp_rank0_io(): + payload = None + if self._is_rank0(): + self._load_rank0_text_encoder() + prompt = self.input_info.prompt + text_encoder_output = self.run_text_encoder(prompt, neg_prompt=self.input_info.negative_prompt) + self._unload_rank0_text_encoder() + payload = { + "inputs": { + "text_encoder_output": text_encoder_output, + "image_encoder_output": None, + }, + "target_shape": self.input_info.target_shape, + } + payload = self._broadcast_rank0_payload(payload) + self.input_info.target_shape = payload["target_shape"] + return payload["inputs"] + prompt = self.input_info.prompt if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): self.text_encoders = self.load_text_encoder() @@ -88,6 +203,22 @@ def _run_input_encoder_local_t2i(self): @ProfilingContext4DebugL2("Run Encoders I2I") def _run_input_encoder_local_i2i(self): + if self._use_tp_rank0_io(): + payload = None + if self._is_rank0(): + self._load_rank0_text_encoder() + payload = { + "inputs": self._run_input_encoder_local_i2i_rank0(), + "target_shape": self.input_info.target_shape, + } + self._unload_rank0_text_encoder() + payload = self._broadcast_rank0_payload(payload) + self.input_info.target_shape = payload["target_shape"] + return payload["inputs"] + + return self._run_input_encoder_local_i2i_rank0() + + def _run_input_encoder_local_i2i_rank0(self): prompt = self.input_info.prompt if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): self.text_encoders = self.load_text_encoder() @@ -330,9 +461,28 @@ def set_img_shapes(self): @ProfilingContext4DebugL1("Run VAE Decoder") def run_vae_decoder(self, latents): + if self._use_tp_rank0_io(): + images = None + if self._is_rank0(): + self._load_rank0_vae() + images = self._decode_latents_with_vae(latents) + self._unload_rank0_vae() + dist.barrier() + return images + if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): self.vae = self.load_vae() + images = self._decode_latents_with_vae(latents) + + if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): + del self.vae + torch_device_module.empty_cache() + gc.collect() + + return images + + def _decode_latents_with_vae(self, latents): B, _, C = latents.shape H = int((self.input_info.latent_image_ids[0, :, 1].max() + 1).item()) @@ -348,14 +498,7 @@ def run_vae_decoder(self, latents): latents = latents.permute(0, 1, 4, 2, 5, 3) latents = latents.reshape(B, C // 4, H * 2, W * 2) - images = self.vae.decode(latents, self.input_info) - - if self.config.get("lazy_load", False) or self.config.get("unload_modules", False): - del self.vae - torch_device_module.empty_cache() - gc.collect() - - return images + return self.vae.decode(latents, self.input_info) @ProfilingContext4DebugL1("RUN pipeline") def run_pipeline(self, input_info): @@ -368,6 +511,7 @@ def run_pipeline(self, input_info): latents, generator = self.run_dit() images = self.run_vae_decoder(latents) + self.end_run() if not input_info.return_result_tensor and is_main_process(): image = images[0]