diff --git a/colossalai/auto_parallel/solver/_utils.py b/colossalai/auto_parallel/solver/_utils.py index 2c545b74c..dd2a09053 100644 --- a/colossalai/auto_parallel/solver/_utils.py +++ b/colossalai/auto_parallel/solver/_utils.py @@ -8,6 +8,7 @@ import warnings from functools import reduce import functools import operator +from .constants import INFINITY_COST def generate_sharding_spec(input_: Union[Node, torch.Tensor], device_mesh: DeviceMesh, @@ -68,19 +69,16 @@ def generate_resharding_costs(nodes: List[Node], for strategy in input_node.strategies_vector: input_sharding_spec = strategy.output_sharding_spec assert isinstance(input_sharding_spec, ShardingSpec), f'The input node should NOT be a tuple of tensor.' - # compute the resharding cost during forward phase - _, _, resharding_cost_forward = shape_consistency_manager.shape_consistency(input_sharding_spec, input_spec) + try: + # compute the resharding cost + _, _, total_resharding_cost = shape_consistency_manager.shape_consistency( + input_sharding_spec, input_spec) - if count_backward: - # In backward phase, we should convert grad with target_spec into input_sharding_spec - _, _, resharding_cost_backward = shape_consistency_manager.shape_consistency( - input_spec, input_sharding_spec) - total_resharding_cost = resharding_cost_forward + resharding_cost_backward - else: - total_resharding_cost = resharding_cost_forward - - # we need multiply the size of elem dtype to get correct communication cost - resharding_cost = total_resharding_cost * size_per_elem_bytes + # we need multiply the size of elem dtype to get correct communication cost + resharding_cost = total_resharding_cost * size_per_elem_bytes + except AssertionError as e: + warnings.warn(f'{e}') + resharding_cost = INFINITY_COST resharding_costs[input_node].append(resharding_cost) return resharding_costs diff --git a/colossalai/auto_parallel/solver/constants.py b/colossalai/auto_parallel/solver/constants.py index 727c3ef35..ecaa74ca7 100644 --- a/colossalai/auto_parallel/solver/constants.py +++ b/colossalai/auto_parallel/solver/constants.py @@ -4,7 +4,7 @@ import operator __all__ = [ 'ELEMENTWISE_MODULE_OP', 'ELEMENTWISE_FUNC_OP', 'RESHAPE_FUNC_OP', 'CONV_MODULE_OP', 'CONV_FUNC_OP', 'LINEAR_MODULE_OP', 'LINEAR_FUNC_OP', 'BATCHNORM_MODULE_OP', 'POOL_MODULE_OP', 'NON_PARAM_FUNC_OP', 'BCAST_FUNC_OP', - 'EMBEDDING_MODULE_OP', 'LAYERNORM_MODULE_OP', 'ELEMENTWISE_METHOD_OP', 'RESHAPE_METHOD_OP' + 'EMBEDDING_MODULE_OP', 'LAYERNORM_MODULE_OP', 'ELEMENTWISE_METHOD_OP', 'RESHAPE_METHOD_OP', 'INFINITY_COST' ] ELEMENTWISE_MODULE_OP = [torch.nn.Dropout, torch.nn.ReLU] diff --git a/colossalai/tensor/shape_consistency.py b/colossalai/tensor/shape_consistency.py index 9da935cd9..3a1f04c8a 100644 --- a/colossalai/tensor/shape_consistency.py +++ b/colossalai/tensor/shape_consistency.py @@ -18,11 +18,225 @@ __all__ = [ ] +def _all_gather(tensor, comm_spec): + ''' + Implement all gather operation on device mesh based on information provided by comm_spec. + ''' + process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] + for rank_list, process_group in process_groups_list: + if dist.get_rank() in rank_list: + tensor_list = [ + torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) + for _ in range(comm_spec.device_mesh.mesh_shape[comm_spec.logical_process_axis]) + ] + tensor = tensor + group = process_group + dist.all_gather(tensor_list, tensor, group=group) + output = torch.cat(tuple(tensor_list), comm_spec.gather_dim).contiguous() + return output + + +def _split(tensor, comm_spec): + ''' + Implement shard operation on device mesh based on information provided by comm_spec. + ''' + process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] + for rank_list, _ in process_groups_list: + if dist.get_rank() in rank_list: + tensor = tensor + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // len(rank_list) + start = length * rank_list.index(dist.get_rank()) + output = torch.narrow(tensor, dim, start, length) + return output + + +def _all_to_all(tensor, comm_spec): + ''' + Implement all to all operation on device mesh based on information provided by comm_spec. + ''' + process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] + for rank_list, process_group in process_groups_list: + if dist.get_rank() in rank_list: + new_shape = list(tensor.shape) + new_shape[comm_spec.shard_dim] = new_shape[comm_spec.shard_dim] // len(rank_list) + new_shape = torch.Size(new_shape) + output_tensor_list = [ + torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list)) + ] + dim = comm_spec.shard_dim + length = tensor.shape[comm_spec.shard_dim] // len(rank_list) + input_tensor_list = [ + torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(len(rank_list)) + ] + group = process_group + dist.all_to_all(output_tensor_list, input_tensor_list, group) + output = torch.cat(tuple(output_tensor_list), comm_spec.gather_dim).contiguous() + return output + + +def _all_reduce(tensor, comm_spec): + ''' + Implement all reduce operation on device mesh based on information provided by comm_spec. + ''' + process_groups_list = comm_spec.device_mesh.process_groups_dict[comm_spec.logical_process_axis] + for rank_list, process_group in process_groups_list: + if dist.get_rank() in rank_list: + dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group) + return tensor + + +class _ReduceGrad(torch.autograd.Function): + """ + A customized communication operation which forward is an identity operation, + backward is all_reduce operation. + + Args: + input_: input matrix. + comm_spec: comm_spec will give information like process group, rank list, etc. + """ + + @staticmethod + def symbolic(graph, input_): + return input_ + + @staticmethod + def forward(ctx, input_, comm_spec): + ctx.comm_spec = comm_spec + return input_ + + @staticmethod + def backward(ctx, grad_output): + return _all_reduce(grad_output, ctx.comm_spec), None + + +class _ReduceInput(torch.autograd.Function): + """ + A customized communication operation which forward is all_reduce operation, + backward is an identity operation. + + Args: + input_: input matrix. + comm_spec: comm_spec will give information like process group, rank list, etc. + """ + + @staticmethod + def symbolic(graph, input_): + return _all_reduce(input_) + + @staticmethod + def forward(ctx, input_, comm_spec): + return _all_reduce(input_, comm_spec) + + @staticmethod + def backward(ctx, grad_output): + return grad_output, None + + +class _SplitForwardGatherBackward(torch.autograd.Function): + """ + A customized communication operation which forward is split operation, + backward is an all gather operation. + + Args: + input_: input matrix. + comm_spec: comm_spec will give information like process group, rank list, etc. + """ + + @staticmethod + def symbolic(graph, input_): + return _split(input_) + + @staticmethod + def forward(ctx, input_, comm_spec): + ctx.comm_spec = comm_spec + return _split(input_, comm_spec) + + @staticmethod + def backward(ctx, grad_output): + return _all_gather(grad_output, ctx.comm_spec), None + + +class _GatherForwardSplitBackward(torch.autograd.Function): + """ + A customized communication operation which forward is an all gather operation, + backward is split operation. + + Args: + input_: input matrix. + comm_spec: comm_spec will give information like process group, rank list, etc. + """ + + @staticmethod + def symbolic(graph, input_): + return _all_gather(input_) + + @staticmethod + def forward(ctx, input_, comm_spec): + ctx.comm_spec = comm_spec + return _all_gather(input_, comm_spec) + + @staticmethod + def backward(ctx, grad_output): + return _split(grad_output, ctx.comm_spec), None + + +class _AllToAll(torch.autograd.Function): + """ + A customized communication operation which forward is an all to all operation, + backward is an all to all operation. + + Args: + input_: input matrix. + comm_spec: comm_spec will give information like process group, rank list, etc. + """ + + @staticmethod + def symbolic(graph, input_): + return _all_to_all(input_) + + @staticmethod + def forward(ctx, input_, comm_spec): + output = _all_to_all(input_, comm_spec) + comm_spec_for_backward = CommSpec(comm_pattern=comm_spec.comm_pattern, + sharding_spec=comm_spec.sharding_spec, + gather_dim=comm_spec.shard_dim, + shard_dim=comm_spec.gather_dim, + logical_process_axis=comm_spec.logical_process_axis) + ctx.comm_spec = comm_spec_for_backward + return output + + @staticmethod + def backward(ctx, grad_outputs): + return _all_to_all(grad_outputs, ctx.comm_spec), None + + +def reduce_grad(input_, comm_spec): + return _ReduceGrad.apply(input_, comm_spec) + + +def reduce_input(input_, comm_spec): + return _ReduceInput.apply(input_, comm_spec) + + +def split_forward_gather_backward(input_, comm_spec): + return _SplitForwardGatherBackward.apply(input_, comm_spec) + + +def gather_forward_split_backward(input_, comm_spec): + return _GatherForwardSplitBackward.apply(input_, comm_spec) + + +def all_to_all(input_, comm_spec): + return _AllToAll.apply(input_, comm_spec) + + class CollectiveCommPattern(Enum): - ALLGATHER = 'all_gather' - ALLTOALL = 'all_to_all' - SHARD = 'shard' - ALLREDUCE = 'all_reduce' + GATHER_FWD_SPLIT_BWD = 'gather_fwd_split_bwd' + ALL2ALL_FWD_ALL2ALL_BWD = 'all2all_fwd_all2all_bwd' + SPLIT_FWD_GATHER_BWD = 'split_fwd_gather_bwd' + REDUCE_FWD_IDENTITY_BWD = 'all_reduce_fwd_identity_bwd' + IDENTITY_FWD_ALLREDUCE_BWD = 'identity_fwd_all_reduce_bwd' class CommSpec: @@ -42,12 +256,19 @@ class CommSpec: logical_process_axis(Union(int, List[int]), Optional): The mesh_dim to implement the communication action. ''' - def __init__(self, comm_pattern, sharding_spec, gather_dim=None, shard_dim=None, logical_process_axis=None): + def __init__(self, + comm_pattern, + sharding_spec, + gather_dim=None, + shard_dim=None, + logical_process_axis=None, + forward_only=False): self.comm_pattern = comm_pattern self.sharding_spec = sharding_spec self.gather_dim = gather_dim self.shard_dim = shard_dim self.logical_process_axis = logical_process_axis + self.forward_only = forward_only if isinstance(self.logical_process_axis, list): self.device_mesh = self.sharding_spec.device_mesh.flatten_device_mesh self.logical_process_axis = 0 @@ -56,21 +277,24 @@ class CommSpec: def __repr__(self): res_list = ["CommSpec:("] - if self.comm_pattern == CollectiveCommPattern.ALLGATHER: - res_list.append(f"comm_pattern:all_gather, ") + if self.comm_pattern == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD: + res_list.append(f"comm_pattern:GATHER_FWD_SPLIT_BWD, ") res_list.append(f"gather_dim:{self.gather_dim}, ") res_list.append(f"logical_process_axis:{self.logical_process_axis})") - elif self.comm_pattern == CollectiveCommPattern.ALLTOALL: - res_list.append(f"comm_pattern:all2all, ") + elif self.comm_pattern == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD: + res_list.append(f"comm_pattern:ALL2ALL_FWD_ALL2ALL_BWD, ") res_list.append(f"gather_dim:{self.gather_dim}, ") res_list.append(f"shard_dim:{self.shard_dim}, ") res_list.append(f"logical_process_axis: {self.logical_process_axis})") - elif self.comm_pattern == CollectiveCommPattern.SHARD: - res_list.append(f"comm_pattern:shard, ") + elif self.comm_pattern == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD: + res_list.append(f"comm_pattern:SPLIT_FWD_GATHER_BWD, ") res_list.append(f"shard_dim:{self.shard_dim}, ") res_list.append(f"logical_process_axis:{self.logical_process_axis})") - elif self.comm_pattern == CollectiveCommPattern.ALLREDUCE: - res_list.append(f"comm_pattern:all_reduce, ") + elif self.comm_pattern == CollectiveCommPattern.REDUCE_FWD_IDENTITY_BWD: + res_list.append(f"comm_pattern:REDUCE_FWD_IDENTITY_BWD, ") + res_list.append(f"logical_process_axis:{self.logical_process_axis})") + elif self.comm_pattern == CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD: + res_list.append(f"comm_pattern:IDENTITY_FWD_ALLREDUCE_BWD, ") res_list.append(f"logical_process_axis:{self.logical_process_axis})") return ''.join(res_list) @@ -82,16 +306,38 @@ class CommSpec: For shard operation, it is an on-chip operation, so the communication cost is zero. ''' comm_size = reduce(operator.mul, self.sharding_spec.get_sharded_shape_per_device(), 1) - if self.comm_pattern == CollectiveCommPattern.ALLGATHER: - return self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis) - if self.comm_pattern == CollectiveCommPattern.ALLTOALL: - return self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis) - if self.comm_pattern == CollectiveCommPattern.ALLREDUCE: - return self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis) - if self.comm_pattern == CollectiveCommPattern.SHARD: + if self.comm_pattern == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD: + forward_communication_cost = self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis) # give a tiny cost to shard - return 10 - raise RuntimeError(f"Could not find a matching CollectiveCommPattern for {self.comm_pattern}.") + backward_communication_cost = 10 + + if self.comm_pattern == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD: + forward_communication_cost = self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis) + # grad should have same shape as input tensor + # all to all operation has same logical process axis as forward. + backward_communication_cost = self.device_mesh.all_to_all_cost(comm_size, self.logical_process_axis) + + if self.comm_pattern == CollectiveCommPattern.REDUCE_FWD_IDENTITY_BWD: + forward_communication_cost = self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis) + backward_communication_cost = 0 + + if self.comm_pattern == CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD: + forward_communication_cost = 0 + backward_communication_cost = self.device_mesh.all_reduce_cost(comm_size, self.logical_process_axis) + + if self.comm_pattern == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD: + # give a tiny cost to shard + forward_communication_cost = 10 + backward_communication_cost = self.device_mesh.all_gather_cost(comm_size, self.logical_process_axis) + try: + if self.forward_only: + total_communication_cost = forward_communication_cost + else: + total_communication_cost = forward_communication_cost + backward_communication_cost + except: + raise RuntimeError(f"Could not find a matching CollectiveCommPattern for {self.comm_pattern}.") + + return total_communication_cost def covert_spec_to_action(self, tensor): ''' @@ -101,64 +347,21 @@ class CommSpec: Argument: tensor(torch.Tensor): Tensor stored in each device, which could be different in different ranks. ''' - process_groups_list = self.device_mesh.process_groups_dict[self.logical_process_axis] - - if self.comm_pattern == CollectiveCommPattern.ALLGATHER: - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - tensor_list = [ - torch.zeros(tensor.shape, dtype=tensor.dtype, device=tensor.device) - for _ in range(self.device_mesh.mesh_shape[self.logical_process_axis]) - ] - tensor = tensor - group = process_group - dist.all_gather(tensor_list, tensor, group=group) - tensor.data = torch.cat(tuple(tensor_list), self.gather_dim) - - elif self.comm_pattern == CollectiveCommPattern.SHARD: - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - tensor = tensor - dim = self.shard_dim - length = tensor.shape[self.shard_dim] // len(rank_list) - start = length * rank_list.index(dist.get_rank()) - tensor.data = torch.narrow(tensor, dim, start, length) - - elif self.comm_pattern == CollectiveCommPattern.ALLTOALL: - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - new_shape = list(tensor.shape) - new_shape[self.shard_dim] = new_shape[self.shard_dim] // len(rank_list) - new_shape = torch.Size(new_shape) - output_tensor_list = [ - torch.zeros(new_shape, dtype=tensor.dtype, device=tensor.device) for _ in range(len(rank_list)) - ] - dim = self.shard_dim - length = tensor.shape[self.shard_dim] // len(rank_list) - input_tensor_list = [ - torch.narrow(tensor, dim, length * i, length).contiguous() for i in range(len(rank_list)) - ] - group = process_group - dist.all_to_all(output_tensor_list, input_tensor_list, group) - tensor.data = torch.cat(tuple(output_tensor_list), self.gather_dim) - - elif self.comm_pattern == CollectiveCommPattern.ALLREDUCE: - # For the consistency of collective communication operation, we temporally do not - # allow all_reduce two different mesh dimensions in the same time. - # e.g.: MatMul[(R, S01), (S01, R)] -> Partial(R, R), - # all_reduce(Partial, logical_pg=(0, 1)) is NOT allowed, instead - # we need to do this in two steps: - # 1. all_reduce(Partial, logical_pg=1) - # 2. all_reduce(Partial, logical_pg=0) - for rank_list, process_group in process_groups_list: - if dist.get_rank() in rank_list: - dist.all_reduce(tensor, op=ReduceOp.SUM, group=process_group) - tensor.data = tensor - + if self.comm_pattern in pattern_to_func_dict: + tensor.data = pattern_to_func_dict[self.comm_pattern](tensor, self) else: tensor.data = tensor +pattern_to_func_dict = { + CollectiveCommPattern.GATHER_FWD_SPLIT_BWD: gather_forward_split_backward, + CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD: all_to_all, + CollectiveCommPattern.SPLIT_FWD_GATHER_BWD: split_forward_gather_backward, + CollectiveCommPattern.REDUCE_FWD_IDENTITY_BWD: reduce_input, + CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD: reduce_grad, +} + + @dataclass class ShapeConsistencyOptions: """ @@ -180,6 +383,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): def __init__(self): self._options = None + self._forward_only = False self.total_communication_cost = 0 self.total_transform_steps = 0 self.cached_spec_pairs_transform_path = {} @@ -193,6 +397,15 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): assert isinstance(options_, ShapeConsistencyOptions) self._options = options_ + @property + def forward_only(self): + return self._forward_only + + @forward_only.setter + def forward_only(self, value): + assert isinstance(value, bool) + self._forward_only = value + def get_all_all_gather_spec(self, source_spec, orig_cost): ''' Get all valid sharding specs from source_spec with single all-gather operation, and @@ -224,7 +437,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): device_mesh_shape: (4, 4): 0} ''' valid_spec_dict = {} - comm_pattern = CollectiveCommPattern.ALLGATHER + comm_pattern = CollectiveCommPattern.GATHER_FWD_SPLIT_BWD for target_pair in source_spec.dim_partition_dict.items(): shard_list = all_gather_simulator(target_pair) index = target_pair[0] @@ -240,10 +453,14 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): # generate the CommSpec to record the action of source_sharding_spec->new_sharding_spec gather_dim = index logical_process_axis = target_pair[1][-1] - comm_spec = CommSpec(comm_pattern, - sharding_spec=source_spec, - gather_dim=gather_dim, - logical_process_axis=logical_process_axis) + comm_spec = CommSpec( + comm_pattern, + sharding_spec=source_spec, + gather_dim=gather_dim, + # shard_dim will be used during backward + shard_dim=gather_dim, + logical_process_axis=logical_process_axis, + forward_only=self.forward_only) # compute the communication cost with CommSpec cost = comm_spec.get_comm_cost() @@ -288,7 +505,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): device_mesh_shape: (4, 4): 0} ''' valid_spec_dict = {} - comm_pattern = CollectiveCommPattern.ALLTOALL + comm_pattern = CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD tensor_dims = len(source_spec.entire_shape) for f_index in range(tensor_dims - 1): for b_index in range(f_index + 1, tensor_dims): @@ -331,7 +548,8 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): sharding_spec=source_spec, gather_dim=gather_dim, shard_dim=shard_dim, - logical_process_axis=logical_process_axis) + logical_process_axis=logical_process_axis, + forward_only=self.forward_only) # compute the communication cost with CommSpec cost = comm_spec.get_comm_cost() @@ -388,7 +606,7 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): device_mesh_shape: (4, 4): 0} ''' valid_spec_dict = {} - comm_pattern = CollectiveCommPattern.SHARD + comm_pattern = CollectiveCommPattern.SPLIT_FWD_GATHER_BWD # legal sharding dims means the mesh_id is still available to use. legal_sharding_dims = [i for i in range(len(source_spec.device_mesh.mesh_shape))] @@ -415,8 +633,10 @@ class ShapeConsistencyManager(metaclass=SingletonMeta): logical_process_axis = shard_list[-1] comm_spec = CommSpec(comm_pattern, sharding_spec=source_spec, + gather_dim=shard_dim, shard_dim=shard_dim, - logical_process_axis=logical_process_axis) + logical_process_axis=logical_process_axis, + forward_only=self.forward_only) # compute the communication cost with CommSpec cost = comm_spec.get_comm_cost() diff --git a/tests/test_tensor/test_comm_spec_apply.py b/tests/test_tensor/test_comm_spec_apply.py index fd843f058..dc51c59d6 100644 --- a/tests/test_tensor/test_comm_spec_apply.py +++ b/tests/test_tensor/test_comm_spec_apply.py @@ -33,7 +33,10 @@ def check_all_gather(device_mesh, rank): sharding_spec = ShardingSpec(device_mesh, tensor_to_check.shape, dim_partition_dict=dim_partition_dict) # CommSpec:(comm_pattern:allgather, gather_dim:1, logical_process_axis:1) - comm_spec = CommSpec(CollectiveCommPattern.ALLGATHER, sharding_spec, gather_dim=1, logical_process_axis=1) + comm_spec = CommSpec(CollectiveCommPattern.GATHER_FWD_SPLIT_BWD, + sharding_spec, + gather_dim=1, + logical_process_axis=1) comm_spec.covert_spec_to_action(sharded_tensor_to_comm) assert sharded_tensor_to_comm.equal(tensor_to_check) @@ -56,7 +59,7 @@ def check_shard(device_mesh, rank): sharding_spec = ShardingSpec(device_mesh, tensor_to_shard.shape, dim_partition_dict=dim_partition_dict) # CommSpec:(comm_pattern:shard, shard_dim:1, logical_process_axis:1) - comm_spec = CommSpec(CollectiveCommPattern.SHARD, sharding_spec, shard_dim=1, logical_process_axis=1) + comm_spec = CommSpec(CollectiveCommPattern.SPLIT_FWD_GATHER_BWD, sharding_spec, shard_dim=1, logical_process_axis=1) comm_spec.covert_spec_to_action(tensor_to_shard) if rank in (0, 2): @@ -102,7 +105,7 @@ def check_all_to_all(device_mesh, rank): sharding_spec = ShardingSpec(device_mesh, torch.Size((4, 2)), dim_partition_dict=dim_partition_dict) # CommSpec:(comm_pattern:shard, shard_dim:1, logical_process_axis:1) - comm_spec = CommSpec(CollectiveCommPattern.ALLTOALL, + comm_spec = CommSpec(CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD, sharding_spec, gather_dim=0, shard_dim=1, @@ -112,7 +115,7 @@ def check_all_to_all(device_mesh, rank): assert tensor_to_comm.equal(tensor_to_check) -def check_all_reduce(device_mesh, rank): +def check_all_reduce_fwd(device_mesh, rank): # tensor to comm tensor_to_comm = torch.ones(2, 2).cuda() * rank @@ -133,8 +136,25 @@ def check_all_reduce(device_mesh, rank): # device_mesh_shape: (2, 2) sharding_spec = ShardingSpec(device_mesh, tensor_to_comm.shape, dim_partition_dict=dim_partition_dict) - # CommSpec:(comm_pattern:all_reduce, logical_process_axis:0) - comm_spec = CommSpec(CollectiveCommPattern.ALLREDUCE, sharding_spec, logical_process_axis=0) + comm_spec = CommSpec(CollectiveCommPattern.REDUCE_FWD_IDENTITY_BWD, sharding_spec, logical_process_axis=0) + comm_spec.covert_spec_to_action(tensor_to_comm) + + assert tensor_to_comm.equal(tensor_to_check) + + +def check_all_reduce_bwd(device_mesh, rank): + # tensor to comm + tensor_to_comm = torch.ones(2, 2).cuda() * rank + + tensor_to_check = torch.ones(2, 2).cuda() * rank + + dim_partition_dict = {} + # DistSpec: + # shard_sequence: R,R + # device_mesh_shape: (2, 2) + sharding_spec = ShardingSpec(device_mesh, tensor_to_comm.shape, dim_partition_dict=dim_partition_dict) + + comm_spec = CommSpec(CollectiveCommPattern.IDENTITY_FWD_ALLREDUCE_BWD, sharding_spec, logical_process_axis=0) comm_spec.covert_spec_to_action(tensor_to_comm) assert tensor_to_comm.equal(tensor_to_check) @@ -157,7 +177,7 @@ def check_all_reduce_in_flatten_device_mesh(device_mesh, rank): sharding_spec = ShardingSpec(device_mesh, tensor_to_comm.shape, dim_partition_dict=dim_partition_dict) # CommSpec:(comm_pattern:all_reduce, logical_process_axis:[0, 1]) - comm_spec = CommSpec(CollectiveCommPattern.ALLREDUCE, sharding_spec, logical_process_axis=[0, 1]) + comm_spec = CommSpec(CollectiveCommPattern.REDUCE_FWD_IDENTITY_BWD, sharding_spec, logical_process_axis=[0, 1]) comm_spec.covert_spec_to_action(tensor_to_comm) assert tensor_to_comm.equal(tensor_to_check) @@ -184,7 +204,8 @@ def check_comm(rank, world_size, port): check_all_to_all(device_mesh, rank) # test all reduce - check_all_reduce(device_mesh, rank) + check_all_reduce_fwd(device_mesh, rank) + check_all_reduce_bwd(device_mesh, rank) # test all reduce in 1D flatten device mesh check_all_reduce_in_flatten_device_mesh(device_mesh, rank) diff --git a/tests/test_tensor/test_shape_consistency.py b/tests/test_tensor/test_shape_consistency.py index 8b1578165..c81bee5e0 100644 --- a/tests/test_tensor/test_shape_consistency.py +++ b/tests/test_tensor/test_shape_consistency.py @@ -106,18 +106,18 @@ def test_shape_consistency(): assert transform_path_str == '[R, S01, R]->[R, S0, R]->[S0, R, R]->[S01, R, R]' # all-gather(S01) -> S0 - assert comm_action_sequence[0].comm_pattern == CollectiveCommPattern.ALLGATHER + assert comm_action_sequence[0].comm_pattern == CollectiveCommPattern.GATHER_FWD_SPLIT_BWD assert comm_action_sequence[0].gather_dim == 1 assert comm_action_sequence[0].logical_process_axis == 1 # all-to-all(R, S0) -> [S0, R] - assert comm_action_sequence[1].comm_pattern == CollectiveCommPattern.ALLTOALL + assert comm_action_sequence[1].comm_pattern == CollectiveCommPattern.ALL2ALL_FWD_ALL2ALL_BWD assert comm_action_sequence[1].gather_dim == 1 assert comm_action_sequence[1].shard_dim == 0 assert comm_action_sequence[1].logical_process_axis == 0 # shard(S0) -> [S01] - assert comm_action_sequence[2].comm_pattern == CollectiveCommPattern.SHARD + assert comm_action_sequence[2].comm_pattern == CollectiveCommPattern.SPLIT_FWD_GATHER_BWD assert comm_action_sequence[2].shard_dim == 0 assert comm_action_sequence[2].logical_process_axis == 1