Fix Context Parallel validation checks by DN6 · Pull Request #12446 · huggingface/diffusers

DN6 · 2025-10-07T12:22:05Z

What does this PR do?

Currently CP inference will run with split hooks even if the attention backend doesn't support it. This can lead to weird results #12443

This PR

Fixes the validation check in the AttentionBackendRegistry to return False for invalid backends
Removes redundant checks in the ContextParallelConfig and adds properties to access device mesh config arguments in a more general way.
Move CP config related validation checks into the CP Config itself, rather than in enable_parallelism

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.

HuggingFaceDocBuilderDev · 2025-10-07T12:30:18Z

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

Left some nits, otherwise looks very nice!

sayakpaul · 2025-10-07T15:11:44Z

src/diffusers/models/_modeling_parallel.py

-            raise NotImplementedError(
-                f"Only rotate_method='allgather' is supported for now, but got {self.rotate_method}."
+
+        if self.ulysses_degree * self.ring_degree > world_size:


Should we hit line as both cannot be set, right?

Both can be set techinically, but currently both can't be > 1. Also this is for cases where you have 3 GPUs available and you set something like ulysses_degree=1 and ring_degree==4 (more GPUs being requested is greater than world_size)

Feels slightly confusing to me but since we're erroring out early for unsupported ulysses_degree and ring_degree value combos, I think it's okay.

sayakpaul · 2025-10-07T15:12:36Z

src/diffusers/models/_modeling_parallel.py

+        self._flattened_mesh = self._mesh._flatten()
+        self._ring_mesh = self._mesh["ring"]
+        self._ulysses_mesh = self._mesh["ulysses"]
+        self._ring_local_rank = self._ring_mesh.get_local_rank()
+        self._ulysses_local_rank = self._ulysses_mesh.get_local_rank()


Can't they be None? Why unguard?

They are internal attributes that are derived from mesh which is set through the setup method. The device mesh object is also only dynamically created when enabled_parallelism is called.

The guards are redundant, they would always be None unless set explicitly for some custom debugging.

sayakpaul · 2025-10-07T15:14:38Z

src/diffusers/models/modeling_utils.py

-                mesh_shape=(cp_config.ring_degree, cp_config.ulysses_degree),
-                mesh_dim_names=("ring", "ulysses"),
-            )
+            from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry


(nit): Would prefer to keep the import at the beginning of the method implementation, if possible (after from .attention_processor import Attention, MochiAttention).

sayakpaul

Just some nits. But not merge-blocker.

sayakpaul · 2025-10-30T03:48:33Z

src/diffusers/models/_modeling_parallel.py

+    @property
+    def mesh_shape(self) -> Tuple[int, int]:
+        """Shape of the device mesh (ring_degree, ulysses_degree)."""
+        return (self.ring_degree, self.ulysses_degree)


Would it be possible to add a small explainer about what it would mean for different values, for example - "(3, 1), (1, 3)", etc.? When both are being set, both cannot be > 1.

sayakpaul · 2025-10-30T03:50:20Z

src/diffusers/models/attention_dispatch.py

-        supports_context_parallel = backend in cls._supports_context_parallel
-        is_degree_greater_than_1 = parallel_config is not None and (
-            parallel_config.context_parallel_config.ring_degree > 1
-            or parallel_config.context_parallel_config.ulysses_degree > 1
-        )
-        return supports_context_parallel and is_degree_greater_than_1
+        supports_context_parallel = backend.value in cls._supports_context_parallel
+        return supports_context_parallel


Very nice cleanup here!

sayakpaul · 2025-10-30T03:53:58Z

src/diffusers/models/modeling_utils.py

-        cp_mesh = None
+        attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
+
+        # Step 1: Validate attention backend supports context parallelism if enabled


(nit): Seems like we are not documenting other steps. So, maybe we can remove this comment or add other steps as comments.

DefTruth · 2025-10-31T07:12:54Z

@sayakpaul It worked for me. When I set the attn backend to native, it threw an error correctly. but i think the native attn backend should be enable for context parallelism as #12563

INFO 10-31 07:05:23 [__init__.py:64] Found attention_backend from config, set attention backend to: native
[rank1]:     transformer.enable_parallelism(
[rank1]:   File "/workspace/dev/vipshop/diffusers/src/diffusers/models/modeling_utils.py", line 1529, in enable_parallelism
[rank1]:     raise ValueError(
[rank1]: ValueError: Context parallelism is enabled but the attention processor 'WanAttnProcessor' is using backend 'native' which does not support context parallelism. Please set a compatible attention backend: ['_native_cudnn', 'flash', 'sage'] using `model.set_attention_backend()` before calling `enable_parallelism()`.

sayakpaul · 2025-11-04T01:53:55Z

@DN6 I think we might have to consider #12563 now that native is supported for Ulysses.

DN6 added 8 commits October 7, 2025 14:42

update

faf61a4

update

428399b

update

1d76322

update

a66787b

update

881e262

update

0845ca0

update

8018a6a

update

f925783

DN6 requested a review from sayakpaul October 7, 2025 12:22

DN6 mentioned this pull request Oct 7, 2025

Fix Flux Context Parallel Bug (Incoherent Image Generation) #12443

Closed

6 tasks

update

5bfc7dd

sayakpaul reviewed Oct 7, 2025

View reviewed changes

DN6 added 2 commits October 8, 2025 14:35

update

fb15ff5

Merge branch 'main' into cp-fix

56114f4

DN6 requested a review from sayakpaul October 30, 2025 02:38

sayakpaul approved these changes Oct 30, 2025

View reviewed changes

DN6 added 2 commits October 30, 2025 22:31

update

4505645

update

3b12a0b

Merge branch 'main' into cp-fix

d65f857

sayakpaul mentioned this pull request Oct 31, 2025

ulysses enabling in native attention path #12563

Merged

Merge branch 'main' into cp-fix

3dcc9ca

DN6 added 2 commits November 6, 2025 17:46

Merge branch 'main' into cp-fix

197dd5f

Merge branch 'main' into cp-fix

e41ca61

DN6 merged commit 8d45f21 into main Nov 10, 2025
31 of 35 checks passed

Conversation

DN6 commented Oct 7, 2025

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Oct 7, 2025

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DefTruth commented Oct 31, 2025

Uh oh!

sayakpaul commented Nov 4, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants