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

[PaulFidika]

Describe the bug

LTX2VideoTransformer3DModel._cp_plan (transformer_ltx2.py:1094) is textually identical to LTXVideoTransformer3DModel._cp_plan (transformer_ltx.py:417), but LTX-2 is a dual-stream audio+video model with a different forward. The plan doesn't cover:

• timestep — per-token on LTX-2, (B, num_video_tokens), so it must be sharded
• the audio stream — audio_hidden_states, audio_timestep, audio_encoder_hidden_states, audio_encoder_attention_mask, audio_proj_out
• 3 of 4 RoPE modules — audio_rope, cross_attn_rope, cross_attn_audio_rope

enable_parallelism() accepts it, then the forward fails at transformer_ltx2.py:634 with size of tensor a (16) must match tensor b (32) — sharded video tokens against the unsharded per-token timestep.

The part I'd flag as more important than the crash: the obvious minimal repair — mirror Wan and shard video only — is silently wrong on LTX-2. Its blocks run bidirectional AV cross-attention; v2a has Q=audio, KV=video, so with video sharded and no collective on that call each rank attends over 1/P of the video keys. Shapes line up, nothing raises, output is wrong. Both streams have to be sharded.

Verified against main @ 60092255.

Reproduction

CPU-only, gloo, 2 ranks, no GPU — spawns its own ranks, so python repro.py is all it takes.

repro.py — --mask and --fixed select the arms

"""
    python repro.py            # A: stock _cp_plan       -> crash
    python repro.py --mask     # B: corrected plan       -> attn_mask head-axis bug
    python repro.py --fixed    # C: corrected plan + fix -> matches single-rank
"""

import argparse, os, traceback
import torch
import torch.distributed as dist
import torch.multiprocessing as mp

import diffusers.models.attention_dispatch as AD
from diffusers.models._modeling_parallel import (
    ContextParallelConfig, ContextParallelInput, ContextParallelOutput,
)
from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel

WORLD = 2

CONFIG = dict(
    in_channels=8, out_channels=8, patch_size=1, patch_size_t=1,
    num_attention_heads=4, attention_head_dim=16, cross_attention_dim=64,
    vae_scale_factors=(8, 32, 32), pos_embed_max_pos=20,
    base_height=2048, base_width=2048, gated_attn=True, cross_attn_mod=True,
    audio_in_channels=8, audio_out_channels=8, audio_patch_size=1,
    audio_patch_size_t=1, audio_num_attention_heads=4, audio_attention_head_dim=8,
    audio_cross_attention_dim=32, audio_scale_factor=4, audio_pos_embed_max_pos=20,
    audio_gated_attn=True, audio_cross_attn_mod=True,
    num_layers=2, caption_channels=32, rope_type="interleaved",
    use_prompt_embeddings=True,
)
B, N_VIDEO, N_AUDIO, N_TEXT = 1, 32, 8, 16   # all divisible by WORLD


def make_inputs():
    g = torch.Generator().manual_seed(0)
    return dict(
        hidden_states=torch.randn(B, N_VIDEO, 8, generator=g),
        audio_hidden_states=torch.randn(B, N_AUDIO, 8, generator=g),
        encoder_hidden_states=torch.randn(B, N_TEXT, 32, generator=g),
        audio_encoder_hidden_states=torch.randn(B, N_TEXT, 32, generator=g),
        timestep=torch.full((B, N_VIDEO), 500.0),        # PER-TOKEN, as the pipelines pass it
        audio_timestep=torch.full((B, N_AUDIO), 500.0),
        sigma=torch.full((B,), 0.5),                     # per-batch, must NOT be split
        audio_sigma=torch.full((B,), 0.5),
        encoder_attention_mask=torch.ones(B, N_TEXT),
        audio_encoder_attention_mask=torch.ones(B, N_TEXT),
        num_frames=2, height=4, width=4,                 # 2 * 4 * 4 = 32 video tokens
        audio_num_frames=N_AUDIO,
        fps=24.0, use_cross_timestep=True, return_dict=False,
    )


def _in(dim=1, nd=3, out=False):
    return ContextParallelInput(split_dim=dim, expected_dims=nd, split_output=out)

_ROPE = {0: _in(out=True), 1: _in(out=True)}

LTX2_CP_PLAN = {
    "": {
        "hidden_states": _in(),
        "timestep": _in(nd=2),
        "audio_hidden_states": _in(),
        "audio_timestep": _in(nd=2),
        "encoder_hidden_states": _in(),
        "encoder_attention_mask": _in(nd=2),
        "audio_encoder_hidden_states": _in(),
        "audio_encoder_attention_mask": _in(nd=2),
    },
    "rope": _ROPE,
    "audio_rope": _ROPE,
    "cross_attn_rope": _ROPE,
    "cross_attn_audio_rope": _ROPE,
    "proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
    "audio_proj_out": ContextParallelOutput(gather_dim=1, expected_dims=3),
}

