diff --git a/hyper_parallel/core/shard/ops/parallel_transpose.py b/hyper_parallel/core/shard/ops/parallel_transpose.py index 31e40330efb8d4fedac80c306513ff7354b95052..c8ccd592d1f1f7b69b5e386775ed85870af3c280 100644 --- a/hyper_parallel/core/shard/ops/parallel_transpose.py +++ b/hyper_parallel/core/shard/ops/parallel_transpose.py @@ -65,3 +65,64 @@ class TransposeDistributedOp(DistributedOp): out_layout = output_layout(*out_tensor_map) return out_layout + + +class TransposeExtViewDistributedOp(DistributedOp): + """Distributed implementation for TransposeExtView operator.""" + + def infer_layout(self, layouts, extra_args): + """ + Infer output layout for TransposeExtView operator. + + Rules: + 1. Output layout is input layout with two tensor-map dimensions swapped. + + Args: + layouts (tuple): Layouts of input tensor. + extra_args (tuple | list): (dim0, dim1) + + Returns: + Layout: Layout for output tensor + """ + if not layouts or layouts[0] is None: + raise ValueError("Input layout is required for TransposeExtView.") + + if not isinstance(extra_args, (tuple, list)) or len(extra_args) < 2: + raise ValueError(f"TransposeExtView expects (dim0, dim1) in extra_args, but got {extra_args}") + + layout = layouts[0] + dim0 = extra_args[0] + dim1 = extra_args[1] + + if not isinstance(dim0, int) or not isinstance(dim1, int): + raise TypeError(f"dim0 and dim1 must be int, but got {type(dim0)} and {type(dim1)}") + + in_tensor_map = layout.alias_tensor_map + if in_tensor_map is None: + raise ValueError("Input layout.alias_tensor_map is None for TransposeExtView.") + + ndim = len(in_tensor_map) + + dim0 = self._normalize_dim(dim0, ndim) + dim1 = self._normalize_dim(dim1, ndim) + + if dim0 == dim1: + return layout + + out_tensor_map = list(in_tensor_map) + out_tensor_map[dim0], out_tensor_map[dim1] = out_tensor_map[dim1], out_tensor_map[dim0] + out_tensor_map = type(in_tensor_map)(out_tensor_map) + + output_layout = Layout( + mesh_shape=layout.mesh_shape, + alias_name=layout.alias_name, + rank_list=layout.rank_list + ) + return output_layout(*out_tensor_map) + + @staticmethod + def _normalize_dim(dim: int, ndim: int) -> int: + """Normalize dim into [0, ndim-1] with MindSpore-style range checks.""" + if dim < -ndim or dim >= ndim: + raise ValueError(f"dim {dim} out of range [-{ndim}, {ndim - 1}]") + return dim + ndim if dim < 0 else dim diff --git a/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml b/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml index e341231ab8b83972d77331fabcc24eeebfaf9631..28067e351ec26b39b4f75b7d7e4f4657a3f84448 100644 --- a/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml +++ b/hyper_parallel/core/shard/ops/yaml/transpose_ops.yaml @@ -1,4 +1,9 @@ Transpose: dist_op_name: _transpose_dist_op distributed_op_class: TransposeDistributedOp + distributed_op_file: parallel_transpose + +TransposeExtView: + dist_op_name: _transpose_ext_view_dist_op + distributed_op_class: TransposeExtViewDistributedOp distributed_op_file: parallel_transpose \ No newline at end of file diff --git a/tests/mindspore/st/shard/test_ops_transpose_ext_view_shell.py b/tests/mindspore/st/shard/test_ops_transpose_ext_view_shell.py new file mode 100644 index 0000000000000000000000000000000000000000..8021d7ee570c7e32b69dbb681efec58de5a522d6 --- /dev/null +++ b/tests/mindspore/st/shard/test_ops_transpose_ext_view_shell.py @@ -0,0 +1,60 @@ +# 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. +# ============================================================================ +"""parallel_transpose_ext_view_shell test""" + +from tests.common.mark_utils import arg_mark +from tests.mindspore.st.utils import msrun_case + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_transpose_ext_view_basic_3d_1(): + ''' + Feature: TransposeExtView operator. + Description: Test TransposeExtView swaps two dims on a 3D tensor in python shard. + Expectation: Run success. + ''' + glog_v = 2 + file_name = "transpose_ext_view_shard_in_python.py" + case_name = "test_transpose_ext_view_basic_3d_1" + master_port = 11294 + msrun_case(glog_v, file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_transpose_ext_view_negative_dims_2(): + ''' + Feature: TransposeExtView operator. + Description: Test TransposeExtView supports negative dims in python shard. + Expectation: Run success. + ''' + glog_v = 2 + file_name = "transpose_ext_view_shard_in_python.py" + case_name = "test_transpose_ext_view_negative_dims_2" + master_port = 11295 + msrun_case(glog_v, file_name, case_name, master_port) + + +@arg_mark(plat_marks=["platform_ascend910b"], level_mark="level0", card_mark="allcards", essential_mark="essential") +def test_transpose_ext_view_same_dims_noop_3(): + ''' + Feature: TransposeExtView operator. + Description: Test TransposeExtView is a no-op when dim0 == dim1 in python shard. + Expectation: Run success. + ''' + glog_v = 2 + file_name = "transpose_ext_view_shard_in_python.py" + case_name = "test_transpose_ext_view_same_dims_noop_3" + master_port = 11296 + msrun_case(glog_v, file_name, case_name, master_port) diff --git a/tests/mindspore/st/shard/transpose_ext_view_shard_in_python.py b/tests/mindspore/st/shard/transpose_ext_view_shard_in_python.py new file mode 100644 index 0000000000000000000000000000000000000000..67a4aeb4d7406061f8f940c32d86adfeaa364f98 --- /dev/null +++ b/tests/mindspore/st/shard/transpose_ext_view_shard_in_python.py @@ -0,0 +1,147 @@ +# 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 transpose ext view shard in python +""" + +import numpy as np + +import mindspore as ms +import mindspore.communication.management as D +from mindspore import nn, Tensor +from hyper_parallel import Layout, shard +from tests.mindspore.st.shard.utils import global_to_local, local_to_global + + +def setup_module(): + ms.set_device("Ascend") + D.init() + + +base_mesh_shape = (2, 2, 2) +base_alias_name = ("dp", "cp", "mp") + + +class TransposeExtViewNet(nn.Cell): + """TransposeExtView composed of transpose and ReLUs""" + + def __init__(self, relu_strategy=None): + super().__init__() + self.transpose = ms.mint.transpose + self.relu = ms.nn.ReLU() + if relu_strategy is not None: + stra = {"forward": {"input": relu_strategy}} + shard(self.relu, stra) + + def construct(self, x, dim0, dim1): + out = self.transpose(x, dim0, dim1) + out = self.relu(out) + out = out + 1 + return out + + +def _standalone_and_parallel_run(x, x_layout, dim0, dim1, relu_in_layout): + """Run standalone and parallel graph and return outputs.""" + # Standalone + standalone_net = TransposeExtViewNet() + standalone_output = standalone_net(x, dim0, dim1) + + # Parallel + x_local = global_to_local(x, x_layout) + parallel_net = TransposeExtViewNet(relu_strategy=(relu_in_layout,)) + parallel_output = parallel_net(x_local, dim0, dim1) + + parallel_output = local_to_global(parallel_output) + return standalone_output, parallel_output + + +def test_transpose_ext_view_basic_3d_1(): + ''' + Feature: TransposeExtView in python shard. + Description: Test TransposeExtView swaps two dims on a 3D tensor. + Expectation: Output matches standalone. + ''' + ms.set_seed(1) + np.random.seed(1) + + d0, d1, d2 = 16, 64, 32 + x = Tensor(np.random.randn(d0, d1, d2).astype(np.float32)) + + layout = Layout(base_mesh_shape, base_alias_name) + + # Shard input to ensure dtensor path is exercised + x_layout = layout("dp", "cp", "mp") + + # After transpose(0, 2): shape becomes (d2, d1, d0). + # Keep ReLU sharding simple: shard last dim by mp. + relu_in_layout = layout("None", "None", "mp") + + standalone_output, parallel_output = _standalone_and_parallel_run( + x, x_layout, dim0=0, dim1=2, relu_in_layout=relu_in_layout + ) + + assert np.allclose(standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3) + + +def test_transpose_ext_view_negative_dims_2(): + ''' + Feature: TransposeExtView in python shard. + Description: Test TransposeExtView supports negative dims (MindSpore semantics). + Expectation: Output matches standalone. + ''' + ms.set_seed(2) + np.random.seed(2) + + d0, d1, d2, d3 = 8, 16, 32, 64 + x = Tensor(np.random.randn(d0, d1, d2, d3).astype(np.float32)) + + layout = Layout(base_mesh_shape, base_alias_name) + + # Shard input: shard dim2 by mp; others replicated + x_layout = layout("None", "None", "mp", "None") + + # swap(-1, -3) => swap dim3 and dim1 + # output shape becomes (d0, d3, d2, d1) + relu_in_layout = layout("None", "mp", "None", "None") + + standalone_output, parallel_output = _standalone_and_parallel_run( + x, x_layout, dim0=-1, dim1=-3, relu_in_layout=relu_in_layout + ) + + assert np.allclose(standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3) + + +def test_transpose_ext_view_same_dims_noop_3(): + ''' + Feature: TransposeExtView in python shard. + Description: Test TransposeExtView is a no-op when dim0 == dim1. + Expectation: Output matches standalone. + ''' + ms.set_seed(3) + np.random.seed(3) + + d0, d1, d2 = 4, 32, 128 + x = Tensor(np.random.randn(d0, d1, d2).astype(np.float32)) + + layout = Layout(base_mesh_shape, base_alias_name) + + x_layout = layout("dp", "cp", "mp") + relu_in_layout = layout("dp", "cp", "mp") + + standalone_output, parallel_output = _standalone_and_parallel_run( + x, x_layout, dim0=1, dim1=1, relu_in_layout=relu_in_layout + ) + + assert np.allclose(standalone_output.asnumpy(), parallel_output.asnumpy(), 1e-3, 1e-3) diff --git a/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose_ext_view.py b/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose_ext_view.py new file mode 100644 index 0000000000000000000000000000000000000000..2c752898015a8954b9a4f745a07fc4930074dd45 --- /dev/null +++ b/tests/mindspore/ut/parallel_ops_infer/test_parallel_transpose_ext_view.py @@ -0,0 +1,187 @@ +# 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. +# ============================================================================ +"""parallel_transpose_ext_view test""" + +import pytest + +from hyper_parallel import Layout +from hyper_parallel.core.shard.ops.parallel_transpose import TransposeExtViewDistributedOp + + +def run_scenario(scenario_name, x_layout, expected_map, extra_args): + """Infer layout of TransposeExtView operator and validate tensor_map.""" + print(f"\n{'=' * 80}") + print(f"Test TransposeExtView, Scenario: {scenario_name}") + print('=' * 80) + + op = TransposeExtViewDistributedOp("TransposeExtView") + output_layout = op.infer_layout((x_layout,), extra_args) + assert output_layout.to_dict()["tensor_map"] == expected_map, \ + f"TransposeExtView failed in scenario '{scenario_name}'. " \ + f"Expected {expected_map}, got {output_layout.to_dict()['tensor_map']}" + + +base_mesh_shape = (2, 2, 2) +base_alias_name = ("dp", "cp", "mp") +base_rank_list = list(range(8)) + + +def test_transpose_ext_view_basic_swap_3d_1(): + """ + Feature: Basic swap. + Description: swap dim0=0 and dim1=2 on 3D tensor map. + Expectation: tensor_map dims swapped. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + run_scenario( + "1. Basic swap (0 <-> 2)", + x_layout, + expected_map=(0,1,2), + extra_args=(0, 2) + ) + + +def test_transpose_ext_view_negative_dims_2(): + """ + Feature: Negative dims. + Description: swap dim0=-1 and dim1=-3 on 3D tensor map. + Expectation: normalized dims swapped. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + # ndim=3: -1 -> 2, -3 -> 0 => swap(2,0) + run_scenario( + "2. Negative dims (-1 <-> -3)", + x_layout, + expected_map=(0,1,2), + extra_args=(-1, -3) + ) + + +def test_transpose_ext_view_noop_same_dims_3(): + """ + Feature: No-op. + Description: dim0 == dim1. + Expectation: output tensor_map unchanged. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + run_scenario( + "3. No-op swap (1 <-> 1)", + x_layout, + expected_map=(2, 1, 0), + extra_args=(1, 1) + ) + + +def test_transpose_ext_view_tuple_alias_dim_4(): + """ + Feature: Tuple alias dim. + Description: swap a normal dim with a tuple-alias dim. + Expectation: tuple moved to the swapped position. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("None", ("dp", "cp"), "mp") + + # tensor_map should be (-1, (2,1), 0) with base_alias_name ("dp","cp","mp") + # swap dim0=1, dim1=2 => (-1, 0, (2,1)) + run_scenario( + "4. Tuple alias swap (1 <-> 2)", + x_layout, + expected_map=(-1, 0, (2, 1)), + extra_args=(1, 2) + ) + + +def test_transpose_ext_view_dim_out_of_range_5(): + """ + Feature: Error handling. + Description: dim0 or dim1 out of range [-ndim, ndim-1]. + Expectation: raise ValueError. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + with pytest.raises(ValueError): + run_scenario( + "5. dim0 out of range", + x_layout, + expected_map=(2, 1, 0), + extra_args=(3, 0) + ) + + with pytest.raises(ValueError): + run_scenario( + "6. dim1 out of range (negative)", + x_layout, + expected_map=(2, 1, 0), + extra_args=(-4, 0) + ) + + +def test_transpose_ext_view_dim_type_error_6(): + """ + Feature: Error handling. + Description: dim0 or dim1 is not int. + Expectation: raise TypeError. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + with pytest.raises(TypeError): + run_scenario( + "7. dim0 type error", + x_layout, + expected_map=(2, 1, 0), + extra_args=("0", 1) + ) + + with pytest.raises(TypeError): + run_scenario( + "8. dim1 type error", + x_layout, + expected_map=(2, 1, 0), + extra_args=(0, None) + ) + + +def test_transpose_ext_view_extra_args_invalid_7(): + """ + Feature: Error handling. + Description: extra_args is not (dim0, dim1). + Expectation: raise ValueError. + """ + x_layout = Layout(base_mesh_shape, base_alias_name, base_rank_list) + x_layout = x_layout("dp", "cp", "mp") + + with pytest.raises(ValueError): + run_scenario( + "9. extra_args missing dim1", + x_layout, + expected_map=(2, 1, 0), + extra_args=(0,) + ) + + with pytest.raises(ValueError): + run_scenario( + "10. extra_args not tuple/list", + x_layout, + expected_map=(2, 1, 0), + extra_args=None + )