diff --git a/mindformers/pynative/distributed/__init__.py b/mindformers/pynative/distributed/__init__.py index 570d69e4d6e95aff99e1358043d565670d71ea74..7eb7d59cb735ebd8dba302b6590f02e4ab738d16 100644 --- a/mindformers/pynative/distributed/__init__.py +++ b/mindformers/pynative/distributed/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Huawei Technologies Co., Ltd +# 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. @@ -13,3 +13,16 @@ # limitations under the License. # ============================================================================ """distributed modules""" +from mindformers.pynative.distributed.style import ( + PrepareModuleInput, + PrepareModuleOutput, + PrepareModuleInputOutput, + ParallelStyle, +) + +__all__ = [ + "PrepareModuleInput", + "PrepareModuleOutput", + "PrepareModuleInputOutput", + "ParallelStyle", +] diff --git a/mindformers/pynative/distributed/style.py b/mindformers/pynative/distributed/style.py new file mode 100644 index 0000000000000000000000000000000000000000..e3ca515bbc0a3d2930cf8aa3df76622fe256caa7 --- /dev/null +++ b/mindformers/pynative/distributed/style.py @@ -0,0 +1,436 @@ +# 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 style for distributed tensor parallel.""" +from typing import Any, Optional, Tuple, Dict, Union +from abc import ABC, abstractmethod + +from hyper_parallel.core.dtensor import DTensor +from hyper_parallel.core.device_mesh import DeviceMesh +from hyper_parallel.core.placement_types import Placement + +import mindspore as ms +from mindspore import nn + + +__all__ = [ + "ParallelStyle", + "PrepareModuleInput", + "PrepareModuleOutput", + "PrepareModuleInputOutput", +] + + +class ParallelStyle(ABC): + """ + The parallel style contract defines how the module or submodule should be parallelized. + + It only defines the ``_apply`` method for parallelization to use, this allows maximum + flexibility for different kind of style implementations. + """ + + src_data_rank: Optional[int] = 0 + + @abstractmethod + def _apply(self, module: nn.Cell, device_mesh: DeviceMesh) -> nn.Cell: + """Apply parallel style to the module.""" + ... + + +class PrepareModuleInput(ParallelStyle): + """ + Configure the nn.Cell's inputs to convert the input tensors of the nn.Cell to DTensors at runtime according to + ``input_layouts``, and perform layout redistribution according to the ``desired_input_layouts``. + + Keyword Args: + input_layouts (Union[Placement, Tuple[Optional[Placement], ...]], optional): + The DTensor layouts of input tensors for the nn.Cell, this is used to convert the input tensors to + DTensors. If some inputs are not Tensor or no need to convert to DTensors, ``None`` need to be specified + as a placeholder. default: None. + desired_input_layouts (Union[Placement, Tuple[Optional[Placement], ...]], optional): + The desired DTensor layout of input tensors for the nn.Cell, this is used to ensure the inputs of the nn.Cell + have the desired DTensor layouts. This argument needs to have the same length with ``input_layouts``. default: None. + input_kwarg_layouts (Dict[str, Placement], optional): + The DTensor layouts of input kwargs for the nn.Cell, this is used to convert the input kwarg tensors to DTensors. + default: None + desired_input_kwarg_layouts (Dict[str, Placement], optional): + The desired DTensor layout of input kwargs for the nn.Cell, this is used to ensure the inputs of the nn.Cell + have the desired DTensor layouts. default: None. + use_local_output (bool, optional): + Whether to use local :class:`Tensor` instead of :class:`DTensor` for the module inputs, default: False. + Returns: + A :class:`ParallelStyle` object that prepares the sharding layouts of the nn.Cell's inputs. + + Example:: + >>> from mindformers.pynative.distributed import PrepareModuleInput + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> from hyper_parallel.core.device_mesh import init_device_mesh + >>> ... + >>> block = TransformerBlock(...) # block is a nn.Cell that contains an "attn" Attention submodule + >>> tp_mesh = init_device_mesh(mesh_shape=(8,), alias_name=("tp",)) + >>> + >>> # According to the style specified below, the first input of attn will be annotated to Sharded DTensor + >>> # and then redistributed to Replicated DTensor. + >>> prepare_input = PrepareModuleInput( + >>> input_layouts=(Shard(0), None, None), + >>> desired_input_layouts=(Replicate(), None, None) + >>> ) + >>> prepare_input._apply(block.attn, tp_mesh) + """ + + def __init__( + self, + *, + input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, + desired_input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, + input_kwarg_layouts: Optional[Dict[str, Placement]] = None, + desired_input_kwarg_layouts: Optional[Dict[str, Placement]] = None, + use_local_output: bool = False, + ): + super().__init__() + self.input_layouts = ( + (input_layouts,) if isinstance(input_layouts, Placement) else input_layouts + ) + self.desired_input_layouts = ( + (desired_input_layouts,) + if isinstance(desired_input_layouts, Placement) + else desired_input_layouts + ) + self.use_local_output = use_local_output + if self.input_layouts is not None: + if self.desired_input_layouts is None: + raise ValueError("desired module inputs should not be None!") + if len(self.input_layouts) != len(self.desired_input_layouts): + raise ValueError( + "input_layouts and desired_input_layouts should have same length!" + ) + self.with_kwargs = input_kwarg_layouts is not None + self.input_kwarg_layouts = input_kwarg_layouts or {} + self.desired_input_kwarg_layouts = desired_input_kwarg_layouts or {} + if self.with_kwargs: + if len(self.input_kwarg_layouts) != len(self.desired_input_kwarg_layouts): + raise ValueError( + "input_kwarg_layouts and desired_input_kwarg_layouts should have same length!" + ) + + def _prepare_input_arg( + self, + input: Any, + mesh: DeviceMesh, + input_layout: Optional[Placement], + desired_layout: Optional[Placement], + ): + """Prepare a single input argument.""" + if input_layout is not None: + if isinstance(input, DTensor): + dt_inp = input + else: + if not isinstance(input, ms.Tensor): + raise ValueError(f"expecting input to be a Tensor, but got {type(input)}") + dt_inp = DTensor.from_local( + input, mesh, (input_layout,) + ) + + if desired_layout is not None and input_layout != desired_layout: + dt_inp = dt_inp.redistribute(mesh, (desired_layout,)) + + return dt_inp.to_local() if self.use_local_output else dt_inp + else: + return input + + def _prepare_input_fn(self, inputs, device_mesh): + """Prepare input arguments.""" + if self.input_layouts is None: + return inputs + prepared_inputs = [] + if not isinstance(inputs, tuple): + inputs = (inputs,) + if len(inputs) != len(self.input_layouts): + raise ValueError("module inputs and input_layouts should have same length!") + + if self.desired_input_layouts is None: + raise ValueError("desired module inputs should not be None!") + + for inp, input_layout, desired_layout in zip( + inputs, self.input_layouts, self.desired_input_layouts + ): + prepared_inputs.append( + self._prepare_input_arg(inp, device_mesh, input_layout, desired_layout) + ) + return tuple(prepared_inputs) + + def _prepare_input_kwarg_fn(self, inputs, kwarg_inputs, device_mesh): + """Prepare input arguments and keyword arguments.""" + prepared_arg_inputs = self._prepare_input_fn(inputs, device_mesh) + prepared_kwarg_inputs = {} + for kwarg_key in kwarg_inputs: + kwarg_val = kwarg_inputs[kwarg_key] + input_layout = self.input_kwarg_layouts.get(kwarg_key) + desired_input_layout = self.desired_input_kwarg_layouts.get(kwarg_key) + + prepared_kwarg_inputs[kwarg_key] = self._prepare_input_arg( + kwarg_val, device_mesh, input_layout, desired_input_layout + ) + + return (prepared_arg_inputs, prepared_kwarg_inputs) + + def _apply(self, module: nn.Cell, device_mesh: DeviceMesh) -> nn.Cell: + """ + Apply PrepareModuleInput style to the module. + + This method registers a forward pre-hook on the module to prepare inputs + according to the specified layouts before the module's construct method is called. + """ + if self.with_kwargs: + module.register_forward_pre_hook( + lambda _, inputs, kwargs: self._prepare_input_kwarg_fn( + inputs, kwargs, device_mesh + ), + with_kwargs=True, + ) + else: + module.register_forward_pre_hook( + lambda _, inputs: self._prepare_input_fn(inputs, device_mesh) + ) + return module + + def __repr__(self) -> str: + tmpstr = self.__class__.__name__ + "(" + tmpstr += f"input_layouts={self.input_layouts}, " + tmpstr += f"desired_input_layouts={self.desired_input_layouts}, " + tmpstr += f"input_kwarg_layouts={self.input_kwarg_layouts}, " + tmpstr += f"desired_input_kwarg_layouts={self.desired_input_kwarg_layouts}, " + tmpstr += f"use_local_output={self.use_local_output}" + tmpstr += ")" + return tmpstr + + +class PrepareModuleOutput(ParallelStyle): + """ + Configure the nn.Cell's outputs to convert the output tensors of the nn.Cell to DTensors at runtime according to + ``output_layouts``, and perform layout redistribution according to the ``desired_output_layouts``. + + Keyword Args: + output_layouts (Union[Placement, Tuple[Optional[Placement], ...]]): + The DTensor layouts of output tensors for the nn.Cell, this is used to convert the output tensors to + DTensors if they are :class:`Tensor`. If some outputs are not Tensor or no need to convert to DTensors, + ``None`` need to be specified as a placeholder. + desired_output_layouts (Union[Placement, Tuple[Placement, ...]]): + The desired DTensor layouts of output tensors for the nn.Cell, this is used to ensure the outputs of the nn.Cell + have the desired DTensor layouts. + use_local_output (bool, optional): + Whether to use local :class:`Tensor` instead of :class:`DTensor` for the module outputs, default: True. + Returns: + A :class:`ParallelStyle` object that prepares the sharding layouts of the nn.Cell's outputs. + + Example:: + >>> from mindformers.pynative.distributed import PrepareModuleOutput + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> from hyper_parallel.core.device_mesh import init_device_mesh + >>> ... + >>> block = TransformerBlock(...) # block is a nn.Cell that contains an "attn" Attention submodule + >>> tp_mesh = init_device_mesh(mesh_shape=(8,), alias_name=("tp",)) + >>> + >>> # According to the style specified below, the output of the TransformerBlock will be converted to Replicated DTensor + >>> # and then redistributed to Sharded DTensor. + >>> prepare_output = PrepareModuleOutput( + >>> output_layouts=Replicate(), + >>> desired_output_layouts=Shard(0) + >>> ) + >>> prepare_output._apply(block, tp_mesh) + """ + + def __init__( + self, + *, + output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], + desired_output_layouts: Union[Placement, Tuple[Placement, ...]], + use_local_output: bool = True, + ): + super().__init__() + self.output_layouts = ( + (output_layouts,) if isinstance(output_layouts, Placement) else output_layouts + ) + self.desired_output_layouts = ( + (desired_output_layouts,) + if isinstance(desired_output_layouts, Placement) + else desired_output_layouts + ) + self.use_local_output = use_local_output + if len(self.output_layouts) != len(self.desired_output_layouts): + raise ValueError( + "output_layouts and desired_output_layouts should have same length!" + ) + + def _prepare_out_fn(self, outputs, device_mesh): + """Prepare output arguments.""" + prepared_outputs = [] + if not isinstance(outputs, tuple): + outputs = (outputs,) + if len(outputs) != len(self.output_layouts): + raise ValueError( + "module outputs and output_layouts should have same length!" + ) + + for out, out_layout, desired_out_layout in zip( + outputs, self.output_layouts, self.desired_output_layouts + ): + if out_layout is not None: + if isinstance(out, DTensor): + dt_out = out + else: + if not isinstance(out, ms.Tensor): + raise ValueError(f"expecting output to be a Tensor, but got {type(out)}") + dt_out = DTensor.from_local( + out, device_mesh, (out_layout,) + ) + + if out_layout != desired_out_layout: + dt_out = dt_out.redistribute(device_mesh, (desired_out_layout,)) + prepared_outputs.append( + dt_out.to_local() if self.use_local_output else dt_out + ) + else: + prepared_outputs.append(out) + + if len(prepared_outputs) == 1: + return prepared_outputs[0] + else: + return tuple(prepared_outputs) + + def _apply(self, module: nn.Cell, device_mesh: DeviceMesh) -> nn.Cell: + """ + Apply PrepareModuleOutput style to the module. + + This method registers a forward hook on the module to prepare outputs + according to the specified layouts after the module's construct method is called. + """ + module.register_forward_hook( + lambda _, inputs, outputs: self._prepare_out_fn(outputs, device_mesh) + ) + return module + + def __repr__(self) -> str: + tmpstr = self.__class__.__name__ + "(" + tmpstr += f"output_layouts={self.output_layouts}, " + tmpstr += f"desired_output_layouts={self.desired_output_layouts}, " + tmpstr += f"use_local_output={self.use_local_output}" + tmpstr += ")" + return tmpstr + + +class PrepareModuleInputOutput(ParallelStyle): + """ + Configure the nn.Cell's inputs (and outputs) to convert the input tensors (and output tensors, respectively) of the nn.Cell + to DTensors at runtime according to ``input_layouts`` (and output_layouts, respectively), and perform layout redistribution + according to the ``desired_input_layouts`` (and ``desired_output_layouts``, respectively). This is a combination of + :class:`PrepareModuleInput` and :class:`PrepareModuleOutput`. + + Keyword Args: + input_layouts (Union[Placement, Tuple[Optional[Placement], ...]], optional): + The DTensor layouts of input tensors for the nn.Cell, this is used to convert the input tensors to + DTensors. If some inputs are not Tensor or no need to convert to DTensors, ``None`` need to be specified + as a placeholder. default: None. + desired_input_layouts (Union[Placement, Tuple[Optional[Placement], ...]], optional): + The desired DTensor layout of input tensors for the nn.Cell, this is used to ensure the inputs of the nn.Cell + have the desired DTensor layouts. This argument needs to have the same length with ``input_layouts``. default: None. + input_kwarg_layouts (Dict[str, Placement], optional): + The DTensor layouts of input kwargs for the nn.Cell, this is used to convert the input kwarg tensors to DTensors. + default: None + desired_input_kwarg_layouts (Dict[str, Placement], optional): + The desired DTensor layout of input kwargs for the nn.Cell, this is used to ensure the inputs of the nn.Cell + have the desired DTensor layouts. default: None. + use_local_input (bool, optional): + Whether to use local :class:`Tensor` instead of :class:`DTensor` for the module inputs, default: False. + output_layouts (Union[Placement, Tuple[Optional[Placement], ...]]): + The DTensor layouts of output tensors for the nn.Cell, this is used to convert the output tensors to + DTensors if they are :class:`Tensor`. If some outputs are not Tensor or no need to convert to DTensors, + ``None`` need to be specified as a placeholder. + desired_output_layouts (Union[Placement, Tuple[Placement, ...]]): + The desired DTensor layouts of output tensors for the nn.Cell, this is used to ensure the outputs of the nn.Cell + have the desired DTensor layouts. + use_local_output (bool, optional): + Whether to use local :class:`Tensor` instead of :class:`DTensor` for the module outputs, default: True. + Returns: + A :class:`ParallelStyle` object that prepares the sharding layouts of the nn.Cell's inputs and outputs. + + Example:: + >>> from mindformers.pynative.distributed import PrepareModuleInputOutput + >>> from hyper_parallel.core.placement_types import Shard, Replicate + >>> from hyper_parallel.core.device_mesh import init_device_mesh + >>> ... + >>> block = TransformerBlock(...) # block is a nn.Cell that contains an "attn" Attention submodule + >>> tp_mesh = init_device_mesh(mesh_shape=(8,), alias_name=("tp",)) + >>> + >>> # According to the style specified below, the first input of attn will be annotated as Sharded DTensor + >>> # and then redistributed to Replicated DTensor, and the output of the TransformerBlock will be annotated + >>> # as Replicated DTensor and then redistributed to Sharded DTensor. + >>> prepare_input_output = PrepareModuleInputOutput( + >>> input_layouts=(Shard(0), None, None, ...), + >>> desired_input_layouts=(Replicate(), None, None, ...), + >>> output_layouts=Replicate(), + >>> desired_output_layouts=Shard(0), + >>> ) + >>> prepare_input_output._apply(block.attn, tp_mesh) + """ + + def __init__( + self, + *, + input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, + desired_input_layouts: Optional[Union[Placement, Tuple[Optional[Placement], ...]]] = None, + input_kwarg_layouts: Optional[Dict[str, Placement]] = None, + desired_input_kwarg_layouts: Optional[Dict[str, Placement]] = None, + use_local_input: bool = False, + output_layouts: Union[Placement, Tuple[Optional[Placement], ...]], + desired_output_layouts: Union[Placement, Tuple[Placement, ...]], + use_local_output: bool = True, + ): + super().__init__() + self.prepare_module_input = PrepareModuleInput( + input_layouts=input_layouts, + desired_input_layouts=desired_input_layouts, + input_kwarg_layouts=input_kwarg_layouts, + desired_input_kwarg_layouts=desired_input_kwarg_layouts, + use_local_output=use_local_input, + ) + self.prepare_module_output = PrepareModuleOutput( + output_layouts=output_layouts, + desired_output_layouts=desired_output_layouts, + use_local_output=use_local_output, + ) + + def _apply(self, module: nn.Cell, device_mesh: DeviceMesh) -> nn.Cell: + """ + Apply PrepareModuleInputOutput style to the module. + + This method applies both PrepareModuleInput and PrepareModuleOutput to the module. + """ + self.prepare_module_input._apply(module, device_mesh) + self.prepare_module_output._apply(module, device_mesh) + return module + + def __repr__(self) -> str: + tmpstr = self.__class__.__name__ + "(" + tmpstr += f"input_layouts={self.prepare_module_input.input_layouts}, " + tmpstr += f"desired_input_layouts={self.prepare_module_input.desired_input_layouts}, " + tmpstr += f"input_kwarg_layouts={self.prepare_module_input.input_kwarg_layouts}, " + tmpstr += f"desired_input_kwarg_layouts={self.prepare_module_input.desired_input_kwarg_layouts}, " + tmpstr += f"use_local_input={self.prepare_module_input.use_local_output}, " + tmpstr += f"output_layouts={self.prepare_module_output.output_layouts}, " + tmpstr += f"desired_output_layouts={self.prepare_module_output.desired_output_layouts}, " + tmpstr += f"use_local_output={self.prepare_module_output.use_local_output}" + tmpstr += ")" + return tmpstr + diff --git a/tests/st/test_ut/test_pynative/test_distributed/__init__.py b/tests/st/test_ut/test_pynative/test_distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dfd3a257a525a4883c54e38c36ccc08b1a7631e3 --- /dev/null +++ b/tests/st/test_ut/test_pynative/test_distributed/__init__.py @@ -0,0 +1,16 @@ +# 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 distributed modules""" + diff --git a/tests/st/test_ut/test_pynative/test_distributed/run_test_style.py b/tests/st/test_ut/test_pynative/test_distributed/run_test_style.py new file mode 100644 index 0000000000000000000000000000000000000000..4011a932d2a60791c91b3a4ca7440cf63472fdd2 --- /dev/null +++ b/tests/st/test_ut/test_pynative/test_distributed/run_test_style.py @@ -0,0 +1,370 @@ +# 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. +# ============================================================================ +"""Run script for testing PrepareModuleInput and PrepareModuleOutput.""" +import argparse +import numpy as np + +from hyper_parallel import init_device_mesh +from hyper_parallel.core.dtensor import DTensor +from hyper_parallel.core.placement_types import Shard, Replicate + +import mindspore as ms +from mindspore import nn, Tensor +from mindformers.pynative.distributed.style import ( + PrepareModuleInput, + PrepareModuleOutput, + PrepareModuleInputOutput, +) +from mindformers.pynative.layers.linear import Linear +from mindformers.parallel_core.utils.init_method import init_method_normal + + +def create_test_tensors(): + """Create test tensors with fixed random seed.""" + np.random.seed(42) + return { + "tensor_1x16": Tensor( + np.random.randn(2, 16).astype(np.float32), dtype=ms.float32 + ), + "tensor_2x8": Tensor( + np.random.randn(2, 8).astype(np.float32), dtype=ms.float32 + ), + "tensor_4x4": Tensor( + np.random.randn(4, 4).astype(np.float32), dtype=ms.float32 + ), + } + + +class TestModule(nn.Cell): + """Test module for PrepareModuleInput and PrepareModuleOutput.""" + + def __init__(self): + super().__init__() + self.linear = Linear( + input_size=8, + output_size=8, + params_dtype="float32", + compute_dtype="float32", + init_method=init_method_normal(), + bias=True, + skip_bias_add=True, + ) + + def construct(self, x, y=None, mask=None): + if mask is not None: + return mask + if y is not None: + return x, y + return self.linear(x) + + +def verify_dtensor_result( + result, expected_dtensor, err_msg, use_local_output=False +): + """Verify DTensor result matches expected DTensor.""" + if use_local_output: + assert isinstance(result, Tensor) + assert not isinstance(result, DTensor) + np.testing.assert_array_equal( + result.asnumpy(), + expected_dtensor.asnumpy(), + err_msg=err_msg, + ) + else: + assert isinstance(result, DTensor) + assert result.placements == expected_dtensor.placements + np.testing.assert_array_equal( + result.to_local().asnumpy(), + expected_dtensor.to_local().asnumpy(), + err_msg=err_msg, + ) + + +def create_reference_module(source_module): + """Create a reference module with same weights as source.""" + ref_module = TestModule() + source_weight = source_module.linear.weight.value() + ref_module.linear.weight.set_data(source_weight) + return ref_module + + +def compute_expected_output( + module, input_tensor, input_layout, desired_layout, device_mesh +): + """Compute expected output after input conversion.""" + dtensor = DTensor.from_local(input_tensor, device_mesh, (input_layout,)) + converted_dtensor = dtensor.redistribute(device_mesh, (desired_layout,)) + local_tensor = converted_dtensor.to_local() + ref_module = create_reference_module(module) + return ref_module(local_tensor) + + +def test_prepare_module_input(device_mesh, test_tensors): + """Test PrepareModuleInput covering multiple scenarios.""" + + # Test: Multiple inputs + module = TestModule() + prepare_inps = PrepareModuleInput( + input_layouts=(Shard(0), None), + desired_input_layouts=(Replicate(), None), + ) + prepare_inps._apply(module, device_mesh) + + tensor1 = test_tensors["tensor_2x8"] + tensor2 = test_tensors["tensor_4x4"] + output_x, output_y = module(tensor1, tensor2) + + sharded_dtensor = DTensor.from_local(tensor1, device_mesh, (Shard(0),)) + expected_dtensor = sharded_dtensor.redistribute( + device_mesh, (Replicate(),) + ) + verify_dtensor_result( + output_x, + expected_dtensor, + "Multiple inputs test: Output_x mismatch", + ) + assert output_y is tensor2 + print("✓ Multiple inputs passed") + + # Test: DTensor input with use_local_output + module_dtensor = TestModule() + prepare_dtensor = PrepareModuleInput( + input_layouts=(Shard(1),), + desired_input_layouts=(Replicate(),), + use_local_output=True, + ) + prepare_dtensor._apply(module_dtensor, device_mesh) + + test_tensor = test_tensors["tensor_4x4"] + test_dtensor = DTensor.from_local(test_tensor, device_mesh, (Shard(1),)) + result = module_dtensor(test_dtensor) + expected_result = compute_expected_output( + module_dtensor, test_tensor, Shard(1), Replicate(), device_mesh + ) + verify_dtensor_result( + result[0], + expected_result[0], + "DTensor input test: Result mismatch", + use_local_output=True, + ) + print("✓ DTensor input passed") + + # Test: use_local_output with shard change + module_local = TestModule() + prepare_local = PrepareModuleInput( + input_layouts=(Shard(0),), + desired_input_layouts=(Shard(1),), + use_local_output=True, + ) + prepare_local._apply(module_local, device_mesh) + + test_tensor = test_tensors["tensor_1x16"] + result = module_local(test_tensor) + expected_result = compute_expected_output( + module_local, test_tensor, Shard(0), Shard(1), device_mesh + ) + verify_dtensor_result( + result[0], + expected_result[0], + "use_local_output test: Result mismatch", + use_local_output=True, + ) + print("✓ use_local_output passed") + + # Test: Keyword arguments + module_kwargs = TestModule() + prepare_kwargs = PrepareModuleInput( + input_kwarg_layouts={"mask": Replicate()}, + desired_input_kwarg_layouts={"mask": Shard(0)}, + ) + prepare_kwargs._apply(module_kwargs, device_mesh) + + x_tensor = test_tensors["tensor_2x8"] + mask_tensor = test_tensors["tensor_2x8"] + output = module_kwargs(x_tensor, mask=mask_tensor) + + replicated_dtensor = DTensor.from_local( + mask_tensor, device_mesh, (Replicate(),) + ) + expected_dtensor = replicated_dtensor.redistribute( + device_mesh, (Shard(0),) + ) + verify_dtensor_result( + output, expected_dtensor, "Kwargs test: Output mismatch" + ) + print("✓ Kwargs passed") + + +def test_prepare_module_output(device_mesh, test_tensors): + """Test PrepareModuleOutput covering multiple scenarios.""" + + # Test: Multiple outputs + module_multi = TestModule() + prepare_output_multi = PrepareModuleOutput( + output_layouts=(Shard(0), None), + desired_output_layouts=(Replicate(), None), + use_local_output=False, + ) + prepare_output_multi._apply(module_multi, device_mesh) + + test_tensor = test_tensors["tensor_4x4"] + output_x, output_y = module_multi(test_tensor, y=test_tensor) + + sharded_dtensor = DTensor.from_local( + test_tensor, device_mesh, (Shard(0),) + ) + expected_dtensor = sharded_dtensor.redistribute( + device_mesh, (Replicate(),) + ) + verify_dtensor_result( + output_x, expected_dtensor, "Multiple outputs test: Output_x mismatch" + ) + assert output_y is test_tensor + print("✓ Multiple outputs passed") + + # Test: Output redistribution + module_dtensor = TestModule() + prepare_dtensor = PrepareModuleOutput( + output_layouts=(Replicate(), None), + desired_output_layouts=(Shard(1), None), + use_local_output=False, + ) + prepare_dtensor._apply(module_dtensor, device_mesh) + + test_tensor = test_tensors["tensor_2x8"] + result = module_dtensor(test_tensor) + ref_module = create_reference_module(module_dtensor) + ref_output = ref_module(test_tensor) + dtensor = DTensor.from_local(ref_output[0], device_mesh, (Replicate(),)) + expected_dtensor = dtensor.redistribute(device_mesh, (Shard(1),)) + verify_dtensor_result( + result[0], + expected_dtensor, + "Output redistribution test: Result mismatch", + ) + print("✓ Output redistribution passed") + + # Test: Shard dimension change with use_local_output + module_shard = TestModule() + prepare_shard = PrepareModuleOutput( + output_layouts=(Shard(0), None), + desired_output_layouts=(Shard(1), None), + use_local_output=True, + ) + prepare_shard._apply(module_shard, device_mesh) + test_tensor = test_tensors["tensor_2x8"] + result = module_shard(test_tensor) + + ref_module = create_reference_module(module_shard) + ref_output = ref_module(test_tensor) + sharded_dtensor = DTensor.from_local( + ref_output[0], device_mesh, (Shard(0),) + ) + expected_dtensor = sharded_dtensor.redistribute( + device_mesh, (Shard(1),) + ) + expected_tensor = expected_dtensor.to_local() + verify_dtensor_result( + result[0], + expected_tensor, + "Shard change test: Result mismatch", + use_local_output=True, + ) + print("✓ Shard change passed") + + +def test_prepare_module_input_output(device_mesh, test_tensors): + """Test PrepareModuleInputOutput covering multiple scenarios.""" + + # Test: Keyword arguments with input/output conversion + module_kwargs = TestModule() + prepare_kwargs = PrepareModuleInputOutput( + input_kwarg_layouts={"mask": Shard(0)}, + desired_input_kwarg_layouts={"mask": Replicate()}, + output_layouts=(Replicate(),), + desired_output_layouts=(Shard(1),), + use_local_output=False, + ) + prepare_kwargs._apply(module_kwargs, device_mesh) + + x_tensor = test_tensors["tensor_2x8"] + mask_tensor = test_tensors["tensor_2x8"] + result = module_kwargs(x_tensor, mask=mask_tensor) + + sharded_mask_dtensor = DTensor.from_local( + mask_tensor, device_mesh, (Shard(0),) + ) + expected_mask_dtensor = sharded_mask_dtensor.redistribute( + device_mesh, (Replicate(),) + ) + expected_output_dtensor = expected_mask_dtensor.redistribute( + device_mesh, (Shard(1),) + ) + verify_dtensor_result( + result, + expected_output_dtensor, + "Kwargs I/O test: Result mismatch", + ) + print("✓ Kwargs I/O passed") + + +def main(): + """Main function to run tests with configurable parameters.""" + from mindspore.communication import init + + parser = argparse.ArgumentParser( + description="Run PrepareModuleInput/Output/InputOutput tests" + ) + parser.add_argument( + "--tp", type=int, default=2, help="Tensor parallel size" + ) + parser.add_argument( + "--test_function", + type=str, + required=True, + choices=[ + "test_prepare_module_input", + "test_prepare_module_output", + "test_prepare_module_input_output", + ], + help="Test function to run", + ) + + args = parser.parse_args() + + # Initialize distributed environment + init() + + # Create device mesh + device_mesh = init_device_mesh( + mesh_shape=(args.tp,), alias_name=("dp",) + ) + + # Create test tensors + test_tensors = create_test_tensors() + + # Run the specified test function + if args.test_function == "test_prepare_module_input": + test_prepare_module_input(device_mesh, test_tensors) + elif args.test_function == "test_prepare_module_output": + test_prepare_module_output(device_mesh, test_tensors) + elif args.test_function == "test_prepare_module_input_output": + test_prepare_module_input_output(device_mesh, test_tensors) + else: + raise ValueError(f"Unknown test function: {args.test_function}") + + +if __name__ == "__main__": + main() diff --git a/tests/st/test_ut/test_pynative/test_distributed/test_style.py b/tests/st/test_ut/test_pynative/test_distributed/test_style.py new file mode 100644 index 0000000000000000000000000000000000000000..2dfa3bf3dca4f5777bf60bb85b3acc4e7a27d65a --- /dev/null +++ b/tests/st/test_ut/test_pynative/test_distributed/test_style.py @@ -0,0 +1,174 @@ +# 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 PrepareModuleInput and PrepareModuleOutput classes.""" +import os +import random +import subprocess +import pytest + +from hyper_parallel.core.placement_types import Shard, Replicate +from mindformers.pynative.distributed.style import ( + PrepareModuleInput, + PrepareModuleOutput, + PrepareModuleInputOutput, +) + + +cur_dir = os.path.dirname(os.path.abspath(__file__)) +run_script_path = os.path.join(cur_dir, "run_test_style.py") +base_port = int( + os.environ.get("ASCEND_PORT_ID", random.randint(50000, 65535)) +) + + +class TestPrepareModule: + """Test cases for PrepareModuleInput, PrepareModuleOutput, and + PrepareModuleInputOutput classes.""" + + @pytest.mark.level0 + @pytest.mark.parametrize( + "test_function,port_offset", + [ + ("test_prepare_module_input", 0), + ("test_prepare_module_output", 1), + ("test_prepare_module_input_output", 2), + ], + ) + def test_multi_card_cases(self, test_function, port_offset): + """Test PrepareModuleInput/Output/InputOutput with multi-card.""" + assert os.path.exists( + run_script_path + ), f"Run script not found: {run_script_path}" + port_id = base_port + port_offset + cmd = ( + f"msrun --worker_num=2 " + f"--local_worker_num=2 " + f"--master_port={port_id} " + f"--log_dir=./msrun_log_test_style --join=True " + f"{run_script_path} --tp 2 --test_function {test_function}" + ) + result = subprocess.run( + cmd, shell=True, capture_output=True, text=True + ) + assert result.returncode == 0, ( + f"Test failed with non-zero exit code: " + f"{result.returncode}.\n" + f"Stdout:\n{result.stdout}\nStderr:\n{result.stderr}" + ) + + def test_prepare_module_input_output_validation(self): + """Test PrepareModuleInputOutput validation, which covers + PrepareModuleInput and PrepareModuleOutput validation.""" + # Test: desired_input_layouts is None when input_layouts is not None + with pytest.raises( + ValueError, match="desired module inputs should not be None" + ): + PrepareModuleInputOutput( + input_layouts=(Shard(0),), + desired_input_layouts=None, + output_layouts=(Shard(0),), + desired_output_layouts=(Replicate(),), + ) + + # Test: input_layouts and desired_input_layouts have different + # lengths + with pytest.raises(ValueError, match="should have same length"): + PrepareModuleInputOutput( + input_layouts=(Shard(0),), + desired_input_layouts=(Replicate(), Shard(0)), + output_layouts=(Shard(0),), + desired_output_layouts=(Replicate(),), + ) + + # Test: input_kwarg_layouts and desired_input_kwarg_layouts have + # different lengths + with pytest.raises(ValueError, match="should have same length"): + PrepareModuleInputOutput( + input_kwarg_layouts={"mask": Shard(0)}, + desired_input_kwarg_layouts={ + "mask": Replicate(), + "other": Replicate(), + }, + output_layouts=(Shard(0),), + desired_output_layouts=(Replicate(),), + ) + + # Test: output_layouts and desired_output_layouts have different + # lengths (single vs multiple) + with pytest.raises(ValueError, match="should have same length"): + PrepareModuleInputOutput( + output_layouts=(Shard(0),), + desired_output_layouts=(Replicate(), Shard(0)), + ) + + # Test: output_layouts and desired_output_layouts have different + # lengths (multiple vs single) + with pytest.raises(ValueError, match="should have same length"): + PrepareModuleInputOutput( + output_layouts=(Shard(0), Replicate()), + desired_output_layouts=(Shard(1),), + ) + + def test_prepare_module_input_repr(self): + """Test PrepareModuleInput __repr__ method.""" + prepare_input = PrepareModuleInput( + input_layouts=(Shard(0),), + desired_input_layouts=(Replicate(),), + use_local_output=True, + ) + + repr_str = repr(prepare_input) + assert "PrepareModuleInput" in repr_str + assert "input_layouts" in repr_str + assert "desired_input_layouts" in repr_str + assert "use_local_output=True" in repr_str + + def test_prepare_module_output_repr(self): + """Test PrepareModuleOutput __repr__ method.""" + prepare_output = PrepareModuleOutput( + output_layouts=(Shard(0),), + desired_output_layouts=(Replicate(),), + use_local_output=True, + ) + + repr_str = repr(prepare_output) + assert "PrepareModuleOutput" in repr_str + assert "output_layouts" in repr_str + assert "desired_output_layouts" in repr_str + assert "use_local_output=True" in repr_str + + def test_prepare_module_input_output_repr(self): + """Test PrepareModuleInputOutput __repr__ method.""" + prepare_input_output = PrepareModuleInputOutput( + input_layouts=(Shard(0),), + desired_input_layouts=(Replicate(),), + output_layouts=(Replicate(),), + desired_output_layouts=(Shard(0),), + use_local_input=True, + use_local_output=True, + ) + + repr_str = repr(prepare_input_output) + assert "PrepareModuleInputOutput" in repr_str + assert "input_layouts" in repr_str + assert "desired_input_layouts" in repr_str + assert "output_layouts" in repr_str + assert "desired_output_layouts" in repr_str + assert "use_local_input=True" in repr_str + assert "use_local_output=True" in repr_str + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])