fix(hooks): Add padding support to context parallel hooks

[Ratish1]

What does this PR do?

This PR now modifies the ContextParallelSplitHook and ContextParallelGatherHook to gracefully handle sequence lengths that are not divisible by the world size.

This PR changes:

1. Generic Padding: The ContextParallelSplitHook now pads any input tensor to a divisible length before sharding.
2. State Management: It temporarily stores the original sequence length on the module instance itself.
3. Generic Trimming: The ContextParallelGatherHook uses this stored length to trim the padding from the final output tensor before returning it.

This ensures that the padding is completely transparent to the model and the end-user, preventing crashes without altering the output shape. The fix is now contained entirely within the hooks and requires no changes to the Qwen transformer or any other model.

I have also added a new unit test in tests/hooks/test_hooks.py that directly tests this new padding and trimming logic,

Fixes #12568

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 @yiyixuxu @DN6

[yiyixuxu]

thanks for the PR! however, we will not want any of these logic go into qwen transformer
would you be interested to how to support this case( not just qwen) from the context parallel hooks
https://github.com/huggingface/diffusers/blob/main/src/diffusers/hooks/context_parallel.py#L204

[Ratish1]

thanks for the PR! however, we will not want any of these logic go into qwen transformer would you be interested to how to support this case( not just qwen) from the context parallel hooks https://github.com/huggingface/diffusers/blob/main/src/diffusers/hooks/context_parallel.py#L204

Hi @yiyixuxu, yes I would be interested to support this change.

I have some follow up questions regarding the model-specific consequences of this padding. The Qwen transformer calculates RoPE based on the original sequence length. When the hook pads the input tensor, the model still needs to be aware of the new, padded length to avoid a shape mismatch inside the RoPE calculation.

Previously, I fixed this by recalculating the sequence length inside the transformer's forward method based on the padded tensor's shape. With the padding logic now in the hook, could you help me what is the preferred way to handle this?

Is it acceptable to keep that small, model-specific part of the logic (recalculating the sequence length for RoPE) inside the Qwen transformer, or is there a more general way to communicate the new padded length from the hook back to the model that I should use instead?. Thanks for your help.

[Ratish1]

Hi @yiyixuxu ,Just wanted to follow up. After looking at the hook implementation as you suggested, I've updated the PR with a new approach that is fully generic and contains all logic within the hooks, with no changes to the transformer.

The solution now involves adding padding in the ContextParallelSplitHook and then trimming it in theContextParallelGatherHook, using the module instance to temporarily store the original sequence length. I've also added a new unit test for this logic in test_hooks.py. Thanks and lmk if you need more changes. I've updated the PR description with the full details.

CC @sayakpaul @DN6