wehub-resource-sync
85 lines
3.2 KiB
Python
85 lines
3.2 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from deepspeed import comm as dist
|
|
from deepspeed.constants import CROSS_RANK, CROSS_SIZE, LOCAL_RANK
|
|
from .data_parallel_writer_factory import DataParallelWriterFactory
|
|
|
|
# TODO: parse socket number from env.
|
|
SOCKETS_PER_MACHINE = 2
|
|
|
|
|
|
@dataclass
|
|
class MPUInfo(object):
|
|
pp_world_size: int
|
|
pp_rank: int
|
|
tp_world_size: int
|
|
tp_rank: int
|
|
dp_world_size: int
|
|
dp_peer_ranks: list
|
|
dp_rank: int
|
|
|
|
|
|
def _create_model_parallel_info(mpu):
|
|
return MPUInfo(pp_world_size=mpu.get_pipeline_model_parallel_world_size(),
|
|
pp_rank=mpu.get_pipeline_model_parallel_rank(),
|
|
tp_world_size=mpu.get_tensor_model_parallel_world_size(),
|
|
tp_rank=mpu.get_tensor_model_parallel_rank(),
|
|
dp_world_size=mpu.get_data_parallel_world_size(),
|
|
dp_peer_ranks=mpu.get_data_parallel_group_ranks(),
|
|
dp_rank=mpu.get_data_parallel_rank())
|
|
|
|
|
|
@dataclass
|
|
class ExpertParallelInfo(object):
|
|
ep_world_size: int
|
|
ep_rank: int
|
|
dp_world_size: int
|
|
dp_peer_ranks: list
|
|
dp_rank: int
|
|
|
|
|
|
def _create_expert_parallel_info(groups):
|
|
group_name = groups._get_max_expert_size_name()
|
|
return ExpertParallelInfo(ep_world_size=groups._get_expert_parallel_world_size(group_name),
|
|
ep_rank=groups._get_expert_parallel_rank(group_name),
|
|
dp_world_size=groups._get_expert_data_parallel_world_size(group_name),
|
|
dp_peer_ranks=groups._get_expert_data_parallel_group_ranks(group_name),
|
|
dp_rank=groups._get_expert_data_parallel_rank(group_name))
|
|
|
|
|
|
@dataclass
|
|
class UniversalParallelInfo(object):
|
|
global_world_size: int
|
|
global_rank: int
|
|
local_rank: int
|
|
mpu_info: MPUInfo
|
|
ep_info: ExpertParallelInfo
|
|
pure_dp: bool
|
|
num_machines: int
|
|
machine_rank: int
|
|
num_sockets: int
|
|
|
|
|
|
def create_universal_parallel_info(groups, has_moe_layers):
|
|
return UniversalParallelInfo(global_world_size=dist.get_world_size(),
|
|
global_rank=dist.get_rank(),
|
|
local_rank=int(os.environ[LOCAL_RANK]),
|
|
mpu_info=None if groups.mpu is None else _create_model_parallel_info(groups.mpu),
|
|
ep_info=_create_expert_parallel_info(groups) if has_moe_layers else None,
|
|
pure_dp=groups.mpu is None
|
|
or groups.mpu.get_data_parallel_world_size() == dist.get_world_size(),
|
|
num_machines=int(os.environ[CROSS_SIZE]),
|
|
machine_rank=int(os.environ[CROSS_RANK]),
|
|
num_sockets=int(os.environ[CROSS_SIZE]) * SOCKETS_PER_MACHINE)
|
|
|
|
|
|
def create_data_parallel_writer_config(groups, parallel_unit, zero_stage, has_moe_layers):
|
|
uni_parallel_info = create_universal_parallel_info(groups, has_moe_layers)
|
|
writer_factory = DataParallelWriterFactory(uni_parallel_info, parallel_unit)
|
|
return writer_factory.create_config(zero_stage, has_moe_layers)
|