Fix Flux Context Parallel Bug (Incoherent Image Generation)

[mali-afridi]

What does this PR do?

Fix Context Parallelism: Implement Ring Attention Pattern for Coherent Multi-GPU Generation

🐛 Problem

I did some testings of the https://huggingface.co/docs/diffusers/main/training/distributed_inference on main branch.
Context parallelism in diffusers was producing fragmented/split images when using multiple GPUs. Instead of generating a single coherent image, each GPU was independently generating its own portion, resulting in visible seams or completely different content in each image segment.

Example: Running with torchrun --nproc-per-node=2 would produce an image that looked like two different images side-by-side rather than one unified image.

🔍 Root Cause Analysis

The issue stems from how attention was computed in context parallel mode:

Before (Broken):

# Each GPU only sees its local chunkGPU0: Q=[B, S/2, H, D], K=[B, S/2, H, D], V=[B, S/2, H, D] # First half onlyGPU1: Q=[B, S/2, H, D], K=[B, S/2, H, D], V=[B, S/2, H, D] # Second half only# Result: Each GPU's attention can only see S/2 context, missing the other half

Each GPU was computing attention using only its local sequence chunk for Q, K, and V. This meant:

• GPU 0's queries could only attend to the first half of the sequence
• GPU 1's queries could only attend to the second half of the sequence
• No cross-GPU attention was happening → broken context → independent generations

✅ Solution: Ring Attention Pattern

This PR implements the Ring Attention pattern where:

• Query (Q) remains local to each GPU (for computational efficiency)
• Key (K) and Value (V) are gathered from all GPUs (for full context)

After (Fixed):

# Each GPU's Q is local, but K and V are globalGPU0: Q=[B, S/2, H, D], K=[B, S_full, H, D], V=[B, S_full, H, D]
GPU1: Q=[B, S/2, H, D], K=[B, S_full, H, D], V=[B, S_full, H, D]
# Result: Each GPU's queries can attend to the FULL sequence context!

📝 Implementation Details

The fix is applied directly in the attention processors after rotary embeddings but before attention computation:

FluxAttnProcessor (transformer_flux.py):

# After rotary embeddings# Gather K and V from all GPUskey_list= [torch.empty_like(key) for_inrange(world_size)]
value_list= [torch.empty_like(value) for_inrange(world_size)]
torch.distributed.all_gather(key_list, key.contiguous())
torch.distributed.all_gather(value_list, value.contiguous())
# Concatenate along sequence dimensionkey=torch.cat(key_list, dim=1) # Now has full sequencevalue=torch.cat(value_list, dim=1) # Now has full sequence

🧪 Testing

For testing, run the following with torchrun --nproc-per-node=2:

import torch, time
import torch.distributed as dist
from diffusers import AutoModel, FluxPipeline, ContextParallelConfig
dist.init_process_group("nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device("cuda", rank)
torch.cuda.set_device(device)
transformer = AutoModel.from_pretrained(
"black-forest-labs/FLUX.1-dev",
subfolder="transformer",
torch_dtype=torch.bfloat16,
parallel_config=ContextParallelConfig(ring_degree=world_size) if world_size > 1 else None
).to(device)
pipeline = FluxPipeline.from_pretrained(
"black-forest-labs/FLUX.1-dev",
transformer=transformer,
torch_dtype=torch.bfloat16,
).to(device)
prompt = "A beautiful sunset over mountains"
generator = torch.Generator(device=device).manual_seed(42)
start_time = time.time()
image = pipeline(
prompt=prompt,
num_inference_steps=28,
generator=generator,
).images[0]
end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
start_time = time.time()
if rank == 0:
image.save(f"flux_{dist.get_world_size()}.png")
print("\n" + "="*60)
print(f"Saved flux_{dist.get_world_size()}.png")
dist.destroy_process_group()

Before Fix:

Result: Two different images side-by-side in output

After Fix:

Result: Single coherent image matching single-GPU output

Summary: This PR fixes context parallelism by ensuring each GPU's attention queries can access the full key-value context from all GPUs, implementing the Ring Attention pattern for coherent multi-GPU image generation.

Note:I have observed that some tensors in QwenImage cannot be divided by world_size (encoder_hidden_states, encoder_hidden_mask etc.) . I am also willing to make a new PR for the QwenImage support for context parallel by padding the tensors to be divisible by world size, similar to chengzeyi/ParaAttention#53 if you guys want to.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul@a-r-r-o-w

[sayakpaul]

Thanks for the PR. From the looks of it, it does seem like it is fully LLM-generated. Also, FWIW, we strive to keep our modeling implementations simple so, I am not sure yet if the changes align with that philosophy. @DN6 WDYT?

[DN6]

Hi @mali-afridi the issue seems to be because an unsupported backend is being used with CP.

This snippet should work

importtorchfromdiffusersimportFluxPipelinefromdiffusersimportContextParallelConfigtry:
torch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
device=torch.device("cuda")
pipe=FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to(device)
pipe.transformer.set_attention_backend("_native_cudnn")
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2))
prompt="A picture of a cat holding a sign that says hello"# Must specify generator so all ranks start with same latents (or pass your own)generator=torch.Generator().manual_seed(42)
image=pipe(prompt, num_inference_steps=28, guidance_scale=4.0, generator=generator).images[0]
ifrank==0:
image.save("output.png")
exceptExceptionase:
raiseefinally:
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

