wehub/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deepspeed/utils/groups.py
+916
View File
@@ -0,0 +1,916 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
# The file has been adapted from https://github.com/NVIDIA/Megatron-LM and retains the following license from the original file
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Support different forms of parallelism in DeepSpeed using multiple process groups.
Given that there are multiple scenarios and use-cases, this file is going to be updated
frequently. For now, the group creation needed for the training scenario is being implemented.
For inference and other new scenarios, the code will be either reused or added to this file.
"""
import os
from deepspeed import comm as dist
from deepspeed.utils import log_dist
from deepspeed.utils.bwc import bwc_tensor_model_parallel_world_size, bwc_pipeline_parallel_world_size
from deepspeed.utils.exceptions import DeprecatedException
from deepspeed.accelerator import get_accelerator
# Expert parallel group that the current rank belongs to.
_EXPERT_PARALLEL_GROUP = {}
# Mapping of expert parallel group to ranks
_EXPERT_PARALLEL_GROUP_RANKS = {}
# Expert data parallel group that the current rank belongs to.
_EXPERT_DATA_PARALLEL_GROUP = {}
# Mapping of expert data parallel group to ranks
_EXPERT_DATA_PARALLEL_GROUP_RANKS = {}
# dist world group needs to be cloned for some cases
_WORLD_GROUP = None
# ZeRO parameter partitioning group that the current rank belongs to.
_ZERO_PARAM_INTRA_PARALLEL_GROUP = None
# global object to maintain mpu object if passed by a Megatron client
mpu = None
# global object that stores tensor parallel world size for experts
expert_tensor_parallel_world_size = 1
# All to All quantized graident communication groups
_ALL_TO_ALL_GROUP = {}
mesh_device = None
_DEVICE_MESH_SPLIT_UNSUPPORTED = "No backend for the parent process group or its backend does not support splitting"
# Deprecated groups initialize function.
def initialize(ep_size=1, mpu=None):
""" Deprecated function. Retained to inform the users."""
raise DeprecatedException(
"Please do not use the groups.initialize() API as it is deprecated. Instead, pass the desired ep_size to deepspeed.moe.layer.MoE(..,ep_size,..)"
)
def _ensure_divisibility(numerator, denominator):
"""Ensure that numerator is divisible by the denominator."""
assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator)
# ======== Start: Tensor Parallel Group Attributes ========
# Intra-layer model parallel group that the current rank belongs to.
_TENSOR_MODEL_PARALLEL_GROUP = None
# Model parallel group (both intra- and pipeline) that the current rank belongs to.
_MODEL_PARALLEL_GROUP = None
# Data parallel group that the current rank belongs to.
_DATA_PARALLEL_GROUP = None
# These values enable us to change the mpu sizes on the fly.
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
_MPU_TENSOR_MODEL_PARALLEL_RANK = None
def _init_tp_groups_with_new_group(tensor_model_parallel_size=1, data_parallel_size=None):
"""Initialize TP/DP groups with explicit rank lists.
This mirrors a 2D DeviceMesh shaped as (data_parallel, tensor_parallel),
while avoiding DeviceMesh's optimized split_group path.
"""
global _DATA_PARALLEL_GROUP
global _MODEL_PARALLEL_GROUP
global _TENSOR_MODEL_PARALLEL_GROUP
world_size = dist.get_world_size()
_ensure_divisibility(world_size, tensor_model_parallel_size)
if data_parallel_size is None:
data_parallel_size = world_size // tensor_model_parallel_size
else:
assert data_parallel_size * tensor_model_parallel_size == world_size, (
f"data_parallel_size ({data_parallel_size}) * tensor_model_parallel_size "
f"({tensor_model_parallel_size}) must equal world_size ({world_size})")
rank = dist.get_rank()
data_parallel_group = None
tensor_model_parallel_group = None
for tensor_rank in range(tensor_model_parallel_size):
ranks = list(range(tensor_rank, world_size, tensor_model_parallel_size))
group = dist.new_group(ranks)
if rank in ranks:
data_parallel_group = group
for data_rank in range(data_parallel_size):
start = data_rank * tensor_model_parallel_size
ranks = list(range(start, start + tensor_model_parallel_size))
group = dist.new_group(ranks)
if rank in ranks:
tensor_model_parallel_group = group
assert data_parallel_group is not None, 'data parallel group is not initialized'
assert tensor_model_parallel_group is not None, 'tensor parallel group is not initialized'
_DATA_PARALLEL_GROUP = data_parallel_group
_TENSOR_MODEL_PARALLEL_GROUP = tensor_model_parallel_group
_MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP
return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
def _init_tp_mesh_device(tensor_model_parallel_size=1, data_parallel_size=None):
"""Initialize model data parallel groups."""
global _DATA_PARALLEL_GROUP
global _MODEL_PARALLEL_GROUP
global _TENSOR_MODEL_PARALLEL_GROUP
if _TENSOR_MODEL_PARALLEL_GROUP is not None:
return
if data_parallel_size is None:
data_parallel_size = dist.get_world_size() // tensor_model_parallel_size
if os.environ.get("TORCH_DISTRIBUTED_DEBUG", "").upper() == "DETAIL":
log_dist("TORCH_DISTRIBUTED_DEBUG=DETAIL detected; initializing TP mesh groups with new_group", ranks=[0])
return _init_tp_groups_with_new_group(tensor_model_parallel_size, data_parallel_size)
try:
mesh_device = dist.initialize_mesh_device((data_parallel_size, tensor_model_parallel_size),
("data_parallel", "tensor_parallel"))
except RuntimeError as exc:
if _DEVICE_MESH_SPLIT_UNSUPPORTED not in str(exc):
raise
log_dist("DeviceMesh process-group splitting is unsupported; falling back to new_group TP mesh groups",
ranks=[0])
return _init_tp_groups_with_new_group(tensor_model_parallel_size, data_parallel_size)
_TENSOR_MODEL_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="tensor_parallel")
_DATA_PARALLEL_GROUP = mesh_device.get_group(mesh_dim="data_parallel")
# They are always equal only in 2D (DP + TP) parallelism.
# _MODEL_PARALLEL_GROUP is assigned the same value as _TENSOR_MODEL_PARALLEL_GROUP
# to allow for future potential changes.
_MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP
return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
def get_tensor_model_parallel_group():
"""Get the tensor model parallel group the caller rank belongs to."""
assert _TENSOR_MODEL_PARALLEL_GROUP is not None, \
'intra_layer_model parallel group is not initialized'
return _TENSOR_MODEL_PARALLEL_GROUP
def get_model_parallel_group():
"""Get the model parallel group the caller rank belongs to."""
assert _MODEL_PARALLEL_GROUP is not None, \
'model parallel group is not initialized'
return _MODEL_PARALLEL_GROUP
def get_data_parallel_group():
"""Get the data parallel group the caller rank belongs to."""
assert _DATA_PARALLEL_GROUP is not None, \
'data parallel group is not initialized'
return _DATA_PARALLEL_GROUP
def set_tensor_model_parallel_world_size(world_size):
"""Set the tensor model parallel size"""
global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
def get_tensor_model_parallel_world_size():
"""Return world size for the tensor model parallel group."""
global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
return dist.get_world_size(group=get_tensor_model_parallel_group())
def get_model_parallel_world_size():
return get_tensor_model_parallel_world_size()
def set_tensor_model_parallel_rank(rank):
"""Set tensor model parallel rank."""
global _MPU_TENSOR_MODEL_PARALLEL_RANK
_MPU_TENSOR_MODEL_PARALLEL_RANK = rank
def get_tensor_model_parallel_rank():
"""Return my rank for the tensor model parallel group."""
global _MPU_TENSOR_MODEL_PARALLEL_RANK
if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
return _MPU_TENSOR_MODEL_PARALLEL_RANK
return dist.get_rank(group=get_tensor_model_parallel_group())
def get_model_parallel_rank():
return get_tensor_model_parallel_rank()
def get_tensor_model_parallel_src_rank():
"""Calculate the global rank corresponding to the first local rank
in the tensor model parallel group."""
global_rank = dist.get_rank()
local_world_size = get_tensor_model_parallel_world_size()
return (global_rank // local_world_size) * local_world_size
def get_data_parallel_world_size():
"""Return world size for the data parallel group."""
return dist.get_world_size(group=get_data_parallel_group())
def get_data_parallel_rank():
"""Return my rank for the data parallel group."""
return dist.get_rank(group=get_data_parallel_group())
# ======== End: Tensor Parallel Group Attributes ========
# Not currently used. Helper function to create a model (tensor) parallel group.
def _create_model_parallel(model_parallel_size_):
"""
Initialize model data parallel groups.
Arguments:
model_parallel_size: number of GPUs used to parallelize model.
Returns:
Tuple of data parallel group and model parallel group
Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we
use 2 GPUs to parallelize the model. The present function will
create 4 model parallel groups and 2 data parallel groups as:
4 model parallel groups:
[g0, g1], [g2, g3], [g4, g5], [g6, g7]
2 data parallel groups:
[g0, g2, g4, g6], [g1, g3, g5, g7]
Note that for efficiency, the caller should make sure adjacent ranks
are on the same DGX box. For example if we are using 2 DGX-1 boxes
with a total of 16 GPUs, rank 0 to 7 belong to the first box and
ranks 8 to 15 belong to the second box.
"""
log_dist(f'Creating model parallel group with size {model_parallel_size_}', ranks=[0])
# Get world size and rank. Ensure some consistencies.
assert dist.is_initialized()
world_size = dist.get_world_size()
model_parallel_size = min(model_parallel_size_, world_size)
_ensure_divisibility(world_size, model_parallel_size)
rank = dist.get_rank()
_DATA_PARALLEL_GROUP = None
_MODEL_PARALLEL_GROUP = None
# Build the data parallel groups.
for i in range(model_parallel_size):
ranks = range(i, world_size, model_parallel_size)
group = dist.new_group(ranks)
if i == (rank % model_parallel_size):
_DATA_PARALLEL_GROUP = group
# Build the model parallel groups.
for i in range(world_size // model_parallel_size):
ranks = range(i * model_parallel_size, (i + 1) * model_parallel_size)
group = dist.new_group(ranks)
if i == (rank // model_parallel_size):
_MODEL_PARALLEL_GROUP = group
return _DATA_PARALLEL_GROUP, _MODEL_PARALLEL_GROUP
def _create_expert_and_data_parallel(expert_parallel_size_,
mp_size=None,
pp_size=None,
mp_mode="tp",
use_data_before_expert_parallel_=False,
folding_spec=None):
"""Create expert and data parallel groups.
When mp_size is None or 1: legacy consecutive ordering (backward compatible).
When mp_size > 1 and folding_spec is not None: AutoEP+AutoTP folding tables.
When mp_size > 1 and mp_mode=="tp": TP-strided rank ordering.
When mp_size > 1 and mp_mode=="sp": consecutive rank ordering.
Note: Caller of this function is responsible to check if the groups already exist.
Example - E + D parallel (legacy path)
world_size = 16
expert_parallel_size = 2 # number of experts in same group
expert_data_parallel_group = [0,2,4,6,8,10,12,14], [1,3,5,7,9,11,13,15] - all reduce is only on MoE params
expert_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - no all reduce, but all to all
data_parallel_group = [0,1,...,15] - all reduce is only on non-MoE
Args:
expert_parallel_size_ (int): Expert parallel group size.
mp_size (int, optional): Model parallel size (TP or SP). None treated as 1.
pp_size (int, optional): Pipeline parallel size. None falls back to mpu.
mp_mode (str): "tp" for TP-strided ordering, "sp" for consecutive ordering.
use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology.
folding_spec: Optional AutoEP+AutoTP folding topology spec.
"""
assert dist.is_initialized()
# Resolve parameters for backward compat
effective_mp_size = 1 if mp_size is None else mp_size
log_dist(f'Creating expert and data parallel groups with size {expert_parallel_size_}', ranks=[0])
world_size = dist.get_world_size()
# Resolve pp_size
if pp_size is not None:
pp_world_size = pp_size
else:
pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)
rank = dist.get_rank()
pp_stride = world_size // pp_world_size
_ensure_divisibility(pp_stride, expert_parallel_size_)
group_name = f"ep_size_{expert_parallel_size_}"
global _EXPERT_DATA_PARALLEL_GROUP
global _EXPERT_DATA_PARALLEL_GROUP_RANKS
global _EXPERT_PARALLEL_GROUP
global _EXPERT_PARALLEL_GROUP_RANKS
# Legacy path: mp_size <= 1 (preserves exact original behavior)
if effective_mp_size <= 1:
ep_stride = pp_stride // expert_parallel_size_
# Build the expert data parallel groups.
# Only create group if it does not already exist
if group_name not in _EXPERT_DATA_PARALLEL_GROUP:
for pp_stage_start in range(0, world_size, pp_stride):
for i in range(expert_parallel_size_):
if use_data_before_expert_parallel_:
ranks = range(pp_stage_start + i * ep_stride, pp_stage_start + (i + 1) * ep_stride)
else:
ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, expert_parallel_size_)
group = dist.new_group(ranks)
log_dist(
f'Creating expert data parallel process group named {group_name} with ranks: {list(ranks)}',
[0])
if rank in ranks:
_EXPERT_DATA_PARALLEL_GROUP[group_name] = group
_EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks
# Build the expert parallel groups.
# Only create group if it does not already exist
if group_name not in _EXPERT_PARALLEL_GROUP:
if use_data_before_expert_parallel_:
for pp_stage_start in range(0, world_size, pp_stride):
for i in range(ep_stride):
ranks = range(pp_stage_start + i, pp_stage_start + pp_stride, ep_stride)
group = dist.new_group(ranks)
log_dist(
f'creating expert parallel process group named {group_name} '
f'with ranks: {list(ranks)}', [0])
if rank in ranks:
_EXPERT_PARALLEL_GROUP[group_name] = group
_EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
else:
for i in range(world_size // expert_parallel_size_):
ranks = range(i * expert_parallel_size_, (i + 1) * expert_parallel_size_)
group = dist.new_group(ranks)
log_dist(
f'creating expert parallel process group named {group_name} '
f'with ranks: {list(ranks)}', [0])
if rank in ranks:
_EXPERT_PARALLEL_GROUP[group_name] = group
_EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
return
# New path: mp_size > 1
if use_data_before_expert_parallel_:
raise NotImplementedError("use_data_before_expert_parallel_ is not supported with mp_size > 1")
if group_name in _EXPERT_PARALLEL_GROUP:
if folding_spec is not None:
from deepspeed.module_inject.auto_ep_folding import assert_group_matches_spec
assert_group_matches_spec(
{
"ep": [_EXPERT_PARALLEL_GROUP_RANKS[group_name]],
"edp": [_EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name]],
},
folding_spec,
)
return # Already created
folding_tables = None
if folding_spec is not None:
from deepspeed.module_inject.auto_ep_folding import expected_folding_group_tables
folding_tables = expected_folding_group_tables(folding_spec)
for pp_stage_start in range(0, world_size, pp_stride):
stage_ranks = list(range(pp_stage_start, pp_stage_start + pp_stride))
stage_rank_set = set(stage_ranks)
if folding_tables is not None:
ep_groups_list = [list(group) for group in folding_tables.ep_groups if set(group).issubset(stage_rank_set)]
edp_groups_list = [
list(group) for group in folding_tables.edp_groups if set(group).issubset(stage_rank_set)
]
else:
# Preserve the existing TP-strided native MoE topology when no
# folding spec was provided by the AutoEP+AutoTP path.
if mp_mode == "tp" and effective_mp_size > 1:
num_tp_groups = len(stage_ranks) // effective_mp_size
ordered_stage_ranks = []
for dp_lane in range(effective_mp_size):
for tp_group_idx in range(num_tp_groups):
ordered_stage_ranks.append(stage_ranks[tp_group_idx * effective_mp_size + dp_lane])
else:
ordered_stage_ranks = stage_ranks
num_ep_groups = len(ordered_stage_ranks) // expert_parallel_size_
ep_groups_list = [
ordered_stage_ranks[g * expert_parallel_size_:(g + 1) * expert_parallel_size_]
for g in range(num_ep_groups)
]
edp_groups_list = [[ep_groups_list[g][pos] for g in range(num_ep_groups)]
for pos in range(expert_parallel_size_)]
for ep_ranks in ep_groups_list:
group = dist.new_group(ep_ranks)
log_dist(f'creating expert parallel process group named {group_name} with ranks: {ep_ranks}', [0])
if rank in ep_ranks:
_EXPERT_PARALLEL_GROUP[group_name] = group
_EXPERT_PARALLEL_GROUP_RANKS[group_name] = ep_ranks
for edp_ranks in edp_groups_list:
group = dist.new_group(edp_ranks)
log_dist(f'Creating expert data parallel process group named {group_name} with ranks: {edp_ranks}', [0])
if rank in edp_ranks:
_EXPERT_DATA_PARALLEL_GROUP[group_name] = group
_EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = edp_ranks
def _get_expert_parallel_ranks(world_size,
tensor_parallel_size_,
expert_parallel_size_,
pipeline_parallel_size_=1,
use_data_before_expert_parallel_=False):
"""Generate expert parallel and expert data parallel group ranks list.
Example - E + M + D parallel
world_size = 16
model_degree = 2
expert_degree = 4 # number of experts in same group
mp_group = [0, 1], [2,3], [4,5] ...
data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15]
expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15]
expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15]
Args:
world_size (int): Distributed world size.
tensor_parallel_size_ (int): Tensor parallel group size.
expert_parallel_size_ (int): Expert parallel group size.
pipeline_parallel_size_ (int): Pipeline parallel group size
use_data_before_expert_parallel_ (bool): Use the D + E instead of E + D topology
Returns:
Expert parallel group ranks and Expert data parallel group ranks list.
"""
_ensure_divisibility(world_size, tensor_parallel_size_ * pipeline_parallel_size_)
dp_world_size = world_size // (tensor_parallel_size_ * pipeline_parallel_size_)
_ensure_divisibility(dp_world_size, expert_parallel_size_)
# Generate data parallel groups
data_parallel_groups = []
dp_group_size = tensor_parallel_size_
pp_stride = world_size // pipeline_parallel_size_
if use_data_before_expert_parallel_:
dp_stride = world_size // expert_parallel_size_ // tensor_parallel_size_ // pipeline_parallel_size_
for pp_stage_start in range(0, world_size, pp_stride):
pp_stage_next = pp_stage_start + pp_stride
for i in range(dp_group_size):
data_parallel_groups.append(list())
for ds in range(dp_stride):
# [0, 4, 8, 12, 16, 20, 24, 28, 2, 6, 10, 14, 18, 22, 26, 30]
# [1, 5, 9, 13, 17, 21, 25, 29, 3, 7, 11, 15, 19, 23, 27, 31]
data_parallel_groups[-1].extend(
list(
range(pp_stage_start + i + ds * tensor_parallel_size_, pp_stage_next,
dp_stride * tensor_parallel_size_)))
else:
for pp_stage_start in range(0, world_size, pp_stride):
pp_stage_next = pp_stage_start + pp_stride
for i in range(dp_group_size):
data_parallel_groups.append(list(range(pp_stage_start + i, pp_stage_next, dp_group_size)))
expert_parallel_groups = []
expert_data_parallel_groups = []
for dp_ranks in data_parallel_groups:
# partition of expert parallel groups, e.g. [0,2,4,6], [8,10,12,14]
part_ep_groups = []
for i in range(0, dp_world_size, expert_parallel_size_):
part_ep_groups.append(dp_ranks[i:i + expert_parallel_size_])
expert_parallel_groups.extend(part_ep_groups)
# zip part_ep_groups get expert data parallel ranks, e.g [0,8],[2,10],[4,12],[6,14]
for expert_dp_ranks in zip(*part_ep_groups):
expert_data_parallel_groups.append(list(expert_dp_ranks))
return expert_parallel_groups, expert_data_parallel_groups
def _create_expert_data_and_model_parallel(expert_parallel_size_, mpu, use_data_before_expert_parallel_=False):
"""
Create expert and data parallel groups based on MPU (model parallel) group.
Note: Caller of this function is responsible to check if the groups already exist.
Example - E + M + D parallel
world_size = 16
model_degree = 2
expert_degree = 4 # number of experts in same group
mp_group = [0, 1], [2,3], [4,5] ...
data_parallel_group =[0,2,4,6,8,10, 12,14], [1,3,5,7,9,11,13,15]
expert_parallel_group = [0,2,4,6], [8,10,12,14] [1,3,5,7], [9,11,13,15]
expert_data_parallel_group = [0,8],[2,10],[4,12],[6,14], [1,9],[3,11],[5,13],[7,15]
"""
assert dist.is_initialized(), "dist is not initialized"
tensor_parallel_size_ = bwc_tensor_model_parallel_world_size(mpu)
global expert_tensor_parallel_world_size
expert_tensor_parallel_world_size = tensor_parallel_size_
world_size = dist.get_world_size()
rank = dist.get_rank()
dp_world_size = _get_data_parallel_world_size()
pp_world_size = 1 if mpu is None else bwc_pipeline_parallel_world_size(mpu)
_ensure_divisibility(world_size, tensor_parallel_size_)
_ensure_divisibility(dp_world_size, expert_parallel_size_)
log_dist(
f"Creating deepspeed groups with model parallel size {tensor_parallel_size_}, "
f"pipeline parallel size {pp_world_size}, expert parallel size {expert_parallel_size_}, "
f"world size {world_size}, dp world size {dp_world_size}", [0])
global _EXPERT_PARALLEL_GROUP, _EXPERT_DATA_PARALLEL_GROUP
global _EXPERT_PARALLEL_GROUP_RANKS, _EXPERT_DATA_PARALLEL_GROUP_RANKS
group_name = f"ep_size_{expert_parallel_size_}"
# Only create groups if they don't already exist
# Need to check conditions outside the group creation loop because of the way torch.dist group creation works
if group_name not in _EXPERT_DATA_PARALLEL_GROUP and group_name not in _EXPERT_PARALLEL_GROUP:
expert_parallel_groups, expert_data_parallel_groups = _get_expert_parallel_ranks(
world_size, tensor_parallel_size_, expert_parallel_size_, pp_world_size, use_data_before_expert_parallel_)
for ranks in expert_parallel_groups:
group = dist.new_group(ranks)
if rank in list(ranks):
_EXPERT_PARALLEL_GROUP[group_name] = group
_EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks
for ranks in expert_data_parallel_groups:
group = dist.new_group(ranks)
if rank in list(ranks):
_EXPERT_DATA_PARALLEL_GROUP[group_name] = group
_EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = ranks
def _get_max_expert_size():
"""Get the maximum ep_size from all the created groups."""
assert _EXPERT_PARALLEL_GROUP is not None, "Warning! Process group not initialized"
keylist = []
for key in _EXPERT_PARALLEL_GROUP.keys():
# index 2 is ep_size in the group name: ep_size_<ep_size>
index = 2
keylist.append(int(key.split('_')[index]))
return max(keylist) if len(keylist) > 0 else None
def _get_max_expert_size_name():
"""Get the name of the group with max. ep_size"""
return f'ep_size_{_get_max_expert_size()}'
def _get_max_expert_parallel_group():
"""Get the max expert parallel size."""
return _get_expert_parallel_group(_get_max_expert_size_name())
def _get_expert_parallel_group(group_name):
"""Get the expert parallel group the caller rank belongs to."""
assert group_name in _EXPERT_PARALLEL_GROUP, \
'expert parallel group is not initialized'
return _EXPERT_PARALLEL_GROUP[group_name]
def _get_expert_parallel_group_ranks(group_name):
"""Get the ranks of the expert parallel group the caller rank belongs to."""
assert group_name in _EXPERT_PARALLEL_GROUP_RANKS, \
'expert parallel group is not initialized'
return _EXPERT_PARALLEL_GROUP_RANKS[group_name]
def _get_expert_parallel_group_dict():
"""Get the expert parallel group dict."""
return _EXPERT_PARALLEL_GROUP
def _get_expert_data_parallel_group(group_name):
"""Get the expert data parallel group the caller rank belongs to."""
assert group_name in _EXPERT_DATA_PARALLEL_GROUP, \
'expert data parallel group is not initialized'
return _EXPERT_DATA_PARALLEL_GROUP[group_name]
def _get_expert_data_parallel_group_ranks(group_name):
"""Get the ranks of the expert data parallel group the caller rank belongs to."""
assert group_name in _EXPERT_DATA_PARALLEL_GROUP_RANKS, \
'expert data parallel group is not initialized'
return _EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name]
def _get_expert_data_parallel_group_dict():
"""Get the expert data parallel group dict."""
return _EXPERT_DATA_PARALLEL_GROUP
def _clone_world_group():
"""Create a clone of the world group
Note: We need to clone the dist world group because we
use dist.get_global_rank() utility function in DeepSpeed at many places.
As that function does not work on dist.group.WORLD, we
need to keep a clone of it.
"""
assert dist.is_initialized(), "dist is not initialized"
global _WORLD_GROUP
if _WORLD_GROUP is None:
# If not cloned already, clone the world group
_WORLD_GROUP = dist.new_group(ranks=range(dist.get_world_size()))
return _WORLD_GROUP
def _get_local_all_to_all_group():
assert dist.is_initialized(), 'dist is not initialized'
global _ALL_TO_ALL_GROUP
device_per_node = get_accelerator().device_count()
num_local = dist.get_world_size() // device_per_node
if num_local == 0 and dist.get_world_size() > 0:
assert dist.get_world_size() >= 1, 'num_gpus must >=1, cannot initialize All-To-All'
cur_rank = []
for i in range(dist.get_world_size()):
cur_rank.append(i)
_ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=cur_rank)
elif num_local == 1:
assert dist.get_world_size(
) == device_per_node, 'num_gpus not equal to device per node, cannot initialize All-To-All'
_ALL_TO_ALL_GROUP['local_0'] = dist.new_group(ranks=[i for i in range(device_per_node)])
else:
assert dist.get_world_size() > device_per_node, 'num_nodes<2 cannot initialize All-To-All'
for i in range(num_local):
local_rank = [j + device_per_node * i for j in range(device_per_node)]
_ALL_TO_ALL_GROUP[f"local_{i}"] = dist.new_group(ranks=local_rank)
for i in range(device_per_node):
cur_rank = []
for j in range(num_local):
cur_rank.append(i + j * device_per_node)
_ALL_TO_ALL_GROUP[f"global_{i}"] = dist.new_group(ranks=cur_rank)
return _ALL_TO_ALL_GROUP
def _get_data_parallel_group():
"""Get the data parallel group the caller rank belongs to."""
assert dist.is_initialized(), 'dist is not initialized'
global mpu
if mesh_device is not None:
return mesh_device.get_group(mesh_dim="data_parallel")
if mpu is not None:
if hasattr(mpu, 'initialize_sequence_parallel'):
return None
else:
return mpu.get_data_parallel_group()
# Return the clone of dist world group
return _clone_world_group()
def _get_data_parallel_group_ranks():
"""Get the ranks of data parallel group the caller rank belongs to."""
assert dist.is_initialized(), \
'dist is not initialized'
global mpu
if mpu is not None:
return mpu.get_data_parallel_group_ranks()
# Return all ranks
return range(dist.get_world_size())
def _get_broadcast_src_rank():
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_global_rank(_get_sequence_data_parallel_group(), 0)
def _get_expert_broadcast_src_rank(group_name):
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_global_rank(_get_expert_data_parallel_group(group_name), 0)
def _get_expert_parallel_world_size(group_name):
"""Return world size for the expert parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_world_size(group=_get_expert_parallel_group(group_name))
def _get_expert_data_parallel_world_size(group_name):
"""Return world size for the expert data parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_world_size(group=_get_expert_data_parallel_group(group_name))
def _get_expert_parallel_rank(group_name):
"""Return my rank for the expert parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_rank(group=_get_expert_parallel_group(group_name))
def _get_expert_parallel_src_rank(group_name):
"""Calculate the global rank corresponding to a local rank zero
in the expert parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
global_rank = dist.get_rank()
local_world_size = _get_expert_parallel_world_size(group_name)
return (global_rank // local_world_size) * local_world_size
def _get_expert_data_parallel_rank(group_name):
"""Return my rank for the expert data parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_rank(group=_get_expert_data_parallel_group(group_name))
def _get_data_parallel_world_size():
"""Return world size for the data parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
if mesh_device is not None:
return dist.get_world_size(mesh_device.get_group(mesh_dim="data_parallel"))
global mpu
if mpu is not None:
if hasattr(mpu, 'initialize_sequence_parallel'):
return None
else:
return mpu.get_data_parallel_world_size()
return dist.get_world_size(group=_get_data_parallel_group())
def _get_model_parallel_world_size():
"""Return world size for the model parallel group."""
global mpu
if mpu is None or hasattr(mpu, 'initialize_sequence_parallel'):
return 1
return mpu.get_model_parallel_world_size()
def _get_data_parallel_rank():
"""Return my rank for the data parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
return dist.get_rank(group=_get_data_parallel_group())
def _get_sequence_parallel_world_size():
"""Return world size for the model parallel group."""
assert dist.is_initialized(), 'dist is not initialized'
global mpu
if mesh_device is not None:
return dist.get_world_size(mesh_device.get_group(mesh_dim="sequence_parallel"))
if mpu is not None and hasattr(mpu, 'get_sequence_parallel_world_size'):
return mpu.get_sequence_parallel_world_size()
return 1
def _get_sequence_parallel_rank():
"""Return my rank for the sequence parallel group."""
global mpu
if mpu is not None and hasattr(mpu, 'get_sequence_parallel_rank'):
return mpu.get_sequence_parallel_rank()
if mesh_device is not None:
return dist.get_rank(mesh_device.get_group(mesh_dim="sequence_parallel"))
return 0
def _get_sequence_parallel_group():
global mpu
if mpu is None or not hasattr(mpu, 'get_sequence_parallel_group'):
if mesh_device is None:
raise KeyError("No sequence parallel group found")
return mesh_device.get_group(mesh_dim="sequence_parallel")
return mpu.get_sequence_parallel_group()
def _get_sequence_data_parallel_world_size():
"""Return world size for the model parallel group."""
global mpu
if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_world_size'):
return mpu.get_sequence_data_parallel_world_size()
return _get_data_parallel_world_size()
def _get_sequence_data_parallel_rank():
"""Return my rank for the data parallel group."""
global mpu
if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_rank'):
return mpu.get_sequence_data_parallel_rank()
return _get_data_parallel_rank()
def _get_sequence_data_parallel_group():
global mpu
# When sequence parallelism is enabled, the process group for zero sharding and
# gradient allreduce must be across both dimensions of data and sequence parallelism.
if mpu is not None and hasattr(mpu, 'get_sequence_data_parallel_group'):
return mpu.get_sequence_data_parallel_group()
return _get_data_parallel_group()
def _get_expert_model_parallel_world_size():
global expert_tensor_parallel_world_size
return expert_tensor_parallel_world_size
def _create_zero_param_parallel_group(group_size):
"""
Create parameter partitioning group within ZeRO data parallel groups.
Example - ZP + D parallel
world_size = 16
zero_hpz_partition_size = 2 # number of ranks with replicated params (dual partitioning)
zero_param_intra_parallel_group = [0, 1], [2,3], [4,5], [6,7], [8,9] - segmented (subgroup) with rep partition
data_parallel_group = [0,1,...,15] - all reduce is on ZeRO model
"""
assert dist.is_initialized()
global _ZERO_PARAM_INTRA_PARALLEL_GROUP
# Only create group if it does not already exist
assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is None, \
'ZeRO parameter intra parallel group is already initialized'
world_size = dist.get_world_size()
rank = dist.get_rank()
zero_param_parallel_size_ = min(group_size, world_size)
_ensure_divisibility(world_size, zero_param_parallel_size_)
# Build the ZeRO param intra parallel groups.
for i in range(world_size // zero_param_parallel_size_):
ranks = range(i * zero_param_parallel_size_, (i + 1) * zero_param_parallel_size_)
group = dist.new_group(ranks)
if i == (rank // zero_param_parallel_size_):
_ZERO_PARAM_INTRA_PARALLEL_GROUP = group
def _get_zero_param_intra_parallel_group():
"""Get the ZeRO parameter partitioning intra parallel group the caller rank belongs to."""
#assert _ZERO_PARAM_INTRA_PARALLEL_GROUP is not None, \
# 'ZeRO parameter partitioning group is not initialized'
#TODO: Add warning
return _ZERO_PARAM_INTRA_PARALLEL_GROUP
def _zero_param_parallel_is_initialized():
"""Check if ZeRO data parallel with parameter partititioning groups are initialized."""
###TODO: assert that MPU is not set
if _ZERO_PARAM_INTRA_PARALLEL_GROUP is None and _DATA_PARALLEL_GROUP is None:
return False
def _get_zero_param_intra_parallel_rank_in_mygroup():
"""Return my rank for the ZeRO parameter inter parallel group."""
return dist.get_rank(group=_get_zero_param_intra_parallel_group())
def _get_zero_param_intra_parallel_group_world_size():
"""Return world size for the ZeRO parameter parallel group."""
return dist.get_world_size(group=_get_zero_param_intra_parallel_group())
def _get_zero_param_intra_parallel_group_ranks():
"""Return all ranks for the ZeRO parameter intra parallel group."""
return dist.get_all_ranks_from_group(group=_get_zero_param_intra_parallel_group())