mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-08 12:30:42 +00:00
[tensor] shape consistency generate transform path and communication cost (#1435)
* [tensor] shape consistency output transform path and communication cost * polish code
This commit is contained in:
YuliangLiu0306
GitHub
@@ -1,4 +1,15 @@
|
||||
import torch
|
||||
from colossalai.device.device_mesh import DeviceMesh
|
||||
from colossalai.tensor.utils import all_gather_simulator, all_to_all_simulator, shard_simulator
|
||||
from copy import deepcopy
|
||||
from enum import Enum
|
||||
from functools import reduce
|
||||
import operator
|
||||
|
||||
ALLGATHER_COST = 20
|
||||
SHARD_COST = 5
|
||||
STEP_PENALTY = 6
|
||||
NAN = 'nan'
|
||||
|
||||
|
||||
class _DimSpec:
|
||||
@@ -15,6 +26,7 @@ class _DimSpec:
|
||||
def __init__(self, shard_list):
|
||||
self.is_replica = len(shard_list) == 0
|
||||
self.shard_list = shard_list
|
||||
self.build_difference_2d_dict()
|
||||
|
||||
def __eq__(self, other):
|
||||
return str(self) == str(other)
|
||||
@@ -27,11 +39,101 @@ class _DimSpec:
|
||||
target += str(dim)
|
||||
return target
|
||||
|
||||
def _convert_str_to_shard_list(self, str_spec):
|
||||
'''
|
||||
Conver str_spec into shard_list.
|
||||
|
||||
Argument:
|
||||
str_spec(str): dim spec in str type.
|
||||
'''
|
||||
|
||||
if str_spec == 'R':
|
||||
return []
|
||||
if str_spec == 'S0':
|
||||
return [0]
|
||||
if str_spec == 'S1':
|
||||
return [1]
|
||||
if str_spec == 'S01':
|
||||
return [0, 1]
|
||||
|
||||
def build_difference_2d_dict(self):
|
||||
'''
|
||||
Build a difference maping for 2D device mesh case. It will be used to
|
||||
compute the difference between DimSpec pairs.
|
||||
'''
|
||||
|
||||
source_spec_list = ['R', 'S0', 'S1', 'S01']
|
||||
target_spec_list = ['R', 'S0', 'S1', 'S01']
|
||||
difference_dict = {}
|
||||
for source_spec in source_spec_list:
|
||||
for target_spec in target_spec_list:
|
||||
legal_sharding_dims = []
|
||||
spec_pair = (deepcopy(source_spec), deepcopy(target_spec))
|
||||
source_shard_list = self._convert_str_to_shard_list(source_spec)
|
||||
target_shard_list = self._convert_str_to_shard_list(target_spec)
|
||||
|
||||
# source same as target
|
||||
if source_shard_list == target_shard_list:
|
||||
difference = 0
|
||||
|
||||
# all_gather(source) -> target
|
||||
elif len(source_shard_list
|
||||
) == len(target_shard_list) + 1 and source_shard_list[:-1] == target_shard_list:
|
||||
difference = ALLGATHER_COST
|
||||
|
||||
# shard(source) -> target
|
||||
elif len(source_shard_list) == len(
|
||||
target_shard_list) - 1 and source_shard_list == target_shard_list[:-1] and target_shard_list[
|
||||
-1] not in source_shard_list:
|
||||
difference = SHARD_COST
|
||||
|
||||
# S1 -> S0 or S0 -> S1
|
||||
elif len(source_shard_list) == len(target_shard_list):
|
||||
# source -> R -> target
|
||||
difference = ALLGATHER_COST + STEP_PENALTY + SHARD_COST
|
||||
|
||||
# R -> S01
|
||||
elif len(source_shard_list) == len(target_shard_list) - 2:
|
||||
difference = SHARD_COST + STEP_PENALTY + SHARD_COST
|
||||
|
||||
# S01 -> R
|
||||
elif len(source_shard_list) == len(target_shard_list) + 2:
|
||||
difference = ALLGATHER_COST + STEP_PENALTY + ALLGATHER_COST
|
||||
|
||||
# S1 -> S01
|
||||
elif len(source_shard_list) == len(target_shard_list) - 1:
|
||||
difference = ALLGATHER_COST + STEP_PENALTY + SHARD_COST + STEP_PENALTY + SHARD_COST
|
||||
|
||||
# S01 -> S1
|
||||
elif len(source_shard_list) == len(target_shard_list) + 1:
|
||||
difference = ALLGATHER_COST + STEP_PENALTY + ALLGATHER_COST + STEP_PENALTY + SHARD_COST
|
||||
|
||||
else:
|
||||
difference = NAN
|
||||
difference_dict[spec_pair] = difference
|
||||
|
||||
self.difference_dict = difference_dict
|
||||
|
||||
def difference(self, other):
|
||||
'''
|
||||
This function is temporarily NOT implemented, it will be codesigned with ShapeConsistency feature.
|
||||
The difference between two _DimSpec.
|
||||
|
||||
Argument:
|
||||
other(_DimSpec): the dim spec to compare with.
|
||||
|
||||
Return:
|
||||
difference(int): the difference between two _DimSpec.
|
||||
|
||||
Example:
|
||||
dim_spec = _DimSpec([0])
|
||||
other_dim_spec = _DimSpec([0, 1])
|
||||
print(dim_spec.difference(other_dim_spec))
|
||||
|
||||
Output:
|
||||
5
|
||||
'''
|
||||
pass
|
||||
difference = self.difference_dict[(str(self), str(other))]
|
||||
return difference
|
||||
|
||||
|
||||
class ShardingSpec:
|
||||
@@ -43,8 +145,9 @@ class ShardingSpec:
|
||||
Argument:
|
||||
device_mesh(DeviceMesh): A logical view of a physical mesh.
|
||||
entire_shape(torch.Size): The entire shape of tensor before sharded.
|
||||
dim_partition_dict(Dict[int, List[int]]): The key is the dimension of tensor to be sharded,
|
||||
dim_partition_dict(Dict[int, List[int]]， optional): The key is the dimension of tensor to be sharded,
|
||||
and the value of the key decribe which logical axis will be sharded in that dimension.
|
||||
sharding_sequence(List[_DimSpec], optional): A straight view of ShardingSpec looks like [R, R, S0, S1].
|
||||
'''
|
||||
|
||||
def __init__(self, device_mesh, entire_shape, dim_partition_dict=None, sharding_sequence=None):
|
||||
@@ -79,12 +182,18 @@ class ShardingSpec:
|
||||
f"find an invalid sharding axis {element} in dim_partition_dict in tensor dimension {dim}.")
|
||||
|
||||
def convert_dict_to_shard_sequence(self):
|
||||
'''
|
||||
Convert dim_partition_dict into list of _DimSpec, and assign it to sharding_sequence.
|
||||
'''
|
||||
sharding_sequence = [_DimSpec([])] * len(self.entire_shape)
|
||||
for dim, shard_list in self.dim_partition_dict.items():
|
||||
sharding_sequence[dim] = _DimSpec(shard_list)
|
||||
self.sharding_sequence = sharding_sequence
|
||||
|
||||
def convert_shard_sequence_to_dict(self):
|
||||
'''
|
||||
Convert sharding_sequence into dim_partition_dict.
|
||||
'''
|
||||
new_dim_partition_dict = {}
|
||||
for index, dim_spec in enumerate(self.sharding_sequence):
|
||||
if not dim_spec.is_replica:
|
||||
@@ -95,6 +204,45 @@ class ShardingSpec:
|
||||
|
||||
def sharding_sequence_difference(self, other):
|
||||
'''
|
||||
This function is temporarily NOT implemented, it will be codesigned with ShapeConsistency feature.
|
||||
This function is a naive version of difference computation. It just simply accumulates difference every dimension between the
|
||||
pair of sharding sequence.
|
||||
|
||||
Example:
|
||||
dim_partition_dict = {0: [0, 1]}
|
||||
# DistSpec:
|
||||
# shard_sequence: S01,R,R
|
||||
# device_mesh_shape: (4, 4)
|
||||
sharding_spec = ShardingSpec(device_mesh, entire_shape, dim_partition_dict)
|
||||
dim_partition_dict_to_compare = {0: [0], 1: [1]}
|
||||
# DistSpec:
|
||||
# shard_sequence: S0,S1,R
|
||||
# device_mesh_shape: (4, 4)
|
||||
sharding_spec_to_compare = ShardingSpec(device_mesh, entire_shape, dim_partition_dict_to_compare)
|
||||
print(sharding_spec.sharding_sequence_difference(sharding_spec_to_compare))
|
||||
|
||||
Output:
|
||||
25
|
||||
|
||||
Argument:
|
||||
other(ShardingSpec): The ShardingSpec to compared with.
|
||||
|
||||
Return:
|
||||
difference(int): Difference between two ShardingSpec.
|
||||
'''
|
||||
pass
|
||||
assert len(self.sharding_sequence) == len(
|
||||
other.sharding_sequence), f'Cannot compare difference for two sharding specs with different length.'
|
||||
difference = 0
|
||||
for orig_dim_spec, other_dim_spec in zip(self.sharding_sequence, other.sharding_sequence):
|
||||
difference += orig_dim_spec.difference(other_dim_spec)
|
||||
return difference
|
||||
|
||||
def get_sharded_shape_per_device(self):
|
||||
|
||||
sharded_shape = list(self.entire_shape)
|
||||
for dim, shard_list in self.dim_partition_dict.items():
|
||||
mesh_list = [self.device_mesh.mesh_shape[mesh_dim] for mesh_dim in shard_list]
|
||||
shard_partitions = reduce(operator.mul, mesh_list, 1)
|
||||
assert sharded_shape[
|
||||
dim] % shard_partitions == 0, f'Cannot shard dimension {dim} into {shard_partitions} partitions.'
|
||||
sharded_shape[dim] //= shard_partitions
|
||||
return torch.Size(sharded_shape)
|
||||
|
||||
Reference in New Issue
Block a user