Wan2.2: Ulysses sequence parallelism on the dit, and CFG with Ulysses on four GPUs (5/5)

configs/wan22_cfg2_sp2_4gpu.yaml

@@ -0,0 +1,28 @@
+model: "wan22"
+
+# See wan22_cfg_parallel_2gpu.yaml. On four GPUs the dit splits both ways: one
+# classifier-free-guidance branch per pair of ranks (cfg_size), and each
+# branch's tokens over its pair (sp_size, Ulysses). The text encoder and both
+# VAE passes run on rank 0.
+max_seq_len: 512
+
+model_kwargs:
+  vae_decode_fast_sizes: ["1280x704", "832x480"]
+
+node_groups:
+  - node_names:
+      - text_encoder
+    ranks:
+      - 0
+
+  - node_names:
+      - vae_encoder
+      - vae_decoder
+    ranks:
+      - 0
+
+  - node_names:
+      - dit
+    ranks: [0, 1, 2, 3]
+    cfg_size: 2
+    sp_size: 2

docs/models.rst

@@ -1556,6 +1556,15 @@ step, correct but no faster. For throughput under concurrency,
 at 4 concurrent vs 4.7 CFG-parallel (whose rank 1 idles through rank 0's VAE
 decode; measured on that earlier version), while one request takes the 1-GPU time.
 
+**Four GPUs.** ``configs/wan22_cfg2_sp2_4gpu.yaml`` also splits each branch's
+tokens over a pair of ranks (``cfg_size: 2, sp_size: 2``, Ulysses on the dit's
+self-attention) to cut one request's latency. Split the guidance branches
+first (``cfg_size: 2``), then add ``sp_size`` on top: on two GPUs,
+``sp_size: 2`` alone is slower than ``cfg_size: 2``.
+``sp_size`` must divide the dit's 24 attention heads. Under concurrency,
+``wan22_dp2.yaml`` stays the throughput layout; the four-GPU layout is unmeasured
+there.
+
 LTX-2.5 (``ltx2_5``)
 --------------------
 

docs/serving.rst

@@ -637,8 +637,8 @@ one whole instance:
    max_seq_len: 512
    node_groups:
      - {node_names: [text_encoder, vae_encoder, vae_decoder], ranks: [0]}
-     # one guidance branch per rank
-     - {node_names: [dit], ranks: [0, 1], cfg_size: 2}
+     # one guidance branch per pair of ranks, each branch's tokens split over its pair
+     - {node_names: [dit], ranks: [0, 1, 2, 3], cfg_size: 2, sp_size: 2}
 
 Several lockstep-parallel nodes can share GPUs when, on each worker, they are in the same
 groups on every axis, each declared parallel on those axes by the model. The ranks then run

mstar/model/components/diffusion/attention.py

@@ -17,11 +17,16 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from typing import Protocol
 
 import torch
 import torch.nn.functional as F
 
+# ``(q, k, v) -> out`` over ``[B, L, H, D]`` tensors (``k`` / ``v`` may have their own
+# ``L``): ``sdpa_attention``, ``joint_attention`` over a ragged resource, or a wrapper.
+Attend = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+
 
 class RaggedAttentionFn(Protocol):
     """``(q, k, v) -> out`` over packed ``[total_tokens, H, D]`` tensors, one segment per

mstar/model/components/distributed/sequence_parallel.py

@@ -19,11 +19,12 @@
 :func:`sp_head_gather` assume ``[tokens, heads, head_dim]``. ``seq_sizes`` is the
 per-rank token count along the sequence (need not be equal — sequences indivisible
 by ``sp`` are handled without padding); the scatter/gather of the residual stream
-at the model boundary uses the same split.
+at the model boundary uses the same split. :class:`SeqParallel` binds the group and
+that split once per forward for a ``[B, L, H, D]`` DiT stream.
 """
 from __future__ import annotations
 
-from typing import Callable
+from typing import Callable, NamedTuple
 
 import torch
 
@@ -171,3 +172,33 @@ def sp_head_gather(sp_group: CommGroup, x: torch.Tensor) -> torch.Tensor:
         return x
     return sp_group.all_gather(x, dim=1)
 
