wehub-resource-sync
217 lines
10 KiB
Python
217 lines
10 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from dataclasses import dataclass
|
|
from deepspeed.checkpoint.reshape_utils import partition_data
|
|
from deepspeed.runtime.zero.config import ZeroStageEnum
|
|
from .constants import *
|
|
|
|
|
|
@dataclass
|
|
class DataParallelWriterConfig(object):
|
|
world_size: int
|
|
rank: int
|
|
global_rank: int
|
|
local_rank: int
|
|
pure_dp: bool
|
|
|
|
|
|
class DataParallelWriterFactory(object):
|
|
|
|
def __init__(self, uni_parallel_info, parallel_unit):
|
|
self._uni_parallel_info = uni_parallel_info
|
|
self._parallel_unit = parallel_unit
|
|
if parallel_unit == CheckpointDataParallel.SOCKET:
|
|
self._num_resources = uni_parallel_info.num_sockets
|
|
else:
|
|
self._num_resources = uni_parallel_info.num_machines
|
|
self._ranks_per_resource = max(1, self._uni_parallel_info.global_world_size // self._num_resources)
|
|
|
|
def create_config(self, zero_stage, has_moe_layers):
|
|
if zero_stage == ZeroStageEnum.weights:
|
|
return self._create_config(1, 0)
|
|
|
|
if has_moe_layers:
|
|
writer_config = self._get_expert_data_parallel_config()
|
|
else:
|
|
writer_config = self._get_data_parallel_config()
|
|
|
|
if writer_config is None and zero_stage >= ZeroStageEnum.optimizer_states:
|
|
return self._create_config(1, 0)
|
|
|
|
return writer_config
|
|
|
|
def _create_config(self, world_size, rank):
|
|
return DataParallelWriterConfig(world_size=world_size,
|
|
rank=rank,
|
|
global_rank=self._uni_parallel_info.global_rank,
|
|
local_rank=self._uni_parallel_info.local_rank,
|
|
pure_dp=self._uni_parallel_info.pure_dp)
|
|
|
|
def _get_expert_data_parallel_config(self):
|
|
ep_info = self._uni_parallel_info.ep_info
|
|
if self._parallel_unit is None:
|
|
dp_rank = ep_info.dp_rank
|
|
return self._create_config(1, 0) if dp_rank == 0 else None
|
|
|
|
assert self._uni_parallel_info.pure_dp, \
|
|
'3D parallelism is not yet supported for data parallel checkpointing.'
|
|
|
|
if self._parallel_unit == CheckpointDataParallel.REPLICA or ep_info.ep_world_size == 1:
|
|
return self._get_parallel_write_for_ddp(ep_info.dp_world_size, ep_info.dp_rank)
|
|
|
|
return self._get_expert_parallel_write_for_2d()
|
|
|
|
def _get_expert_parallel_write_for_2d(self):
|
|
ep_info = self._uni_parallel_info.ep_info
|
|
|
|
def _get_expert_slice_resources(expert_resources, resource_name):
|
|
ep_world_size = ep_info.ep_world_size
|
|
slices_per_resource = min(self._ranks_per_resource, ep_world_size)
|
|
assert slices_per_resource <= len(expert_resources)
|
|
|
|
ep_num_resources = len(expert_resources)
|
|
assert ep_num_resources % slices_per_resource == 0, f'{resource_name}: Expected ep_num_resources={ep_num_resources} to multiple of slices_per_resource={slices_per_resource} for ep_world_size={ep_world_size}'
|
|
|
|
slice_partitions = partition_data(expert_resources, slices_per_resource)
|
|
# print(
|
|
# f'edp_resource_partition: self._uni_parallel_info.global_rank={self._uni_parallel_info.global_rank} expert_resources={expert_resources} slices_per_resource={slices_per_resource} ep_world_size={ep_world_size} slice_partitions={slice_partitions}'
|
|
# )
|
|
resource_index = ep_info.ep_rank % slice_resources
|
|
return slice_partitions[resource_index]
|
|
|
|
dp_ranks = ep_info.dp_peer_ranks
|
|
expert_resources = [r // self._ranks_per_resource for r in dp_ranks]
|
|
slice_resources = _get_expert_slice_resources(expert_resources, self._parallel_unit)
|
|
assert all([idx < self._num_resources for idx in expert_resources]), \
|
|
f'Detected invalid resource index in expert_resources={expert_resources}, self._num_resources={self._num_resources}'
|
|
return self._assign_resources_to_tensor_slice(slice_resources, ep_info.ep_rank, dp_ranks)
|
|
|
|
def _get_data_parallel_config(self):
|
|
mpu_info = self._uni_parallel_info.mpu_info
|
|
if self._parallel_unit is None:
|
|
dp_rank = self._uni_parallel_info.dp_rank if mpu_info is None else mpu_info.dp_rank
|
|
return self._create_config(1, 0) if dp_rank == 0 else None
|
|
|
|
if self._uni_parallel_info.pure_dp:
|
|
return self._get_parallel_write_for_ddp(self._uni_parallel_info.global_world_size,
|
|
self._uni_parallel_info.global_rank)
|
|
|
|
if self._parallel_unit == CheckpointDataParallel.REPLICA:
|
|
return self._create_config(mpu_info.dp_world_size, mpu_info.dp_rank)
|
|
|
|
return self._get_parallel_write_for_3d()
|
|
|
|
def _get_parallel_write_for_3d(self):
|
|
mpu_info = self._uni_parallel_info.mpu_info
|
|
my_global_rank = self._uni_parallel_info.global_rank
|
|
|
|
def _expand_resources(resource_list, new_size):
|
|
old_size = len(resource_list)
|
|
if old_size >= new_size:
|
|
return resource_list
|
|
|
|
assert new_size % old_size == 0, f'Expect new_size={new_size} to be multiple of old_size={old_size}'
|
|
multiplier = new_size // old_size
|
|
new_resource_list = []
|
|
for r in resource_list:
|
|
new_resource_list += [r] * multiplier
|
|
# print(f'expand_resources: {my_global_rank=} {old_size=} {new_size=} {resource_list=} {new_resource_list=}')
|
|
return new_resource_list
|
|
|
|
# Getting resource partition for a tensor slice is a 2-step process
|
|
# 1. Get resource partitions for all pipeline stages. A pipeline stage is a 2D grid of size TP x DP
|
|
def _get_pipeline_stage_resources(resource_indices):
|
|
num_resources = len(resource_indices)
|
|
pp_world_size = mpu_info.pp_world_size
|
|
if num_resources < pp_world_size:
|
|
resource_indices = _expand_resources(resource_indices, pp_world_size)
|
|
num_resources = pp_world_size
|
|
global_resource_partitions = partition_data(resource_indices, pp_world_size)
|
|
pp_rank = mpu_info.pp_rank
|
|
return global_resource_partitions[pp_rank]
|
|
|
|
# 2. Get resource partition for tensor slice. A tensor slice is a 1D vector of size DP
|
|
def _get_tensor_slice_resources(resource_indices, resource_name):
|
|
pipe_stage_resources = _get_pipeline_stage_resources(resource_indices)
|
|
tp_world_size = mpu_info.tp_world_size
|
|
if len(pipe_stage_resources) < tp_world_size:
|
|
pipe_stage_resources = _expand_resources(pipe_stage_resources, tp_world_size)
|
|
tp_num_resources = len(pipe_stage_resources)
|
|
assert tp_num_resources % tp_world_size == 0, \
|
|
f'{resource_name}: Expected tp_num_resources={tp_num_resources} to multiple of tp_world_size={tp_world_size}'
|
|
|
|
pipe_stage_resource_partitions = partition_data(pipe_stage_resources, tp_world_size)
|
|
tp_rank = mpu_info.tp_rank
|
|
return pipe_stage_resource_partitions[tp_rank]
|
|
|
|
def _get_model_parallel_slice_resources():
|
|
# Get resources of my dp peer ranks
|
|
resources = [(r // self._ranks_per_resource) for r in mpu_info.dp_peer_ranks]
|
|
if len(resources) < self._ranks_per_resource:
|
|
resources = _expand_resources(resources, self._ranks_per_resource)
|
|
|
|
resource_partitions = partition_data(resources, self._ranks_per_resource)
|
|
mp_rank = (mpu_info.pp_rank * mpu_info.tp_world_size) + mpu_info.tp_rank
|
|
slice_rank = mp_rank % self._ranks_per_resource
|
|
return resource_partitions[slice_rank]
|
|
|
|
num_slices = mpu_info.tp_world_size * mpu_info.pp_world_size
|
|
if num_slices > self._ranks_per_resource:
|
|
slice_resources = _get_model_parallel_slice_resources()
|
|
else:
|
|
all_resources = list(range(self._num_resources))
|
|
slice_resources = _get_tensor_slice_resources(all_resources, self._parallel_unit)
|
|
|
|
return self._assign_resources_to_tensor_slice(slice_resources, mpu_info.tp_rank, mpu_info.dp_peer_ranks)
|
|
|
|
def _get_slice_writers(self, slice_resources, my_dp_ranks):
|
|
resource_map = {}
|
|
for res in slice_resources:
|
|
resource_map[res] = [r for r in my_dp_ranks if (r // self._ranks_per_resource) == res]
|
|
|
|
# Only one writer per resource, and we conventionally pick the first rank as writer.
|
|
return [ranks[0] for ranks in resource_map.values()]
|
|
|
|
def _assign_resources_to_tensor_slice(self, slice_resources, my_slice_index, my_dp_ranks):
|
|
my_global_rank = self._uni_parallel_info.global_rank
|
|
slice_writer_ranks = self._get_slice_writers(slice_resources, my_dp_ranks)
|
|
my_resource_index = my_global_rank // self._ranks_per_resource
|
|
print(
|
|
f'resource_assign: my_global_rank={my_global_rank} my_slice_index={my_slice_index} my_dp_ranks={my_dp_ranks} slice_resources={slice_resources} slice_writer_ranks={slice_writer_ranks}'
|
|
)
|
|
if my_resource_index in slice_resources and my_global_rank in slice_writer_ranks:
|
|
my_writer_index = (my_global_rank - slice_writer_ranks[0]) // self._ranks_per_resource
|
|
num_slice_writers = len(slice_writer_ranks)
|
|
print(
|
|
f'slice_writer: my_global_rank={my_global_rank} my_writer_index={my_writer_index} num_slice_writers={num_slice_writers}'
|
|
)
|
|
return self._create_config(num_slice_writers, my_writer_index)
|
|
|
|
return None
|
|
|
|
def _get_parallel_write_for_ddp(self, dp_world_size, dp_rank):
|
|
if self._parallel_unit == CheckpointDataParallel.REPLICA:
|
|
return self._create_config(dp_world_size, dp_rank)
|
|
|
|
num_machines = self._uni_parallel_info.num_machines
|
|
if self._parallel_unit == CheckpointDataParallel.SOCKET:
|
|
if dp_world_size == num_machines:
|
|
# There is one rank per machine
|
|
return self._create_config(num_machines, dp_rank)
|
|
|
|
num_sockets = self._uni_parallel_info.num_sockets
|
|
ranks_per_socket = dp_world_size // num_sockets
|
|
if dp_rank % ranks_per_socket == 0:
|
|
return self._create_config(num_sockets, dp_rank // ranks_per_socket)
|
|
else:
|
|
return None
|
|
|
|
ranks_per_machine = dp_world_size // num_machines
|
|
if dp_rank % ranks_per_machine == 0:
|
|
return self._create_config(num_machines, self._uni_parallel_info.machine_rank)
|
|
|
|
return None
|