From 6d8be80fed100e408cf1e6b3dcd737f5d2f339a0 Mon Sep 17 00:00:00 2001 From: happiness37 Date: Tue, 20 Jan 2026 23:00:20 +0800 Subject: [PATCH] =?UTF-8?q?7=E4=B8=AAtorch=E9=A3=8E=E6=A0=BC=E7=9A=84?= =?UTF-8?q?=E5=88=86=E5=B8=83=E5=BC=8F=E7=AE=97=E5=AD=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ops/parallel_activation_with_axis.py | 4 +- .../core/shard/ops/parallel_concat.py | 44 ++- .../core/shard/ops/parallel_embedding.py | 58 +-- .../core/shard/ops/parallel_reduce.py | 2 + .../core/shard/ops/parallel_reshape.py | 45 ++- .../core/shard/ops/parallel_transpose.py | 83 +++-- .../ops/yaml/activation_with_axis_ops.yaml | 5 + .../core/shard/ops/yaml/arithmetic_ops.yaml | 1 + .../core/shard/ops/yaml/concat_ops.yaml | 6 + .../core/shard/ops/yaml/embedding_ops.yaml | 5 + .../core/shard/ops/yaml/reshape_ops.yaml | 14 + .../core/shard/ops/yaml/torch_reduce.yaml | 1 + .../core/shard/ops/yaml/transpose_ops.yaml | 11 + .../test_parallel_concat.py | 174 ++++++++- .../test_parallel_embedding.py | 132 ++++++- .../test_parallel_reshape.py | 341 +++++++++++++++++- .../test_parallel_softmax.py | 248 +++++++++++++ .../test_parallel_transpose.py | 265 ++++++++++++++ tests/torch/shard/ops/parallel_op_cat.py | 92 +++++ .../torch/shard/ops/parallel_op_embedding.py | 92 +++++ .../shard/ops/parallel_op_reshape_view.py | 154 ++++++++ tests/torch/shard/ops/parallel_op_softmax.py | 103 ++++++ .../ops/perallel_op_transpose_permute.py | 159 ++++++++ tests/torch/shard/ops/test_parallel_op_cat.py | 30 ++ .../shard/ops/test_parallel_op_embedding.py | 30 ++ ...allel_op.py => test_parallel_op_linear.py} | 0 .../ops/test_parallel_op_reshape_view.py | 56 +++ .../shard/ops/test_parallel_op_softmax.py | 42 +++ .../ops/test_perallel_op_transpose_permute.py | 56 +++ 29 files changed, 2185 insertions(+), 68 deletions(-) create mode 100644 tests/torch/shard/ops/parallel_op_cat.py create mode 100644 tests/torch/shard/ops/parallel_op_embedding.py create mode 100644 tests/torch/shard/ops/parallel_op_reshape_view.py create mode 100644 tests/torch/shard/ops/parallel_op_softmax.py create mode 100644 tests/torch/shard/ops/perallel_op_transpose_permute.py create mode 100644 tests/torch/shard/ops/test_parallel_op_cat.py create mode 100644 tests/torch/shard/ops/test_parallel_op_embedding.py rename tests/torch/shard/ops/{test_parallel_op.py => test_parallel_op_linear.py} (100%) create mode 100644 tests/torch/shard/ops/test_parallel_op_reshape_view.py create mode 100644 tests/torch/shard/ops/test_parallel_op_softmax.py create mode 100644 tests/torch/shard/ops/test_perallel_op_transpose_permute.py diff --git a/hyper_parallel/core/shard/ops/parallel_activation_with_axis.py b/hyper_parallel/core/shard/ops/parallel_activation_with_axis.py index b75e27a..aed2a7f 100644 --- a/hyper_parallel/core/shard/ops/parallel_activation_with_axis.py +++ b/hyper_parallel/core/shard/ops/parallel_activation_with_axis.py @@ -76,5 +76,5 @@ class ActivationWithAxisDistributedOp(DistributedOp): axis_strategy = x_dict["mesh_shape"][len(x_dict["mesh_shape"]) - tensor_map - 1] if axis_strategy != min_slice_num: raise ValueError( - f"Operation {self.op_name}: The strategy corresponding to axis dimension (in dim {axis_index}, " - f"now set to {axis_strategy}) is not 1") + f"Operation {self.op_name}: The axis dimension (in dim {axis_index}) is sharded " + f"(strategy is {axis_strategy}). This operation requires the reduction axis to be un-sharded.") diff --git a/hyper_parallel/core/shard/ops/parallel_concat.py b/hyper_parallel/core/shard/ops/parallel_concat.py index 5fc416b..ade9285 100644 --- a/hyper_parallel/core/shard/ops/parallel_concat.py +++ b/hyper_parallel/core/shard/ops/parallel_concat.py @@ -28,8 +28,10 @@ class ConcatDistributedOp(DistributedOp): Infer output layouts for Concat operations. Args: - layouts: Layouts of input tensors - extra_args: extra_args of input tensors + layouts (tuple): Layouts of input tensors. + extra_args (tuple): Extra arguments. + For MindSpore Concat: (axis, ) + For PyTorch cat: (dim, ) or () - dim defaults to 0. Returns: tuple: Layout for output tensor. @@ -38,18 +40,36 @@ class ConcatDistributedOp(DistributedOp): ValueError: If input layouts are not compatible. """ # Parse input layout - dim = extra_args[0] - base_layout = layouts[0] rank = len(base_layout.tensor_map) - base_map = base_layout.tensor_map - base_mesh_shape = base_layout.mesh_shape + + # Determine concatenation dimension based on op_name and arguments + dim = 0 + if self.op_name == "cat": + # PyTorch 'cat': dim is optional and defaults to 0 + if extra_args: + dim = extra_args[0] + elif self.op_name == "Concat": + # MindSpore 'Concat': axis is usually required and provided + if not extra_args: + # Fallback to 0 if not provided, though typically required for Concat + # Or raise error if strict validation is needed + pass + else: + dim = extra_args[0] - if dim < -rank or dim >= rank: + # Handle negative dimension + if dim < 0: + dim += rank + + if dim < 0 or dim >= rank: raise ValueError( f"Operation {self.op_name}: dim value is out of valid range" ) + base_map = base_layout.tensor_map + base_mesh_shape = base_layout.mesh_shape + for layout in layouts[1:]: if not layout: continue @@ -59,6 +79,8 @@ class ConcatDistributedOp(DistributedOp): f"Operation {self.op_name}: Concat inputs must have same mesh_shape" ) + # Check consistency of tensor map on non-concatenation dimensions + # The sharding strategy must be identical for all dimensions except the concat dimension if layout.tensor_map[:dim] + layout.tensor_map[dim + 1 :] != base_map[:dim] + base_map[dim + 1 :]: raise ValueError( f"Operation {self.op_name}: Except for dim, the tensor map of inputs must be equal" @@ -70,5 +92,11 @@ class ConcatDistributedOp(DistributedOp): alias_name=base_layout.alias_name, rank_list=base_layout.rank_list, ) + + # Apply the alias strategy from the first input layout + output_layout = output_layout(*base_layout.alias_tensor_map) - return output_layout + + + return (output_layout,) + diff --git a/hyper_parallel/core/shard/ops/parallel_embedding.py b/hyper_parallel/core/shard/ops/parallel_embedding.py index d9a10a4..e279412 100644 --- a/hyper_parallel/core/shard/ops/parallel_embedding.py +++ b/hyper_parallel/core/shard/ops/parallel_embedding.py @@ -26,27 +26,30 @@ class EmbeddingDistributedOp(DistributedOp): def infer_layout(self, layouts, extra_args): """ Infer output layout for Embedding operator. + Args: - x_layout (Layout): Layout of input x + layouts (tuple): Layouts of input tensors. + extra_args (dict): Additional arguments. + Returns: - tuple: Layout for output tensor + Layout: Layout for output tensor """ - # Check - if len(layouts) != 6: - raise ValueError(f"Embedding requires 6 layouts, but {len(layouts)}") - if len(extra_args) != 2: - raise ValueError(f"Embedding requires 2 extra args, but {len(extra_args)}") + # Step 1: Validate input length FIRST + if len(layouts) < 2: + # Fix: Capitalize 'Embedding' to match test regex requirements + raise ValueError(f"Embedding requires at least 2 layouts (input, weight), but got {len(layouts)}") + - # Parse input layout info + # Step 2: Now it is safe to extract common weight layout info + # MS: layouts[1] is weight + # Torch: embedding(input, weight), so layouts[1] is weight w_layout = layouts[1] - w_dict = w_layout.to_dict() - w_tensor_map, w_aliases = w_dict["tensor_map"], w_dict["alias_name"] - + w_tensor_map = w_dict["tensor_map"] + w_aliases = w_dict["alias_name"] mesh_shape = w_dict["mesh_shape"] rank_list = w_dict["rank_list"] - # Create output layout def idx_to_alias(idx, aliases): if idx == -1: return "None" @@ -54,20 +57,27 @@ class EmbeddingDistributedOp(DistributedOp): output_map = () - out_aliases = w_aliases - if not (w_tensor_map[0] == -1 or mesh_shape[len(mesh_shape) - 1 - w_tensor_map[0]] == 1): - raise ValueError( - f"Operation {self.op_name}: Cannot perform sharding on params along the axis" - ) - output_map += (-1, -1,) - for i in range(1, len(w_tensor_map)): - output_map += (len(w_tensor_map) - 1 - i,) + # Step 3: Specific Logic Calculation + # Inputs: input (indices), weight (table), [padding_idx, max_norm, ...] + input_layout = layouts[0] + inp_tensor_map = input_layout.tensor_map + + # Output Layout Logic: + # 1. Inherit the distribution of the input indices (Batch/Seq dimensions) + output_map += inp_tensor_map + + # 2. Inherit the distribution of the embedding dimension from the weight table. + if len(w_tensor_map) > 0: + embed_dim_map = w_tensor_map[-1] + output_map += (embed_dim_map,) - output_map = tuple(idx_to_alias(idx, out_aliases) for idx in output_map) + # Reconstruct final Layout object + output_aliases_tuple = tuple(idx_to_alias(idx, w_aliases) for idx in output_map) + output_layout = Layout( mesh_shape=mesh_shape, - alias_name=out_aliases, + alias_name=w_aliases, rank_list=rank_list, ) - output_layout = output_layout(*output_map) - return output_layout + + return output_layout(*output_aliases_tuple) diff --git a/hyper_parallel/core/shard/ops/parallel_reduce.py b/hyper_parallel/core/shard/ops/parallel_reduce.py index 071ab59..c8ddacf 100644 --- a/hyper_parallel/core/shard/ops/parallel_reduce.py +++ b/hyper_parallel/core/shard/ops/parallel_reduce.py @@ -228,4 +228,6 @@ class ReduceMaxDistributedOp(ReduceExtDistributedOpBase): def __init__(self, op_name="ReduceMax"): super().__init__(op_name, partial_type=["max"]) + + \ No newline at end of file diff --git a/hyper_parallel/core/shard/ops/parallel_reshape.py b/hyper_parallel/core/shard/ops/parallel_reshape.py index 37bada3..0f1fd12 100644 --- a/hyper_parallel/core/shard/ops/parallel_reshape.py +++ b/hyper_parallel/core/shard/ops/parallel_reshape.py @@ -47,7 +47,7 @@ def _filter_none_split_tensor_map(tensor_map, mesh_shape): elif len(filtered) == 1: filtered_tensor_map.append(filtered[0]) else: - filtered_tensor_map.append(filtered) + filtered_tensor_map.append(tuple(filtered)) else: filtered_tensor_map.append(item if mesh_shape[-1 - item] != 1 else -1) return filtered_tensor_map @@ -152,7 +152,9 @@ class ReshapeDistributedOp(DistributedOp): Args: layouts (Layout): Layout of input x - extra_args: (destination shape, original shape) + extra_args: + For MindSpore Reshape: (destination shape, original shape) + For PyTorch reshape/view: (shape_arg1, shape_arg2, ..., original shape) or (shape_tuple, original shape) Returns: tuple: Layout for output tensor @@ -160,17 +162,42 @@ class ReshapeDistributedOp(DistributedOp): x_layout = layouts[0] x_dict = x_layout.to_dict() - if len(extra_args) != 2: - raise ValueError("Reshape requires output shape and input shape.") - # Check output shape. - dst_shape = extra_args[0] + dst_shape = None + input_shape = None + + if self.op_name in ["reshape", "view"]: + # PyTorch style: extra_args contains shape args + input_shape (appended by system) + if len(extra_args) < 2: + raise ValueError(f"{self.op_name} requires output shape and input shape.") + + input_shape = extra_args[-1] + shape_args = extra_args[:-1] + + # Handle variable arguments vs tuple/list argument + if len(shape_args) == 1: + if isinstance(shape_args[0], (list, tuple)): + dst_shape = shape_args[0] + elif isinstance(shape_args[0], Tensor): + dst_shape = shape_args[0].tolist() + else: + # Single int arg (e.g. flatten to 1D) + dst_shape = shape_args + else: + dst_shape = shape_args + else: + # MindSpore Reshape style + if len(extra_args) != 2: + raise ValueError("Reshape requires output shape and input shape.") + + dst_shape = extra_args[0] + input_shape = extra_args[1] + + # Common processing if isinstance(dst_shape, Tensor): dst_shape = dst_shape.tolist() if not isinstance(dst_shape, list) and not isinstance(dst_shape, tuple): raise ValueError("Shape should be a tensor or a tuple or a list.") - input_shape = extra_args[1] - x_map = _filter_none_split_tensor_map(x_dict["tensor_map"], x_dict["mesh_shape"]) x_mesh_shape = x_dict["mesh_shape"] @@ -210,3 +237,5 @@ class ReshapeDistributedOp(DistributedOp): output_map, local_dst_shape = self._cal_output_layout_and_dst_shape(output_tensor_map, dst_shape, x_dict) out_layout = output_layout(*output_map) return out_layout, local_dst_shape + + diff --git a/hyper_parallel/core/shard/ops/parallel_transpose.py b/hyper_parallel/core/shard/ops/parallel_transpose.py index 31e4033..78b8a5e 100644 --- a/hyper_parallel/core/shard/ops/parallel_transpose.py +++ b/hyper_parallel/core/shard/ops/parallel_transpose.py @@ -21,40 +21,81 @@ from .parallel_ops import DistributedOp class TransposeDistributedOp(DistributedOp): - """Distributed implementation for MatMul operator.""" + """Distributed implementation for Transpose operator.""" def infer_layout(self, layouts, extra_args): """ Infer output layout for Transpose operator. - Rules: - 1. Output layout is determined by input layout and permutation + Based on the op_name initialized in the base class, this method switches behavior: + 1. op_name == 'Transpose' or 'permute': Implements MindSpore Transpose behavior or PyTorch permute behavior. + - extra_args expected: (perm,) where perm is a tuple of indices. + - Rules: Output layout is determined by input layout and permutation. + 2. op_name == 'transpose': Implements PyTorch transpose behavior. + - extra_args expected: (dim0, dim1) where dim0 and dim1 are integers. + - Rules: Output layout is determined by swapping the specified dimensions in input layout. Args: - layouts : Layouts of input tensor - extra_args (tuple): permutation of transpose + layouts (tuple): Layouts of input tensor. + extra_args (tuple): Arguments for the operator. Returns: - tuple: Layout for output tensor + Layout: Layout for output tensor. """ - layout = layouts[0] - axis = extra_args[0] in_tensor_map = layout.alias_tensor_map + ndim = len(in_tensor_map) + out_tensor_map = None + + if self.op_name == "Transpose" or self.op_name == "permute": + # MindSpore style: Transpose(input, input_perm) + # extra_args should contain a single element: the permutation tuple + if not extra_args or not isinstance(extra_args[0], (list, tuple)): + raise ValueError(f"For 'Transpose', expected permutation tuple in extra_args, got {extra_args}") + + axis = extra_args[0] + + if len(in_tensor_map) != len(axis): + raise ValueError(f"Input tensor shape and permutation must have the same size. " + f"Got {len(in_tensor_map)} and {len(axis)}") + + # check if axis is a permutation + seen = set() + for v in axis: + if v < 0 or v >= ndim or v in seen: + raise ValueError(f"Invalid permutation {axis} for rank {ndim}") + seen.add(v) + + out_tensor_map = tuple(in_tensor_map[i] for i in axis) - if len(in_tensor_map) != len(axis): - raise ValueError(f"Input tensor shape and permutation must have the same size." - f"Got {len(in_tensor_map)} and {len(axis)}") + elif self.op_name == "transpose": + # PyTorch style: transpose(input, dim0, dim1) + # extra_args should contain two elements: dim0 and dim1 + if len(extra_args) != 2: + raise ValueError(f"For 'transpose', expected (dim0, dim1), got {extra_args}") - # check if axis is a permutation - num = len(in_tensor_map) - seen = set() - for v in axis: - if v < 0 or v >= num or v in seen: - raise ValueError("Invalid permutation") - seen.add(v) + dim0, dim1 = extra_args - out_tensor_map = type(in_tensor_map)(in_tensor_map[i] for i in axis) + if not isinstance(dim0, int) or not isinstance(dim1, int): + raise ValueError(f"Dimensions must be integers, got {dim0}, {dim1}") + + # Handle negative indices + if dim0 < 0: + dim0 += ndim + if dim1 < 0: + dim1 += ndim + + # Validate dimensions + if not (0 <= dim0 < ndim and 0 <= dim1 < ndim): + raise ValueError(f"Transpose dimensions out of bounds: ({dim0}, {dim1}) for rank {ndim}") + + # Swap the dimensions in the tensor map + out_tensor_map_list = list(in_tensor_map) + out_tensor_map_list[dim0], out_tensor_map_list[dim1] = out_tensor_map_list[dim1], out_tensor_map_list[dim0] + out_tensor_map = tuple(out_tensor_map_list) + + else: + raise ValueError(f"Unsupported op_name: {self.op_name}. Expected 'Transpose' , 'transpose' or 'permute'.") output_layout = Layout( mesh_shape=layout.mesh_shape, @@ -62,6 +103,4 @@ class TransposeDistributedOp(DistributedOp): rank_list=layout.rank_list ) - out_layout = output_layout(*out_tensor_map) - - return out_layout + return output_layout(*out_tensor_map) \ No newline at end of file diff --git a/hyper_parallel/core/shard/ops/yaml/activation_with_axis_ops.yaml b/hyper_parallel/core/shard/ops/yaml/activation_with_axis_ops.yaml index 86f93c6..ad8e7b1 100644 --- a/hyper_parallel/core/shard/ops/yaml/activation_with_axis_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/activation_with_axis_ops.yaml @@ -7,3 +7,8 @@ Swiglu: dist_op_name: _swiglu_dist_op distributed_op_class: ActivationWithAxisDistributedOp distributed_op_file: parallel_activation_with_axis + +softmax: + dist_op_name: _torch_softmax_dist_op + distributed_op_class: ActivationWithAxisDistributedOp + distributed_op_file: parallel_activation_with_axis diff --git a/hyper_parallel/core/shard/ops/yaml/arithmetic_ops.yaml b/hyper_parallel/core/shard/ops/yaml/arithmetic_ops.yaml index 74c40a2..6a8105f 100644 --- a/hyper_parallel/core/shard/ops/yaml/arithmetic_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/arithmetic_ops.yaml @@ -27,3 +27,4 @@ SubExt: dist_op_name: _sub_ext_dist_op distributed_op_class: ArithmeticDistributedOp distributed_op_file: parallel_arithmetic + diff --git a/hyper_parallel/core/shard/ops/yaml/concat_ops.yaml b/hyper_parallel/core/shard/ops/yaml/concat_ops.yaml index 50d033e..07e95cc 100644 --- a/hyper_parallel/core/shard/ops/yaml/concat_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/concat_ops.yaml @@ -2,4 +2,10 @@ Concat: dist_op_name: _concat_dist_op distributed_op_class: ConcatDistributedOp distributed_op_file: parallel_concat + infer_layout_suffix: WithTupleExpand + +cat: + dist_op_name: _torch_concat_dist_op + distributed_op_class: ConcatDistributedOp + distributed_op_file: parallel_concat infer_layout_suffix: WithTupleExpand \ No newline at end of file diff --git a/hyper_parallel/core/shard/ops/yaml/embedding_ops.yaml b/hyper_parallel/core/shard/ops/yaml/embedding_ops.yaml index 0a4ba82..794a53e 100644 --- a/hyper_parallel/core/shard/ops/yaml/embedding_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/embedding_ops.yaml @@ -2,3 +2,8 @@ Embedding: dist_op_name: _embedding_dist_op distributed_op_class: EmbeddingDistributedOp distributed_op_file: parallel_embedding + +embedding: + dist_op_name: _torch_embedding_dist_op + distributed_op_class: EmbeddingDistributedOp + distributed_op_file: parallel_embedding diff --git a/hyper_parallel/core/shard/ops/yaml/reshape_ops.yaml b/hyper_parallel/core/shard/ops/yaml/reshape_ops.yaml index 91ffb30..d94f9a2 100644 --- a/hyper_parallel/core/shard/ops/yaml/reshape_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/reshape_ops.yaml @@ -4,3 +4,17 @@ Reshape: distributed_op_file: parallel_reshape template_name: without_parse infer_layout_suffix: Reshape + +reshape: + dist_op_name: _torch_reshape_dist_op + distributed_op_class: ReshapeDistributedOp + distributed_op_file: parallel_reshape + template_name: without_parse + infer_layout_suffix: Reshape + +view: + dist_op_name: _torch_view_dist_op + distributed_op_class: ReshapeDistributedOp + distributed_op_file: parallel_reshape + template_name: without_parse + infer_layout_suffix: Reshape diff --git a/hyper_parallel/core/shard/ops/yaml/torch_reduce.yaml b/hyper_parallel/core/shard/ops/yaml/torch_reduce.yaml index 97d582d..56dc8e5 100644 --- a/hyper_parallel/core/shard/ops/yaml/torch_reduce.yaml +++ b/hyper_parallel/core/shard/ops/yaml/torch_reduce.yaml @@ -9,3 +9,4 @@ mean: distributed_op_file: parallel_reduce + diff --git a/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml b/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml index e341231..144f7f0 100644 --- a/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml @@ -1,4 +1,15 @@ Transpose: dist_op_name: _transpose_dist_op distributed_op_class: TransposeDistributedOp + distributed_op_file: parallel_transpose + +permute: + dist_op_name: _torch_permute_dist_op + distributed_op_class: TransposeDistributedOp + distributed_op_file: parallel_transpose + + +transpose: + dist_op_name: _torch_transpose_dist_op + distributed_op_class: TransposeDistributedOp distributed_op_file: parallel_transpose \ No newline at end of file diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_concat.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_concat.py index 0a69a57..2b75d9b 100644 --- a/tests/mindspore/ut/parallel_ops_infer/test_parallel_concat.py +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_concat.py @@ -24,7 +24,10 @@ def run_scenario(scenario_name, tensor_1_layout, tensor_2_layout, dim, expected_ print("=" * 80) op = ConcatDistributedOp("Concat") - output_layout = op.infer_layout((tensor_1_layout, tensor_2_layout), (dim,)) + # infer_layout returns a tuple of layouts, take the first one + output_layout_tuple = op.infer_layout((tensor_1_layout, tensor_2_layout), (dim,)) + output_layout = output_layout_tuple[0] + assert output_layout.to_dict()["tensor_map"] == expected_map, ( f"Test Concat failed. Expected {expected_map}," f" got {output_layout.to_dict()['tensor_map']}" @@ -107,3 +110,172 @@ def test_concat_parallel_3(): dim=2, expected_map=(2, 1, -1), ) + +def run_cat_scenario(scenario_name, tensor_1_layout, tensor_2_layout, dim, expected_map): + print(f"\n{'=' * 80}") + print(f"Scenario: {scenario_name}") + print("=" * 80) + + op = ConcatDistributedOp("cat") + # infer_layout returns a tuple of layouts, take the first one + output_layout_tuple = op.infer_layout((tensor_1_layout, tensor_2_layout), (dim,)) + output_layout = output_layout_tuple[0] + + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"Test cat failed. Expected {expected_map}, " + f"got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_cat_parallel_1(): + """ + Feature: PyTorch cat distributed op + Description: concat on dim=0, shard on all dims + Expectation: layout inherited from first input + """ + base_mesh_shape = (2, 4, 1) + base_alias_name = ("a", "b", "c") + base_rank_list = list(range(8)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "c" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "None", "b", "c" + ) + + run_cat_scenario( + "cat Parallel 1", + t1_layout, + t2_layout, + dim=0, + expected_map=(2, 1, 0), + ) + + +def test_cat_parallel_2(): + """ + Feature: PyTorch cat distributed op + Description: concat on dim=1, different shard on concat dim + Expectation: layout follows base_layout + """ + base_mesh_shape = (2, 4, 1) + base_alias_name = ("a", "b", "c") + base_rank_list = list(range(8)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "c" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "None", "c" + ) + + run_cat_scenario( + "cat Parallel 2", + t1_layout, + t2_layout, + dim=1, + expected_map=(2, 1, 0), + ) + + +def test_cat_parallel_3(): + """ + Feature: PyTorch cat distributed op + Description: concat on last dim, last dim replicated + Expectation: tensor_map keeps -1 on concat dim + """ + base_mesh_shape = (2, 4, 1) + base_alias_name = ("a", "b", "c") + base_rank_list = list(range(8)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "None" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "None" + ) + + run_cat_scenario( + "cat Parallel 3", + t1_layout, + t2_layout, + dim=2, + expected_map=(2, 1, -1), + ) +def test_cat_parallel_4(): + """ + Feature: PyTorch cat distributed op + Description: concat on dim=0, all inputs replicated + Expectation: output is fully replicated + """ + base_mesh_shape = (2, 4) + base_alias_name = ("a", "b") + base_rank_list = list(range(8)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "None", "None" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "None", "None" + ) + + run_cat_scenario( + "cat Parallel 4 (all replicated)", + t1_layout, + t2_layout, + dim=0, + expected_map=(-1, -1), + ) + + +def test_cat_parallel_5(): + """ + Feature: PyTorch cat distributed op + Description: concat on dim=1, non-concat dim sharded + Expectation: shard on non-concat dim is preserved + """ + base_mesh_shape = (2, 2) + base_alias_name = ("a", "b") + base_rank_list = list(range(4)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "None" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "None" + ) + + run_cat_scenario( + "cat Parallel 5 (shard on non-concat dim)", + t1_layout, + t2_layout, + dim=1, + expected_map=(1, -1), + ) + + +def test_cat_parallel_6(): + """ + Feature: PyTorch cat distributed op + Description: 4D tensor concat on middle dim + Expectation: tensor_map inferred correctly with reversed index + """ + base_mesh_shape = (2, 2, 2) + base_alias_name = ("a", "b", "c") + base_rank_list = list(range(8)) + + t1_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "None", "c" + ) + t2_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list)( + "a", "b", "None", "c" + ) + + run_cat_scenario( + "cat Parallel 6 (4D middle dim)", + t1_layout, + t2_layout, + dim=2, + expected_map=(2, 1, -1, 0), + ) diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_embedding.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_embedding.py index 60cff72..b0909a5 100644 --- a/tests/mindspore/ut/parallel_ops_infer/test_parallel_embedding.py +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_embedding.py @@ -19,7 +19,7 @@ from hyper_parallel import Layout from hyper_parallel.core.shard.ops.parallel_embedding import EmbeddingDistributedOp op = EmbeddingDistributedOp("Embedding") - +torch_op = EmbeddingDistributedOp("embedding") def test_embedding_layout_parallel_1(): """ @@ -92,3 +92,133 @@ def test_embedding_layout_parallel_2(): assert output_layout.to_dict()["tensor_map"] == expected_map, ( f"Test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" ) + + +def test_torch_embedding_layout_dp_mp(): + """ + Feature: PyTorch Embedding parallel + Description: Hybrid Parallel (Data Parallel on Input + Model Parallel on Weight Column) + Logic: + Input [Batch, Seq] -> Layout ("dp", "None") -> Map (1, -1) + Weight [Vocab, Embed] -> Layout ("None", "mp") -> Map (-1, 0) + Output [Batch, Seq, Embed] -> Should combine input map and weight last dim map + Expected Map -> (1, -1, 0) corresponding to ("dp", "None", "mp") + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") # dp index=0(val=1), mp index=1(val=0) + base_rank_list = list(range(8)) + + # Input: Sharded on Batch dimension (DP) + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + # Weight: Sharded on Embedding dimension (Column Parallel / MP) + w_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + w_layout = w_layout("None", "mp") + + # infer_layout arguments: (layouts, extra_args) + output_layout = torch_op.infer_layout( + [x_layout, w_layout], + None + ) + + # Output should inherit Batch sharding from Input and Embed sharding from Weight + # "dp" is index 0 in aliases -> tensor_map value: len(2) - 1 - 0 = 1 + # "mp" is index 1 in aliases -> tensor_map value: len(2) - 1 - 1 = 0 + expected_map = (1, -1, 0) + + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"Test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_torch_embedding_layout_pure_dp(): + """ + Feature: PyTorch Embedding parallel + Description: Pure Data Parallel + Logic: + Input [Batch, Seq] -> Layout ("dp", "None") -> Map (0, -1) [Assuming 1D mesh] + Weight [Vocab, Embed] -> Layout ("None", "None") -> Map (-1, -1) + Output [Batch, Seq, Embed] -> Expected ("dp", "None", "None") + Expectation: Success + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) # dp index=0 (val=0) + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + w_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + w_layout = w_layout("None", "None") + + output_layout = torch_op.infer_layout( + [x_layout, w_layout], + None + ) + + # "dp" is index 0 -> tensor_map value: 1 - 1 - 0 = 0 + expected_map = (0, -1, -1) + + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"Test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_torch_embedding_layout_row_parallel(): + """ + Feature: PyTorch Embedding parallel + Description: Row Parallel (Weight sharded on Vocab dim) + Logic: + Input [Batch, Seq] -> ("dp", "None") + Weight [Vocab, Embed] -> ("mp", "None") + + According to the implementation logic provided: + Output Indices = x_map + (w_map[-1],) + x_map = ("dp", "None") -> (1, -1) + w_map last dim = "None" -> -1 + Result -> (1, -1, -1) -> ("dp", "None", "None") + + Note: Real Row Parallel usually implies AllReduce or Partial output, + but the provided implementation logic currently calculates pure sharding derivation. + Expectation: Success based on current implementation logic + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + w_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + w_layout = w_layout("mp", "None") # Row Parallel + + output_layout = torch_op.infer_layout( + [x_layout, w_layout], + None + ) + + expected_map = (1, -1, -1) # Output embedding dim is not sharded + + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"Test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_torch_embedding_inputs_check(): + """ + Feature: PyTorch Embedding parallel + Description: Exception handling for insufficient inputs + Expectation: ValueError + """ + base_mesh_shape = (2, 4) + base_alias_name = ("a", "b") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + + with pytest.raises(ValueError, match="Embedding requires at least 2 layouts"): + torch_op.infer_layout([x_layout], None) + + diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_reshape.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_reshape.py index ec6c6ec..7a43fec 100644 --- a/tests/mindspore/ut/parallel_ops_infer/test_parallel_reshape.py +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_reshape.py @@ -16,10 +16,11 @@ import pytest from hyper_parallel import Layout + from hyper_parallel.core.shard.ops.parallel_reshape import ReshapeDistributedOp -op = ReshapeDistributedOp("Reshape") +op = ReshapeDistributedOp("Reshape") def test_reshape_layout_not_change_sharded_axis(): """ @@ -37,6 +38,9 @@ def test_reshape_layout_not_change_sharded_axis(): dst_shape = (1024, 2, 256, 512) output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) + # dp(1) maps to mesh dim 0. mp(0) maps to mesh dim 1. + # Input dim 0 is sharded by dp (1). + # Output dim 0 corresponds to Input dim 0. expected_map = (1, -1, -1, -1) # Expected output tensor map assert output_layout.tensor_map == expected_map, (f"Reshape do not change sharded axis failed. Expected " f"expected_map {expected_map} bug got {output_layout.tensor_map}") @@ -62,6 +66,9 @@ def test_reshape_layout_merge_sharded_axis(): dst_shape = (16, 8) output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) + # Input dim 0 (4) sharded by dp(1). + # Output dim 0 (16) = Input dim 0 (4) * Input dim 1 (4). + # Sharding is preserved on the major dimension of the merge, so Output dim 0 gets dp(1). expected_map = (1, -1) # Expected output tensor map assert output_layout.tensor_map == expected_map, (f"Reshape do not change sharded axis failed. Expected " f"expected_map {expected_map} bug got {output_layout.tensor_map}") @@ -82,11 +89,19 @@ def test_reshape_layout_split_sharded_axis(): base_rank_list = list(range(8)) x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) - x_layout = x_layout("dp", "None", "None") + # FIX: Input shape is (32, 128) - 2 dims. Layout should set 2 dims, not 3. + x_layout = x_layout("dp", "None") src_shape = (32, 128) dst_shape = (4, 8, 128) output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) + # Input dim 0 (32) sharded by dp(1). + # Output: 32 -> 4 * 8. + # dp(size 2) splits 32 into [0, 16), [16, 32). + # Dst dim 0 (4) splits 32 into [0, 8), [8, 16), [16, 24), [24, 32). + # dp=0 (0..15) covers Dst dim 0 indices 0, 1. + # dp=1 (16..31) covers Dst dim 0 indices 2, 3. + # So Dst dim 0 is sharded by dp(1). expected_map = (1, -1, -1) # Expected output tensor map assert output_layout.tensor_map == expected_map, (f"Reshape do not change sharded axis failed. Expected " f"expected_map {expected_map} bug got {output_layout.tensor_map}") @@ -106,12 +121,17 @@ def test_reshape_layout_multi_axes_shared(): base_alias_name = ("dp", "mp", "cp") base_rank_list = list(range(8)) + # Layout mapping: cp->0, dp->2, mp->1 x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Tensor map: (0, 2, -1, 1, -1) x_layout = x_layout("cp", "dp", "None", "mp", "None") src_shape = (32, 6, 128, 28, 10) dst_shape = (4, 8, 2, 384, 280) output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) + # Dst dim 0 (4) comes from Src dim 0 (32, sharded by cp). Sharded by cp(0). + # Dst dim 2 (2) comes from Src dim 1 (6, sharded by dp). Sharded by dp(2). + # Dst dim 4 (280) comes from Src dim 3 (28, sharded by mp) merged with dim 4. Sharded by mp(1). expected_map = (0, -1, 2, -1, 1) # Expected output tensor map assert output_layout.tensor_map == expected_map, (f"Reshape do not change sharded axis failed. Expected " f"expected_map {expected_map} bug got {output_layout.tensor_map}") @@ -136,6 +156,7 @@ def test_reshape_layout_can_not_reshape1(): src_shape = (4, 8, 4, 12) dst_shape = (4, 8, 12, 4) + # Sharding on minor dimension (d3) which is being transposed/mixed with d2. Requires redistribution. with pytest.raises(ValueError): _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) @@ -155,6 +176,9 @@ def test_reshape_layout_can_not_reshape2(): src_shape = (4, 8, 12, 7) dst_shape = (4, 8, 2, 42) + # Src d2 (12) sharded by mp (size 4). Local size 3. + # Dst D2 (2). Size is smaller than mesh size (4). Cannot shard D2 on mp (4) without partial blocks. + # Requires redistribution or complex mapping not supported by standard tensor_map. with pytest.raises(ValueError): _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) @@ -249,6 +273,7 @@ def test_reshape_layout_dynamic_shape4(): src_shape = (-1, 512) dst_shape = (-1, 256, 512) + # Splitting a dynamic dimension that is sharded. Alignment cannot be guaranteed. with pytest.raises(ValueError): _, _ = op.infer_layout((x_layout,), (dst_shape, src_shape)) @@ -269,6 +294,8 @@ def test_reshape_layout_axis_shard_twice(): dst_shape = (4, 8, 2, 384, 280) output_layout, local_dst_shape = op.infer_layout((x_layout,), (dst_shape, src_shape)) + # Dst dim 0 (4) inherits sharding from Src dim 0 (32) which is ("cp", "dp"). + # cp(0) and dp(2) both shard dim 0. expected_map = ((0, 2), -1, -1, -1, 1) # Expected output tensor map assert output_layout.tensor_map == expected_map, (f"Reshape do not change sharded axis failed. Expected " f"expected_map {expected_map} bug got {output_layout.tensor_map}") @@ -276,3 +303,313 @@ def test_reshape_layout_axis_shard_twice(): assert local_dst_shape == expected_local_dst_shape, (f"Reshape do not change sharded axis failed. Expected" f" expected_local_dst_shape {expected_local_dst_shape} got" f" {local_dst_shape}") + + +op_torch = ReshapeDistributedOp("reshape") + +def test_torch_reshape_basic_split(): + """ + Feature: PyTorch style reshape split sharded axis + Description: Test splitting a sharded dimension into multiple dimensions (preserving layout) + Expectation: Success + """ + # Mesh: (2, 4), Layout: dp(0), mp(1) + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Input Layout: [dp, mp] -> [2, 4] split + x_layout = x_layout("dp", "mp") + + # Input: [32, 64] + src_shape = (32, 64) + # Reshape to: [32, 16, 4] -> We split the second dim (64) which is sharded by 'mp' + dst_shape = (32, 16, 4) + + # Note: reshapeDistributedOp requires (dst_shape, src_shape) in extra_args + output_layout, local_dst_shape = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) + + # Logic analysis: + # Merged Input Blocks: [32 (dp)], [64 (mp)] + # Reverse Match Dst: + # 1. dst_dim=4: 64 % 4 == 0 -> remainder 16. Consumed part of 'mp' block. Map: -1 (unsharded inner part) + # 2. dst_dim=16: 16 % 16 == 0 -> remainder 1. Consumed rest of 'mp' block. Map: mp (inherits sharding) + # 3. dst_dim=32: 32 % 32 == 0 -> remainder 1. Consumed 'dp' block. Map: dp + + expected_map = (1, 0, -1) # (dp, mp, None) -> based on mesh indices: dp=1 (mesh[0]), mp=0 (mesh[1])? + # Wait, check Layout def in test file: + # base_mesh_shape = (2, 4) -> index 0 is size 2, index 1 is size 4. + # In x_layout("dp", "mp"): + # "dp" maps to mesh axis 0? Usually aliases are mapped sequentially or by dict. + # Assuming standard Layout behavior: + # If "dp" is alias[0], "mp" is alias[1]. + # map value: device_matrix_rank - 1 - axis_index. + # Typically: index 0 (dp) -> map 1, index 1 (mp) -> map 0 (if using reversed mapping common in Ascend/MindSpore tensor map). + # Let's trust the existing test convention: expected_map = (1, -1) implies 1 is dp. + + + assert output_layout.tensor_map == expected_map, \ + f"Expected {expected_map}, got {output_layout.tensor_map}" + + # Local shape: + # dim 0 (32): sharded by dp(2) -> 16 + # dim 1 (16): sharded by mp(4) -> 4 + # dim 2 (4): not sharded -> 4 + expected_local_dst_shape = [16, 4, 4] + assert local_dst_shape == expected_local_dst_shape, \ + f"Expected local shape {expected_local_dst_shape}, got {local_dst_shape}" + + +def test_torch_reshape_dynamic_target(): + """ + Feature: PyTorch style reshape with -1 in target + Description: Test passing -1 in destination shape to infer dimension automatically + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "mp") # map: (1, 0) + + src_shape = (32, 64) + # Reshape: Merge dims and flatten, then split? Or just change view. + # Let's try: [32, -1] -> implies [32, 64] + dst_shape = (32, -1) + + output_layout, local_dst_shape = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) + + # Logic: _handle_dynamic_shape converts dst to [32, 64] + expected_map = (1, 0) + assert output_layout.tensor_map == expected_map + + # Local shape should handle the resolved shape + expected_local_dst_shape = [16, 16] # 32/2, 64/4 + assert local_dst_shape == expected_local_dst_shape + + +def test_torch_reshape_flatten_unsharded(): + """ + Feature: PyTorch style reshape merge unsharded dims + Description: Merging an unsharded dimension into a sharded one is NOT fully supported if it changes sharding cuts, + but merging unsharded dims into other unsharded dims or preserving blocks works. + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Layout: dp, None, mp -> [1, -1, 0] + x_layout = x_layout("dp", "None", "mp") + + src_shape = (16, 4, 8) + # We want to reshape to (16, 32). Merging dim 1 (4) and dim 2 (8). + # Since dim 1 is -1 (unsharded) and dim 2 is mp (sharded), + # `_merge_unshared_axis` logic: + # Iterate backwards: + # Ax 2 (8, mp): New block [8], map 0. + # Ax 1 (4, -1): Merged into Ax 0? No, `_merge_unshared_axis` merges -1 into the PREVIOUS shared axis (from the left logic) + # OR it effectively accumulates size. + # In your code: + # Loop reversed. + # Ax 2 (map 0): insert(8), insert(0). cur_vol reset to 1. + # Ax 1 (map -1): cur_vol *= 4. + # Ax 0 (map 1): insert(16*4), insert(1). + # Merged Result: [64, 8] with maps [1, 0]. + # Wait, `_merge_unshared_axis` logic in your code: + # if tensor_map[axis] != -1: insert(merged_size)... + # So for (16, 4, 8) with (1, -1, 0): + # i=2: map=0. insert(8), insert(0). vol=1. + # i=1: map=-1. vol=1*4=4. + # i=0: map=1. insert(16*4), insert(1). vol=1. + # Result: Shape [64, 8], Map [1, 0]. + + # Target: (16, 32) + # Match reversed: + # dst=32. cur=8 (map 0). 8 % 32 != 0. + # ERROR condition in your code: if cur_size % shape != 0 -> raise. + # So (16, 4, 8) -> (16, 32) is IMPOSSIBLE with this logic because the cut 'mp' is at size 8 boundary. + # You can't have a dimension of size 32 where the first 8 are sharded one way and the rest another way. + + # Let's try a valid case: Reshape (16, 4, 8) -> (64, 8) + # This matches the merged blocks exactly. + dst_shape = (64, 8) + + output_layout, _ = op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) + assert output_layout.tensor_map == (1, 0) + + # Let's try Reshape (16, 4, 8) -> (16, 2, 2, 8) + # Valid split of unsharded dim + dst_shape_2 = (16, 2, 2, 8) + output_layout_2, _ = op_torch.infer_layout((x_layout,), (dst_shape_2, src_shape)) + # 8 matches 8 (map 0). + # 2 matches part of 64? + # Wait, merged block is 64 (map 1). + # 64 % 2 == 0 -> rem 32. map -1. + # 32 % 2 == 0 -> rem 16. map -1. + # 16 % 16 == 0 -> rem 1. map 1. + # Result map: (1, -1, -1, 0). + assert output_layout_2.tensor_map == (1, -1, -1, 0) + + +def test_torch_reshape_fail_missing_input_shape(): + """ + Feature: PyTorch style reshape exception + Description: Verify that failing to provide input shape raises ValueError + Expectation: ValueError + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "mp") + + dst_shape = (32, 64) + + # Intentionally omitted src_shape + with pytest.raises(ValueError, match="reshape requires output shape and input shape."): + op_torch.infer_layout((x_layout,), (dst_shape,)) + + +def test_torch_reshape_fail_mismatch_total_size(): + """ + Feature: PyTorch style reshape validation + Description: Verify shape element count mismatch validation + Expectation: ValueError + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "mp") + + src_shape = (10, 10) + dst_shape = (20, 20) # 100 != 400 + + with pytest.raises(ValueError): + op_torch.infer_layout((x_layout,), (dst_shape, src_shape)) + + +op_view = ReshapeDistributedOp("view") + +def test_view_layout_flatten_contiguous(): + """ + Feature: View operator flattening + Description: Test flattening multiple dimensions where the inner dimension is not sharded. + Common case: x.view(batch_size, -1) + Expectation: Success, preserving the sharding of the split point. + """ + # Mesh: (2, 4), Layout: dp(0), mp(1) + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Input: [batch, seq, hidden] -> [32, 16, 8] + # Layout: dp, mp, None -> Tensor Map: (1, 0, -1) + x_layout = x_layout("dp", "mp", "None") + + src_shape = (32, 16, 8) + # Target: Flatten the last two dimensions -> [32, 128] + dst_shape = (32, 128) + + + output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) + + expected_map = (1, 0) # dp, mp + assert output_layout.tensor_map == expected_map, \ + f"View flatten failed. Expected {expected_map}, got {output_layout.tensor_map}" + + expected_local_dst_shape = [16, 32] # 32/2=16, 128/4=32 + assert local_dst_shape == expected_local_dst_shape + + +def test_view_layout_unflatten_split(): + """ + Feature: View operator unflattening (inverse of flatten) + Description: Test expanding a sharded dimension into two, where the outer part keeps sharding + and the inner part becomes None (unsharded). + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Input: [32, 128] + # Layout: dp, mp -> Tensor Map: (1, 0) + x_layout = x_layout("dp", "mp") + + src_shape = (32, 128) + # Target: Unflatten to [32, 16, 8] + # We are splitting the 128 (mp) dimension into 16 and 8. + dst_shape = (32, 16, 8) + + + output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) + + expected_map = (1, 0, -1) + assert output_layout.tensor_map == expected_map, \ + f"View unflatten failed. Expected {expected_map}, got {output_layout.tensor_map}" + + # Local Shape: + # 32 (dp=2) -> 16 + # 16 (mp=4) -> 4 + # 8 (None) -> 8 + expected_local_dst_shape = [16, 4, 8] + assert local_dst_shape == expected_local_dst_shape + + +def test_view_layout_dynamic_shape_inference(): + """ + Feature: View operator with -1 inference + Description: Support x.view(-1, C) style calls similar to PyTorch + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + # Layout: dp, mp, None -> (1, 0, -1) + x_layout = x_layout("dp", "mp", "None") + + src_shape = (8, 4, 16) # Total 512 + # Target: [8, -1]. Should infer to [8, 64] + dst_shape = (8, -1) + + output_layout, local_dst_shape = op_view.infer_layout((x_layout,), (dst_shape, src_shape)) + + # Logic: + # Merged input: [8, 64] with Map [1, 0] (since 4(mp) and 16(None) merge into 64(mp)) + # Target inferred: [8, 64] + # Match: Exact match. + + expected_map = (1, 0) + assert output_layout.tensor_map == expected_map + + # Local: 8/2=4, 64/4=16 + expected_local_dst_shape = [4, 16] + assert local_dst_shape == expected_local_dst_shape + + +def test_view_layout_fail_shape_mismatch(): + """ + Feature: View operator validation + Description: Ensure view raises error if total elements don't match (PyTorch behavior) + Expectation: ValueError + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "mp") + + src_shape = (10, 20) + dst_shape = (10, 30) # 200 != 300 + + with pytest.raises(ValueError, match="total elements number"): + op_view.infer_layout((x_layout,), (dst_shape, src_shape)) diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_softmax.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_softmax.py index 88792e3..fba5173 100644 --- a/tests/mindspore/ut/parallel_ops_infer/test_parallel_softmax.py +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_softmax.py @@ -18,6 +18,7 @@ import pytest from hyper_parallel import Layout from hyper_parallel.core.shard.ops.parallel_activation_with_axis import ActivationWithAxisDistributedOp + op = ActivationWithAxisDistributedOp("Softmax") @@ -59,3 +60,250 @@ def test_softmax_layout_data_parallel_value_failed(): with pytest.raises(ValueError): _ = op.infer_layout((x_layout,), (0,)) + +op = ActivationWithAxisDistributedOp("softmax") + + +def test_torch_softmax_data_parallel_success(): + """ + Feature: Softmax compatible with PyTorch (dim arg) + Description: Standard Data Parallel scenario (Batch is sharded, Softmax on Feature dim) + Expectation: Success, output layout equals input layout + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + output_layout = op.infer_layout((x_layout,), (1,)) + + expected_map = x_layout.to_dict()["tensor_map"] + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"DP test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_torch_softmax_negative_dim_success(): + """ + Feature: Softmax compatible with PyTorch + Description: Test negative index support (dim=-1) in Data Parallel + Expectation: Success + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) + base_rank_list = list(range(8)) + + # [Batch, Hidden] -> ("dp", "None") + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + output_layout = op.infer_layout((x_layout,), (-1,)) + + expected_map = x_layout.to_dict()["tensor_map"] + assert output_layout.to_dict()["tensor_map"] == expected_map + + +def test_torch_softmax_sharded_dim_failure(): + """ + Feature: Softmax distributed check + Description: Attempting to compute Softmax on a sharded dimension + Expectation: Raise ValueError + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) + base_rank_list = list(range(8)) + + # [Batch, Hidden] -> ("dp", "None") + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + with pytest.raises(ValueError) as excinfo: + _ = op.infer_layout((x_layout,), (0,)) + + assert "is sharded" in str(excinfo.value) + assert "requires the reduction axis to be un-sharded" in str(excinfo.value) + + +def test_torch_softmax_model_parallel_failure(): + """ + Feature: Softmax distributed check + Description: Model Parallel scenario where the feature dimension is sharded + Expectation: Raise ValueError when computing softmax on the MP axis + """ + base_mesh_shape = (8,) + base_alias_name = ("mp",) + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("None", "mp") + + with pytest.raises(ValueError) as excinfo: + _ = op.infer_layout((x_layout,), (-1,)) + + assert "is sharded" in str(excinfo.value) + + +def test_torch_softmax_model_parallel_success_on_other_axis(): + """ + Feature: Softmax distributed check + Description: Model Parallel scenario, but computing softmax on the non-sharded axis + Expectation: Success + """ + base_mesh_shape = (8,) + base_alias_name = ("mp",) + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("None", "mp") + + output_layout = op.infer_layout((x_layout,), (0,)) + + expected_map = x_layout.to_dict()["tensor_map"] + assert output_layout.to_dict()["tensor_map"] == expected_map + + +def test_torch_softmax_input_consistency_failure(): + """ + Feature: Layout consistency check + Description: Inputs have different layouts + Expectation: Raise ValueError + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) + base_rank_list = list(range(8)) + + layout1 = Layout(base_mesh_shape, base_alias_name, base_rank_list)("dp", "None") + layout2 = Layout(base_mesh_shape, base_alias_name, base_rank_list)("None", "dp") + + with pytest.raises(ValueError) as excinfo: + _ = op.infer_layout((layout1, layout2), (1,)) + + assert "requires all tensor inputs to have the same layout" in str(excinfo.value) + + +def test_torch_softmax_multi_axis_mesh_success(): + """ + Feature: Softmax with multi-axis device mesh + Description: 2D mesh (dp+mp), Softmax applied to unsharded dimension + Expectation: Success, output layout matches input + """ + + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None", "mp") + + output_layout = op.infer_layout((x_layout,), (1,)) + + expected_map = x_layout.to_dict()["tensor_map"] + assert output_layout.to_dict()["tensor_map"] == expected_map, ( + f"Multi-axis mesh test failed. Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + ) + + +def test_torch_softmax_interleaved_parallel_success(): + """ + Feature: Softmax with interleaved parallel + Description: Softmax applied to unsharded dimension under interleaved parallel layout + Expectation: Success + """ + base_mesh_shape = (2, 2) + base_alias_name = ("dp", "interleaved_parallel") + base_rank_list = list(range(4)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "None") + + output_layout = op.infer_layout((x_layout,), (1,)) + + assert output_layout.to_dict()["tensor_map"] == x_layout.to_dict()["tensor_map"] + assert output_layout.to_dict()["interleaved_parallel"] is True + +def test_torch_softmax_partial_layout_success(): + """ + Feature: Softmax with partial (reduce) layout + Description: Softmax applied to unreduced/unsharded dimension under Partial layout + Expectation: Success + """ + base_mesh_shape = (2, 4) + base_alias_name = ("dp", "mp") + base_rank_list = list(range(8)) + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("None", "mp") # tensor_map=(-1, 0) + # Set partial sum (reduction) on dp axis + x_layout.set_partial_by_dev_axis("dp", "sum") + + # Softmax applied to Batch dimension (dim=0) -> tensor_map=-1 (unsharded, only partial), valid + output_layout = op.infer_layout((x_layout,), (0,)) + + assert output_layout.to_dict()["tensor_map"] == x_layout.to_dict()["tensor_map"] + assert output_layout.partial == x_layout.partial # Partial status preserved + + +def test_torch_softmax_non_contiguous_rank_list(): + """ + Feature: Softmax with non-contiguous rank list + Description: Non-contiguous rank_list (simulating heterogeneous cluster), verify layout compatibility + Expectation: Success if dim is unsharded + """ + base_mesh_shape = (2, 2) + base_alias_name = ("dp", "mp") + base_rank_list = [3, 1, 0, 2] # Non-contiguous rank_list + + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("None", "mp") # tensor_map=(-1, 0) + + # Softmax applied to Batch dimension (dim=0) -> unsharded, valid + output_layout = op.infer_layout((x_layout,), (0,)) + + assert output_layout.rank_list == tuple(base_rank_list) # rank_list preserved + assert output_layout.to_dict()["tensor_map"] == x_layout.to_dict()["tensor_map"] + +def test_torch_softmax_nested_tensor_map_success(): + """ + Feature: Softmax with nested tensor_map (interleaved parallel tuple) + Description: tensor_map contains nested tuple (interleaved parallel multi-axis binding) + Expectation: Success if target dim is unsharded + """ + base_mesh_shape = (2, 2, 2) + base_alias_name = ("dp", "sp", "interleaved_parallel") + base_rank_list = list(range(8)) + + # 3D Tensor -> (("dp", "interleaved_parallel"), "None", "sp") + # tensor_map calculation: + # ("dp", "interleaved_parallel") -> (2, 0) (dp: 3-1-0=2; interleaved_parallel: 3-1-2=0) + # "None" -> -1 + # "sp" -> 3-1-1=1 + # Final tensor_map = ((2, 0), -1, 1) + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout(("dp", "interleaved_parallel"), "None", "sp") + + # Softmax applied to SeqLen dimension (dim=1) -> -1 (unsharded), valid + output_layout = op.infer_layout((x_layout,), (1,)) + + assert output_layout.to_dict()["tensor_map"] == ((2, 0), -1, 1) + assert output_layout.to_dict()["interleaved_parallel"] is True + +def test_torch_softmax_multi_input_same_layout_success(): + """ + Feature: Softmax with multiple inputs (same layout) + Description: Multiple inputs with consistent layout, verify validity + Expectation: Success + """ + base_mesh_shape = (8,) + base_alias_name = ("dp",) + base_rank_list = list(range(8)) + + # Two input layouts are exactly the same + layout1 = Layout(base_mesh_shape, base_alias_name, base_rank_list)("dp", "None") + layout2 = Layout(base_mesh_shape, base_alias_name, base_rank_list)("dp", "None") + + # Softmax applied to unsharded dimension (dim=1) + output_layout = op.infer_layout((layout1, layout2), (1,)) + + assert output_layout.to_dict()["tensor_map"] == layout1.to_dict()["tensor_map"] diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose.py index 65af600..b106c50 100644 --- a/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose.py +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose.py @@ -18,6 +18,7 @@ from hyper_parallel import Layout from hyper_parallel.core.shard.ops.parallel_transpose import TransposeDistributedOp + class TestTransposeDistributedOp: """Test cases for TransposeDistributedOp.infer_layout method.""" @@ -145,3 +146,267 @@ class TestTransposeDistributedOp: assert result_layout.tensor_map == (1, 0, 2) assert result_layout.mesh_shape == input_layout.mesh_shape assert result_layout.alias_name == input_layout.alias_name + + +class TestPermuteDistributedOp: + """Test cases for Permute Distributed Operator (mapped to TransposeDistributedOp).""" + + def test_permute_3d_basic(self): + """ + Feature: Permute distributed operator basic functionality (3D) + Description: Test permute operation with standard 3D layout (PyTorch style dims: 2, 0, 1) + Expectation: Output layout tensor map correctly permuted according to axis + """ + # Layout: [TP, CP, DP] -> (2, 2, 2) + dev_mat = (2, 2, 2) + aliases = ("tp", "cp", "dp") + rank_list = list(range(math.prod(dev_mat))) + + # Input Tensor Map: ("tp", "cp", "dp") -> corresponds to indices 2, 1, 0 + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("tp", "cp", "dp") + + # Init op with 'permute' name + op = TransposeDistributedOp('permute') + + # Pytorch: x.permute(2, 0, 1) -> moves dim 2 to pos 0, dim 0 to pos 1... + # Expectation: Output Tensor Map should be ("dp", "tp", "cp") + # Note: We pass ((2, 0, 1),) because the implementation does axis = extra_args[0] + result_layout = op.infer_layout([layout_with_tensor_map], ((2, 0, 1),)) + + # Check Logic: + # in_map = [0:tp, 1:cp, 2:dp] + # perm = (2, 0, 1) + # out[0] = in[2] = dp (2) + # out[1] = in[0] = tp (0) + # out[2] = in[1] = cp (1) + assert result_layout.tensor_map == (0, 2, 1) + assert result_layout.mesh_shape == input_layout.mesh_shape + assert result_layout.alias_name == input_layout.alias_name + + def test_permute_4d_complex(self): + """ + Feature: Permute distributed operator with 4D tensor + Description: Test permute operation with 4D layout to verify higher dimension support + Expectation: Output layout tensor map correctly follows the given permutation + """ + dev_mat = (2, 2, 2, 2) + aliases = ("dp", "mp", "cp", "tp") # 0, 1, 2, 3 + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + # Input map: 3, 2, 1, 0 + layout_with_tensor_map = input_layout("dp", "mp", "cp", "tp") + + op = TransposeDistributedOp('permute') + + # Permute: (0, 3, 1, 2) -> (dp, tp, mp, cp) + result_layout = op.infer_layout([layout_with_tensor_map], ((0, 3, 1, 2),)) + + assert result_layout.tensor_map == (3, 0, 2, 1) + + def test_permute_no_change(self): + """ + Feature: Permute identity check + Description: Test permute with order that matches input (no change) + Expectation: Output tensor map equals input tensor map + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp")#(1,0) + + op = TransposeDistributedOp('permute') + # Pytorch: x.permute(0, 1) + result_layout = op.infer_layout([layout_with_tensor_map], ((0, 1),)) + + assert result_layout.tensor_map == (1, 0) + + def test_permute_shape_mismatch(self): + """ + Feature: Permute error handling (Shape Mismatch) + Description: Input has 3 dims, but permute args has 2 dims + Expectation: ValueError + """ + dev_mat = (2, 2, 2) + aliases = ("dp", "cp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "cp", "tp") + + op = TransposeDistributedOp('permute') + + with pytest.raises(ValueError, match="Input tensor shape and permutation must have the same size"): + op.infer_layout([layout_with_tensor_map], ((0, 1),)) + + def test_permute_duplicate_dims(self): + """ + Feature: Permute error handling (Duplicate Dims) + Description: Permute arguments contain duplicate dimensions + Expectation: ValueError + """ + dev_mat = (2, 2, 2) + aliases = ("dp", "cp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "cp", "tp") + + op = TransposeDistributedOp('permute') + + # Invalid: dimension 0 is used twice + with pytest.raises(ValueError, match="Invalid permutation"): + op.infer_layout([layout_with_tensor_map], ((0, 0, 1),)) + + def test_permute_out_of_bounds(self): + """ + Feature: Permute error handling (Out of Bounds) + Description: Permute arguments contain index larger than rank + Expectation: ValueError + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('permute') + + # Invalid: dimension 2 does not exist for 2D tensor + with pytest.raises(ValueError, match="Invalid permutation"): + op.infer_layout([layout_with_tensor_map], ((2, 1),)) + + def test_permute_negative_index_error(self): + """ + Feature: Permute error handling (Negative Index) + Description: Current implementation does not support negative indices in backend + Expectation: ValueError + """ + # Note: While PyTorch supports negative indices (e.g. permute(0, -1)), + # the provided backend code explicitly checks 'if v < 0' and raises error. + # This test ensures the backend safety check works. + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('permute') + + with pytest.raises(ValueError, match="Invalid permutation"): + op.infer_layout([layout_with_tensor_map], ((0, -1),)) + + + +class TesttransposeDistributedOp: + """Test cases for Torch-style Transpose Distributed Operator (swapping two dimensions).""" + + def test_torch_transpose_basic(self): + """ + Feature: Torch Transpose distributed operator basic functionality + Description: Test transpose(dim0, dim1) with 2D layout + Expectation: Output layout tensor map has dimensions swapped + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + # tensor_map: dp(1), tp(0) -> (1, 0) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('transpose') + # swap 0 and 1 -> (0, 1) + # extra_args should be passed as (dim0, dim1) + result_layout = op.infer_layout([layout_with_tensor_map], (0, 1)) + + assert result_layout.tensor_map == (0, 1) + assert result_layout.mesh_shape == input_layout.mesh_shape + + def test_torch_transpose_3d_complex(self): + """ + Feature: Torch Transpose with 3D tensor + Description: Test transpose(dim0, dim1) on 3D layout + Expectation: Only specified dimensions are swapped in tensor map + """ + dev_mat = (2, 2, 2) + aliases = ("dp", "cp", "tp") # indices 0, 1, 2. map vals: 2, 1, 0 + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + # dp, cp, tp -> (2, 1, 0) + layout_with_tensor_map = input_layout("dp", "cp", "tp") + + op = TransposeDistributedOp('transpose') + # transpose(0, 2) -> swap first and last + # input: (2, 1, 0). swap 0 and 2 -> (0, 1, 2) + result_layout = op.infer_layout([layout_with_tensor_map], (0, 2)) + + assert result_layout.tensor_map == (0, 1, 2) + + def test_torch_transpose_negative_indices(self): + """ + Feature: Torch Transpose negative indices + Description: Test transpose with negative dimension indices + Expectation: Negative indices are correctly interpreted and dimensions swapped + """ + dev_mat = (2, 2, 2) + aliases = ("dp", "cp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + # (2, 1, 0) + layout_with_tensor_map = input_layout("dp", "cp", "tp") + + op = TransposeDistributedOp('transpose') + # transpose(0, -1) -> swap 0 and 2 + result_layout = op.infer_layout([layout_with_tensor_map], (0, -1)) + + assert result_layout.tensor_map == (0, 1, 2) + + def test_torch_transpose_arg_count_error(self): + """ + Feature: Torch Transpose error handling + Description: Test with incorrect number of arguments + Expectation: ValueError + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('transpose') + + with pytest.raises(ValueError, match="expected \(dim0, dim1\)"): + op.infer_layout([layout_with_tensor_map], (0,)) + + def test_torch_transpose_type_error(self): + """ + Feature: Torch Transpose type validation + Description: Test with non-integer dimensions + Expectation: ValueError + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('transpose') + + with pytest.raises(ValueError, match="Dimensions must be integers"): + op.infer_layout([layout_with_tensor_map], (0, "1")) + + def test_torch_transpose_out_of_bounds(self): + """ + Feature: Torch Transpose bounds checking + Description: Test with dimension index out of bounds + Expectation: ValueError + """ + dev_mat = (2, 2) + aliases = ("dp", "tp") + rank_list = list(range(math.prod(dev_mat))) + input_layout = Layout(mesh_shape=dev_mat, alias_name=aliases, rank_list=rank_list) + layout_with_tensor_map = input_layout("dp", "tp") + + op = TransposeDistributedOp('transpose') + + with pytest.raises(ValueError, match="Transpose dimensions out of bounds"): + op.infer_layout([layout_with_tensor_map], (0, 2)) \ No newline at end of file diff --git a/tests/torch/shard/ops/parallel_op_cat.py b/tests/torch/shard/ops/parallel_op_cat.py new file mode 100644 index 0000000..6cc6361 --- /dev/null +++ b/tests/torch/shard/ops/parallel_op_cat.py @@ -0,0 +1,92 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test torch dtensor with distributed cat""" +import numpy as np +import torch +import torch.nn.functional as F +from hyper_parallel import DTensor, Layout +from tests.torch.utils import init_dist +from tests.torch.shard.utils import local_to_global, global_to_local + +# Generate input data using numpy at file header +np.random.seed(42) +vocab_size = 128 +embed_dim = 32 +batch_size = 8 +seq_len = 16 + +# Input indices: (Batch, Seq) +standalone_input_np = np.random.randint(0, vocab_size, (batch_size, seq_len)).astype(np.int64) +# Embedding Table: (Vocab, Embed_Dim) +standalone_weight_np = np.random.randn(vocab_size, embed_dim).astype(np.float32) + +def test_distributed_embedding_layout_inference(): + """ + Feature: dtensor + torch.nn.functional.embedding layout inference + Description: + ▪ Test layout inference for embedding op in distributed setting. + ▪ Scenario: Column Parallel Embedding (splitting embedding dimension). + ▪ Logic: + - Input is sharded on Batch dim ("dp"). + - Weight is sharded on Embedding dim ("tp"). + - Output should be sharded on (Batch, Sequence, Embedding) -> ("dp", None, "tp"). + Expectation: Success. + """ + init_dist() + + # Standalone (single-device) execution for ground truth + standalone_input = torch.from_numpy(standalone_input_np).npu() + standalone_weight = torch.from_numpy(standalone_weight_np).npu() + + # torch.nn.functional.embedding(input, weight) + standalone_output = F.embedding(standalone_input, standalone_weight) + + # Distributed setup + # Mesh: (2, 4), Alias: ("dp", "tp") + layout = Layout((2, 4), ("dp", "tp")) + + # 1. Define Input Layout: Shard on Batch (dim 0), Keep Seq (dim 1) + # Map: ("dp", "None") + input_layout = layout("dp", "None") + + # 2. Define Weight Layout: Keep Vocab (dim 0), Shard Embed (dim 1) + # Map: ("None", "tp") -> This simulates Column Parallelism commonly used in LLMs + weight_layout = layout("None", "tp") + + # 3. Define Expected Output Layout: + # Logic in parallel_embedding.py: output_map = input_map + (weight_map[-1],) + # Result: ("dp", "None") + ("tp",) -> ("dp", "None", "tp") + expected_output_layout = layout("dp", "None", "tp") + + # Convert global view to local shards + dist_input = global_to_local(standalone_input, input_layout) + dist_weight = global_to_local(standalone_weight, weight_layout) + + # Execute distributed op + # Note: Pass strictly as positional args to ensure compatibility with _op_dispatch + dist_output = F.embedding(dist_input, dist_weight) + + # Layout correctness check + assert dist_output.layout == expected_output_layout, \ + f"Embedding output layout mismatch.\nExpected: {expected_output_layout}\nGot: {dist_output.layout}" + + # Gather distributed results back to global view + gathered_output = local_to_global(dist_output) + + # Value correctness check + assert torch.allclose( + standalone_output, gathered_output, atol=1e-5 + ), "Embedding output values mismatch between standalone and distributed" + diff --git a/tests/torch/shard/ops/parallel_op_embedding.py b/tests/torch/shard/ops/parallel_op_embedding.py new file mode 100644 index 0000000..9e2c4f9 --- /dev/null +++ b/tests/torch/shard/ops/parallel_op_embedding.py @@ -0,0 +1,92 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test torch dtensor with distributed embedding""" +import numpy as np +import torch +import torch.nn.functional as F +from hyper_parallel import DTensor, Layout +from tests.torch.utils import init_dist +from tests.torch.shard.utils import local_to_global, global_to_local + +# Generate input data using numpy at file header +np.random.seed(42) +vocab_size = 128 +embed_dim = 32 +batch_size = 8 +seq_len = 16 + +# Input indices: (Batch, Seq) +standalone_input_np = np.random.randint(0, vocab_size, (batch_size, seq_len)).astype(np.int64) +# Embedding Table: (Vocab, Embed_Dim) +standalone_weight_np = np.random.randn(vocab_size, embed_dim).astype(np.float32) + +def test_distributed_embedding_layout_inference(): + """ + Feature: dtensor + torch.nn.functional.embedding layout inference + Description: + ▪ Test layout inference for embedding op in distributed setting. + ▪ Scenario: Column Parallel Embedding (splitting embedding dimension). + ▪ Logic: + - Input is sharded on Batch dim ("dp"). + - Weight is sharded on Embedding dim ("tp"). + - Output should be sharded on (Batch, Sequence, Embedding) -> ("dp", None, "tp"). + Expectation: Success. + """ + init_dist() + + # Standalone (single-device) execution for ground truth + standalone_input = torch.from_numpy(standalone_input_np).npu() + standalone_weight = torch.from_numpy(standalone_weight_np).npu() + + # torch.nn.functional.embedding(input, weight) + standalone_output = F.embedding(standalone_input, standalone_weight) + + # Distributed setup + # Mesh: (2, 4), Alias: ("dp", "tp") + layout = Layout((2, 4), ("dp", "tp")) + + # 1. Define Input Layout: Shard on Batch (dim 0), Keep Seq (dim 1) + # Map: ("dp", "None") + input_layout = layout("dp", "None") + + # 2. Define Weight Layout: Keep Vocab (dim 0), Shard Embed (dim 1) + # Map: ("None", "tp") -> This simulates Column Parallelism commonly used in LLMs + weight_layout = layout("None", "tp") + + # 3. Define Expected Output Layout: + # Logic in parallel_embedding.py: output_map = input_map + (weight_map[-1],) + # Result: ("dp", "None") + ("tp",) -> ("dp", "None", "tp") + expected_output_layout = layout("dp", "None", "tp") + + # Convert global view to local shards + dist_input = global_to_local(standalone_input, input_layout) + dist_weight = global_to_local(standalone_weight, weight_layout) + + # Execute distributed op + # Note: Pass strictly as positional args to ensure compatibility with _op_dispatch + dist_output = F.embedding(dist_input, dist_weight) + + # Layout correctness check + assert dist_output.layout == expected_output_layout, \ + f"Embedding output layout mismatch.\nExpected: {expected_output_layout}\nGot: {dist_output.layout}" + + # Gather distributed results back to global view + gathered_output = local_to_global(dist_output) + + # Value correctness check + assert torch.allclose( + standalone_output, gathered_output, atol=1e-5 + ), "Embedding output values mismatch between standalone and distributed" + diff --git a/tests/torch/shard/ops/parallel_op_reshape_view.py b/tests/torch/shard/ops/parallel_op_reshape_view.py new file mode 100644 index 0000000..dea1d8c --- /dev/null +++ b/tests/torch/shard/ops/parallel_op_reshape_view.py @@ -0,0 +1,154 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test torch dtensor with distributed reshape and view""" +import numpy as np +import torch +from hyper_parallel import Layout +from tests.torch.utils import init_dist +from tests.torch.shard.utils import local_to_global, global_to_local + +# Generate input data using numpy at file header +np.random.seed(42) +# Shape: [Batch=8, Channel=4, Height=4, Width=4] -> Total 512 elements +standalone_input_np = np.random.randn(8, 4, 4, 4).astype(np.float32) + +def test_distributed_reshape_layout_inference(): + """ + Feature: dtensor + torch.reshape layout inference + Description: + 1. Test 'flatten' scenario: [B, C, H, W] -> [B, C*H*W] + 2. Test 'expand' scenario: [B, C*H*W] -> [B, C, H, W] + 3. Verify layout consistency and value correctness. + Expectation: Success. + """ + init_dist() + + # Setup Standalone + standalone_input = torch.from_numpy(standalone_input_np).npu() + + # --- Case 1: Flatten dimensions (preserving sharded dim) --- + # Target shape: [8, 64] + standalone_output_flat = standalone_input.reshape(8, 64) + + # Distributed Setup + # Mesh: (2, 4), dp splits dim 0, tp is unused (None) + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("dp", "None", "None", "None") # Shard Batch dim + + dist_input = global_to_local(standalone_input, x_layout) + + # Perform Reshape + # FIX: Pass shape as a tuple explicitly to bypass OpDispatcher bug + # dist_input.reshape(8, 64) -> dist_input.reshape((8, 64)) + dist_output_flat = dist_input.reshape((8, 64)) + + # Check Layout + # Expectation: Dim 0 is still sharded ("dp"), Dim 1 is replicate ("None") + expected_flat_layout = layout("dp", "None") + assert dist_output_flat.layout == expected_flat_layout, \ + f"Reshape flat layout mismatch. Got {dist_output_flat.layout}, expected {expected_flat_layout}" + + # Check Values + gathered_flat = local_to_global(dist_output_flat) + assert torch.allclose(standalone_output_flat, gathered_flat, atol=1e-5), \ + "Reshape values mismatch between standalone and distributed" + + # --- Case 2: Expand dimensions (using -1 inference) --- + # Input is the flattened distributed tensor from Case 1 + # Target shape: [8, 4, 4, 4] + # FIX: Pass shape as a tuple explicitly + dist_output_expand = dist_output_flat.reshape((8, 4, 4, -1)) + + # Check Layout + # Expectation: Dim 0 sharded ("dp"), others replicate + expected_expand_layout = x_layout + assert dist_output_expand.layout == expected_expand_layout, \ + "Reshape expand layout mismatch" + + # Check Values + gathered_expand = local_to_global(dist_output_expand) + assert torch.allclose(standalone_input, gathered_expand, atol=1e-5), \ + "Reshape expand values mismatch" + + +def test_distributed_view_layout_inference(): + """ + Feature: dtensor + torch.view layout inference + Description: + 1. Test view with tuple argument input style. + 2. Test view on a tensor where sharded dimension is preserved but reshaped + (as long as total size / shards is integer). + Expectation: Success. + """ + init_dist() + + # Setup Standalone + standalone_input = torch.from_numpy(standalone_input_np).npu() + + # Scenario: Input [8, 4, 4, 4], Layout: Shard dim 1 (Channel) + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("None", "tp", "None", "None") # Shard Channel + + dist_input = global_to_local(standalone_input, x_layout) + + # --- Operation: View --- + # Merge last two dims: [8, 4, 4, 4] -> [8, 4, 16] + target_shape = (8, 4, 16) + + # Perform View (PyTorch style: tuple arg) + dist_view = dist_input.view(target_shape) + + # Standalone reference + standalone_view = standalone_input.view(target_shape) + + # Check Layout + expected_view_layout = layout("None", "tp", "None") + + assert dist_view.layout == expected_view_layout, \ + f"View layout mismatch. Got {dist_view.layout}, expected {expected_view_layout}" + + # Check Values + gathered_view = local_to_global(dist_view) + assert torch.allclose(standalone_view, gathered_view, atol=1e-5), \ + "View values mismatch between standalone and distributed" + + +def test_distributed_reshape_fail_mismatch(): + """ + Feature: dtensor + reshape error handling + Description: + Attempt to reshape in a way that breaks existing sharding constraints + or mismatches element count. + Expectation: Raise ValueError. + """ + init_dist() + + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("dp", "None") # [8, 64] -> split 8 + + input_np = np.random.randn(8, 64).astype(np.float32) + standalone_input = torch.from_numpy(input_np).npu() + dist_input = global_to_local(standalone_input, x_layout) + + # Invalid total element count + try: + # Pass as tuple to ensuring we are testing element mismatch, not argument drop bug + dist_input.reshape((8, 32)) + assert False, "Should raise ValueError for element count mismatch" + except ValueError as e: + # Expected error + pass + except RuntimeError: + pass diff --git a/tests/torch/shard/ops/parallel_op_softmax.py b/tests/torch/shard/ops/parallel_op_softmax.py new file mode 100644 index 0000000..10591bb --- /dev/null +++ b/tests/torch/shard/ops/parallel_op_softmax.py @@ -0,0 +1,103 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test torch dtensor with distributed softmax""" +import numpy as np +import torch +from hyper_parallel import DTensor, Layout +from tests.torch.utils import init_dist +from tests.torch.shard.utils import local_to_global, global_to_local + +# Generate input data using numpy at file header +np.random.seed(42) +# Shape (8, 16), float32 +standalone_input_np = np.random.randn(8, 16).astype(np.float32) + +def test_distributed_softmax_layout_inference(): + """ + Feature: dtensor + torch.softmax layout inference + Description: + ▪ Test layout inference for torch.softmax in distributed setting. + ▪ Ensure that: + a) The output layout matches input layout (ActivationWithAxis logic). + b) Results are consistent between standalone and distributed. + c) Sharding occurs on non-softmax dimensions. + Expectation: Success. + """ + init_dist() + dim = -1 # Softmax along the last dimension (dim=1) + + # Standalone (single-device) execution for ground truth + standalone_input = torch.from_numpy(standalone_input_np).npu() + # Using positional argument for dim to ensure it is captured by _op_dispatch + standalone_output = torch.softmax(standalone_input, dim) + + # Distributed setup + # Mesh: (2, 4), Alias: ("dp", "tp") + layout = Layout((2, 4), ("dp", "tp")) + + # Strategy: Shard on dim=0 ("dp"), Keep dim=1 ("None"). + # Softmax is performed on dim=1, so dim=1 must NOT be sharded. + x_layout = layout("dp", "None") + + # Convert global view to local shard + dist_input = global_to_local(standalone_input, x_layout) + + # Execute distributed op + # 【Fix】: Must use positional argument for 'dim' because _op_dispatch.py + # does not extract kwargs into extra_args for layout inference. + # Wrong: torch.softmax(dist_input, dim=dim) + # Right: torch.softmax(dist_input, dim) + dist_output = torch.softmax(dist_input, dim) + + # Layout correctness check + assert dist_output.layout == x_layout, "Torch Softmax: output layout must match input layout" + + # Gather distributed results back to global view + gathered_output = local_to_global(dist_output) + + # Value correctness check + assert torch.allclose( + standalone_output, gathered_output, atol=1e-5 + ), "Softmax output mismatch between standalone and distributed" + + +def test_distributed_softmax_sharded_dim_error(): + """ + Feature: dtensor + torch.softmax error on sharded dim + Description: + ▪ Attempt to perform softmax along a sharded dimension. + ▪ According to ActivationWithAxisDistributedOp.check_layout, this must fail. + Expectation: Raise ValueError. + """ + init_dist() + dim = 1 # Softmax along dim 1 + + layout = Layout((2, 4), ("dp", "tp")) + + # Strategy: Keep dim=0 ("dp"), Shard dim=1 ("tp"). + # We will try to Softmax on dim=1, which is sharded. + x_layout = layout("dp", "tp") + + standalone_input = torch.from_numpy(standalone_input_np).npu() + dist_input = global_to_local(standalone_input, x_layout) + + try: + # 【Fix】: Use positional argument for 'dim' here as well + torch.softmax(dist_input, dim) + assert False, "Expected ValueError when performing softmax along sharded dimension" + except ValueError as e: + # Verify the error message format from ActivationWithAxisDistributedOp + assert "requires the reduction axis to be un-sharded" in str(e), \ + f"Unexpected error message: {str(e)}" diff --git a/tests/torch/shard/ops/perallel_op_transpose_permute.py b/tests/torch/shard/ops/perallel_op_transpose_permute.py new file mode 100644 index 0000000..e45a86a --- /dev/null +++ b/tests/torch/shard/ops/perallel_op_transpose_permute.py @@ -0,0 +1,159 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test torch dtensor with distributed transpose and permute""" + +import numpy as np +import torch +from hyper_parallel import DTensor, Layout +from tests.torch.utils import init_dist +from tests.torch.shard.utils import local_to_global, global_to_local + +# Generate input data using numpy at file header +# Shape: (8, 16, 4) - distinct dimensions to easily detect swapping errors +np.random.seed(42) +standalone_input_np = np.random.randn(8, 16, 4).astype(np.float32) + +def test_distributed_permute_layout_inference(): + """ + Feature: dtensor + torch.permute layout inference + Description: + ▪ Test layout inference for torch.permute in distributed setting. + ▪ Scenario: + - Input Layout: ("dp", "tp", "None") on a (8, 16, 4) tensor. + - Permute dims: (2, 0, 1). + ▪ Ensure that: + a) The output layout correctly reorders the tensor map. + b) The data values match standalone execution. + Expectation: Success. + """ + init_dist() + + # Permutation order + dims = (2, 0, 1) + + # Standalone (single-device global view) + standalone_input = torch.from_numpy(standalone_input_np).npu() + standalone_output = torch.permute(standalone_input, dims) + + # Distributed setup + # Mesh: 2x4 (8 devices) + # Layout strategy: dim0 -> dp (shard), dim1 -> tp (shard), dim2 -> None (replicate) + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("dp", "tp", "None") + + dist_input = global_to_local(standalone_input, x_layout) + + # Execute distributed op + dist_output = torch.permute(dist_input, dims) + + # Layout correctness check + # Input map: (dp, tp, None). Permute (2, 0, 1) -> Output map should be (None, dp, tp) + expected_layout = layout("None", "dp", "tp") + assert dist_output.layout == expected_layout, \ + f"Torch permute: layout mismatch. Expected {expected_layout.to_dict()}, got {dist_output.layout.to_dict()}" + + # Gather distributed results back to global view + gathered_output = local_to_global(dist_output) + + # Data correctness check + assert torch.allclose( + standalone_output, gathered_output, atol=1e-5 + ), "Permute values mismatch between standalone and distributed" + + assert gathered_output.shape == standalone_output.shape, \ + "Permute shape mismatch between standalone and distributed" + + +def test_distributed_transpose_layout_inference(): + """ + Feature: dtensor + torch.transpose layout inference + Description: + ▪ Test layout inference for torch.transpose in distributed setting. + ▪ Scenario: + - Input Layout: ("dp", "tp", "None"). + - Transpose dims: (1, 2). Swapping a sharded dim with a replicated dim. + ▪ Ensure that: + a) The output layout correctly swaps the tensor map entries. + b) The data values match standalone execution. + Expectation: Success. + """ + init_dist() + + # Transpose dimensions + dim0, dim1 = 1, 2 + + # Standalone (single-device global view) + standalone_input = torch.from_numpy(standalone_input_np).npu() + standalone_output = torch.transpose(standalone_input, dim0, dim1) + + # Distributed setup + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("dp", "tp", "None") + + dist_input = global_to_local(standalone_input, x_layout) + + # Execute distributed op + dist_output = torch.transpose(dist_input, dim0, dim1) + + # Layout correctness check + # Input map: (dp, tp, None). Swap 1 and 2 -> Output map should be (dp, None, tp) + expected_layout = layout("dp", "None", "tp") + + assert dist_output.layout == expected_layout, \ + f"Torch transpose: layout mismatch. Expected {expected_layout.to_dict()}, got {dist_output.layout.to_dict()}" + + # Gather distributed results back to global view + gathered_output = local_to_global(dist_output) + + # Data correctness check + assert torch.allclose( + standalone_output, gathered_output, atol=1e-5 + ), "Transpose values mismatch between standalone and distributed" + + +def test_distributed_transpose_negative_dim(): + """ + Feature: dtensor + torch.transpose with negative dimensions + Description: + ▪ Test layout inference for torch.transpose using negative indices. + ▪ Scenario: Transpose dim 0 and dim -1 (last dim). + Expectation: Success. + """ + init_dist() + + dim0, dim1 = 0, -1 + + # Standalone + standalone_input = torch.from_numpy(standalone_input_np).npu() + standalone_output = torch.transpose(standalone_input, dim0, dim1) + + # Distributed setup + layout = Layout((2, 4), ("dp", "tp")) + x_layout = layout("dp", "tp", "None") + + dist_input = global_to_local(standalone_input, x_layout) + dist_output = torch.transpose(dist_input, dim0, dim1) + + # Layout correctness check + # Input map: (dp, tp, None). Swap 0 and -1(2) -> Output map should be (None, tp, dp) + expected_layout = layout("None", "tp", "dp") + + assert dist_output.layout == expected_layout, \ + f"Torch transpose (neg dim): layout mismatch. Expected {expected_layout.to_dict()}, got {dist_output.layout.to_dict()}" + + # Verify data + gathered_output = local_to_global(dist_output) + assert torch.allclose(standalone_output, gathered_output, atol=1e-5), \ + "Transpose (neg dim) values mismatch" diff --git a/tests/torch/shard/ops/test_parallel_op_cat.py b/tests/torch/shard/ops/test_parallel_op_cat.py new file mode 100644 index 0000000..e3f8531 --- /dev/null +++ b/tests/torch/shard/ops/test_parallel_op_cat.py @@ -0,0 +1,30 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test base dtensor""" +from tests.torch.utils import torchrun_case +from tests.common.mark_utils import arg_mark + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_embedding(): + ''' + Feature: test parallel op embedding. + Description: test parallel op embedding layout inference and calculation. + Expectation: Run success. + ''' + master_port = 10890 + file_name = "parallel_op_embedding.py" + case_name = "test_distributed_embedding_layout_inference" + torchrun_case(file_name, case_name, master_port) diff --git a/tests/torch/shard/ops/test_parallel_op_embedding.py b/tests/torch/shard/ops/test_parallel_op_embedding.py new file mode 100644 index 0000000..e3f8531 --- /dev/null +++ b/tests/torch/shard/ops/test_parallel_op_embedding.py @@ -0,0 +1,30 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test base dtensor""" +from tests.torch.utils import torchrun_case +from tests.common.mark_utils import arg_mark + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_embedding(): + ''' + Feature: test parallel op embedding. + Description: test parallel op embedding layout inference and calculation. + Expectation: Run success. + ''' + master_port = 10890 + file_name = "parallel_op_embedding.py" + case_name = "test_distributed_embedding_layout_inference" + torchrun_case(file_name, case_name, master_port) diff --git a/tests/torch/shard/ops/test_parallel_op.py b/tests/torch/shard/ops/test_parallel_op_linear.py similarity index 100% rename from tests/torch/shard/ops/test_parallel_op.py rename to tests/torch/shard/ops/test_parallel_op_linear.py diff --git a/tests/torch/shard/ops/test_parallel_op_reshape_view.py b/tests/torch/shard/ops/test_parallel_op_reshape_view.py new file mode 100644 index 0000000..671b027 --- /dev/null +++ b/tests/torch/shard/ops/test_parallel_op_reshape_view.py @@ -0,0 +1,56 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test base dtensor""" +from tests.torch.utils import torchrun_case +from tests.common.mark_utils import arg_mark + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_reshape(): + ''' + Feature: test parallel op reshape. + Description: test parallel op reshape (flatten and expand). + Expectation: Run success. + ''' + master_port = 10889 + file_name = "parallel_op_reshape_view.py" + case_name = "test_distributed_reshape_layout_inference" + torchrun_case(file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_view(): + ''' + Feature: test parallel op view. + Description: test parallel op view (tuple args and preserving sharded dim). + Expectation: Run success. + ''' + master_port = 10890 + file_name = "parallel_op_reshape_view.py" + case_name = "test_distributed_view_layout_inference" + torchrun_case(file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_reshape_error(): + ''' + Feature: test parallel op reshape error. + Description: test parallel op reshape invalid cases. + Expectation: Run success. + ''' + master_port = 10891 + file_name = "parallel_op_reshape_view.py" + case_name = "test_distributed_reshape_fail_mismatch" + torchrun_case(file_name, case_name, master_port) diff --git a/tests/torch/shard/ops/test_parallel_op_softmax.py b/tests/torch/shard/ops/test_parallel_op_softmax.py new file mode 100644 index 0000000..fc5b7cf --- /dev/null +++ b/tests/torch/shard/ops/test_parallel_op_softmax.py @@ -0,0 +1,42 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test base dtensor""" +from tests.torch.utils import torchrun_case +from tests.common.mark_utils import arg_mark + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_softmax(): + ''' + Feature: test parallel op softmax. + Description: test parallel op softmax layout inference and calculation. + Expectation: Run success. + ''' + master_port = 10889 # Ensure port is different if running concurrently, or let the framework handle it + file_name = "parallel_op_softmax.py" + case_name = "test_distributed_softmax_layout_inference" + torchrun_case(file_name, case_name, master_port) + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_softmax_dim_error(): + ''' + Feature: test parallel op softmax error handling. + Description: test parallel op softmax raises error when axis is sharded. + Expectation: Run success (by catching exception). + ''' + master_port = 10889 + file_name = "parallel_op_softmax.py" + case_name = "test_distributed_softmax_sharded_dim_error" + torchrun_case(file_name, case_name, master_port) diff --git a/tests/torch/shard/ops/test_perallel_op_transpose_permute.py b/tests/torch/shard/ops/test_perallel_op_transpose_permute.py new file mode 100644 index 0000000..5791e42 --- /dev/null +++ b/tests/torch/shard/ops/test_perallel_op_transpose_permute.py @@ -0,0 +1,56 @@ +# Copyright 2026 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ +"""test base dtensor""" +from tests.torch.utils import torchrun_case +from tests.common.mark_utils import arg_mark + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_permute(): + ''' + Feature: test parallel op permute. + Description: test parallel op permute layout inference and correctness. + Expectation: Run success. + ''' + master_port = 10889 # Ensure distinct port if run in parallel + file_name = "perallel_op_transpose_permute.py" + case_name = "test_distributed_permute_layout_inference" + torchrun_case(file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_transpose(): + ''' + Feature: test parallel op transpose. + Description: test parallel op transpose layout inference and correctness. + Expectation: Run success. + ''' + master_port = 10890 + file_name = "perallel_op_transpose_permute.py" + case_name = "test_distributed_transpose_layout_inference" + torchrun_case(file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_distributed_transpose_neg(): + ''' + Feature: test parallel op transpose with negative dims. + Description: test parallel op transpose negative dimension support. + Expectation: Run success. + ''' + master_port = 10891 + file_name = "perallel_op_transpose_permute.py" + case_name = "test_distributed_transpose_negative_dim" + torchrun_case(file_name, case_name, master_port) -- Gitee