Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT#1627
Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT#1627pkisfaludi-nv wants to merge 1 commit into
Conversation
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) <noreply@anthropic.com>
|
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews, 🤖 CI and 🙋 you
🔴 PR merge requirementsWaiting for
This rule is failing.
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Ulysses sequence-parallel (SP) support for causal attention in causal_wanvideo.py. Instead of partitioning across the sequence dimension, the causal pipeline replicates the full sequence on every rank and parallelizes attention across the head dimension. Per-rank head outputs are then gathered back using sequence_model_parallel_all_gather. Additionally, the multiplication of sequence length by the sequence-parallel world size when computing rotary position embeddings (RoPE) has been removed to prevent crashes, as the sequence is unsharded. There are no review comments to address, and I have no further feedback to provide.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
1 participant
Problem
Running a self-forcing causal Wan model (e.g.
SFWan2.1-T2V-A14B,SFWan2.1-T2V-1.3B) with Ulysses sequence parallelism (num_gpus>=2) crashes:Root cause
causal_wanvideo.pybuilds RoPE withrope_num_frames = post_patch_num_frames * get_sp_world_size()(sp× 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). The attention then applies that over-long RoPE to the full local query → shape mismatch (9360 = 2 × 4680at sp=2).Fix
* get_sp_world_size()from the RoPE length in all branches (relativistic / standard / train) so RoPE matches the full local sequence (no-op at sp=1).[rank*H/sp : (rank+1)*H/sp]over the full sequence, thensequence_model_parallel_all_gather(dim=2)reconstructs all heads. Applied uniformly to all three attention paths (flex_attention, KV-cacheself.attn, relativistic window). A sequence all-to-all (à lawanvideo.py) is incompatible with the causal KV-cache windowing (expressed in global token coords), so head-parallel is the exact-matching approach; heads are independent → numerically identical to single-GPU.Verification
wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers, 480×832, 21 frames (7 autoregressive KV-cache blocks), seed 1234:torch.equal=True,max_abs_diff=0.0— byte-identical; 2-GPU no longer crashes.Note: verified on the standard KV-cache path (this model uses
rope_cache_policy="absolute"); therelativisticpath gets the identical head-slice+all-gather transform for consistency.