wehub-resource-sync
204 lines
7.4 KiB
Python
204 lines
7.4 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
from torch import Tensor
|
|
|
|
from deepspeed import comm as dist
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.utils import logger
|
|
from deepspeed.utils.torch import jit_script_compat
|
|
|
|
|
|
def _log_rank0(msg):
|
|
if dist.get_rank() == 0:
|
|
logger.info(msg)
|
|
|
|
|
|
@jit_script_compat
|
|
def scale_tensors(tensors: List[Tensor], scale: int):
|
|
for t in tensors:
|
|
t.div_(scale)
|
|
|
|
|
|
@dataclass
|
|
class MiCS_CommGroups:
|
|
""""""
|
|
param_shard_group = None
|
|
param_shard_size = -1
|
|
param_shard_rank = -1
|
|
|
|
param_repli_group = None
|
|
param_repli_size = -1
|
|
param_repli_rank = -1
|
|
|
|
param_intra_node_group = None
|
|
param_inter_node_shard_group = None
|
|
|
|
|
|
def create_mics_comm_groups(
|
|
shard_size,
|
|
dp_group,
|
|
hierarchical_allgather=False,
|
|
mpu=None,
|
|
):
|
|
"""
|
|
create shard-group, replicate-group from config_file
|
|
TODO: consider broadcast the config from rank0
|
|
|
|
Returns:
|
|
MiCS_CommGroups
|
|
"""
|
|
# env var for debugging purpose
|
|
ndevices_per_node = int(os.environ.get("NDEV_PER_NODE", get_accelerator().device_count()))
|
|
_log_rank0(f'creating MiCS communication groups with per node device size {ndevices_per_node}')
|
|
groups = MiCS_CommGroups()
|
|
|
|
if mpu is not None:
|
|
assert dp_group == mpu.get_data_parallel_group()
|
|
|
|
# full size of the world
|
|
world_size = dist.get_world_size()
|
|
# global rank
|
|
global_rank = dist.get_rank()
|
|
|
|
config = _generate_mics_config(world_size, ndevices_per_node, shard_size, 1)
|
|
ranks_of_shard_group = config['shard_groups']
|
|
ranks_of_repli_group = config['replicate_groups']
|
|
if len(ranks_of_repli_group) == 0:
|
|
assert len(ranks_of_shard_group) == 1, "replicate groups are empty only for single shard group"
|
|
for r in ranks_of_shard_group[0]:
|
|
ranks_of_repli_group.append([r])
|
|
|
|
# for simplicity
|
|
assert _sizes_all_same(ranks_of_repli_group), "replicate groups must have the same size"
|
|
assert _sizes_all_same(ranks_of_shard_group), "shard groups must have the same size"
|
|
|
|
assert sum([len(g) for g in ranks_of_shard_group]) == dist.get_world_size(), "all sharded ranks "
|
|
if len(ranks_of_shard_group) > 1: # if only shard on one group then no need for replicate groups
|
|
assert len(ranks_of_shard_group) == len(
|
|
ranks_of_repli_group[0]), "number of shard groups must equal to the size of each replicate group"
|
|
|
|
global_rank = dist.get_rank()
|
|
# create shard groups
|
|
for shard_ranks in ranks_of_shard_group:
|
|
_group = dist.new_group(shard_ranks)
|
|
if global_rank in shard_ranks:
|
|
groups.param_shard_group = _group
|
|
groups.param_shard_size = len(shard_ranks)
|
|
groups.param_shard_rank = dist.get_rank(_group)
|
|
logger.info(f'rank {global_rank}, shard group'
|
|
f' {groups.param_shard_rank}/{dist.get_world_size(group=_group)}')
|
|
|
|
# create replicate groups
|
|
for repli_ranks in ranks_of_repli_group:
|
|
if len(repli_ranks) > 1:
|
|
_group = dist.new_group(repli_ranks)
|
|
if global_rank in repli_ranks:
|
|
groups.param_repli_group = _group
|
|
groups.param_repli_size = len(repli_ranks)
|
|
groups.param_repli_rank = dist.get_rank(group=_group)
|
|
logger.info(f'rank {global_rank} '
|
|
f'replicate group {groups.param_repli_rank}/{dist.get_world_size(group=_group)}')
|
|
else:
|
|
groups.param_repli_group = None
|
|
groups.param_repli_size = 1
|
|
groups.param_repli_rank = 0
|
|
logger.info(f'rank {global_rank} replicate group 0/1')
|
|
|
|
# assign shard group size as world size
|
|
assert groups.param_shard_size == len(ranks_of_shard_group[0])
|
|
|
|
if hierarchical_allgather:
|
|
# create hierarchy inter-node, intra-node groups
|
|
# n_span_nodes = config['shard_span']
|
|
n_span_nodes = config['span_nodes']
|
|
assert n_span_nodes > 1, "sharding spans on single node, no need for hierarchy allgather"
|
|
assert len(ranks_of_shard_group[0]) % n_span_nodes == 0
|
|
|
|
n_gpu_per_node = len(ranks_of_shard_group[0]) // n_span_nodes
|
|
intra_node_ranks_group = []
|
|
inter_node_ranks_group = []
|
|
for shard_group in ranks_of_shard_group:
|
|
_intra_node_ranks = []
|
|
for i in range(0, len(shard_group), n_gpu_per_node):
|
|
_intra_node_ranks.append(shard_group[i:i + n_gpu_per_node])
|
|
_inter_node_ranks = []
|
|
for i in range(n_gpu_per_node):
|
|
_ranks = [_g[i] for _g in _intra_node_ranks]
|
|
_inter_node_ranks.append(_ranks)
|
|
|
|
intra_node_ranks_group.append(_intra_node_ranks)
|
|
inter_node_ranks_group.append(_inter_node_ranks)
|
|
|
|
_log_rank0(f"create for hierarchy all-gather groups: intra nodes {intra_node_ranks_group}")
|
|
_log_rank0(f"create for hierarchy all-gather groups: inter nodes {inter_node_ranks_group}")
|
|
|
|
# create communicators
|
|
for shard_group in intra_node_ranks_group:
|
|
for intra_node_ranks in shard_group:
|
|
_group = dist.new_group(intra_node_ranks)
|
|
if global_rank in intra_node_ranks:
|
|
groups.param_intra_node_group = _group
|
|
_log_rank0(f'create group for intra node ranks {intra_node_ranks}')
|
|
|
|
for shard_group in inter_node_ranks_group:
|
|
for inter_node_ranks in shard_group:
|
|
_group = dist.new_group(inter_node_ranks)
|
|
if global_rank in inter_node_ranks:
|
|
groups.param_inter_node_shard_group = _group
|
|
_log_rank0(f'create group for inter node ranks {inter_node_ranks}')
|
|
return groups
|
|
|
|
|
|
def _generate_mics_config(world_size, ndev_per_node, shard_size, pp_size=1):
|
|
"""Generating the configuration for sharding This shard config generation assume
|
|
that the pipeline stages are partitioned in order, i.e., first ranks
|
|
hold the stage0, etc.
|
|
|
|
Args:
|
|
|
|
shard_size (int): zero3 data-parallel shard size, FIXME:
|
|
change the name later
|
|
|
|
pp_size (int): pipeline parallel size, currently, only work with
|
|
pipeline parallelism + zero
|
|
|
|
"""
|
|
assert world_size % pp_size == 0
|
|
assert (world_size // pp_size) % shard_size == 0, \
|
|
f"dp group size is not dividable by dp_shard_size, "\
|
|
f" (world_size {world_size}, pp_size {pp_size}, dp_shard_size {shard_size})"
|
|
|
|
config = {}
|
|
shard_groups = np.arange(world_size).reshape(-1, shard_size)
|
|
replicate_groups = []
|
|
for i in range(shard_size):
|
|
same_shard_ranks = shard_groups[:, i].tolist()
|
|
n_ranks = len(same_shard_ranks)
|
|
replicate_size = n_ranks // pp_size
|
|
replicate_groups.extend([same_shard_ranks[j:j + replicate_size] for j in range(0, n_ranks, replicate_size)])
|
|
|
|
config['replicate_groups'] = replicate_groups
|
|
config['shard_groups'] = shard_groups.tolist()
|
|
config["span_nodes"] = len(shard_groups[0]) // ndev_per_node
|
|
return config
|
|
|
|
|
|
def _sizes_all_same(groups):
|
|
"""all groups have same length"""
|
|
all_same = True
|
|
for g in groups:
|
|
if len(g) != len(groups[0]):
|
|
return False
|
|
return all_same
|