Skip to content

Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT #1627

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pkisfaludi-nv wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix Ulysses sequence-parallel crash in self-forcing causal Wan DiT #1627

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 53 additions & 8 deletions fastvideo/models/dits/causal_wanvideo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@
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,
Expand Down Expand Up @@ -103,6 +106,29 @@ def forward(self,
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]
Expand All @@ -126,11 +152,14 @@ def forward(self,
)

x = flex_attention(
query=padded_roped_query.transpose(2, 1),
key=padded_roped_key.transpose(2, 1),
value=padded_v.transpose(2, 1),
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
Expand Down Expand Up @@ -196,7 +225,12 @@ def forward(self,
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:
Expand Down Expand Up @@ -608,10 +642,16 @@ def _forward_inference(
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),
Expand Down Expand Up @@ -718,7 +758,12 @@ def _forward_train(self,
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,
# 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,
Expand Down
Loading