From 5a06b460b2ac1de307d7bb11bb5a9acdd591cb2f Mon Sep 17 00:00:00 2001 From: Peter Kisfaludi Date: Mon, 20 Jul 2026 18:04:16 -0700 Subject: [PATCH] Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT Under Ulysses SP (num_gpus>=2), self-forcing causal Wan models (causal_wanvideo.py) crash in `_apply_rotary_emb`: RuntimeError: The size of tensor a (4680) must match tensor b (9360) at non-singleton dimension 1 Cause: the DiT builds RoPE for `post_patch_num_frames * get_sp_world_size()` (sp x too long) but the causal pipeline *replicates* the full token sequence and a full-size KV cache on every rank (it does not shard the sequence), so the attention applies that over-long RoPE to the full local query. Fix: - Remove `* get_sp_world_size()` from the RoPE length (relativistic, standard, and train branches) so RoPE matches the full local sequence. No-op at sp=1. - Parallelize the causal attention across HEADS (Ulysses over heads): each rank computes attention for its head slice `[rank*H/sp:(rank+1)*H/sp]` over the full sequence, then `sequence_model_parallel_all_gather(dim=2)` reconstructs all heads. Applied uniformly to all three attention paths (flex_attention, KV-cache `self.attn`, and the relativistic window). A sequence all-to-all a la wanvideo.py is incompatible with the causal KV-cache windowing (global token coords), so head-parallel is the exact-matching approach; heads are independent so it is numerically identical to single-GPU. Verified: `wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers`, 480x832, 21 frames (7 autoregressive KV-cache blocks), seed 1234, on 2x GPU vs 1x GPU: torch.equal=True, max_abs_diff=0.0 (byte-identical). 2-GPU previously crashed. Co-Authored-By: Claude Opus 4.8 (1M context) --- fastvideo/models/dits/causal_wanvideo.py | 60 ++++++++++++++++++++---- 1 file changed, 52 insertions(+), 8 deletions(-) diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 2875edec1..83ce09ec1 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -18,7 +18,10 @@ import torch.distributed as dist import fastvideo.envs as envs from fastvideo.attention import (DistributedAttention, LocalAttention) from fastvideo.configs.models.dits import WanVideoConfig -from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather) +from fastvideo.distributed.parallel_state import (get_sp_parallel_rank, + get_sp_world_size) from fastvideo.forward_context import get_forward_context from fastvideo.layers.layernorm import (FP32LayerNorm, LayerNormScaleShift, RMSNorm, ScaleResidual, ScaleResidualLayerNormScaleShift) @@ -98,6 +101,29 @@ class CausalWanSelfAttention(nn.Module): 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) + # Ulysses sequence-parallel (SP): the causal pipeline replicates the full + # (unsharded) token sequence on every rank, so RoPE, the block mask and + # the KV cache are all full-length here. We parallelize the attention + # across the head dimension (à la Ulysses): each rank computes attention + # for its own slice of heads over the full sequence, then all-gathers the + # per-head outputs to reconstruct all heads. This is numerically identical + # to single-GPU attention (heads are independent) and avoids a sequence + # all-to-all that would be incompatible with the causal KV-cache windowing + # (expressed in global token coordinates). Applied uniformly to all three + # attention paths below (flex, relativistic window, and the standard + # KV-cache path). + sp_world_size = get_sp_world_size() + sp_rank = get_sp_parallel_rank() + if sp_world_size > 1: + assert self.num_heads % sp_world_size == 0, ( + f"num_heads ({self.num_heads}) must be divisible by sequence " + f"parallel size ({sp_world_size})") + heads_per_rank = self.num_heads // sp_world_size + head_slice = slice(sp_rank * heads_per_rank, + (sp_rank + 1) * heads_per_rank) + else: + head_slice = slice(None) + if kv_cache is None: # Padding for flex attention padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] @@ -117,10 +143,13 @@ class CausalWanSelfAttention(nn.Module): [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], device=v.device, dtype=v.dtype)], dim=1) - x = flex_attention(query=padded_roped_query.transpose(2, 1), - key=padded_roped_key.transpose(2, 1), - value=padded_v.transpose(2, 1), + x = flex_attention(query=padded_roped_query[:, :, head_slice].transpose(2, 1), + key=padded_roped_key[:, :, head_slice].transpose(2, 1), + value=padded_v[:, :, head_slice].transpose(2, 1), block_mask=block_mask)[:, :, :-padded_length].transpose(2, 1) + if sp_world_size > 1: + # Gather per-rank head outputs back to the full head dimension. + x = sequence_model_parallel_all_gather(x.contiguous(), dim=2) else: current_end = current_start + q.shape[1] sink_tokens = self.sink_size * frame_seqlen @@ -178,7 +207,11 @@ class CausalWanSelfAttention(nn.Module): is_neox_style=False).type_as(v) key_window = _apply_rotary_emb(key_window, cos[:window_len], sin[:window_len], is_neox_style=False).type_as(v) - x = self.attn(roped_query, key_window, value_window) + x = self.attn(roped_query[:, :, head_slice], key_window[:, :, head_slice], + value_window[:, :, head_slice]) + if sp_world_size > 1: + # Gather per-rank head outputs back to the full head dimension. + x = sequence_model_parallel_all_gather(x.contiguous(), dim=2) if isinstance(kv_cache["global_end_index"], torch.Tensor): kv_cache["global_end_index"].fill_(current_end) else: @@ -581,10 +614,16 @@ class CausalWanTransformer3DModel(BaseDiT): # fixed table over [0, max_attention_frames); attention slices it per step max_attention_frames = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES if self.local_attn_size == -1 else self.local_attn_size) - rope_num_frames = max_attention_frames * get_sp_world_size() + # No multiplication by the sequence-parallel world size: the causal + # pipeline replicates the full (unsharded) sequence on every rank, so + # RoPE must cover the full local sequence length. Multiplying by the + # SP world size produced a 2x-too-long RoPE and crashed + # _apply_rotary_emb under SP. Attention is parallelized across heads + # instead (see CausalWanSelfAttention.forward). + rope_num_frames = max_attention_frames rope_start_frame = 0 else: - rope_num_frames = post_patch_num_frames * get_sp_world_size() + rope_num_frames = post_patch_num_frames rope_start_frame = start_frame # 0 when kv_cache is None freqs_cos, freqs_sin = get_rotary_pos_embed((rope_num_frames, post_patch_height, post_patch_width), self.hidden_size, @@ -682,7 +721,12 @@ class CausalWanTransformer3DModel(BaseDiT): d = self.hidden_size // self.num_attention_heads rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] freqs_cos, freqs_sin = get_rotary_pos_embed( - (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + # No multiplication by the sequence-parallel world size: the causal + # pipeline replicates the full (unsharded) sequence on every rank, so + # RoPE must cover the full local sequence length. Attention is + # parallelized across heads instead (see + # CausalWanSelfAttention.forward). + (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size, self.num_attention_heads, rope_dim_list,