_ORIG = AD.TemplatedUlyssesAttention.forward

def _patched(ctx, query, key, value, attn_mask, *args, **kwargs):
    parallel_config = kwargs.get("_parallel_config") or args[-1]
    if attn_mask is not None and attn_mask.ndim == 4:
        cp = parallel_config.context_parallel_config
        world_size, local_rank = cp.ulysses_degree, cp._ulysses_local_rank
        num_heads = query.shape[2]
        if attn_mask.shape[1] == num_heads and num_heads % world_size == 0:
            attn_mask = attn_mask.chunk(world_size, dim=1)[local_rank].contiguous()
    return _ORIG(ctx, query, key, value, attn_mask, *args, **kwargs)


def worker(rank, mode):
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "29591")
    dist.init_process_group("gloo", rank=rank, world_size=WORLD)
    torch.manual_seed(0)

    model = LTX2VideoTransformer3DModel(**CONFIG).eval()
    with torch.no_grad():
        ref_video, ref_audio = model(**make_inputs())

    if mode == "fixed":
        AD.TemplatedUlyssesAttention.forward = staticmethod(_patched)
    kwargs = {} if mode == "stock" else {"cp_plan": LTX2_CP_PLAN}
    model.enable_parallelism(config=ContextParallelConfig(ulysses_degree=WORLD), **kwargs)

    try:
        with torch.no_grad():
            video, audio = model(**make_inputs())
        if rank == 0:
            print(f"[{mode}] ran. max|delta| vs single-rank: "
                  f"video={(video - ref_video).abs().max():.3e} "
                  f"audio={(audio - ref_audio).abs().max():.3e}")
    except Exception as exc:
        if rank == 0:
            frames = [ln.strip() for ln in traceback.format_exc().splitlines()
                      if "transformer_ltx2.py" in ln or "attention_dispatch.py" in ln]
            print(f"[{mode}] FAILED: {type(exc).__name__}: {exc}")
            print(f"[{mode}]   at {frames[-1] if frames else '?'}")
    dist.destroy_process_group()


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--fixed", action="store_true")
    ap.add_argument("--mask", action="store_true")
    a = ap.parse_args()
    mp.spawn(worker, args=("fixed" if a.fixed else "mask" if a.mask else "stock",),
             nprocs=WORLD, join=True)

$ python repro.py
[stock] FAILED: RuntimeError: The size of tensor a (16) must match the size of tensor b (32) at non-singleton dimension 1
[stock]   at File ".../transformer_ltx2.py", line 634, in forward

$ python repro.py --mask
[mask] FAILED: RuntimeError: The size of tensor a (2) must match the size of tensor b (4) at non-singleton dimension 1
[mask]   at File ".../attention_dispatch.py", line 851, in _native_attention_forward_op

$ python repro.py --fixed
[fixed] ran. max|delta| vs single-rank: video=4.768e-07 audio=2.384e-07

The --fixed arm is the 12-entry plan in the repro, which matches single-rank to fp32 accumulation error through the a2v/v2a path.

Secondary: TemplatedUlyssesAttention doesn't shard attn_mask's head axis

Independent of LTX-2, but LTX-2 is the first model to reach it (it applies CP to masked cross-attention; Wan guards that at transformer_wan.py:152). attention_dispatch.py:2393 all-gathers a sharded mask on the KV axis, but q/k/v come back from the all-to-all with H // world_size heads while the mask keeps all H → broadcast failure in SDPA. Distinct from #13696/#13756, which fixed the sequence axis. Fix is a chunk(world_size, dim=1)[ulysses_local_rank] alongside the KV gather; same gap in the ulysses_anything branch at :2609.

Two smaller notes: _prepare_cp_input warns and skips on a dims mismatch rather than raising, so expected_dims=3 for "rope" silently degrades to unsharded at rope_type="split" (4-D output); and video_self_attention_mask is (B, T_v, T_v), a two-axis shard ContextParallelInput can't express, so the IC-LoRA path can't be fixed by a plan alone.

System Info

diffusers main @ 60092255 (also on 0.39.0) · torch 2.13.0 · Python 3.12 · Linux · CPU-only, gloo, 2 ranks

Who can help?

@sayakpaul @DN6 @yiyixuxu · cc @zhtmike re the Ulysses mask handling

[sayakpaul]

Would it be possible to trim down the issue description a little bit in a simpler manner? Also, would you be open to opening a PR to fix this?

[PaulFidika]

Short version: LTX-2's _cp_plan is LTX-1's plan, so the per-token timestep, the audio stream, and 3 of 4 RoPE modules are never sharded — enable_parallelism() crashes, and the obvious video-only repair silently corrupts a2v/v2a cross-attention. There is also a small independent bug where TemplatedUlyssesAttention all-gathers the attn mask on the KV axis but never shards its head axis. The repro in the body runs on CPU. PR with both fixes, validated on 2xH100 with the released checkpoint: #14317