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

[zhihengy]

Summary

USP (Ulysses × 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 ×sp, while training stays numerically equivalent to sp=1 up to bounded floating-point rounding. The engine is model-agnostic — a family opts in by declaring a SequenceParallelPlan; Wan is enabled with zero family-specific attention code. Default --sequence-parallel-size 1 is a no-op.

Design

Coordinates. global rank = dp_rank · sp + sp_rank, with sp = ulysses × ring. Ulysses ranks are contiguous (its all-to-all is bandwidth-hungry and unoverlappable → keep it on NVLink); ring ranks are strided (its p2p overlaps with compute → fit for cross-node). Layout is pure functions (sp_mesh.py), no distributed init required.

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

Sequence lifecycle. The full latent enters the model; boundary hooks — declared by the model's own diffusers _cp_plan — slice hidden_states after patchify (and rope to the local chunk), and gather (a differentiable all-gather) after proj_out. Model outputs are therefore full-sequence and bitwise identical across sp ranks: loss / log_prob / metrics code is untouched.

Attention. The default installer intercepts diffusers' dispatch_attention_fn — the single choke point every modern diffusers processor routes attention through, where the model itself already flags self- vs cross-attention per call site. Self-attention is rerouted to sp_ops.usp_attention (ulysses a2a → local SDPA, or ring via torch's _templated_ring_attention → a2a back); cross-attention and the model's own processors run untouched. No per-family processor copies to maintain.

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: the sequence gather's backward all-reduces(SUM) over the sp group (×sp on every partial grad), and FSDP's mean over dp·sp divides it back — (1/(dp·sp)) · Σ sp·g = (1/dp) · Σ_dp g. No per-param cross-sp sync, no optimizer divide factor. Measured at dp2×sp2: 54.4 → 27.2 GB/rank.

Model seam. ModelBackend.sequence_parallel_plan(model) assembles {boundaries, attention installer, head count}; the diffusers default reads the model's _cp_plan. Native-modeling backends override the method wholesale. The engine references no model class.

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

Why a 2D mesh + hand-built subgroups (not a 3D mesh)

The device mesh is 2D (dp, sp); ulysses/ring groups are built with dist.new_group from the pure-function layout rather than as mesh dimensions:

1. Consumers decide the shape. Mesh dims matter to DTensor/FSDP machinery; ulysses/ring groups feed raw collectives (all_to_all_single, ring send/recv), which take ProcessGroup handles. A 3D mesh would create dims with zero consumers.
2. One _flatten instead of two overlapping ones. FSDP needs flatten(dp, sp); a 3D mesh would additionally need flatten(ring, ulysses) for the sp group, and the two flattens share dims — fragile territory on a private API.
3. Degree-1 honesty. u/r vary per config and frequently hit size 1; the convention is "degree-1 → None → the dimension doesn't exist", which composes badly with size-1 mesh dims.
4. Every mesh dim costs a NCCL communicator whether used or not.

The 2D-mesh row-major order and the pure-function layout are the same implicit contract; it is pinned by tests (NCCL group members asserted equal to the pure-function lists).

Limitations — everything outside the validated envelope fails loudly

Condition	Where it raises
world % sp != 0	validate_sp_config, driver-side at startup
sp % ulysses != 0	resolve_sp_degrees
num_heads % ulysses != 0	apply_sequence_parallel (the plan carries the real head count)
S % sp != 0 (padding unimplemented)	shard_sequence, first forward
--fsdp-attention-backend with sp>1	validate_sp_support, driver-side, before any weight load
Model backend without sequence_parallel_plan	validate_sp_support
Model without _cp_plan, or wildcard plan keys (e.g. transformer_blocks.*)	plan construction
Model whose attention doesn't route through dispatch_attention_fn	attention install
attention mask / dropout>0 / is_causal / custom scale under USP	dispatch wrapper, at call time
Boundary keys unmappable to the module's forward signature (peft-unwrapped)	hook install

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

Validation (selected)

Test	What it locks	Result
E2E topology study: 4 bands × 10 optimizer steps (sp=1 / u4 / u2r2 / u1r4, each 5 real rollouts + PickScore + weight sync)	rounding does not accumulate; training health under every topology	per-step train-vs-rollout canary flat in all 4 bands (tails == heads, ~9e-6); loss/grad_norm/clipfrac in the sp=1 band; rollout-0 rewards bitwise identical across all 4 topologies
Determinism smoke (u2r1, u1r2)	run-to-run repeatability under --deterministic-mode: one tree, same inputs, run twice	forward + all grads bitwise across double-runs, ring included
Attention parity vs the stock diffusers processor (sp2·u2 / sp4·u4 / sp4·u2r2, each with and without grad-ckpt, plus an fp32 band)	operator correctness	ulysses forward & input grads bitwise (rel=0.00e+00); weight grads 3–8e-3 bf16 → 3–8e-7 fp32 (scaling law ⇒ rounding, not logic)
Grad-sync parity under real FSDP2 (LoRA / distinct-dp / grad-accum variants)	the full gradient path the optimizer sees: the sequence gather's backward sp-sum paired with FSDP's mean must reproduce the single-process gradient exactly — cross-checked by a second, simpler parameter placement that computes the same gradients with an explicit per-param sum	all 69 params == single-process reference; sp-rank outputs assert bitwise-equal; clip total norms == reference
Weight-sync parity (sp2·u2 / sp4·u4 / sp4·u2r2)	train→rollout weight reconstruction	checksums bitwise == single-process reference on every rank

Numerics: deterministic vs random deviation, and what --deterministic-mode buys

Ulysses — deterministic deviation, in exactly one place. Its 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 (one long GEMM accumulation vs per-rank partials + NCCL sum). This is deterministic reordering, not noise: 3–8e-3 rel in bf16 (1–2 ulp of bf16's 2⁻⁸ epsilon, the floor for any reordered sum), collapsing to 3–8e-7 in fp32 — the error scales with mantissa width, the fingerprint of rounding rather than a logic error.

Ring — adds deterministic deviation to forward too, plus one genuinely random source. Ring replaces the single-kernel softmax with r online-softmax chunk merges (the same algebra flash attention uses across SRAM tiles, moved to rank boundaries): forward deviates from sp=1 by ~3.6e-3 — frozen, reproduced identically across reruns and grad-ckpt variants. The merge chain accumulates in fp32 inside torch's template, so this does not grow with ring degree. The random component lives in the aten flash backward it calls: gradient accumulation via atomicAdd, whose thread arrival order varies run to run.

--deterministic-mode eliminates the random component — verified. use_deterministic_algorithms(True, warn_only=False) engages the aten deterministic gate, switching the flash backward to ordered accumulation; the determinism smoke shows double-runs bitwise identical (outputs and every gradient) on the pure-ring path, and warn_only=False doubles as a detector — any op on the path lacking a deterministic implementation would raise, and none does. Note the precise claim: deterministic mode buys repeatability (run-to-run bitwise), not sp=1-equality — the rounding deviation above remains, frozen and certified as rounding by the fp32 scaling law. The 10-step topology study confirms it does not accumulate in any topology.