Skip to content

Fix LTX-2 context parallel plan and Ulysses attn mask head sharding - #14317

Open
PaulFidika wants to merge 1 commit into
huggingface:mainfrom
PaulFidika:fix-ltx2-cp-plan
Open

Fix LTX-2 context parallel plan and Ulysses attn mask head sharding#14317
PaulFidika wants to merge 1 commit into
huggingface:mainfrom
PaulFidika:fix-ltx2-cp-plan

Conversation

@PaulFidika

Copy link
Copy Markdown

LTX-2's _cp_plan was a copy of LTX-1's, so enable_parallelism() crashed on the per-token timestep and never sharded the audio stream or 3 of the 4 RoPE modules — and sharding video only would be silently wrong for the bidirectional a2v/v2a cross-attention. The plan is now set in __init__ because the RoPE sequence dim depends on config: (B, T, D) for rope_type="interleaved" vs (B, H, T, D//2) for "split" (what the released checkpoint uses). Two adjacent fixes this surfaced: TemplatedUlyssesAttention now shards a 4-D attn mask's head axis to match the post-all-to-all QKV heads, and the CP hook now skips planned params the caller didn't pass (LTX-2's optional audio_timestep) instead of raising.

CPU repro from #14316 on this branch (2 ranks, gloo, built-in plan, both rope types):

[interleaved] ran. max|delta| vs single-rank: video=4.768e-07 audio=2.384e-07
[split] ran. max|delta| vs single-rank: video=3.576e-07 audio=2.980e-07

Also validated on 2xH100 (NVLink) with the released Lightricks/LTX-2 checkpoint (LTX2Pipeline, 1280x704, 121 frames, bf16): a single transformer forward under ulysses_degree=2 matches single-GPU exactly (fp32 max|delta| <= 4.2e-7, bf16 bitwise on a small config), and the full 20-step run produces the same video+audio take with bf16 accumulation drift 4-9x below the delta between two seeds (latent rel-RMS 0.21 vs 0.92 video, 0.07 vs 0.62 audio).

The two remaining gaps stay tracked in #14316: the (B, T_v, T_v) IC-LoRA self-attention mask (not expressible as a plan entry) and the ulysses_anything mask path. Fixes#14316.

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

fixes-issuehooksmodelssize/MPR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LTX-2 _cp_plan appears copied from LTX-1 and fails under enable_parallelism() (crash + one silent-correctness hazard)

1 participant

@PaulFidika