I've opened a PR to raise an error when an incompatible backend is used: #12446

[mali-afridi]

Interesting, yeah the _native_cudnn worked. What should we do about the QwenImage though? The tensors don't get divided by world size. Padding seems to work but it will cause SSIM change as mentioned in the link in description.

[mali-afridi]

For Qwen Reproducibility:

import torch, time
from diffusers import QwenImagePipeline, ContextParallelConfig
try:
torch.distributed.init_process_group("nccl")
rank = torch.distributed.get_rank()
device = torch.device("cuda", rank % torch.cuda.device_count())
torch.cuda.set_device(device)
pipeline = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16)
pipeline.to(device)
pipeline.transformer.set_attention_backend("_native_cudnn")
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2))
# pipeline.transformer.set_attention_backend("flash")
# positive_magic = {
# "en": "Ultra HD, 4K, cinematic composition.", # for english prompt,
# "zh": "超清，4K，电影级构图" # for chinese prompt,
# }
prompt = '''A coffee shop entrance features a chalkboard sign reading "Qwen Coffee 😊 $2 per cup," with a neon light beside it displaying "通义千问". Next to it hangs a poster showing a beautiful Chinese woman, and beneath the poster is written "π≈3.1415926-53589793-23846264-33832795-02384197". Ultra HD, 4K, cinematic composition'''
negative_prompt = " " # using an empty string if you do not have specific concept to remove
# Generate with different aspect ratios
aspect_ratios = {
"1:1": (1328, 1328),
"16:9": (1664, 928),
"9:16": (928, 1664),
"4:3": (1472, 1140),
"3:4": (1140, 1472),
"3:2": (1584, 1056),
"2:3": (1056, 1584),
}
width, height = aspect_ratios["16:9"]
# Must specify generator so all ranks start with same latents (or pass your own)
generator = torch.Generator().manual_seed(42)
start_time = time.time()
image = pipeline(prompt=prompt,
# negative_prompt=negative_prompt,
# width=width,
# height=height,
num_inference_steps=50,
# true_cfg_scale=4.0,
generator=generator).images[0]
end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
if rank == 0:
image.save("output1.png")
except Exception as e:
import traceback
print(f"An error occurred: {e}")
print("Traceback (most recent call last):")
traceback.print_exc()
raise
finally:
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()

Error: https://pastebin.com/NhcEhb3s

[DefTruth]

i got bad result w/o this fix for FLUX.1-dev 2 GPU:

[leeguandong]

For Qwen Reproducibility:对于 Qwen 可复现性：

