github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 09:07:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Feature] MoE Ulysses Support (#5918)

Browse Source 
* moe sp support

* moe sp bug solve

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Haze188
2024-07-18 11:37:56 +08:00
committed by Hongxin Liu
parent 3e2b6132b7
commit 404b16faf3
6 changed files with 571 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/cluster/process_group_mesh.py
View File

@@ -209,7 +209,7 @@ class ProcessGroupMesh:
axis: Union[int, List[int]],
indices_at_axis: Optional[Union[List[int], List[List[int]]]] = None,
backend: Optional[str] = None,
return_ranks_by_group: bool = False
return_ranks_by_group: bool = False,
) -> Union[ProcessGroup, List[Tuple[int, ...]]]:
"""Create all process groups along the given axis, and return the one which the current process belongs to.
@@ -257,7 +257,11 @@ class ProcessGroupMesh:
return target_group
def get_group_along_axis(
self, axis: Union[int, List[int]], indices_at_axis: Optional[List[int]] = None, backend: Optional[str] = None, return_ranks_by_group: bool = False
self,
axis: Union[int, List[int]],
indices_at_axis: Optional[List[int]] = None,
backend: Optional[str] = None,
return_ranks_by_group: bool = False,
) -> Union[ProcessGroup, List[Tuple[int, ...]]]:
"""Get the process group along the given axis which the current process belongs to. If the process group doesn't exist, it will be created.