support CP in native flash attention

[sywangyi]

What does this PR do?

• General functionalities: @sayakpaul @yiyixuxu @DN6

native flash attention could support both Ulysses and ring attention

[sywangyi]

since native attention only support Ulysses Attention, we need an attention worked for Ring Attention. xpu enable _scaled_dot_product_flash_attention in torch. so we could use it for ring attention

[sywangyi]

@yao-matrix

[sayakpaul]

Cool work! Could you also supplement a fully working code snippet?

[sywangyi]

yes, the PR also works for cuda.

import random

import numpy as np
import torch
from torch import distributed as dist

from diffusers import AutoencoderKLWan, ContextParallelConfig, WanPipeline
from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video


model_id="Wan-AI/Wan2.2-T2V-A14B-Diffusers"

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    device = torch.device(f"cuda:{dist.get_rank()}")
    torch.cuda.set_device(device)
    return device


def set_seed_for_all_ranks(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    generator = torch.Generator(device="cuda")
    generator.manual_seed(seed)
    return generator


device = setup_distributed()
generator = set_seed_for_all_ranks(42)
onload_device = device
offload_device = torch.device("cpu")

vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
# group-offloading
pipe = WanPipeline.from_pretrained(
    model_id,
    vae=vae,
    torch_dtype=torch.bfloat16,
)
ring_degree = torch.distributed.get_world_size()
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=ring_degree))
pipe.transformer_2.enable_parallelism(config=ContextParallelConfig(ring_degree=ring_degree))
pipe.transformer.set_attention_backend("_native_flash")
pipe.transformer_2.set_attention_backend("_native_flash")
pipe.to("cuda")

pipe.vae.enable_tiling(tile_sample_min_height=480,tile_sample_min_width=960,tile_sample_stride_height=352,tile_sample_stride_width=640)
height = 704
width = 1280
num_frames = 24
num_inference_steps = 50
guidance_scale = 5.0


prompt = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩>残留，丑陋的，残缺的>，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画>面，杂>乱的背景，三条腿，背>景人很多，倒着走"

output = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=height,
    width=width,
    num_frames=num_frames,
    guidance_scale=guidance_scale,
    num_inference_steps=num_inference_steps,
    generator=generator,
).frames[0]
if torch.distributed.get_rank() == 0:
    export_to_video(output, "5bit2v_output.mp4", fps=24)
if dist.is_initialized():
    torch.distributed.destroy_process_group()

[sywangyi]

you could run it using torchrun --nproc-per-node 4 test.py, without the PR, the output is corrupted.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

without the PR, the output is corrupted.

Hmm, it should raise an error no?

diffusers/src/diffusers/models/attention_dispatch.py

Line 1935 in 1567243

def _native_flash_attention(

On main we don't have parallel support.

[sywangyi]

previous if the attention does not support CP, but application use CP. no error is raised and corrupted video is generated.
see the output wo the PR

with the PR

[sayakpaul]

Weird. Will check and fix this. Cc: @DN6

[sayakpaul]

Okay I tracked it down.

The order in which we're calling things matters. For example, if we do:

pipeline.transformer.set_attention_backend("_native_flash")
pipeline.transformer.enable_parallelism(config=cp_config)

It rightfully errors out:

[rank0]: ValueError: Context parallelism is enabled but the attention processor 'FluxAttnProcessor' is using backend '_native_flash' which does not support context parallelism. Please set a compatible attention backend: ['_native_cudnn', 'flash', 'native', 'sage'] using `model.set_attention_backend()` before calling `enable_parallelism()`.

But for any other combinations, it silently passes through. Will fix.