Feat: USP sequence parallelism (Ulysses x Ring) via SequenceParallelPlan; Wan enabled

[zhihengy]

Summary

USP (Ulysses x Ring) sequence parallelism for DiT GRPO training: the latent sequence of a single sample is sharded across GPUs, so activation memory scales /sp and trainable video length x sp, while training stays numerically equivalent to sp=1 up to bounded floating-point rounding. Default --sequence-parallel-size 1 is a no-op.

The Wan call chain under SP:

actor (sp > 1)
├─ boundary hooks, declared by the model's diffusers _cp_plan:
│    shard hidden_states after patchify → gather (differentiable all-gather) after proj_out.
│    The gather's backward all-reduces(SUM) over the sp group; FSDP's 1/(dp·sp) mean then
│    restores (1/dp)·Σ_dp exactly. Model outputs stay full-sequence and bitwise identical
│    across sp ranks, so loss/log_prob code is untouched.
└─ wrapped dispatch_attention_fn (the modeling module's, found through the MRO under
   FSDP's class swizzle), marker-gated by _USPDispatchConfig set as the processors'
   _parallel_config:
     self-attention        → usp_attention: Ulysses all-to-all (heads sharded, sequence
                             gathered), then Ring via torch's ring templates with an aten
                             kernel selected from RING_KERNELS (ring_backend = the call
                             site's backend kwarg, i.e. --fsdp-attention-backend)
     cross-attention / I2V → original dispatch_attention_fn, stock path (text K/V are
                             replicated, local attention is exact; the configured
                             backend stays live)

Calls without the marker pass through the original dispatch untouched; the model's own processors are never replaced.

Design

Coordinates. global rank = dp_rank · sp + sp_rank, with sp = ulysses × ring. Ulysses ranks are contiguous (its all-to-all is bandwidth-hungry → keep it on NVLink); ring ranks are strided (its transfers overlap with compute → fit for cross-node). Layout is pure functions (sequence_parallel/topology.py).

Data. Dispatch is by dp_rank, so sp peers automatically share samples — the foundation for every "identical loss across the sp group" argument, and it makes dual-DiT expert routing consistent for free.

Parameter placement. FSDP always shards over the flattened dp×sp mesh — its memory account is identical to sp=1; SP is invisible to FSDP. Gradient correctness is one exact identity: (1/(dp·sp)) · Σ sp·g = (1/dp) · Σ_dp g. Measured at dp2×sp2: 54.4 → 27.2 GB/rank.

Attention install. A family declares a SequenceParallelPlan{boundaries, attention_installer, num_attention_heads}; DiffusersModelBackend.sequence_parallel_plan returns the model's pre-attached _miles_sp_plan if present, else assembles the default from the model's _cp_plan boundaries with apply_dispatch_sp_attention as installer. Native backends override sequence_parallel_plan wholesale. set_attention_backend runs independently of SP install; the backend reaches the wrapped dispatch as the ordinary call-site kwarg.

Ring kernel selection. --fsdp-attention-backend also picks the ring kernel via RING_KERNELS (_native_flash → aten flash, _native_cudnn → aten cudnn; both return LSE and pair with a real backward kernel). Everything else is rejected for ring at argument validation. The cudnn kernel measured 1.8x faster than aten flash for Wan shapes on H200 (fwd+bwd, B1 H40 S32768 D128 bf16: 129ms vs 235ms).

Why not diffusers' own context parallelism. Its training backward is broken on 0.37.0 in four distinct places (native op recomputes SDPA without enable_grad, fixed upstream in 0.38; cudnn/flash backward ops re-transpose already-transposed saved tensors; cudnn forward computes LSE only when return_lse; TemplatedRingAttention.backward feeds every rotation step from chunk-0 saved tensors — the last three still present on upstream main). The wrapped dispatch passes no diffusers parallel_config downstream, so that branch is unreachable.

Config surface: --sequence-parallel-size + --ulysses-degree (ring = sp / ulysses, derived) + --fsdp-attention-backend.

Limitations — everything outside the validated envelope fails loudly

Condition	Where it raises
sp < 1, ulysses < 0, world % sp, sp % ulysses	topology validation, driver-side at startup
num_attention_heads % ulysses != 0	apply_sequence_parallel
S % sp != 0 (padding unimplemented)	shard_sequence, first forward
ring > 1 on torch < 2.11 (torch's private ring templates; the CI image pins torch==2.11; 2.9's template merge is additionally flash-only)	ModuleNotFoundError at first ring call
ring > 1 with a backend outside RING_KERNELS	validate_sp_args, driver-side
modeling module without dispatch_attention_fn and no family _miles_sp_plan	apply_dispatch_sp_attention, at install
model backend without sequence_parallel_plan	validate_sp_args, driver-side
model without _cp_plan, or wildcard plan keys	plan construction
attention mask / dropout / is_causal / explicit scale on self-attention under SP	wrapped dispatch, at call time

Known non-guarded boundaries: sp>1 weight grads differ from sp=1 by bf16 summation rounding (inherent, see numerics); multi-node is untested; checkpoint resume across topologies is untested.

Validation

Offline suite (4x H200; ulysses bands on torch 2.9.1, ring bands on torch 2.11):

Test	What it locks	Result
CPU fast tests (108)	validation rules, plan ownership, ring kernel rejection	PASS
Op-level parity, 5 bands (u4r1 / u2r2 / u1r4 flash; u2r2 / u1r4 cudnn)	usp_attention vs full-sequence SDPA reference	PASS — u4r1 forward AND dQ/dK/dV bitwise; ring forward 5.6e-3, grads 6.1–7.2e-3 (rtol 8e-3); cudnn deviation identical to flash to 3 digits
Deterministic double-run (u4r1, u2r2, u1r4)	--deterministic-mode under SP incl. ring	PASS — forward and all grads bitwise in every band
Grad-sync parity under FSDP2 (3 bands incl. u2r2)	the dp×sp gradient identity on the production path	PASS (07-21 revalidation at f76da39; code path unchanged)
Weight-sync parity (3 bands)	train→rollout weight reconstruction	PASS (07-21 revalidation; checksums bitwise, dtype/shape-sensitive)

E2E (Wan2.2 PickScore GRPO, 4 GPUs colocate, 5 rollouts = 10 optimizer steps per band; ring imports re-homed to torch 2.9 for the box's conda env, one-line delta):

Band	Result
sp1 / u4r1 / u2r2 / u1r4 / u2r2+_native_cudnn	sp1 / u4r1 / u2r2 / u1r4 PASS (10 optimizer steps each); u2r2+cudnn: rollout + reward identical, train step unrunnable on the box's torch-2.9 ring-template re-home (2.9's merge is flash-only, cudnn LSE layout mismatches) — the cudnn ring kernel is locked by the op-level band on torch 2.11
rollout-0 reward stats across all 5 bands	identical to every printed digit across all 5 bands (mean=0.7977 std=0.0511 min=0.6587 max=0.8907)
train-vs-rollout canary (log_prob_mean_abs_diff), 10 steps per band	flat in every completed band, tails ~ heads (sp1 8.3e-6->1.0e-5; u4r1 8.9e-6->1.3e-5; u2r2 8.9e-6->9.1e-6; u1r4 8.5e-6->1.2e-5); no accumulation

Per-band mean optimizer-step wall time (same recipe, 8 prompts x 8 samples, 5 frames 480x480):

Band	dp x sp	opt-step wall
sp1	dp4	6.6s
u4r1	dp1 x sp4	14.2s
u2r2	dp1 x sp4	15.8s
u1r4	dp1 x sp4	15.2s

SP at dp1 pays ~2.2-2.4x per optimizer step at this scale (serialized grad accumulation + per-layer collectives); the same ratio reproduces at the production recipe (48x16, 480x832).

Numerics

Ulysses — deterministic deviation, in exactly one place. The all-to-all is a pure permutation and each head's attention runs whole on one rank with the unchanged kernel, so forward and input grads are bitwise equal to sp=1. Only weight grads differ: dW = xᵀ·dy reduces over the token axis — precisely the axis SP splits — so the summation tree changes. 3–8e-3 rel in bf16, collapsing to 3–8e-7 in fp32: the error scales with mantissa width, the fingerprint of rounding.

Ring — adds deterministic deviation to forward, plus one random source. Ring replaces the single-kernel softmax with r online-softmax chunk merges; forward deviates from sp=1 by ~5.6e-3, frozen across reruns. The random component is the fused kernel's backward atomicAdd ordering. The cudnn ring kernel shows the same deviation class as flash (op-level bands: normalized deltas equal to 3 digits, ≤7.25e-3 vs the 8e-3 band).

--deterministic-mode buys repeatability, not sp=1-equality. use_deterministic_algorithms(True, warn_only=False) switches the aten backward to ordered accumulation; double-runs are bitwise identical (outputs and every gradient) in all three topologies, and warn_only=False doubles as a detector for any uncovered op on the path. The rounding deviation above remains, frozen. The 10-step topology study confirms it does not accumulate.

[Rockdu]

Made a few edits, others are good for now