+
+class SeqParallel(NamedTuple):
+    """Ulysses sequence parallelism over a ``[B, L, H, D]`` token stream: ``sizes`` is
+    each rank's share of its ``L``. ``scatter`` cuts a replicated per-token tensor to
+    this rank's tokens and ``gather`` reassembles them, both a no-op on one rank.
+
+    Only an attention whose queries and keys are both the sharded stream needs the
+    exchange: it attends every token over this rank's head-group (``attend``). A
+    query shard attending a replicated stream needs nothing, and a replicated query
+    attending the sharded keys merges the ranks' partial attentions by their
+    log-sum-exps (``attend_shards``; ``run`` must take ``return_lse``)."""
+
+    group: CommGroup
+    sizes: list[int]
+
+    def scatter(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor:
+        return scatter_sequence(self.group, x, self.sizes, dim=dim)
+
+    def gather(self, x: torch.Tensor, dim: int = 1) -> torch.Tensor:
+        return gather_sequence(self.group, x, self.sizes, dim=dim)
+
+    def attend(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, run: Callable[..., torch.Tensor],
+    ) -> torch.Tensor:
+        return ulysses_attention(self.group, q, k, v, run, self.sizes, seq_dim=1, head_dim=2)
+
+    def attend_shards(
+        self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, run: Callable[..., torch.Tensor],
+    ) -> torch.Tensor:
+        return attention_over_sharded_keys(self.group, q, k, v, run)

mstar/model/lingbot_world/components/dit.py

@@ -41,7 +41,6 @@
 """
 
 import math
-from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import NamedTuple
 
@@ -50,6 +49,7 @@
 from einops import rearrange
 from torch import nn
 
+from mstar.model.components.diffusion.attention import Attend
 from mstar.model.lingbot_world.config import LingBotWorldConfig
 
 
@@ -90,10 +90,6 @@ def _flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch
     return x.type(out_dtype)
 
 
-# ``(q, k, v) -> out`` on ``[1, L, H, D]`` tensors; ``k`` / ``v`` have their own length.
-Attend = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
-
-
 class DitAttention(NamedTuple):
     """The DiT's two attentions: the chunk's queries over its rolling K/V
     window, and over the text K/V."""

mstar/model/ltx2_5/components/connectors.py

@@ -20,9 +20,8 @@
 from torch import nn
 
 from mstar.distributed.communication import CommGroup
-from mstar.model.components.diffusion.attention import sdpa_attention
+from mstar.model.components.diffusion.attention import Attend, sdpa_attention
 from mstar.model.ltx2_5.components.layers import (
-    Attend,
     GatedAttention,
     GeluFeedForward,
     RotaryTable,

mstar/model/ltx2_5/components/layers.py

@@ -13,25 +13,21 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
 from typing import NamedTuple
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
 from mstar.distributed.communication import CommGroup
+from mstar.model.components.diffusion.attention import Attend
 from mstar.model.components.distributed.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
 )
-from mstar.model.components.distributed.sequence_parallel import attention_over_sharded_keys, ulysses_attention
-
-# ``(q, k, v) -> out`` over ``[B, L, H, D]`` tensors (``k``/``v`` may have their own
-# ``L``); see ``mstar.model.components.diffusion.attention``.
-Attend = Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]
+from mstar.model.components.distributed.sequence_parallel import SeqParallel
 
 
 class RotaryTable(NamedTuple):
@@ -45,26 +41,6 @@ def heads(self, start: int, count: int) -> "RotaryTable":
         return RotaryTable(self.cos[:, :, start:start + count], self.sin[:, :, start:start + count])
 
 
-class SeqParallel(NamedTuple):
-    """Ulysses sequence parallelism over a ``[B, L, H, D]`` stream, the video: ``sizes``
-    is each rank's share of its ``L``.
-
-    Only an attention whose queries and keys are both the sharded stream needs the
-    exchange: it attends every token over this rank's head-group (``attend``). A
-    query shard attending a replicated stream needs nothing, and a replicated query
-    attending the sharded keys merges the ranks' partial attentions by their
-    log-sum-exps (``attend_shards``; ``run`` must take ``return_lse``)."""
-
-    group: CommGroup
-    sizes: list[int]
-
-    def attend(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, run: Attend) -> torch.Tensor:
-        return ulysses_attention(self.group, q, k, v, run, self.sizes, seq_dim=1, head_dim=2)
-
-    def attend_shards(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, run: Attend) -> torch.Tensor:
-        return attention_over_sharded_keys(self.group, q, k, v, run)
-
-
 def rms_norm_no_weight(x: torch.Tensor, eps: float) -> torch.Tensor:
     """``diffusers.models.normalization.RMSNorm(elementwise_affine=False)``: the
     variance in fp32, the product promoted to fp32, then back to the input dtype."""

mstar/model/ltx2_5/components/transformer.py

@@ -27,13 +27,12 @@
 from torch import nn
 
 from mstar.distributed.communication import CommGroup
-from mstar.model.components.distributed.sequence_parallel import gather_sequence, scatter_sequence, sp_seq_split
+from mstar.model.components.diffusion.attention import Attend
+from mstar.model.components.distributed.sequence_parallel import SeqParallel, sp_seq_split
 from mstar.model.ltx2_5.components.layers import (
-    Attend,
     GatedAttention,
     GeluFeedForward,
     RotaryTable,
-    SeqParallel,
     rms_norm_no_weight,
     split_rope_table,
 )
@@ -302,7 +301,7 @@ def forward(
         sp = None
         if self.sp_group.world_size > 1:
             sp = SeqParallel(self.sp_group, sp_seq_split(latents.shape[1], self.sp_group.world_size))
-            latents = scatter_sequence(sp.group, latents, sp.sizes, dim=1)
+            latents = sp.scatter(latents)
             rope = rope._replace(video=_shard_rows(sp, rope.video), cross_video=_shard_rows(sp, rope.cross_video))
         x = self.proj_in(latents)
         a = self.audio_proj_in(audio_latents)
@@ -316,7 +315,7 @@ def forward(
         a = nn.functional.layer_norm(a, (a.shape[-1],), eps=1e-6) * (1 + a_scale) + a_shift
         video = self.proj_out(x)
         if sp is not None:
-            video = gather_sequence(sp.group, video, sp.sizes, dim=1)
+            video = sp.gather(video)
         return video, self.audio_proj_out(a)
 
 

mstar/model/wan22/components/dit.py

@@ -33,16 +33,19 @@
 two-linear MLPs ARE the shared ``components.mlp.MLP``.
 """
 
+import functools
 import math
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-from mstar.model.components.diffusion.attention import RaggedAttentionFn, joint_attention
+from mstar.distributed.communication import CommGroup
+from mstar.model.components.diffusion.attention import Attend, RaggedAttentionFn, joint_attention
 from mstar.model.components.diffusion.ffn import gelu_tanh_ffn
 from mstar.model.components.diffusion.modulation import ada_layer_norm
 from mstar.model.components.diffusion.qk_norm_rope import rms_norm_rope
+from mstar.model.components.distributed.sequence_parallel import SeqParallel, sp_seq_split
 from mstar.model.components.mlp import MLP
 from mstar.model.wan22.config import Wan22Config
 
@@ -163,7 +166,7 @@ def forward(
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor | None = None,
         rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
-        attn_fn: RaggedAttentionFn | None = None,
+        attend: Attend = joint_attention,
     ) -> torch.Tensor:
         kv_source = hidden_states if encoder_hidden_states is None else encoder_hidden_states
         query = self.q_proj(hidden_states)
@@ -178,8 +181,9 @@ def forward(
             key = self.k_norm(key).unflatten(2, (self.num_heads, -1))
 
         # Bidirectional, no mask: SDPA (the reference's dispatch_attention_fn
-        # native path) without a ragged callable, one segment per row with one.
-        out = joint_attention(query, key, value, ragged=attn_fn)
+        # native path) by default; the DiT passes self-attention its ragged and
+        # Ulysses ``attend``.
+        out = attend(query, key, value)
         return self.o_proj(out.flatten(2, 3))
 
 
@@ -240,7 +244,7 @@ def forward(
         encoder_hidden_states: torch.Tensor,
         timestep_proj: torch.Tensor,
         rotary_emb: tuple[torch.Tensor, torch.Tensor],
-        self_attn_fn: RaggedAttentionFn | None = None,
+        self_attend: Attend = joint_attention,
     ) -> torch.Tensor:
         # [1, 1, 6, dim] + [B, frames, 6, dim] -> six [B, frames, dim] fp32
         # tensors, each applied to all of a frame's tokens.
@@ -249,7 +253,7 @@ def forward(
         ).unbind(2)
 
         _, norm_hidden = ada_layer_norm(hidden_states, shift_msa, scale_msa, eps=self.norm1.eps)
-        attn_output = self.self_attn(norm_hidden, rotary_emb=rotary_emb, attn_fn=self_attn_fn)
+        attn_output = self.self_attn(norm_hidden, rotary_emb=rotary_emb, attend=self_attend)
         hidden_states, norm_hidden = ada_layer_norm(
             hidden_states, None, None, eps=self.norm2.eps, weight=self.norm2.weight, bias=self.norm2.bias,
             residual=attn_output, gate=gate_msa,
@@ -271,10 +275,21 @@ class Wan22DiT(nn.Module):
     """The dense 5B TI2V video DiT: Conv3d patchify -> 30 blocks -> adaLN
     output head + unpatchify. Built on the meta device and materialized by
     ``weight_loader.build_wan22_dit`` (meta -> ``cast_serving_dtypes`` ->
-    ``to_empty(device)`` -> checkpoint load)."""
+    ``to_empty(device)`` -> checkpoint load).
+
+    ``sp_group`` runs it under Ulysses sequence parallelism: each rank keeps a
+    contiguous slice of the post-patchify tokens, self-attention all-to-alls
+    them around the kernel, and the slices are gathered before unpatchify."""
 
-    def __init__(self, config: Wan22Config):
+    def __init__(self, config: Wan22Config, sp_group: CommGroup | None = None):
         super().__init__()
+        self.sp_group = sp_group or CommGroup.trivial()
+        if config.num_attention_heads % self.sp_group.world_size:
+            # Ulysses gives each rank an equal share of the heads
+            raise ValueError(
+                f"sp_size {self.sp_group.world_size} must divide the dit's "
+                f"{config.num_attention_heads} attention heads"
+            )
         if config.qk_norm != "rms_norm_across_heads" or not config.cross_attn_norm:
             # These two config facts are baked into Wan22DiTAttention /
             # Wan22DiTBlock (across-heads q/k RMSNorm, affine norm2 before
@@ -343,18 +358,27 @@ def forward(
         hidden_states = self.patch_embedding(hidden_states)
         hidden_states = hidden_states.flatten(2).transpose(1, 2)
 
+        # every per-token tensor is cut to this rank's tokens (a no-op on one rank)
+        sp = SeqParallel(self.sp_group, sp_seq_split(hidden_states.shape[1], self.sp_group.world_size))
+        if self.sp_group.world_size > 1:
+            # a rank's tokens need not be whole frames, so it modulates per token
+            timestep = timestep.repeat_interleave(hidden_states.shape[1] // grid[0], dim=1)
+        hidden_states, timestep = sp.scatter(hidden_states), sp.scatter(timestep)
+        rotary_emb = tuple(sp.scatter(t, dim=0) for t in rotary_emb)  # [S, D] tables
+        self_attend = functools.partial(sp.attend, run=functools.partial(joint_attention, ragged=self_attn_fn))
+
         temb, timestep_proj, encoder_hidden_states = self.condition_embedder(
             timestep, encoder_hidden_states
         )
 
         for block in self.blocks:
-            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, self_attn_fn)
+            hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, self_attend)
 
         # Output head: 2-chunk adaLN in fp32 ([1, 1, 2, dim] island table +
         # bf16 temb promotes, per frame), then project and unpatchify.
         shift, scale = (self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2)).unbind(2)
         _, hidden_states = ada_layer_norm(hidden_states, shift, scale, eps=self.norm_out.eps)
-        hidden_states = self.proj_out(hidden_states)
+        hidden_states = sp.gather(self.proj_out(hidden_states))
 
         hidden_states = hidden_states.reshape(batch_size, *grid, p_t, p_h, p_w, -1)
         hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)