import torch, time
from diffusers import QwenImagePipeline, ContextParallelConfig
try:
torch.distributed.init_process_group("nccl")
rank = torch.distributed.get_rank()
device = torch.device("cuda", rank % torch.cuda.device_count())
torch.cuda.set_device(device)
pipeline = QwenImagePipeline.from_pretrained("Qwen/Qwen-Image", torch_dtype=torch.bfloat16)
pipeline.to(device)
pipeline.transformer.set_attention_backend("_native_cudnn")
pipeline.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2))
# pipeline.transformer.set_attention_backend("flash")
# positive_magic = {
# "en": "Ultra HD, 4K, cinematic composition.", # for english prompt,
# "zh": "超清，4K，电影级构图" # for chinese prompt,
# }
prompt = '''A coffee shop entrance features a chalkboard sign reading "Qwen Coffee 😊 $2 per cup," with a neon light beside it displaying "通义千问". Next to it hangs a poster showing a beautiful Chinese woman, and beneath the poster is written "π≈3.1415926-53589793-23846264-33832795-02384197". Ultra HD, 4K, cinematic composition'''
negative_prompt = " " # using an empty string if you do not have specific concept to remove
# Generate with different aspect ratios
aspect_ratios = {
"1:1": (1328, 1328),
"16:9": (1664, 928),
"9:16": (928, 1664),
"4:3": (1472, 1140),
"3:4": (1140, 1472),
"3:2": (1584, 1056),
"2:3": (1056, 1584),
}
width, height = aspect_ratios["16:9"]
# Must specify generator so all ranks start with same latents (or pass your own)
generator = torch.Generator().manual_seed(42)
start_time = time.time()
image = pipeline(prompt=prompt,
# negative_prompt=negative_prompt,
# width=width,
# height=height,
num_inference_steps=50,
# true_cfg_scale=4.0,
generator=generator).images[0]
end_time = time.time()
print(f"Time taken: {end_time - start_time} seconds")
if rank == 0:
image.save("output1.png")
except Exception as e:
import traceback
print(f"An error occurred: {e}")
print("Traceback (most recent call last):")
traceback.print_exc()
raise
finally:
if torch.distributed.is_initialized():
torch.distributed.destroy_process_group()

Error: https://pastebin.com/NhcEhb3s错误：https://pastebin.com/NhcEhb3s

I also encountered AssertionError: Tensor size along dimension to be sharded must be divisible by mesh size when testing the parallelism of qwen-image-edit. How can I fix this? I look forward to your solution.

importtorchfromPILimportImagefromloguruimportloggerfromtransformersimportQwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessorfromdiffusersimportQwenImagePipeline, ContextParallelConfig, QwenImageEditPlusPipeline, \
AutoencoderKLQwenImage, QwenImageTransformer2DModelfromoptimum.quantoimportfreeze, qfloat8_e4m3fn, quantizetorch.distributed.init_process_group("nccl")
rank=torch.distributed.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
torch.cuda.set_device(device)
model_id="/share/gdli7/models/checkpoints/Qwen/Qwen-Image-Edit-2509"dtype=torch.bfloat16quant_model=Truetransformer=QwenImageTransformer2DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=dtype)
text_encoder=Qwen2_5_VLForConditionalGeneration.from_pretrained(model_id, subfolder="text_encoder",
torch_dtype=dtype)
ifquant_model:
quantize(text_encoder, weights=qfloat8_e4m3fn)
freeze(text_encoder)
quantize(transformer, weights=qfloat8_e4m3fn)
freeze(transformer)
pipe=QwenImageEditPlusPipeline.from_pretrained(model_id,
# vae=vae,# image_encoder=image_encoder,transformer=transformer,
text_encoder=text_encoder,
# tokenizer=tokenizer,torch_dtype=dtype)
print(f"Current scheduler: {type(pipe.scheduler).__name__}")
print(f"Scheduler config: {pipe.scheduler.config}")
pipe.to("cuda")
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ring_degree=2))
pipe.transformer.set_attention_backend("_native_cudnn")
image=Image.open("/share/gdli7/common/AIPhoto/yarn-art-pikachu.png").resize((512,512)).convert("RGB")
prompt="Make Pikachu hold a sign that says 'Qwen is awesome', yarn art style, detailed, vibrant colors"generator=torch.Generator().manual_seed(42)
image=pipe(image, prompt, num_inference_steps=50, generator=generator).images[0]
ifrank==0:
image.save("output.png")
iftorch.distributed.is_initialized():
torch.distributed.destroy_process_group()

[sayakpaul]

@DefTruth could you try out #12446?

[sayakpaul]

@mali-afridi what's your recommended fix for Qwen?

[mali-afridi]

@sayakpaul I would suggest padding to make the dimension divisible, like I did in chengzeyi/ParaAttention#53
But it comes with a cost to slightly change the image since it also attends to those padded tokens

[sayakpaul]

Hmm, I see. That would be my preliminary approach, too. Can't we use masks to avoid that?

[DefTruth]

@DefTruth could you try out #12446?

sorry for late reply, i will take a try today.

[sayakpaul]

@DefTruth LMK how it went.

[DefTruth]

@DefTruth LMK how it went.

By using the _native_cudnn backend, I have currently avoided this problem

[sayakpaul]

Okay so it works. How about #12563?

[DefTruth]

Okay so it works. How about #12563?

also work

[sayakpaul]

Closing this since it has been resolved.