From a46f5ea380ee1398c8286ec7ff721a3068d9ab92 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 23 Nov 2025 05:33:37 +0000 Subject: [PATCH] working --- .../basic/basic_self_forcing_causal.py | 3 +- fastvideo/attention/layer.py | 31 ++----- fastvideo/layers/rotary_embedding.py | 12 +++ fastvideo/models/dits/causal_wanvideo.py | 31 ++++++- .../pipelines/stages/causal_denoising.py | 84 +++++++------------ 5 files changed, 82 insertions(+), 79 deletions(-) diff --git a/examples/inference/basic/basic_self_forcing_causal.py b/examples/inference/basic/basic_self_forcing_causal.py index 1a0c91668..903bdceb3 100644 --- a/examples/inference/basic/basic_self_forcing_causal.py +++ b/examples/inference/basic/basic_self_forcing_causal.py @@ -20,7 +20,8 @@ def main(): ) sampling_param = SamplingParam.from_pretrained(model_name) - sampling_param.num_frames = 13 + # sampling_param.num_frames = 13 + sampling_param.num_frames = 77 prompt = ( "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " "wide with interest. The playful yet serene atmosphere is complemented by soft " diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index 4ca94e31e..f2d850ca7 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -70,7 +70,8 @@ class DistributedAttention(nn.Module): k_cache: torch.Tensor | None = None, v_cache: torch.Tensor | None = None, return_current_kv: bool = False, - ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, + torch.Tensor | None]: """Forward pass for distributed attention. Args: @@ -97,32 +98,18 @@ class DistributedAttention(nn.Module): ctx_attn_metadata = forward_context.attn_metadata # Stack QKV - logger.info( - f"rank: {local_rank}, q: {q.shape}, k: {k.shape}, v: {v.shape}", - local_main_process_only=False) qkv = torch.cat([q, k, v], dim=0) # [3, seq_len, num_heads, head_dim] # Redistribute heads across sequence dimension qkv = sequence_model_parallel_all_to_all_4D(qkv, scatter_dim=2, gather_dim=1) - # q = sequence_model_parallel_all_to_all_4D(q, - # scatter_dim=2, - # gather_dim=1) - # k = sequence_model_parallel_all_to_all_4D(k, - # scatter_dim=2, - # gather_dim=1) - # v = sequence_model_parallel_all_to_all_4D(v, - # scatter_dim=2, - # gather_dim=1) - # Apply backend-specific preprocess_qkv qkv = self.attn_impl.preprocess_qkv(qkv, ctx_attn_metadata) # Concatenate with replicated QKV if provided if replicated_q is not None: - assert False, "should not be used" assert replicated_k is not None and replicated_v is not None replicated_qkv = torch.cat( [replicated_q, replicated_k, replicated_v], @@ -141,21 +128,19 @@ class DistributedAttention(nn.Module): v_total = v_new if k_cache is not None or v_cache is not None: - assert False, "KV cache is not supported" + # assert False, "KV cache is not supported" assert k_cache is not None and v_cache is not None assert k_cache.shape == v_cache.shape logger.info( f"rank: {local_rank}, k_cache: {k_cache.shape}, v_cache: {v_cache.shape}", local_main_process_only=False) - assert k_cache.shape[2] == v_cache.shape[2], "Number of heads must be the same" + assert k_cache.shape[2] == v_cache.shape[ + 2], "Number of heads must be the same" if k_cache.shape[1] > 0: k_total = torch.cat([k_cache, k_new], dim=1) v_total = torch.cat([v_cache, v_new], dim=1) assert k_total.shape == v_total.shape, "Key and value shapes must be the same" - logger.info( - f"rank: {local_rank}, k_total: {k_total.shape}, v_total: {v_total.shape}", - local_main_process_only=False) output = self.attn_impl.forward(q, k_total, v_total, ctx_attn_metadata) # Redistribute back if using sequence parallelism @@ -166,6 +151,7 @@ class DistributedAttention(nn.Module): # TODO: make this asynchronous replicated_output = sequence_model_parallel_all_gather( replicated_output.contiguous(), dim=2) + # Apply backend-specific postprocess_output output = self.attn_impl.postprocess_output(output, ctx_attn_metadata) @@ -173,7 +159,7 @@ class DistributedAttention(nn.Module): scatter_dim=1, gather_dim=2) if return_current_kv: - return output, replicated_output, k_new, v_new # [batch_size, seq_len, num_heads/sp_world_size, head_dim] + return output, replicated_output, k_new, v_new # [batch_size, seq_len, num_heads/sp_world_size, head_dim] else: return output, replicated_output, None, None @@ -193,7 +179,8 @@ class DistributedAttention_VSA(DistributedAttention): replicated_v: torch.Tensor | None = None, gate_compress: torch.Tensor | None = None, **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, + torch.Tensor | None]: """Forward pass for distributed attention. Args: diff --git a/fastvideo/layers/rotary_embedding.py b/fastvideo/layers/rotary_embedding.py index 0a1c1dc28..6f6048708 100644 --- a/fastvideo/layers/rotary_embedding.py +++ b/fastvideo/layers/rotary_embedding.py @@ -431,6 +431,7 @@ def get_rotary_pos_embed( shard_dim: int = 0, dtype: torch.dtype = torch.float32, start_frame: int = 0, + use_sp_shard: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """ Generate rotary positional embeddings for the given sizes. @@ -444,6 +445,8 @@ def get_rotary_pos_embed( theta_rescale_factor: Rescale factor for theta. Defaults to 1.0 interpolation_factor: Factor to scale positions. Defaults to 1.0 shard_dim: Which dimension to shard for sequence parallelism. Defaults to 0. + use_sp_shard: Whether to shard rotary embeddings for SP. Defaults to True. + Set to False for Ulysses SP where all-to-all gives each rank the full sequence. Returns: Tuple of (cos, sin) tensors for rotary embeddings @@ -464,6 +467,15 @@ def get_rotary_pos_embed( sp_rank = sp_group.rank_in_group sp_world_size = sp_group.world_size + # For Ulysses SP, don't shard rotary embeddings + if not use_sp_shard: + sp_rank = 0 + sp_world_size = 1 + + logger.info( + f"sp_rank: {sp_rank}, sp_world_size: {sp_world_size}, use_sp_shard: {use_sp_shard}", + local_main_process_only=False) + freqs_cos, freqs_sin = get_nd_rotary_pos_embed( rope_dim_list, rope_sizes, diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index fac623f80..dc3dcf256 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -89,6 +89,25 @@ class CausalWanSelfAttention(nn.Module): cache_start = current_start cos, sin = freqs_cis + + # For Ulysses SP, slice rotary embeddings to match local sequence portion + # Each rank has a portion of the sequence, but freqs_cis is for the full sequence + from fastvideo.distributed.parallel_state import get_sp_group + sp_group = get_sp_group() + sp_world_size = sp_group.world_size + if sp_world_size > 1: + # Calculate this rank's slice of the sequence + local_seq_len = q.shape[1] # sequence length on this rank + total_seq_len = local_seq_len * sp_world_size + rank_in_sp_group = sp_group.rank_in_group + start_idx = rank_in_sp_group * local_seq_len + end_idx = start_idx + local_seq_len + + # Slice the rotary embeddings to match local sequence + cos = cos[start_idx:end_idx] + sin = sin[start_idx:end_idx] + logger.info(f'rank: {rank_in_sp_group}, slicing RoPE from {start_idx}:{end_idx}, local_seq_len: {local_seq_len}', local_main_process_only=False) + roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v) roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v) @@ -132,7 +151,7 @@ class CausalWanSelfAttention(nn.Module): sp_rank = sp_group.rank_in_group logger.info('rank: %s, num_new_tokens: %s', sp_rank, num_new_tokens, local_main_process_only=False) - prev_global_end = kv_cache["global_end_index"].item() + # prev_global_end = kv_cache["global_end_index"].item() # prev_local_end = kv_cache["local_end_index"].item() prev_local_end = current_start # prev_local_end = @@ -512,6 +531,8 @@ class CausalWanTransformer3DModel(BaseDiT): # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + # For Ulysses SP, we need full rotary embeddings (no temporal sharding) + # because after all-to-all, each rank will have the full sequence freqs_cos, freqs_sin = get_rotary_pos_embed( (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), @@ -520,7 +541,8 @@ class CausalWanTransformer3DModel(BaseDiT): rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - start_frame=start_frame # Assume that start_frame is 0 when kv_cache is None + start_frame=start_frame, # Assume that start_frame is 0 when kv_cache is None + use_sp_shard=False # Don't shard rotary embeddings for Ulysses SP ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) @@ -611,6 +633,8 @@ class CausalWanTransformer3DModel(BaseDiT): # Get rotary embeddings d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + # For Ulysses SP, we need full rotary embeddings (no temporal sharding) + # because after all-to-all, each rank will have the full sequence freqs_cos, freqs_sin = get_rotary_pos_embed( (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), @@ -619,7 +643,8 @@ class CausalWanTransformer3DModel(BaseDiT): rope_dim_list, dtype=torch.float32 if current_platform.is_mps() else torch.float64, rope_theta=10000, - start_frame=start_frame + start_frame=start_frame, + use_sp_shard=False # Don't shard rotary embeddings for Ulysses SP ) freqs_cos = freqs_cos.to(hidden_states.device) freqs_sin = freqs_sin.to(hidden_states.device) diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 2907f6912..e3f45f6f6 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/fastvideo/pipelines/stages/causal_denoising.py @@ -222,46 +222,38 @@ class CausalDMDDenosingStage(DenoisingStage): block_sizes.pop(0) latents[:, :, :1, :, :] = first_frame_latent + # Setup noise generator for SP + sp_world_size = get_sp_world_size() + sp_rank = get_sp_parallel_rank() + sp_group = sp_world_size > 1 + noise_generator = batch.generator[0] if isinstance( + batch.generator, list) else batch.generator + + if sp_group: + try: + # Use initial_seed + rank to ensure diversity across ranks + seed = noise_generator.initial_seed() + noise_generator = torch.Generator( + device=noise_generator.device).manual_seed(seed + sp_rank) + except Exception as e: + logger.warning(f"Could not create rank-specific generator: {e}") + # DMD loop in causal blocks with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: for idx, current_num_frames in enumerate(block_sizes): - logger.info("================================================") - logger.info(f"Block {idx}") - logger.info("================================================") - logger.info(f"Current num frames: {current_num_frames}") - logger.info( - f"Processing block {start_index} to {start_index + current_num_frames}" - ) - logger.info(f"Latents shape: {latents.shape}") current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] sp_world_size, rank_in_sp_group = get_sp_world_size( ), get_sp_parallel_rank() sp_group = sp_world_size > 1 - # logger.info(f"SP group: {sp_group}") - logger.info( - f"before SP: Current latents shape: {current_latents.shape}" - ) if sp_group: current_latents = rearrange(current_latents, "b c (n t) h w -> b c n t h w", n=sp_world_size).contiguous() - logger.info(f"rearranged latents shape: {current_latents.shape}") current_latents = current_latents[:, :, rank_in_sp_group, :, :, :] - logger.info( - f"after SP: Current latents shape: {current_latents.shape}" - ) - # Per-rank temporal offset (in frames) within this block when using SP - frames_per_rank = ( - current_num_frames // - sp_world_size) if sp_group else current_num_frames - # logger.info(f"frames_per_rank: {frames_per_rank}") - rank_offset_frames = (rank_in_sp_group * - frames_per_rank) if sp_group else 0 - # logger.info(f"rank {rank_in_sp_group} rank_offset_frames: {rank_offset_frames}", local_main_process_only=False) # use BTCHW for DMD conversion routines noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) video_raw_latent_shape = noise_latents_btchw.shape @@ -360,12 +352,10 @@ class CausalDMDDenosingStage(DenoisingStage): [1], dtype=torch.long, device=pred_video_btchw.device) - noise = torch.randn( - video_raw_latent_shape, - dtype=pred_video_btchw.dtype, - generator=(batch.generator[0] if isinstance( - batch.generator, list) else - batch.generator)).to(self.device) + noise = torch.randn(video_raw_latent_shape, + dtype=pred_video_btchw.dtype, + generator=noise_generator).to( + self.device) noise_btchw = noise if boundary_timestep is not None and i < len( high_noise_timesteps) - 1: @@ -393,25 +383,6 @@ class CausalDMDDenosingStage(DenoisingStage): if progress_bar is not None: progress_bar.update() - # Write back and advance - logger.info( - f"before all gather: Current latents shape: {current_latents.shape}" - ) - context_bcthw = current_latents.clone().to(target_dtype) - if sp_group: - current_latents = sequence_model_parallel_all_gather( - current_latents, dim=2) - logger.info( - f"after all gather: Current latents shape: {current_latents.shape}" - ) - logger.info( - f"start_index: {start_index}, current_num_frames: {current_num_frames}" - ) - logger.info( - f"before write back: Latents shape: {latents.shape}") - latents[:, :, start_index:start_index + - current_num_frames, :, :] = current_latents - logger.info(f"after write back: Latents shape: {latents.shape}") # Re-run with context timestep to update KV cache using clean context context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) @@ -419,8 +390,6 @@ class CausalDMDDenosingStage(DenoisingStage): device=latents.device, dtype=torch.long) * int(context_noise) context_bcthw = current_latents.to(target_dtype) - # context_frame_start = (pos_start_base + start_index) - # context_token_start = context_frame_start * self.frame_seq_length with torch.autocast(device_type="cuda", dtype=target_dtype, @@ -458,6 +427,13 @@ class CausalDMDDenosingStage(DenoisingStage): **pos_cond_kwargs, ) + if sp_group: + current_latents = sequence_model_parallel_all_gather( + current_latents, dim=2) + + # Write back and advance + latents[:, :, start_index:start_index + + current_num_frames, :, :] = current_latents start_index += current_num_frames if boundary_timestep is not None: @@ -474,7 +450,9 @@ class CausalDMDDenosingStage(DenoisingStage): kv_cache1 = [] num_attention_heads = self.transformer.num_attention_heads sp_world_size = get_sp_world_size() - assert num_attention_heads % max(sp_world_size, 1) == 0, "num_attention_heads must be divisible by sp_world_size" + assert num_attention_heads % max( + sp_world_size, + 1) == 0, "num_attention_heads must be divisible by sp_world_size" heads_per_rank = num_attention_heads // max(sp_world_size, 1) attention_head_dim = self.transformer.attention_head_dim if self.local_attn_size != -1: @@ -557,4 +535,4 @@ class CausalDMDDenosingStage(DenoisingStage): result.add_check( "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) - return result \ No newline at end of file + return result