fix: disable dynamo on Ulysses all-to-all to prevent FSDP×compile NCCL deadlock

[samsja]

Summary

Fixes the NCCL collective ordering deadlock that occurs when torch.compile is used together with Ulysses context parallelism and FSDP (reported in PR #3186, Slurm job 1578).

Root cause

At NCCL sequence number 3750, rank 0 was stuck on FSDP's reduce_scatter_tensor (from post_backward) while rank 1 was stuck on Ulysses' all_to_all_single (from dist_nn.functional backward). Both ranks had completed all collectives up to 3749, then diverged — a classic collective ordering mismatch that deadlocks NCCL after the 1-hour watchdog timeout.

The issue is that torch.compile traces the Ulysses all-to-all collectives into the compiled backward graph for each decoder layer, while FSDP's post_backward reduce-scatter hooks fire eagerly via the autograd engine outside the compiled graph. The compiler may reorder the all-to-all backward differently across ranks (especially for MoE models with data-dependent expert routing or variable-length sequences), causing the two collective systems to interleave in inconsistent orders across ranks.

Fix

Wrap _all_to_all_seq_to_head and _all_to_all_head_to_seq with @torch._dynamo.disable. This forces a graph break around the collectives so they execute eagerly in the autograd engine — the same scheduling path as FSDP hooks — restoring a single consistent backward ordering between the two collective systems. This mirrors the existing torch._dynamo.disable on the FA4 flash kernel itself.

---

Note

Medium Risk
Touches distributed training collectives under compile+FSDP; the change is small and aligns scheduling with existing FA4 dynamo-disable behavior, but incorrect ordering could still deadlock or regress multi-GPU training.

Overview
Fixes an NCCL collective-ordering deadlock when torch.compile, Ulysses context parallelism, and FSDP run together (PR #3186 / job 1578).

_all_to_all_seq_to_head and _all_to_all_head_to_seq in ulysses_attn.py are now wrapped with @torch._dynamo.disable, with docstrings noting that Ulysses all_to_all_single must run eagerly in the autograd engine so its backward ordering matches FSDP post_backward reduce-scatter hooks instead of being reordered inside compiled backward graphs (which can differ across ranks).

This follows the same pattern already used for the FA4 flash kernel in that module.

^{Reviewed by Cursor Bugbot for commit a8e305f. Bugbot is set up for automated code reviews on this repo. Configure here.}

[rasdani]

this unblocks my laguna training

should me merge?