diff --git a/lightx2v/common/distributed/__init__.py b/lightx2v/common/distributed/__init__.py new file mode 100644 index 000000000..0f979d8f9 --- /dev/null +++ b/lightx2v/common/distributed/__init__.py @@ -0,0 +1,11 @@ +from .pipeline_comm import PipelineComm +from .pipeline_state import ( + PipelineRuntimeState, + get_pipeline_parallel_rank, + get_pipeline_parallel_world_size, + get_pipeline_runtime_state, + get_pp_group, + init_pipeline_parallel_state, + is_pipeline_first_stage, + is_pipeline_last_stage, +) diff --git a/lightx2v/common/distributed/pipeline_comm.py b/lightx2v/common/distributed/pipeline_comm.py new file mode 100644 index 000000000..521913272 --- /dev/null +++ b/lightx2v/common/distributed/pipeline_comm.py @@ -0,0 +1,115 @@ +"""P2P communication manager for pipeline parallelism. + +Uses the default CUDA stream for all NCCL operations. NCCL's internal +stream handles the actual async data transfer, so a separate comm_stream +is unnecessary and would only introduce cross-stream sync overhead. + +isend requests are returned to the caller who MUST store them to +prevent tensor GC before the send completes. +""" + +from typing import Dict, List, Tuple + +import torch +import torch.distributed as dist + + +class PipelineComm: + """P2P communication between adjacent pipeline stages.""" + + def __init__(self, pp_group: dist.ProcessGroup): + self.pp_group = pp_group + self.rank = dist.get_rank(pp_group) + self.world_size = dist.get_world_size(pp_group) + + self.ranks = list(dist.get_process_group_ranks(pp_group)) + self.prev_rank = self.ranks[(self.rank - 1) % self.world_size] + self.next_rank = self.ranks[(self.rank + 1) % self.world_size] + self._device_group = pp_group + + self._recv_tasks_queue: List[Tuple[str, int]] = [] + self._receiving_tasks: List[Tuple[object, str, int]] = [] + self._recv_buffers: Dict[Tuple[str, int], torch.Tensor] = {} + + # ------------------------------------------------------------------ + # Synchronous send / recv (sync pipeline) + # ------------------------------------------------------------------ + + def pipeline_send(self, tensor: torch.Tensor, name: str = "latent", skip_shape: bool = False): + tensor = tensor.contiguous() + if not skip_shape: + shape_info = torch.tensor( + [tensor.ndim] + list(tensor.shape), + device=tensor.device, + dtype=torch.int64, + ) + padded = torch.zeros(9, device=tensor.device, dtype=torch.int64) + padded[: len(shape_info)] = shape_info + dist.send(padded, dst=self.next_rank, group=self._device_group) + dist.send(tensor, dst=self.next_rank, group=self._device_group) + + def pipeline_recv(self, name: str = "latent", shape=None, dtype=None) -> torch.Tensor: + if dtype is None: + dtype = torch.bfloat16 + if shape is not None: + buf = torch.empty(shape, dtype=dtype, device=torch.cuda.current_device()) + dist.recv(buf, src=self.prev_rank, group=self._device_group) + return buf + shape_info = torch.zeros(9, device=torch.cuda.current_device(), dtype=torch.int64) + dist.recv(shape_info, src=self.prev_rank, group=self._device_group) + ndim = shape_info[0].item() + recv_shape = tuple(shape_info[1 : 1 + ndim].tolist()) + buf = torch.empty(recv_shape, dtype=dtype, device=torch.cuda.current_device()) + dist.recv(buf, src=self.prev_rank, group=self._device_group) + return buf + + # ------------------------------------------------------------------ + # Asynchronous send / recv (async pipeline) + # All on default stream — NCCL's internal stream handles the real + # async transfer, and req.wait() just inserts a stream-side + # dependency (non-blocking on CPU in the same-stream case). + # ------------------------------------------------------------------ + + def pipeline_isend(self, tensor: torch.Tensor, name: str = "latent", segment_idx: int = 0): + """Non-blocking send on the current (default) stream. + + Returns a Work object that the caller SHOULD store to prevent + the tensor from being garbage-collected before the send + completes. The Work's wait() only inserts a stream-side + dependency — it does not block the CPU thread. + """ + tensor = tensor.contiguous() + return dist.isend(tensor, dst=self.next_rank, group=self._device_group) + + def add_pipeline_recv_task(self, idx: int = 0, name: str = "latent", shape=None, dtype=None): + self._recv_tasks_queue.append((name, idx)) + if (name, idx) not in self._recv_buffers: + assert shape is not None and dtype is not None + self._recv_buffers[(name, idx)] = torch.empty(shape, dtype=dtype, device=torch.cuda.current_device()) + + def recv_next(self): + """Post next irecv on the current (default) stream. + + Non-blocking on CPU: dist.irecv enqueues the work on NCCL's + internal stream and returns immediately. + """ + if not self._recv_tasks_queue: + raise ValueError("No more tasks to receive") + name, idx = self._recv_tasks_queue.pop(0) + buf = self._recv_buffers.get((name, idx)) + assert buf is not None + req = dist.irecv(buf, src=self.prev_rank, group=self._device_group) + self._receiving_tasks.append((req, name, idx)) + + def get_pipeline_recv_data(self, idx: int = 0, name: str = "latent") -> torch.Tensor: + """Wait for and return a previously posted async receive. + + In the single-stream model, req.wait() inserts a stream-side + wait for the NCCL op's completion on the current stream. It + does NOT block the CPU thread unless a timeout is set. + """ + assert self._receiving_tasks + req, rname, ridx = self._receiving_tasks.pop(0) + assert rname == name and ridx == idx + req.wait() + return self._recv_buffers[(name, idx)] diff --git a/lightx2v/common/distributed/pipeline_state.py b/lightx2v/common/distributed/pipeline_state.py new file mode 100644 index 000000000..e84ba4b85 --- /dev/null +++ b/lightx2v/common/distributed/pipeline_state.py @@ -0,0 +1,152 @@ +"""Pipeline parallel runtime state and stage helpers for LightX2V. + +Manages patch metadata (how the latent image is split across pipeline patches) +and provides stage-identification utilities (is_pipeline_first_stage, etc.). +""" + +from typing import List, Optional + +import torch.distributed as dist + +# --------------------------------------------------------------------------- +# Global state +# --------------------------------------------------------------------------- + +_pp_group: Optional[dist.ProcessGroup] = None +_runtime_state: Optional["PipelineRuntimeState"] = None + + +def init_pipeline_parallel_state(pp_group: dist.ProcessGroup): + """Register the pipeline-parallel process group. Called once during + ``set_parallel_config`` when ``pp_size > 1``.""" + global _pp_group, _runtime_state + _pp_group = pp_group + _runtime_state = PipelineRuntimeState() + + +# --------------------------------------------------------------------------- +# Stage helpers +# --------------------------------------------------------------------------- + + +def get_pp_group() -> dist.ProcessGroup: + assert _pp_group is not None, "pipeline parallel group is not initialised" + return _pp_group + + +def get_pipeline_parallel_rank() -> int: + if _pp_group is None: + return 0 + return dist.get_rank(_pp_group) + + +def get_pipeline_parallel_world_size() -> int: + if _pp_group is None: + return 1 + return dist.get_world_size(_pp_group) + + +def is_pipeline_first_stage() -> bool: + return get_pipeline_parallel_rank() == 0 + + +def is_pipeline_last_stage() -> bool: + return get_pipeline_parallel_rank() == get_pipeline_parallel_world_size() - 1 + + +def get_pipeline_runtime_state() -> "PipelineRuntimeState": + assert _runtime_state is not None, "PipelineRuntimeState not initialised" + return _runtime_state + + +# --------------------------------------------------------------------------- +# Runtime state +# --------------------------------------------------------------------------- + + +class PipelineRuntimeState: + """Runtime metadata for patch-level pipeline parallelism (PipeFusion). + + Computes how the latent token sequence is split into *patches* so that + each pipeline stage processes a subset of patches in async mode. + """ + + def __init__(self): + self.num_pipeline_patch: int = 1 + self.pipeline_patch_idx: int = 0 + self.patch_mode: bool = False # True = async, False = sync + self.warmup_steps: int = 1 + + # Patch metadata (along the latent token / sequence dimension) + self.pp_patches_token_num: List[int] = [0] + self.pp_patches_token_start_end_idx_global: List[List[int]] = [[0, 0]] + + # Input parameters + self.height: int = 0 + self.width: int = 0 + self.batch_size: int = 1 + self.packed_h: int = 0 + self.packed_w: int = 0 + self.vae_scale_factor: int = 16 + self.patch_size: int = 1 + + # -- configuration ------------------------------------------------------- + + def set_input_parameters( + self, + height: int, + width: int, + batch_size: int = 1, + num_pipeline_patch: Optional[int] = None, + warmup_steps: int = 1, + vae_scale_factor: int = 16, + patch_size: int = 1, + total_tokens: Optional[int] = None, + ): + self.height = height + self.width = width + self.batch_size = batch_size + self.vae_scale_factor = vae_scale_factor + self.patch_size = patch_size + self.warmup_steps = warmup_steps + if num_pipeline_patch is not None: + self.num_pipeline_patch = num_pipeline_patch + + if total_tokens is not None: + self.packed_h = 0 + self.packed_w = 0 + tok_count = total_tokens + else: + # Compute packed dimensions + multiple_of = vae_scale_factor * 2 + self.packed_h = height // multiple_of + self.packed_w = width // multiple_of + tok_count = self.packed_h * self.packed_w + + # Split tokens evenly across patches + base = tok_count // self.num_pipeline_patch + remainder = tok_count % self.num_pipeline_patch + self.pp_patches_token_num = [] + self.pp_patches_token_start_end_idx_global = [] + start = 0 + for i in range(self.num_pipeline_patch): + n = base + (1 if i < remainder else 0) + self.pp_patches_token_num.append(n) + self.pp_patches_token_start_end_idx_global.append([start, start + n]) + start += n + + # -- patch mode ---------------------------------------------------------- + + def set_patched_mode(self, patch_mode: bool): + self.patch_mode = patch_mode + self.pipeline_patch_idx = 0 + + def next_patch(self): + if self.patch_mode: + self.pipeline_patch_idx += 1 + if self.pipeline_patch_idx >= self.num_pipeline_patch: + self.pipeline_patch_idx = 0 + + @property + def current_patch_token_start_end(self) -> List[int]: + return self.pp_patches_token_start_end_idx_global[self.pipeline_patch_idx] diff --git a/lightx2v/common/ops/attn/flash_attn.py b/lightx2v/common/ops/attn/flash_attn.py index 8ed341848..bde3945aa 100755 --- a/lightx2v/common/ops/attn/flash_attn.py +++ b/lightx2v/common/ops/attn/flash_attn.py @@ -52,9 +52,10 @@ def apply( softmax_scale = kwargs.get("softmax_scale", None) if len(q.shape) == 3: bs = 1 + total_seqlen = q.shape[0] elif len(q.shape) == 4: bs = q.shape[0] - total_seqlen = bs * max_seqlen_q + total_seqlen = bs * q.shape[1] if bs == 1: if len(q.shape) == 3: @@ -127,9 +128,10 @@ def apply( softmax_scale = kwargs.get("softmax_scale", None) if len(q.shape) == 3: bs = 1 + total_seqlen = q.shape[0] elif len(q.shape) == 4: bs = q.shape[0] - total_seqlen = bs * max_seqlen_q + total_seqlen = bs * q.shape[1] if bs == 1: if len(q.shape) == 3: @@ -209,7 +211,7 @@ def apply( k, v, ) - x = x.reshape(bs * max_seqlen_q, -1) + x = x.reshape(q.shape[1], -1) return x diff --git a/lightx2v/models/networks/base_model.py b/lightx2v/models/networks/base_model.py index 44d98124c..8b5f0b4e3 100644 --- a/lightx2v/models/networks/base_model.py +++ b/lightx2v/models/networks/base_model.py @@ -399,7 +399,11 @@ def _load_safetensor_to_dict(self, file_path, unified_dtype, sensitive_layer): remove_keys = self.remove_keys if hasattr(self, "remove_keys") else [] preserve_keys = self.preserved_keys if hasattr(self, "preserved_keys") else None # None means all keys are preserved, otherwise only keys in preserve_keys are preserved - if self.device.type != "cpu" and dist.is_initialized(): + # In PipeFusion mode, load weights to CPU first to avoid OOM — each + # stage only needs a subset of block weights on GPU. + if self.config.get("pipefusion_parallel", False): + device = "cpu" + elif self.device.type != "cpu" and dist.is_initialized(): device = dist.get_rank() else: device = str(self.device) diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py b/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py new file mode 100644 index 000000000..fbd3b4694 --- /dev/null +++ b/lightx2v/models/networks/flux2/infer/pipefusion/__init__.py @@ -0,0 +1,2 @@ +from .pipeline_driver import Flux2PipelineDriver +from .transformer_infer import Flux2PipeFusionTransformerInfer diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py new file mode 100644 index 000000000..a512db1d3 --- /dev/null +++ b/lightx2v/models/networks/flux2/infer/pipefusion/pipeline_driver.py @@ -0,0 +1,395 @@ +"""Sync/Async pipeline driver for Flux2 PipeFusion. + +Orchestrates the denoising loop across pipeline stages: +- **Sync pipeline** (warmup): each timestep, all stages process the full latent + sequentially (stage 0 -> stage 1 -> ... -> last stage). +- **Async pipeline** (main loop): each timestep, the latent is split into + patches; stages process different patches concurrently, overlapping compute + and P2P communication. +""" + +import torch +import torch.distributed as dist + +from lightx2v.common.distributed import ( + PipelineComm, + get_pipeline_parallel_world_size, + get_pipeline_runtime_state, + get_pp_group, + is_pipeline_first_stage, + is_pipeline_last_stage, +) + + +class Flux2PipelineDriver: + """Drives the PipeFusion denoising loop for Flux2.""" + + def __init__(self, model, config): + self.model = model + self.config = config + self.state = get_pipeline_runtime_state() + self.pp_comm = PipelineComm(get_pp_group()) + self._is_first = is_pipeline_first_stage() + self._is_last = is_pipeline_last_stage() + self._pp_world_size = get_pipeline_parallel_world_size() + self._dtype = config.get("dtype", torch.bfloat16) + if isinstance(self._dtype, str): + self._dtype = getattr(torch, self._dtype) + + # ================================================================== + # Public entry point + # ================================================================== + + def run_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None): + """Run the full denoising loop with PipeFusion. + + Returns final latents on the last stage, ``None`` on other stages. + """ + warmup_steps = self.state.warmup_steps + + if self._pp_world_size > 1 and len(timesteps) > warmup_steps: + latents = self._sync_pipeline( + latents, + prompt_embeds, + text_ids, + latent_image_ids, + timesteps[:warmup_steps], + scheduler, + do_cfg=do_cfg, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + ) + latents = self._async_pipeline( + latents, + prompt_embeds, + text_ids, + latent_image_ids, + timesteps[warmup_steps:], + scheduler, + do_cfg=do_cfg, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + ) + else: + latents = self._sync_pipeline( + latents, + prompt_embeds, + text_ids, + latent_image_ids, + timesteps, + scheduler, + do_cfg=do_cfg, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + ) + return latents + + # ================================================================== + # Sync pipeline (warmup) + # ================================================================== + + def _sync_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None): + self.state.set_patched_mode(patch_mode=False) + + for step_idx, t in enumerate(timesteps): + scheduler.step_index = step_idx + scheduler.step_pre(step_idx) + + if do_cfg: + # Conditional pass + cond_result = self._sync_pass( + latents, + prompt_embeds, + text_ids, + latent_image_ids, + t, + scheduler, + ) + # Unconditional pass + uncond_result = self._sync_pass( + latents, + negative_prompt_embeds, + negative_text_ids or text_ids, + latent_image_ids, + t, + scheduler, + ) + if self._is_last: + noise_pred_cond = cond_result + noise_pred_uncond = uncond_result + guidance_scale = self.config.get("sample_guide_scale", 1.0) + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + scheduler.noise_pred = noise_pred + scheduler.latents = latents + scheduler.step_post() + latents = scheduler.latents + else: + noise_pred = self._sync_pass( + latents, + prompt_embeds, + text_ids, + latent_image_ids, + t, + scheduler, + ) + if self._is_last: + scheduler.noise_pred = noise_pred + scheduler.latents = latents + scheduler.step_post() + latents = scheduler.latents + + # P2P: last stage sends updated latents to first stage (circular) + # Only rank 0 needs updated latents (for x_embedder in next step). + # Ranks 1-6 don't participate — no global sync barrier. + if self._pp_world_size > 1: + if self._is_last: + # Last stage sends to first stage + dist.send(latents.contiguous(), dst=self.pp_comm.ranks[0], group=self.pp_comm.pp_group) + elif self._is_first: + # First stage receives from last stage + latents = torch.empty_like(latents) + dist.recv(latents, src=self.pp_comm.ranks[-1], group=self.pp_comm.pp_group) + # Ranks 1-6: no op (don't need updated latents) + + return latents + + def _sync_pass(self, latents, prompt_embeds, text_ids, latent_image_ids, t, scheduler): + """Single sync forward pass through all stages. + + Returns ``noise_pred`` on last stage, ``None`` on other stages. + P2P always carries separate (image, text) streams. + """ + # NOTE: do NOT clear KV cache here. Sync mode populates per-patch + # slots so async mode can use them as "stale" KV for global attention. + + if self._is_first: + # First stage: run pre_infer + blocks + pre_infer_out = self.model.pre_infer.infer( + weights=self.model.pre_weight, + hidden_states=latents, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + ) + hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out) + + if self._is_last: + return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep) + else: + # Always send both latent and encoder_hidden_state (skip_shape + # to avoid .item() CPU-GPU sync on receiver) + self.pp_comm.pipeline_send(hidden_states, name="latent", skip_shape=True) + self.pp_comm.pipeline_send(enc_hidden, name="encoder_hidden_state", skip_shape=True) + return None + else: + # Non-first stage: always receive both streams + # Pass pre-computed shapes to avoid .item() CPU-GPU sync + inner_dim = self.config.get("num_attention_heads", 24) * self.config.get("attention_head_dim", 64) + if latent_image_ids.ndim == 3: + img_len = latent_image_ids.shape[1] + else: + img_len = latent_image_ids.shape[0] + if prompt_embeds is not None: + txt_len = prompt_embeds.shape[1] if prompt_embeds.ndim == 3 else prompt_embeds.shape[0] + else: + txt_len = 0 + hidden_states = self.pp_comm.pipeline_recv(name="latent", shape=(img_len, inner_dim), dtype=self._dtype) + enc_hidden = self.pp_comm.pipeline_recv(name="encoder_hidden_state", shape=(txt_len, inner_dim), dtype=self._dtype) + + pre_infer_out = self.model.pre_infer.infer_partial( + weights=self.model.pre_weight, + hidden_states=hidden_states, + encoder_hidden_states=enc_hidden, + txt_ids=text_ids, + img_ids=latent_image_ids, + ) + hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out) + + if self._is_last: + return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep) + else: + self.pp_comm.pipeline_send(hidden_states, name="latent", skip_shape=True) + self.pp_comm.pipeline_send(enc_hidden, name="encoder_hidden_state", skip_shape=True) + return None + + # ================================================================== + # Async pipeline (main loop) + # ================================================================== + + def _async_pipeline(self, latents, prompt_embeds, text_ids, latent_image_ids, timesteps, scheduler, do_cfg=False, negative_prompt_embeds=None, negative_text_ids=None): + self.state.set_patched_mode(patch_mode=True) + num_patch = self.state.num_pipeline_patch + patch_token_nums = self.state.pp_patches_token_num + inner_dim = self.config.get("num_attention_heads", 24) * self.config.get("attention_head_dim", 64) + # Raw latent channels (before x_embedder): rank=0 recv from last stage + # gets raw latents [1, L, C_in]; other stages recv embedded [L, D] + raw_channels = getattr(self.model, "in_channels", self.config.get("transformer_in_channels", self.config.get("in_channels", 128))) + + # Split latents into patches (dim=1 for [B, L, C]) + if self._is_first or self._is_last: + patch_latents = list(latents.split(patch_token_nums, dim=1)) + else: + patch_latents = [None] * num_patch + + # Split image ids by patch + patch_latent_image_ids = [] + for start, end in self.state.pp_patches_token_start_end_idx_global: + if latent_image_ids.ndim == 3: + patch_latent_image_ids.append(latent_image_ids[:, start:end, :]) + else: + patch_latent_image_ids.append(latent_image_ids[start:end, :]) + + # Compute txt_len for buffer allocation + if prompt_embeds is not None: + txt_len = prompt_embeds.shape[1] if prompt_embeds.ndim == 3 else prompt_embeds.shape[0] + else: + txt_len = 0 + + # Pre-allocate recv buffers and pre-post all receives + # First stage: receives raw latents [1, L, C_in] from last stage (circular) + # Non-first stages: receives embedded encoder + latent [L, D] from previous stage + recv_timesteps = len(timesteps) - 1 if self._is_first else len(timesteps) + for _ in range(recv_timesteps): + if not self._is_first: + self.pp_comm.add_pipeline_recv_task( + 0, + "encoder_hidden_state", + shape=(txt_len, inner_dim), + dtype=self._dtype, + ) + for patch_idx in range(num_patch): + # First stage (rank=0) receives raw latents [1, L, C_in] from + # last stage; other stages receive embedded [L, D] + if self._is_first: + latent_shape = (1, patch_token_nums[patch_idx], raw_channels) + else: + latent_shape = (patch_token_nums[patch_idx], inner_dim) + self.pp_comm.add_pipeline_recv_task( + patch_idx, + "latent", + shape=latent_shape, + dtype=self._dtype, + ) + + last_patch_latents = [None] * num_patch if self._is_last else None + first_async_recv = True + total_steps = len(timesteps) + + # Track pending isend requests to prevent tensor GC before send completes + pending_isends = [] + + for i, t in enumerate(timesteps): + scheduler.step_index = i + self.state.warmup_steps + scheduler.step_pre(scheduler.step_index) + + for patch_idx in range(num_patch): + if self._is_last: + last_patch_latents[patch_idx] = patch_latents[patch_idx] + + # ---- 1. Receive current patch's data ---- + if self._is_first and i == 0: + pass # first stage, first step: has initial latents + else: + if first_async_recv: + if not self._is_first and patch_idx == 0: + self.pp_comm.recv_next() + self.pp_comm.recv_next() + first_async_recv = False + if not self._is_first and patch_idx == 0: + last_encoder_hidden_states = self.pp_comm.get_pipeline_recv_data(0, "encoder_hidden_state") + if not (self._is_first and i == 0): + patch_latents[patch_idx] = self.pp_comm.get_pipeline_recv_data(patch_idx, "latent") + + # ---- 2. Compute (default stream) ---- + cur_enc = prompt_embeds if self._is_first else last_encoder_hidden_states + result = self._async_backbone( + patch_latents[patch_idx], + cur_enc, + text_ids, + patch_latent_image_ids[patch_idx], + scheduler, + ) + + # ---- 3. Send result (default stream, after compute) ---- + # Store isend request to prevent tensor GC before send completes + if self._is_last: + noise_pred = result + scheduler.scheduler._step_index = i + self.state.warmup_steps + patch_latents[patch_idx] = scheduler.step_post_patch(noise_pred, last_patch_latents[patch_idx], t) + if i != total_steps - 1: + req = self.pp_comm.pipeline_isend(patch_latents[patch_idx], name="latent", segment_idx=patch_idx) + pending_isends.append((req, patch_latents[patch_idx])) + else: + hidden_states, next_enc = result + if patch_idx == 0: + req = self.pp_comm.pipeline_isend(next_enc, name="encoder_hidden_state", segment_idx=0) + pending_isends.append((req, next_enc)) + req = self.pp_comm.pipeline_isend(hidden_states, name="latent", segment_idx=patch_idx) + pending_isends.append((req, hidden_states)) + + # ---- 4. Post next irecv (default stream — NCCL internal + # stream handles the actual async transfer; no cross-stream + # sync needed.) ---- + if not (self._is_first and i == 0): + is_last_step = i == total_steps - 1 + is_last_patch = patch_idx == num_patch - 1 + if not (is_last_step and is_last_patch): + if self._is_first: + self.pp_comm.recv_next() + else: + if is_last_patch: + self.pp_comm.recv_next() + self.pp_comm.recv_next() + + # ---- 5. Wait for old isends (limit pending to prevent GC issues) ---- + while len(pending_isends) > num_patch * 2: + old_req, _ = pending_isends.pop(0) + old_req.wait() + + self.state.next_patch() + + # Wait for all remaining isends before returning + for req, _ in pending_isends: + req.wait() + pending_isends.clear() + + if self._is_last: + return torch.cat(patch_latents, dim=1) + return None + + def _async_backbone(self, patch_latent, encoder_hidden_states, text_ids, patch_img_ids, scheduler): + """Backbone forward for a single patch in async mode.""" + if self._is_first: + pre_infer_out = self.model.pre_infer.infer( + weights=self.model.pre_weight, + hidden_states=patch_latent, + encoder_hidden_states=encoder_hidden_states, + txt_ids=text_ids, + img_ids=patch_img_ids, + ) + else: + pre_infer_out = self.model.pre_infer.infer_partial( + weights=self.model.pre_weight, + hidden_states=patch_latent, + encoder_hidden_states=encoder_hidden_states, + txt_ids=text_ids, + img_ids=patch_img_ids, + ) + + hidden_states, enc_hidden, num_txt = self.model.transformer_infer.infer(self.model.transformer_weights, pre_infer_out) + + if self._is_last: + return self._run_post_infer(hidden_states, enc_hidden, num_txt, pre_infer_out.timestep) + else: + return (hidden_states, enc_hidden) + + # ================================================================== + # Shared helpers + # ================================================================== + + def _run_post_infer(self, hidden_states, enc_hidden, num_txt, timestep): + """Run post_infer on the last stage.""" + if enc_hidden is None and num_txt > 0: + hidden_states = hidden_states[num_txt:, ...] + noise_pred = self.model.post_infer.infer(self.model.post_weight, hidden_states, timestep) + return noise_pred diff --git a/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py new file mode 100644 index 000000000..faad09e18 --- /dev/null +++ b/lightx2v/models/networks/flux2/infer/pipefusion/transformer_infer.py @@ -0,0 +1,203 @@ +"""PipeFusion-enabled transformer infer for Flux2. + +Subclasses ``Flux2TransformerInfer`` to: +1. Run only the current pipeline stage's block subset. +2. Apply stale-KV caching in async (patched) mode: image KV is cached across + patches while text KV stays fresh. +3. Return ``(hidden_states, encoder_hidden_states, num_txt_tokens)`` so the + pipeline driver can P2P-pass intermediate activations between stages. +""" + +import torch +import torch.nn.functional as F + +from ..transformer_infer import Flux2TransformerInfer + + +class Flux2PipeFusionTransformerInfer(Flux2TransformerInfer): + """Transformer infer with PipeFusion block splitting and stale-KV cache.""" + + def __init__(self, config): + super().__init__(config) + from lightx2v.common.distributed import ( + get_pipeline_runtime_state, + is_pipeline_first_stage, + is_pipeline_last_stage, + ) + + self.pipeline_state = get_pipeline_runtime_state() + self._is_first_stage = is_pipeline_first_stage() + self._is_last_stage = is_pipeline_last_stage() + + # Stale-KV cache: block_idx -> [ [k, v] per patch slot ] + self._kv_cache: dict = {} + + # Pre-allocated full K/V buffers per block (lazily created on first async use). + # Avoids repeated torch.cat allocations per timestep. + self._full_k_bufs: dict = {} + self._full_v_bufs: dict = {} + + # ------------------------------------------------------------------ + # Stale-KV hook (overrides base class no-op) + # ------------------------------------------------------------------ + + def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx): + """Per-patch-slot KV cache for PipeFusion. + + Semantics: + - Cache is indexed by (block_idx, patch_slot). Each slot stores the + image K/V computed for that patch when it was last processed. + - SYNC mode: split full image K/V by patch, populate ALL slots. + Return input unchanged (full attention runs normally). + - ASYNC mode: update current patch's slot with fresh K/V; use full + cache (fresh + stale from prior timestep) for attention. + + Optimization: pre-allocated buffers + copy_ instead of torch.cat + to avoid memory allocations per generation. + """ + num_patch = self.pipeline_state.num_pipeline_patch + if num_patch <= 1 or num_txt_tokens <= 0: + return key, value + + # Split text / image along sequence dim + text_key, img_key = key.split([num_txt_tokens, key.shape[0] - num_txt_tokens], dim=0) + text_value, img_value = value.split([num_txt_tokens, value.shape[0] - num_txt_tokens], dim=0) + + patch_token_nums = self.pipeline_state.pp_patches_token_num + + if not self.pipeline_state.patch_mode: + # Sync mode: split full image K/V by patch, populate all slots. + # .clone() ensures cached tensors own their storage (views of + # transient QKV would become invalid after this timestep). + if block_idx not in self._kv_cache: + self._kv_cache[block_idx] = [None] * num_patch + split_ks = img_key.split(patch_token_nums, dim=0) + split_vs = img_value.split(patch_token_nums, dim=0) + for i in range(num_patch): + self._kv_cache[block_idx][i] = [ + split_ks[i].clone(), + split_vs[i].clone(), + ] + return key, value + + # ---- Async mode ---- + + cur_slot = self.pipeline_state.pipeline_patch_idx + if block_idx not in self._kv_cache: + self._kv_cache[block_idx] = [None] * num_patch + + # Store fresh K/V in cache (clone for persistence across timesteps) + self._kv_cache[block_idx][cur_slot] = [img_key.clone(), img_value.clone()] + + # Build full K/V using pre-allocated buffer + copy_ (avoids torch.cat) + total_img = sum(patch_token_nums) + full_len = num_txt_tokens + total_img + + if block_idx not in self._full_k_bufs or self._full_k_bufs[block_idx].shape[0] != full_len or self._full_k_bufs[block_idx].dtype != key.dtype: + self._full_k_bufs[block_idx] = torch.empty(full_len, *key.shape[1:], dtype=key.dtype, device=key.device) + self._full_v_bufs[block_idx] = torch.empty(full_len, *value.shape[1:], dtype=value.dtype, device=value.device) + + buf_k = self._full_k_bufs[block_idx] + buf_v = self._full_v_bufs[block_idx] + + # Copy text K/V (fresh, from current patch's computation) + buf_k[:num_txt_tokens].copy_(text_key) + buf_v[:num_txt_tokens].copy_(text_value) + + # Copy each slot's image K/V into buffer + offset = num_txt_tokens + for slot in range(num_patch): + n = patch_token_nums[slot] + if slot == cur_slot: + # Fresh from this patch (copy from img_key, already cloned to cache) + buf_k[offset : offset + n].copy_(img_key) + buf_v[offset : offset + n].copy_(img_value) + else: + # Stale from cache (previous timestep) + cached = self._kv_cache[block_idx][slot] + buf_k[offset : offset + n].copy_(cached[0]) + buf_v[offset : offset + n].copy_(cached[1]) + offset += n + + return buf_k[:full_len], buf_v[:full_len] + + def clear_kv_cache(self): + """Clear stale-KV cache. + + NOTE: stale-KV cache persists ACROSS timesteps by design — that's the + whole point of "stale" KV. This method is provided for defensive + cleanup only and should NOT be called between timesteps in async mode. + """ + self._kv_cache.clear() + self._full_k_bufs.clear() + self._full_v_bufs.clear() + + # ------------------------------------------------------------------ + # PipeFusion forward + # ------------------------------------------------------------------ + + def infer(self, block_weights, pre_infer_out): + """Run this stage's blocks only. + + Returns ``(hidden_states, encoder_hidden_states, num_txt_tokens)``. + + For non-last stages, streams are ALWAYS split back to (image, text) + before returning, so P2P always carries separate streams with + consistent shapes. + """ + hidden_states = pre_infer_out.hidden_states + encoder_hidden_states = pre_infer_out.encoder_hidden_states + timestep = pre_infer_out.timestep + image_rotary_emb = pre_infer_out.image_rotary_emb + image_rotary_positions = pre_infer_out.image_rotary_positions + + # Compute num_txt_tokens + if encoder_hidden_states is not None: + num_txt_tokens = encoder_hidden_states.shape[0] + else: + # Streams already concatenated by previous stage — split them + txt_ids = pre_infer_out.txt_ids + num_txt_tokens = txt_ids.shape[0] if txt_ids is not None else 0 + if num_txt_tokens > 0: + encoder_hidden_states = hidden_states[:num_txt_tokens, ...] + hidden_states = hidden_states[num_txt_tokens:, ...] + + # Modulation embeddings (computed on every stage) + timestep_act = F.silu(timestep) + double_stream_mod_img = block_weights.double_stream_modulation_img_linear.apply(timestep_act) + double_stream_mod_txt = block_weights.double_stream_modulation_txt_linear.apply(timestep_act) + single_stream_mod = block_weights.single_stream_modulation_linear.apply(timestep_act) + + # Double-stream blocks (this stage's subset) + for block in block_weights.double_blocks: + encoder_hidden_states, hidden_states = self.infer_double_stream_block( + block, + hidden_states, + encoder_hidden_states, + double_stream_mod_img, + double_stream_mod_txt, + image_rotary_emb, + image_rotary_positions, + ) + + # Single-stream blocks: cat [text, image], run, then split back + has_single = len(block_weights.single_blocks) > 0 + if has_single: + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=0) + + for block in block_weights.single_blocks: + hidden_states = self.infer_single_stream_block( + block, + hidden_states, + None, + single_stream_mod, + image_rotary_emb, + image_rotary_positions, + num_txt_tokens=num_txt_tokens, + ) + + # Split back to (text, image) + encoder_hidden_states = hidden_states[:num_txt_tokens, ...] + hidden_states = hidden_states[num_txt_tokens:, ...] + + return hidden_states, encoder_hidden_states, num_txt_tokens diff --git a/lightx2v/models/networks/flux2/infer/pre_infer.py b/lightx2v/models/networks/flux2/infer/pre_infer.py index 651c6a95f..a55cb71b8 100644 --- a/lightx2v/models/networks/flux2/infer/pre_infer.py +++ b/lightx2v/models/networks/flux2/infer/pre_infer.py @@ -135,6 +135,37 @@ def infer(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img image_rotary_positions=image_rotary_positions, ) + def infer_partial(self, weights, hidden_states, encoder_hidden_states, txt_ids=None, img_ids=None): + """Compute timestep embedding and RoPE only (skip x_embedder / context_embedder). + + Used by non-first pipeline stages that receive already-embedded + hidden_states and encoder_hidden_states via P2P. + """ + timesteps_proj = self.scheduler.timesteps_proj + timestep_embed = weights.timestep_embedder_linear_1.apply(timesteps_proj) + timestep_embed = F.silu(timestep_embed) + timestep_embed = weights.timestep_embedder_linear_2.apply(timestep_embed) + + txt_ids_final = txt_ids if txt_ids is not None else getattr(self.scheduler, "txt_ids", None) + img_ids_final = img_ids if img_ids is not None else getattr(self.scheduler, "latent_image_ids", None) + + num_txt_tokens = encoder_hidden_states.shape[0] if encoder_hidden_states is not None else 0 + image_rotary_emb, image_rotary_positions = self.get_rope_cache(txt_ids_final, img_ids_final, num_txt_tokens) + if img_ids_final is not None and img_ids_final.ndim == 3: + img_ids_final = img_ids_final[0] + if txt_ids_final is not None and txt_ids_final.ndim == 3: + txt_ids_final = txt_ids_final[0] + + return Flux2PreInferModuleOutput( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep_embed, + txt_ids=txt_ids_final, + img_ids=img_ids_final, + image_rotary_emb=image_rotary_emb, + image_rotary_positions=image_rotary_positions, + ) + class Flux2DevPreInfer(Flux2PreInfer): """Pre-processing inference for Flux2 Dev. diff --git a/lightx2v/models/networks/flux2/infer/transformer_infer.py b/lightx2v/models/networks/flux2/infer/transformer_infer.py index ba8a754eb..a37695a17 100644 --- a/lightx2v/models/networks/flux2/infer/transformer_infer.py +++ b/lightx2v/models/networks/flux2/infer/transformer_infer.py @@ -32,6 +32,14 @@ def __init__(self, config): self.seq_p_fp4_comm = False self.enable_head_parallel = False + def _maybe_apply_stale_kv(self, key, value, num_txt_tokens, block_idx): + """Hook for stale-KV cache in PipeFusion mode. No-op in base class. + + Subclasses (PipeFusion) override this to cache image KV across patches + while keeping text KV fresh. + """ + return key, value + def set_scheduler(self, scheduler): self.scheduler = scheduler @@ -95,8 +103,14 @@ def infer_double_stream_block( query, key = block_weights.rope.apply(query, key, image_rotary_emb, positions=image_rotary_positions) + # Stale-KV hook (no-op in base class; PipeFusion subclass overrides) + num_txt_tokens = encoder_hidden_states.shape[0] + key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx) + total_len = query.shape[0] - cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32) + kv_len = key.shape[0] # may differ from total_len in PipeFusion (stale-KV) + cu_seqlens_q = torch.tensor([0, total_len], dtype=torch.int32) + cu_seqlens_kv = torch.tensor([0, kv_len], dtype=torch.int32) model_cls = self.config.get("model_cls", "flux2_klein") @@ -107,7 +121,7 @@ def infer_double_stream_block( k=key, v=value, slice_qkv_len=txt_len, - cu_seqlens_qkv=cu_seqlens, + cu_seqlens_qkv=cu_seqlens_q, attention_module=block_weights.calculate, seq_p_group=self.seq_p_group, use_fp8_comm=self.seq_p_fp8_comm, @@ -121,10 +135,10 @@ def infer_double_stream_block( q=query, k=key, v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=total_len, - max_seqlen_kv=total_len, + max_seqlen_kv=kv_len, model_cls=model_cls, ) @@ -197,8 +211,13 @@ def infer_single_stream_block( query, key = block_weights.rope.apply(query, key, image_rotary_emb, positions=image_rotary_positions) + # Stale-KV hook (no-op in base class; PipeFusion subclass overrides) + key, value = self._maybe_apply_stale_kv(key, value, num_txt_tokens, block_weights.block_idx) + total_len = query.shape[0] - cu_seqlens = torch.tensor([0, total_len], dtype=torch.int32) + kv_len = key.shape[0] # may differ from total_len in PipeFusion (stale-KV) + cu_seqlens_q = torch.tensor([0, total_len], dtype=torch.int32) + cu_seqlens_kv = torch.tensor([0, kv_len], dtype=torch.int32) model_cls = self.config.get("model_cls", "flux2_klein") @@ -208,7 +227,7 @@ def infer_single_stream_block( k=key, v=value, slice_qkv_len=num_txt_tokens, - cu_seqlens_qkv=cu_seqlens, + cu_seqlens_qkv=cu_seqlens_q, attention_module=block_weights.calculate, seq_p_group=self.seq_p_group, use_fp8_comm=self.seq_p_fp8_comm, @@ -222,10 +241,10 @@ def infer_single_stream_block( q=query, k=key, v=value, - cu_seqlens_q=cu_seqlens, - cu_seqlens_kv=cu_seqlens, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=total_len, - max_seqlen_kv=total_len, + max_seqlen_kv=kv_len, model_cls=model_cls, ) diff --git a/lightx2v/models/networks/flux2/model.py b/lightx2v/models/networks/flux2/model.py index 9f9e15c5a..88d18e568 100644 --- a/lightx2v/models/networks/flux2/model.py +++ b/lightx2v/models/networks/flux2/model.py @@ -29,6 +29,10 @@ def __init__(self, config, model_path, device): self._init_tensor_parallel() self._init_infer_class() self._init_weights() + # In PipeFusion mode, weights were loaded to CPU to avoid OOM; + # move only this stage's subset to GPU. + if self.config.get("pipefusion_parallel", False): + self.to_cuda() self._init_infer() def _init_tensor_parallel(self): @@ -272,7 +276,8 @@ def _init_infer(self): self.transformer_infer = self.transformer_infer_class(self.config) self.pre_infer = self.pre_infer_class(self.config) self.post_infer = self.post_infer_class(self.config) - self.pre_infer.set_rope(self.transformer_weights.double_blocks[0].rope) + blocks = self.transformer_weights.double_blocks or self.transformer_weights.single_blocks + self.pre_infer.set_rope(blocks[0].rope) if hasattr(self.transformer_infer, "offload_manager_double") and hasattr(self.transformer_infer, "offload_manager_single"): self._init_offload_manager() @@ -357,7 +362,13 @@ class Flux2KleinTransformerModel(_Flux2TransformerModelBase): def _init_infer_class(self): feature_caching = self.config.get("feature_caching", "NoCaching") - if feature_caching in ("NoCaching", "None"): + if self.config.get("pipefusion_parallel", False): + from lightx2v.models.networks.flux2.infer.pipefusion.transformer_infer import ( + Flux2PipeFusionTransformerInfer, + ) + + self.transformer_infer_class = Flux2PipeFusionTransformerInfer + elif feature_caching in ("NoCaching", "None"): if self.cpu_offload and self.offload_granularity == "block": self.transformer_infer_class = Flux2OffloadTransformerInfer else: diff --git a/lightx2v/models/networks/flux2/weights/transformer_weights.py b/lightx2v/models/networks/flux2/weights/transformer_weights.py index 5a3790c8c..d9e15b298 100644 --- a/lightx2v/models/networks/flux2/weights/transformer_weights.py +++ b/lightx2v/models/networks/flux2/weights/transformer_weights.py @@ -163,8 +163,51 @@ def __init__(self, config): self.num_single_layers = config.get("num_single_layers", 20) self.mm_type = config.get("dit_quant_scheme", "Default") - 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)]) + # -- Pipeline-parallel block splitting -------------------------------- + pp_size = config.get("pipefusion_parallel", False) + if pp_size: + from lightx2v.common.distributed import ( + get_pipeline_parallel_rank, + get_pipeline_parallel_world_size, + ) + + pp_rank = get_pipeline_parallel_rank() + pp_world_size = get_pipeline_parallel_world_size() + else: + pp_rank = 0 + pp_world_size = 1 + + if pp_world_size > 1: + # Split double_blocks + single_blocks across pipeline stages. + # Blocks are assigned contiguously: stage 0 gets the first chunk, + # stage 1 the next, etc. A stage may span the double→single + # boundary (it will then have both types). + total_blocks = self.num_layers + self.num_single_layers + base = total_blocks // pp_world_size + remainder = total_blocks % pp_world_size + stage_start = pp_rank * base + min(pp_rank, remainder) + stage_end = stage_start + base + (1 if pp_rank < remainder else 0) + + double_start = min(stage_start, self.num_layers) + double_end = min(stage_end, self.num_layers) + single_start = max(0, stage_start - self.num_layers) + single_end = max(0, stage_end - self.num_layers) + + self.double_blocks = WeightModuleList([Flux2DoubleBlockWeights(config, i) for i in range(double_start, double_end)]) + self.single_blocks = WeightModuleList([Flux2SingleBlockWeights(config, i) for i in range(single_start, single_end)]) + # Track whether this stage crosses the double→single boundary + self._has_double = double_end > double_start + self._has_single = single_end > single_start + self._stage_start = stage_start + self._stage_end = stage_end + else: + 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._has_double = True + self._has_single = True + self._stage_start = 0 + self._stage_end = self.num_layers + self.num_single_layers + self.register_offload_buffers(config) self.add_module("double_blocks", self.double_blocks) diff --git a/lightx2v/models/runners/flux2/flux2_runner.py b/lightx2v/models/runners/flux2/flux2_runner.py index b5d405853..55a8f291a 100644 --- a/lightx2v/models/runners/flux2/flux2_runner.py +++ b/lightx2v/models/runners/flux2/flux2_runner.py @@ -243,6 +243,12 @@ def _run_dit_local_i2i(self, total_steps=None): return latents, generator def run(self, total_steps=None): + if self.config.get("pipefusion_parallel", False): + return self._run_pipefusion(total_steps) + return self._run_sequential(total_steps) + + def _run_sequential(self, total_steps=None): + """Existing synchronous denoising loop (single-GPU or non-PipeFusion).""" if total_steps is None: total_steps = self.model.scheduler.infer_steps for step_index in range(total_steps): @@ -262,6 +268,75 @@ def run(self, total_steps=None): return self.model.scheduler.latents, self.model.scheduler.generator + def _run_pipefusion(self, total_steps=None): + """PipeFusion denoising loop: pipeline driver controls all timesteps.""" + from lightx2v.common.distributed import ( + get_pipeline_runtime_state, + is_pipeline_last_stage, + ) + + if total_steps is None: + total_steps = self.model.scheduler.infer_steps + + # Initialize pipeline runtime state with image dimensions + pipeline_state = get_pipeline_runtime_state() + height = self.input_info.latent_shape[1] # packed_h * packed_w tokens + # Reconstruct actual height/width from latent_image_ids + latent_image_ids = self.model.scheduler.latent_image_ids + if self.input_info.target_shape is not None: + actual_height, actual_width = self.input_info.target_shape + else: + actual_height = actual_width = 1024 + + num_pipeline_patch = self.config.get("parallel", {}).get("num_pipeline_patch", 4) + warmup_steps = self.config.get("parallel", {}).get("pipeline_warmup_steps", 1) + + pipeline_state.set_input_parameters( + height=actual_height, + width=actual_width, + batch_size=1, + num_pipeline_patch=num_pipeline_patch, + warmup_steps=warmup_steps, + vae_scale_factor=self.config.get("vae_scale_factor", 16), + total_tokens=self.input_info.latent_shape[1], + ) + + # Prepare inputs + latents = self.model.scheduler.latents + text_encoder_output = self.inputs["text_encoder_output"] + prompt_embeds = text_encoder_output["prompt_embeds"] + text_ids = text_encoder_output.get("text_ids") + latent_image_ids = self.model.scheduler.latent_image_ids + + do_cfg = self.config.get("enable_cfg", True) and self.config.get("sample_guide_scale", 1.0) > 1.0 + negative_prompt_embeds = text_encoder_output.get("negative_prompt_embeds") if do_cfg else None + negative_text_ids = text_encoder_output.get("negative_text_ids") if do_cfg else None + + timesteps = self.model.scheduler.timesteps + + # Run pipeline + from lightx2v.models.networks.flux2.infer.pipefusion.pipeline_driver import ( + Flux2PipelineDriver, + ) + + driver = Flux2PipelineDriver(self.model, self.config) + latents = driver.run_pipeline( + latents=latents, + prompt_embeds=prompt_embeds, + text_ids=text_ids, + latent_image_ids=latent_image_ids, + timesteps=timesteps, + scheduler=self.model.scheduler, + do_cfg=do_cfg, + negative_prompt_embeds=negative_prompt_embeds, + negative_text_ids=negative_text_ids, + ) + + if latents is not None and is_pipeline_last_stage(): + self.model.scheduler.latents = latents + + return self.model.scheduler.latents, self.model.scheduler.generator + def get_custom_shape(self): default_aspect_ratios = { "16:9": [1344, 768], @@ -370,13 +445,33 @@ def run_pipeline(self, input_info): self.set_img_shapes() latents, generator = self.run_dit() - images = self.run_vae_decoder(latents) + + # In PipeFusion mode, only the last stage has final latents + if self.config.get("pipefusion_parallel", False): + from lightx2v.common.distributed import is_pipeline_last_stage + + if is_pipeline_last_stage(): + # Offload transformer weights and clear KV cache before VAE decode to avoid OOM + self.model.transformer_weights.to_cpu() + if hasattr(self.model.transformer_infer, "clear_kv_cache"): + self.model.transformer_infer.clear_kv_cache() + torch_device_module.empty_cache() + gc.collect() + images = self.run_vae_decoder(latents) + else: + images = None + else: + images = self.run_vae_decoder(latents) self.end_run() - if not input_info.return_result_tensor and is_main_process(): - image = images[0] - image.save(input_info.save_result_path) - logger.info(f"Image saved: {input_info.save_result_path}") + # Save image: in PipeFusion mode, last stage has the image; + # in normal mode, main process (rank 0) has it. + if not input_info.return_result_tensor: + should_save = is_pipeline_last_stage() if self.config.get("pipefusion_parallel", False) else is_main_process() + if should_save and images is not None: + image = images[0] + image.save(input_info.save_result_path) + logger.info(f"Image saved: {input_info.save_result_path}") del latents, generator torch_device_module.empty_cache() diff --git a/lightx2v/models/schedulers/flux2/scheduler.py b/lightx2v/models/schedulers/flux2/scheduler.py index 32e007cde..432bf4321 100755 --- a/lightx2v/models/schedulers/flux2/scheduler.py +++ b/lightx2v/models/schedulers/flux2/scheduler.py @@ -146,6 +146,15 @@ def step_post(self): ) self.latents = latents + def step_post_patch(self, noise_pred, latents, t): + """Patch-level scheduler step for async PipeFusion mode. + + Unlike ``step_post``, this operates on a single patch's latents and + noise_pred, and does not apply FLS enhancement (which requires the + full latent). + """ + return self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + def _encode_image(self, image): image = image.to(device=AI_DEVICE, dtype=GET_DTYPE()) encoder_output = self.vae.encode_vae_image(image) diff --git a/lightx2v/pipeline.py b/lightx2v/pipeline.py index 41ab93634..0eb125d44 100755 --- a/lightx2v/pipeline.py +++ b/lightx2v/pipeline.py @@ -422,11 +422,14 @@ def enable_cache( self.magcache_retention_ratio = magcache_retention_ratio self.magcache_ratios = magcache_ratios - def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses"): + def enable_parallel(self, cfg_p_size=1, seq_p_size=1, seq_p_attn_type="ulysses", pp_size=1, num_pipeline_patch=4, pipeline_warmup_steps=1): self.parallel = { "cfg_p_size": cfg_p_size, "seq_p_size": seq_p_size, "seq_p_attn_type": seq_p_attn_type, + "pp_size": pp_size, + "num_pipeline_patch": num_pipeline_patch, + "pipeline_warmup_steps": pipeline_warmup_steps, } @torch.no_grad() diff --git a/lightx2v/utils/set_config.py b/lightx2v/utils/set_config.py index 8497133b1..516a527ed 100755 --- a/lightx2v/utils/set_config.py +++ b/lightx2v/utils/set_config.py @@ -28,6 +28,7 @@ def get_default_config(): "parallel": False, "seq_parallel": False, "cfg_parallel": False, + "pipefusion_parallel": False, "enable_cfg": False, "use_image_encoder": True, } @@ -396,11 +397,12 @@ def set_parallel_config(config): tensor_p_size = int(config["parallel"].get("tensor_p_size", 1)) cfg_p_size = int(config["parallel"].get("cfg_p_size", 1)) seq_p_size = int(config["parallel"].get("seq_p_size", 1)) + pp_size = int(config["parallel"].get("pp_size", 1)) world_size = dist.get_world_size() - expected_world_size = tensor_p_size * cfg_p_size * seq_p_size + expected_world_size = tensor_p_size * cfg_p_size * seq_p_size * pp_size if expected_world_size != world_size: raise ValueError( - f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) != world_size ({world_size})." + f"Parallel sizes must match the distributed world size: tensor_p_size ({tensor_p_size}) * cfg_p_size ({cfg_p_size}) * seq_p_size ({seq_p_size}) * pp_size ({pp_size}) != world_size ({world_size})." ) if tensor_p_size > 1: @@ -424,15 +426,29 @@ def set_parallel_config(config): mesh_dim_names=tuple(mesh_dim_names), ) config["tensor_parallel"] = True - config["seq_parallel"] = seq_p_size > 1 - config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) + config["seq_parallel"] = False + config["cfg_parallel"] = False + config["pipefusion_parallel"] = False else: - # Original 2D mesh for cfg_p and seq_p - config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, seq_p_size), mesh_dim_names=("cfg_p", "seq_p")) + # Multi-dimensional mesh: (cfg_p, pp, seq_p) + cfg_p_size = config["parallel"].get("cfg_p_size", 1) + pp_size = config["parallel"].get("pp_size", 1) + seq_p_size = config["parallel"].get("seq_p_size", 1) + assert cfg_p_size * pp_size * seq_p_size == dist.get_world_size(), ( + f"cfg_p_size ({cfg_p_size}) * pp_size ({pp_size}) * seq_p_size ({seq_p_size}) must be equal to world_size ({dist.get_world_size()})" + ) + config["device_mesh"] = init_device_mesh(AI_DEVICE, (cfg_p_size, pp_size, seq_p_size), mesh_dim_names=("cfg_p", "pp", "seq_p")) config["tensor_parallel"] = False config["seq_parallel"] = seq_p_size > 1 config["cfg_parallel"] = bool(config.get("enable_cfg", False) and cfg_p_size > 1) + config["pipefusion_parallel"] = pp_size > 1 + if pp_size > 1: + from lightx2v.common.distributed import init_pipeline_parallel_state + + pp_group = config["device_mesh"].get_group(mesh_dim="pp") + init_pipeline_parallel_state(pp_group) + # warmup dist if AI_DEVICE == "cuda": warmup_device = f"{AI_DEVICE}:{torch.cuda.current_device()}" @@ -440,6 +456,8 @@ def set_parallel_config(config): warmup_device = AI_DEVICE _a = torch.zeros([1], device=warmup_device) dist.all_reduce(_a) + else: + config["pipefusion_parallel"] = False def print_config(config):