Fix Context Parallel validation checks #12446
Conversation
|
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. |
| 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: |
There was a problem hiding this comment.
Should we hit line as both cannot be set, right?
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
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.
| 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() |
There was a problem hiding this comment.
Can't they be None? Why unguard?
There was a problem hiding this comment.
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.
| mesh_shape=(cp_config.ring_degree, cp_config.ulysses_degree), | ||
| mesh_dim_names=("ring", "ulysses"), | ||
| ) | ||
| from .attention_dispatch import AttentionBackendName, _AttentionBackendRegistry |
There was a problem hiding this comment.
(nit): Would prefer to keep the import at the beginning of the method implementation, if possible (after from .attention_processor import Attention, MochiAttention).
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
enable_parallelismFixes # (issue)
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
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.