mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
220 lines
7.7 KiB
Python
220 lines
7.7 KiB
Python
'''Copyright The Microsoft DeepSpeed Team'''
|
|
|
|
from .reshape_utils import partition_data
|
|
|
|
|
|
class meg_2d_parallel_map(object):
|
|
|
|
def __init__(self, pp_degree, tp_degree):
|
|
self.pp_degree = pp_degree
|
|
self.tp_degree = tp_degree
|
|
self.map = {}
|
|
|
|
def simple_init(self):
|
|
self.map = {
|
|
self._make_key(i // self.tp_degree, i % self.tp_degree): [i]
|
|
for i in range(self.pp_degree * self.tp_degree)
|
|
}
|
|
|
|
def add_data(self, pp_index, tp_index, data):
|
|
self._validate_indices(pp_index, tp_index)
|
|
assert type(data) is list
|
|
|
|
key = self._make_key(pp_index, tp_index)
|
|
if not key in self.map.keys():
|
|
self.map[key] = []
|
|
self.map[key] += data
|
|
|
|
def get_data(self, pp_index=None, tp_index=None):
|
|
self._validate_indices(pp_index, tp_index)
|
|
pp_indices = list(range(self.pp_degree)) if pp_index is None else [pp_index]
|
|
tp_indices = list(range(self.tp_degree)) if tp_index is None else [tp_index]
|
|
|
|
result = []
|
|
for i in pp_indices:
|
|
for j in tp_indices:
|
|
result += self.map[self._make_key(i, j)]
|
|
|
|
return result
|
|
|
|
def print_data(self, tag):
|
|
print(f'{tag}')
|
|
for key, value in self.map.items():
|
|
print(f'{key} = {value}')
|
|
|
|
def _validate_indices(self, pp_index, tp_index):
|
|
assert pp_index is None or pp_index < self.pp_degree
|
|
assert tp_index is None or tp_index < self.tp_degree
|
|
|
|
def _make_key(self, i, j):
|
|
return f'{i},{j}'
|
|
|
|
|
|
def _reshape_tp_dimension(old_2d_map, new_tp_degree):
|
|
old_pp_degree = old_2d_map.pp_degree
|
|
new_2d_map = meg_2d_parallel_map(old_pp_degree, new_tp_degree)
|
|
for i in range(old_pp_degree):
|
|
ranks_for_pp_index = old_2d_map.get_data(pp_index=i, tp_index=None)
|
|
split_ranks = partition_data(ranks_for_pp_index, new_tp_degree)
|
|
for j in range(new_tp_degree):
|
|
new_2d_map.add_data(i, j, split_ranks[j])
|
|
|
|
return new_2d_map
|
|
|
|
|
|
def _reshape_pp_dimension(old_2d_map, new_pp_degree):
|
|
old_tp_degree = old_2d_map.tp_degree
|
|
new_2d_map = meg_2d_parallel_map(new_pp_degree, old_tp_degree)
|
|
for i in range(old_tp_degree):
|
|
ranks_for_tp_index = old_2d_map.get_data(pp_index=None, tp_index=i)
|
|
split_ranks = partition_data(ranks_for_tp_index, new_pp_degree)
|
|
for j in range(new_pp_degree):
|
|
new_2d_map.add_data(j, i, split_ranks[j])
|
|
|
|
return new_2d_map
|
|
|
|
|
|
def reshape_meg_2d_parallel(old_pp_degree, old_tp_degree, new_pp_degree, new_tp_degree, verbose=False):
|
|
assert new_pp_degree <= old_pp_degree
|
|
assert new_tp_degree <= old_tp_degree
|
|
|
|
old_2d_map = meg_2d_parallel_map(old_pp_degree, old_tp_degree)
|
|
old_2d_map.simple_init()
|
|
if verbose:
|
|
old_2d_map.print_data(f'original_2d_map:')
|
|
|
|
if old_tp_degree != new_tp_degree:
|
|
new_tp_map = _reshape_tp_dimension(old_2d_map, new_tp_degree)
|
|
else:
|
|
new_tp_map = old_2d_map
|
|
if verbose:
|
|
new_tp_map.print_data(f'after_tp_reshape:')
|
|
|
|
if old_pp_degree != new_pp_degree:
|
|
final_map = _reshape_pp_dimension(new_tp_map, new_pp_degree)
|
|
else:
|
|
final_map = new_tp_map
|
|
|
|
if verbose:
|
|
final_map.print_data(f'final_2d_map:')
|
|
|
|
return final_map
|
|
|
|
|
|
def get_mpu_ranks(tp_size=1, pp_size=1, dp_size=1, virtual_pp_size=None):
|
|
"""
|
|
Initialize model data parallel groups.
|
|
|
|
Arguments:
|
|
tp_size: number of GPUs used to parallelize model tensor.
|
|
pp_size: number of GPUs used to parallelize model pipeline.
|
|
dp_size: number of GPUs used to parallelize model data.
|
|
|
|
Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
|
|
use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
|
|
the model pipeline. The present function will
|
|
create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
|
|
and 8 data-parallel groups as:
|
|
8 data_parallel groups:
|
|
[g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
|
|
8 tensor model-parallel groups:
|
|
[g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
|
|
4 pipeline model-parallel groups:
|
|
[g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
|
|
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.
|
|
"""
|
|
|
|
world_size = tp_size * pp_size * dp_size
|
|
|
|
print(f"\n\n*** tp={tp_size}, pp={pp_size}, dp={dp_size}, world={world_size}")
|
|
|
|
tensor_model_parallel_size = min(tp_size, world_size)
|
|
pipeline_model_parallel_size = min(pp_size, world_size)
|
|
data_parallel_size = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size)
|
|
|
|
num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size
|
|
num_pipeline_model_parallel_groups = world_size // pipeline_model_parallel_size
|
|
num_data_parallel_groups = world_size // data_parallel_size
|
|
|
|
# Build the data-parallel groups.
|
|
all_dp_group_ranks = []
|
|
for i in range(pipeline_model_parallel_size):
|
|
start_rank = i * num_pipeline_model_parallel_groups
|
|
end_rank = (i + 1) * num_pipeline_model_parallel_groups
|
|
for j in range(tensor_model_parallel_size):
|
|
ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
|
|
all_dp_group_ranks.append(list(ranks))
|
|
|
|
print("DP", all_dp_group_ranks)
|
|
|
|
# Build the model-parallel groups.
|
|
all_pp_group_ranks = []
|
|
for i in range(data_parallel_size):
|
|
ranks = [data_parallel_group_ranks[i] for data_parallel_group_ranks in all_dp_group_ranks]
|
|
all_pp_group_ranks.append(list(ranks))
|
|
|
|
print(f"PP", all_pp_group_ranks)
|
|
|
|
# Build the tensor model-parallel groups.
|
|
all_tp_group_ranks = []
|
|
for i in range(num_tensor_model_parallel_groups):
|
|
ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)
|
|
all_tp_group_ranks.append(list(ranks))
|
|
|
|
print(f"TP", all_tp_group_ranks)
|
|
|
|
return all_tp_group_ranks, all_pp_group_ranks, all_dp_group_ranks
|
|
|
|
# # Build the pipeline model-parallel groups and embedding groups
|
|
# # (first and last rank in each pipeline model-parallel group).
|
|
# for i in range(num_pipeline_model_parallel_groups):
|
|
# ranks = range(i, world_size,
|
|
# num_pipeline_model_parallel_groups)
|
|
# print(f"EMB{i}", list(ranks))
|
|
|
|
|
|
def reshape(src, tgt):
|
|
"""
|
|
reshape([tp_size_src, pp_size_src, dp_size_src],
|
|
[tp_size_tgt, pp_size_tgt, dp_size_tgt])
|
|
"""
|
|
|
|
print(f"\n\n*** Reshaping: {src} => {tgt}")
|
|
|
|
tp_size_src, pp_size_src, dp_size_src = src
|
|
tp_size_tgt, pp_size_tgt, dp_size_tgt = tgt
|
|
|
|
tp_ranks1, pp_ranks1, dp_ranks1 = get_mpu_ranks(tp_size=tp_size_src, pp_size=pp_size_src, dp_size=dp_size_src)
|
|
tp_ranks2, pp_ranks2, dp_ranks2 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_src, dp_size=dp_size_src)
|
|
tp_ranks3, pp_ranks3, dp_ranks3 = get_mpu_ranks(tp_size=tp_size_tgt, pp_size=pp_size_tgt, dp_size=dp_size_src)
|
|
|
|
# handle tp contraction first
|
|
print("\n*** TP contraction:")
|
|
|
|
for i, r in enumerate(tp_ranks1):
|
|
print(f'{tp_ranks1[i]} => {tp_ranks2[i]}')
|
|
|
|
# handle pp contraction next
|
|
|
|
print("\n*** PP contraction:")
|
|
|
|
for i, r in enumerate(pp_ranks1):
|
|
print(f'{pp_ranks2[i]} => {pp_ranks3[i]}')
|
|
|
|
|
|
# easy
|
|
#reshape([2,2,1],[1,1,1])
|
|
|
|
# probably need more logic to suggest how to pack
|
|
#reshape([4,4,1],[2,2,1])
|
|
|
|
#reshape([2,4,2], [8,32,1])
|
|
|
|
# get_mpu_ranks(2,2,2)
|
|
# get_mpu_ranks(4,2,1)
|
|
# get_mpu_ranks(2,4,1)
|
|
# get_mpu_ranks(1,1,8)
|