Context parallel bug when using QwenImage

[sayakpaul]

Refer more here: #12443

More context under "Result: Single coherent image matching single-GPU output" here.

Before proceeding with a PR, please reply to this issue mentioning the proposal you have in mind, and we can take it from there.

Please mention @DN6 @yiyixuxu and myself (@sayakpaul).

[Ratish1]

Hi @DN6 , @yiyixuxu and @sayakpaul , I'd like to contribute to this if possible. Proposal: Add padding to encoder_hidden_states and attention_mask in the QwenImage context parallelism setup to make them divisible by world_size, resolving the AssertionError. Adjust the masks to ignore padded tokens, ensuring coherent outputs matching single-GPU without quality changes. Thanks.

[sayakpaul]

Thanks for showing your interest!

Would it be similar to #12443 (comment)?

[Ratish1]

Thanks for the reply @sayakpaul . Yes, my proposal is similar to the padding approach suggested in that comment from #12443. To address the potential image change from attending to padded tokens (as noted there), I plan to explicitly adjust the attention masks to ignore those pads e.g., by setting the mask values for padded positions to prevent any attention flow to them. This should ensure the outputs remain coherent and identical to single-GPU results without quality degradation. Let me know if you prefer a different approach.

[Ratish1]

Hello @sayakpaul , I just wanted to follow up on this incase you were busy. thanks.

[sayakpaul]

Cool, let's do it then :)

[Ratish1]

Hello @sayakpaul , I have created the PR for this and I also made the necessary changes to not actually change the Qwen transformer model but support it for a more generic case from the context parallel hooks. Please do let me know if you need any changes or I need to change my approach for this issue. Thanks.

[sayakpaul]

You're all set. I will mention your name beside the issue. Thank you!

[Ratish1]

Thank You