[attention] Add Ulysses Anything backward

[NancyFyong]

What does this PR do?

Adds output backward to TemplatedUlyssesAnythingAttention, which currently raises NotImplementedError.

• Saves the backend context and separate query/key partition metadata, then reuses the existing all-to-all helpers to restore Q/K/V gradients.
• Fixes uneven local key-padding mask assembly by trimming each gathered shard before concatenation.
• Preserves supported padded-head layouts and marks auxiliary LSE outputs non-differentiable.
• Fails fast on per-head attention masks, which Ulysses Anything cannot shard.
• Fails fast when backpropagating through Helios with context parallelism.
• Adds shared tiny-model backward suites for seven models with a built-in CP plan: Flux, Flux2, Qwen-Image, Wan, ChronoEdit, LTX-Video and MiniMax H3.

This follows the inference support in #12996. Searches for open Ulysses/backward and Ulysses-Anything work found no direct duplicate. Related #14341 and #14349 address backend layout and cuDNN mask handling rather than this missing backward path.

Draft: GPU validation remains pending. This does not implement Ring-Anything backward or extend causal attention, GQA, or hybrid Ring + Ulysses-Anything support.

Unsupported cases

Both cases below fail identically with the existing ordinary Ulysses backward and on main, so they predate this change. This PR only makes the failures explicit.

Helios training. Its plan gathers activations after attention and FFN inside every block. The residual stream, norms and modulation between them run replicated on each rank. The split backward only pads the local shard's gradient, so each rank combines the full residual gradient with a partial attention/FFN gradient. On two CPU/Gloo ranks:

• Forward matched exactly.
• Parameter gradients were wrong for 1097 of 1101 parameters under SUM, 1098 under MEAN, and 1098 with no reduction.

Making the split backward all-gather recovered every parameter, but only with region-specific reductions. Supporting that would require changing the hook semantics, so Helios now raises on backward instead.

This check only triggers on backward. Inference and grad-enabled forward passes are unchanged.

Per-head masks (LTX, LTX2). Their processors expand the text mask to (B, heads, 1, S). Ulysses Anything scatters heads but not the mask's head dimension. Previously this crashed inside the attention backend with a shape mismatch; it now raises a clear ValueError before any communication. Head-broadcast masks are unaffected.

• The LTX-Video suite runs without its mask. With a head-broadcast (or no) mask, LTX works: on the same CPU/Gloo check, its SUM-reduced gradients matched exactly.
• LTX2 is not covered. Its CP plan is also incomplete (unsharded RoPE and audio stream), which Fix LTX-2 context parallel plan and Ulysses attn mask head sharding #14317 fixes for ordinary Ulysses. Fix LTX-2 context parallel plan and Ulysses attn mask head sharding #14317 leaves the Ulysses-Anything mask path to LTX-2 _cp_plan appears copied from LTX-1 and fails under enable_parallelism() (crash + one silent-correctness hazard) #14316.

Model test coverage

Each model uses its existing real class and tiny tester configuration, with model-specific inputs for uneven sequences. H3 retains T2VA/FL2VA/Ref2VA-style packed layouts.

The shared worker compares every output, loss, and parameter gradient against an identically initialized serial model. Parameter gradients are SUM-reduced across CP ranks before comparison. On two CPU/Gloo ranks with even-length fixtures, Flux, Flux2, Qwen-Image, Wan, ChronoEdit and LTX matched exactly under SUM, and failed on every parameter under MEAN.

Coverage	Result
Seven models x checkpointing off/on: serial fixture forward/backward	14 passed
Seven models x SP=2/4 x checkpointing off/on: distributed numerical parity	28 collected, pending GPU execution
Generic UAA output/Q/K/V-gradient parity and per-head-mask rejection on CPU/Gloo, 2/4 ranks	2 passed
Generic FA3-varlen parity, 2/4 ranks	2 collected, pending GPU execution
Helios backward guard	1 passed

Models without a built-in CP plan are outside this matrix; their SP support is not claimed.

Validation

Python 3.11.13, PyTorch 2.13.0+cu129, CPU-only. The source checkout and isolated Hugging Face Hub 1.31.0 dependency were selected through PYTHONPATH.

export CUDA_VISIBLE_DEVICES=''
export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1
export TORCH_COMPILE_DISABLE=1 TORCHINDUCTOR_DISABLE=1
python -m pytest -q -p no:cacheprovider \
  tests/models/test_attention_dispatch.py \
  tests/models/transformers/test_models_transformer_{flux,flux2,qwenimage,wan,chronoedit,ltx,minimax_h3}.py \
  -k UlyssesAnything
python -m pytest -q -p no:cacheprovider \
  tests/models/transformers/test_models_transformer_{helios,ltx,ltx2}.py

Results:

• 16 passed, 30 skipped, 1010 deselected. The model-level passes validate serial fixtures, not distributed gradients. Hardware skips are not passes.
• 117 passed, 90 skipped, 3 xpassed for the touched model files, including the Helios guard.

Through the real enable_parallelism path on two CPU/Gloo ranks:

• Helios raised on backward under both Ulysses Anything and ordinary Ulysses, before any parameter received a gradient. Inference and grad-enabled forward passes still ran.
• Masked LTX2 raised the new per-head-mask error.

The original attention regression reproduced the missing-backward exception. After adding backward, the four-rank masked case exposed the mask-assembly defect and passed after trimming individual shards.

make style PYTHONPATH="$PYTHONPATH"
make fix-copies PYTHONPATH="$PYTHONPATH"
make quality PYTHONPATH="$PYTHONPATH"
git diff --check

All passed. No new runtime dependency is introduced.

Final self-review notes

AI assistance (pi coding agent) was used. The complete diff was reviewed against the repository rubric. The human submitter reviewed every changed line and approved the public wording.

Verdict: NEEDS CHANGES before merge.

• Remaining validation: the 28 model-level accelerator parity cases and two FA3-varlen cases, including auxiliary LSE checks. No real-weight or performance claim is made.
• Scope retained for review: training is documented for zero dropout and a backward-capable backend. Replicated parameter gradients must be summed, not averaged, across CP ranks. Existing global/local mask inference and padded-head helper restrictions are not redesigned.
• Dead-code review: the shared worker is used by all seven model suites, and saved query/key metadata is consumed by the production backward path. The Helios hook and mask check each have a test. No unused helper was found.

Before submitting

• Did you use an AI agent to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read the philosophy doc?
• Was this discussed/approved via a GitHub issue or the forum?
• Did you update the documentation?
• Did you write new necessary tests?

Who can review?

Feedback on the Ulysses-Anything backward contract and distributed test coverage is welcome. Accelerator validation is required before marking this PR ready.