The backward pass of QwenImageTransformer failed with Ulysses SP.

[zhtmike]

Describe the bug

I am not sure whether the backward pass of Ulysses SP is formally supported, but I found that backward ops like _native_attention_backward_op is implemented in the codebase. When I try to run QwenImageTransformer with the backward pass, I encounter errors related to shape mismatches.

Reproduction

We can reproduce the results with the following code snippets (relies on the PR #13278 to fix the forward pass first)

import argparse

import torch
import torch.distributed as dist
import torch.nn.functional as F
from diffusers.models import QwenImageTransformer2DModel
from diffusers.models._modeling_parallel import ContextParallelConfig


def init_model(device, enable_sp: bool):
 model = QwenImageTransformer2DModel(
 num_layers=2,
 num_attention_heads=4,
 attention_head_dim=32,
 joint_attention_dim=3584,
 axes_dims_rope=(8, 12, 12),
 ).to(device, dtype=torch.bfloat16)

 if enable_sp:
 model.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))

 return model


def make_batch(device):
 dtype = torch.bfloat16
 torch.manual_seed(0)
 hidden_states = torch.randn(2, 256, 64, device=device, dtype=dtype)
 encoder_hidden_states = torch.randn(2, 32, 3584, device=device, dtype=dtype)
 encoder_hidden_states_mask = torch.ones(2, 32, device=device, dtype=torch.bool)
 timestep = torch.rand(2, device=device, dtype=dtype)
 img_shapes = [[(1, 16, 16)]] * 2
 target = torch.randn(2, 256, 64, device=device, dtype=dtype)
 return (
 hidden_states,
 encoder_hidden_states,
 encoder_hidden_states_mask,
 timestep,
 img_shapes,
 target,
 )


def train(enable_sp: bool):
 dist.init_process_group(backend="nccl")
 rank = dist.get_rank()
 device = torch.device("cuda", rank)
 torch.cuda.set_device(device)

 model = init_model(device, enable_sp)
 model.train()

 (
 hidden_states,
 encoder_hidden_states,
 encoder_hidden_states_mask,
 timestep,
 img_shapes,
 target,
 ) = make_batch(device)

 pred = model(
 hidden_states=hidden_states,
 encoder_hidden_states=encoder_hidden_states,
 encoder_hidden_states_mask=encoder_hidden_states_mask,
 timestep=timestep,
 img_shapes=img_shapes,
 return_dict=False,
 )[0]

 loss = F.mse_loss(pred.float(), target.float())
 loss.backward()

 if rank == 0:
 print(f"loss={loss.item():.6f}")

 dist.destroy_process_group()


if __name__ == "__main__":
 parser = argparse.ArgumentParser()
 parser.add_argument("--enable-sp", action="store_true")
 args = parser.parse_args()
 train(enable_sp=args.enable_sp)

Logs

• Without Ulysses SP enabled,

torchrun --nproc-per-node 2 toy_train.py

the script runs and produces the expected output:

loss=1.351188

• With Ulysses SP enabled

torchrun --nproc-per-node 2 toy_train.py --enable-sp

it runs with error

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([2, 2, 288, 32]) and output[0] has a shape of torch.Size([2, 288, 2, 32]).

System Info

• 🤗 Diffusers version: 0.38.0.dev0
• Platform: Linux-5.15.0-1053-nvidia-x86_64-with-glibc2.35
• Running on Google Colab?: No
• Python version: 3.12.12
• PyTorch version (GPU?): 2.9.1+cu128 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 0.36.2
• Transformers version: 4.57.6
• Accelerate version: 1.12.0
• PEFT version: 0.18.1
• Bitsandbytes version: not installed
• Safetensors version: 0.7.0
• xFormers version: not installed
• Accelerator: NVIDIA H800, 81559 MiB
  NVIDIA H800, 81559 MiB
  NVIDIA H800, 81559 MiB
  NVIDIA H800, 81559 MiB
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: yes

Who can help?

@sayakpaul

[zhtmike]

@sayakpaul I can help to fix the backward pass of SDPA and add test coverage

[sayakpaul]

That would be awesome. Let's do it.