Ring attention backward uses stale iteration-0 KV from ctx.saved_tensors — silent gradient corruption when training with context parallelism

[pthombre]

Describe the bug

TemplatedRingAttention (in src/diffusers/models/attention_dispatch.py) produces incorrect gradients whenever ContextParallelConfig(ring_degree > 1) is used for training. The forward pass is correct; only the backward is broken, and with the flash backend it fails silently — no error, no NaN, just wrong gradients — so a training run looks perfectly healthy while learning from corrupted gradients.

Root cause

The ring backward rotates key/value chunks between ranks each ring iteration through local Python variables, but the per-iteration gradient computation is delegated to ctx.backward_op(ctx, grad_out) — and every backward op (_native_flash_attention_backward_op, _cudnn_attention_backward_op, the flash-attn op, …) reads Q/K/V from ctx.saved_tensors.

Those tensors were saved once, during the forward pass at ring iteration 0 (the forward passes _save_ctx=i == 0 to the attention op), and therefore contain only the local KV chunk. The rotated remote KV chunks never reach the backward computation: every ring iteration of the backward re-computes gradients against the same iteration-0 KV, and the contributions that should come from remote KV are computed against the wrong tensors.

How it surfaces per backend

Backend	Symptom with ring_degree > 1
flash (flash-attn 2)	Silent corruption. Forward loss matches a single-GPU reference to ~4e-4; per-tensor gradient relative errors are 2.2 / 1.7 / 1.5 (dq/dk/dv, ring=2) and 7.1 / 3.7 / 3.4 (ring=4) — no error raised.
_native_flash	Backward crashes: Number of heads in key/value must divide number of heads in query (compounded by the separate K/V double-transpose bug in that op's backward).
_native_cudnn	Backward crashes with a cuDNN GQA shape error (same compounding).
native	Forward raises Native attention does not support return_lse=True — by design (ring needs LSE for partial-result merging), so this backend never reaches the bug.

The hybrid path (_templated_unified_attention, ring × ulysses) wraps TemplatedRingAttention, so it inherits the bug: ring=2 × ulysses=2 with flash shows the same 2.2/1.7/1.5 gradient errors. Pure Ulysses (ring_degree=1) is unaffected and gradient-correct.

Why this matters

CP is documented for inference, where this is invisible. But the templated attention path is autograd-capable and is the natural building block for context-parallel training of video/image diffusion transformers, where sequence lengths make CP most valuable. Because the flash failure mode is silent, anyone training with ring or hybrid CP today gets a model that trains "successfully" on wrong gradients. (Downstream in NeMo AutoModel we currently hard-reject ring_degree > 1 for training because of this.)

Reproduction

Op-level gradient parity test, 2 ranks (torchrun --nproc-per-node 2 repro.py):

# repro.py — compare dispatch_attention_fn(ring=2) gradients against a
# single-process fp32 SDPA reference over the full sequence.
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from diffusers.models.attention_dispatch import dispatch_attention_fn
from diffusers.models._modeling_parallel import ParallelConfig, ContextParallelConfig

RING, ULYSSES = 2, 1  # flip to (1, 2) to see Ulysses pass the same check
B, S, H, D = 2, 512, 8, 64

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)

# identical inputs on every rank
torch.manual_seed(777)
q = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16, requires_grad=True)
grad_out = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16)

# reference: full-sequence fp32 SDPA on one process
ref_out = F.scaled_dot_product_attention(
    q.float().transpose(1, 2), k.float().transpose(1, 2), v.float().transpose(1, 2)
).transpose(1, 2)
ref_dq, ref_dk, ref_dv = torch.autograd.grad(ref_out, (q, k, v), grad_out.float())

# CP: each rank runs its sequence shard through the templated ring path
mesh = dist.device_mesh.init_device_mesh("cuda", (RING, ULYSSES), mesh_dim_names=("ring", "ulysses"))
cp_config = ContextParallelConfig(ring_degree=RING, ulysses_degree=ULYSSES, mesh=mesh)
parallel_config = ParallelConfig(context_parallel_config=cp_config)
# NOTE: setup mirrors what ModelMixin.enable_parallelism does internally;
# adjust to the private API of the installed version if it differs.

shard = slice(rank * S // world, (rank + 1) * S // world)
qs, ks, vs = (t.detach()[:, shard].clone().requires_grad_(True) for t in (q, k, v))
out = dispatch_attention_fn(qs, ks, vs, backend="flash", parallel_config=parallel_config)
out.backward(grad_out[:, shard])

for name, got, ref in (("dq", qs.grad, ref_dq), ("dk", ks.grad, ref_dk), ("dv", vs.grad, ref_dv)):
    ref_shard = ref[:, shard].to(got.dtype)
    rel = (got - ref_shard).norm() / ref_shard.norm()
    print(f"rank{rank} {name} rel_err={rel.item():.3e}")

Observed (bf16, tolerance 3e-2; Ulysses passes at ~4.5e-3):

ring=2 ulysses=1 backend=flash:  dq rel_err=2.2e+00  dk rel_err=1.7e+00  dv rel_err=1.5e+00   # no exception raised
ring=1 ulysses=2 backend=flash:  dq rel_err=4.5e-03  dk rel_err=4.5e-03  dv rel_err=3.8e-03   # PASS

Forward outputs match the reference in both cases (~3.7e-3 bf16), confirming the defect is confined to backward.

Expected behavior

Ring-attention backward gradients should match the full-sequence reference to the same tolerance as Ulysses. That requires the backward to see the same rotated KV chunk per ring iteration as the forward did — e.g. by having the templated backward pass the current-iteration KV into the backward op explicitly (rather than the op reading iteration-0 tensors from ctx.saved_tensors), or by saving per-iteration tensors, or by re-communicating KV during backward as ring-flash-attention implementations do (recompute-communication tradeoff).

System info

• diffusers: 0.39.0 (bug present in ≤ 0.39; TemplatedRingAttention in models/attention_dispatch.py)
• torch: 2.x + CUDA, NCCL backend
• GPUs: 8×H100 (repro uses 2)
• Related but distinct bugs we can file separately if useful:
  1. _native_flash_attention_backward_op re-transposes K/V that the forward already saved in kernel layout (crashes on the CP training path even with ring_degree=1... i.e., under Ulysses too when this op drives backward).
  2. _cudnn_attention_forward_op calls the kernel with compute_log_sumexp=return_lse, which is False during training, saving lse=None for backward → garbage gradients; its backward also has the same K/V double-transpose as (1). Net effect: _native_cudnn silently produces NaN/garbage gradients under CP training.
  3. Minor: the flash backend's backward op on the CP path advances the CUDA philox RNG state on every call even with dropout 0 (determinism paper-cut; other backends don't).