Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_norm.py: 76%
84 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""
16Distributed implementation for Norm operators (RmsNorm, layer_norm).
17"""
19from typing import Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from .parallel_ops import DistributedOp
25def _normalize_rmsnorm_args(x, gamma, epsilon=1e-6):
26 """Normalize RmsNorm args to positional form.
28 MindSpore Primitive RmsNorm receives (x, gamma, epsilon) as positional arguments.
29 """
30 return (x, gamma, epsilon), {}
33def _normalize_layernorm_args(input_tensor, normalized_shape, weight=None, bias=None, eps=1e-5):
34 """Normalize layer_norm args to positional form.
36 torch.nn.functional.layer_norm(input_tensor, normalized_shape, weight=None, bias=None, eps=1e-5)
37 has no keyword-only parameters, so everything stays positional.
38 """
39 return (input_tensor, normalized_shape, weight, bias, eps), {}
42class NormDistributedOp(DistributedOp):
43 """Distributed implementation for RmsNorm operator."""
45 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
46 """
47 Preprocess arguments for RmsNorm operator.
49 Args:
50 args (tuple): Positional arguments (x, gamma) where both are DTensors.
51 kwargs (dict): Keyword arguments (none expected).
53 Returns:
54 tuple: (local_args, local_kwargs, cache_values)
55 """
56 args, kwargs = _normalize_rmsnorm_args(*args, **kwargs)
57 x, gamma, epsilon = args
58 local_args = (x.to_local(), gamma.to_local(), epsilon)
59 local_kwargs = {}
60 cache_values = [x.layout, gamma.layout]
61 return local_args, local_kwargs, cache_values
63 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
64 """
65 Infer output layouts for RmsNorm operator.
67 Rules:
68 1. Inputs must not have Partial status.
69 2. x and gamma must share the same mesh_shape.
70 3. Dimensions being normalized (the last len(gamma_tensor_map) dims of x)
71 must not be sharded.
72 4. The sharding of the normalized dimensions of x must match gamma's sharding.
73 5. Output layout keeps sharding on non-normalized dims and replicates
74 on normalized dims.
76 Args:
77 cache_values (list): [x_layout, gamma_layout]
79 Returns:
80 tuple: ((x_layout, out_layout), None)
82 Raises:
83 ValueError: If any rule above is violated.
84 """
85 if len(cache_values) < 2:
86 raise ValueError(
87 f"For {self.op_name}, cache_values size {len(cache_values)} is less than 2."
88 )
89 x_layout = cache_values[0]
90 gamma_layout = cache_values[1]
91 # Check partial inputs
92 if not self._allow_partial_inputs:
93 self._check_partial_inputs([x_layout, gamma_layout])
94 x_mesh_shape = x_layout.mesh_shape
95 gamma_mesh_shape = gamma_layout.mesh_shape
96 if x_mesh_shape != gamma_mesh_shape:
97 raise ValueError(f"{self.op_name} inputs must have same mesh_shape")
98 x_alias_map = x_layout.alias_tensor_map
99 gamma_alias_map = gamma_layout.alias_tensor_map
100 if len(gamma_alias_map) > len(x_alias_map):
101 raise ValueError(
102 f"For {self.op_name}, gamma ndim {len(gamma_alias_map)} cannot exceed "
103 f"input ndim {len(x_alias_map)}."
104 )
105 begin_norm_axis = len(x_alias_map) - len(gamma_alias_map)
106 for alias_entry in x_alias_map[begin_norm_axis:]:
107 entries = alias_entry if isinstance(alias_entry, tuple) else (alias_entry,)
108 for name in entries:
109 if name == "None":
110 continue
111 axis_idx = x_layout.alias_name.index(name)
112 if x_mesh_shape[axis_idx] > 1:
113 raise ValueError(f"{self.op_name} is disabled to support the splitting after "
114 f"begin_norm_axis {begin_norm_axis} for input 0.")
115 if x_alias_map[begin_norm_axis:] != gamma_alias_map:
116 raise ValueError(f"For {self.op_name}, input sharding from begin_norm_axis "
117 f"{begin_norm_axis}, {x_alias_map[begin_norm_axis:]}, should equal "
118 f"gamma sharding {gamma_alias_map}.")
119 output_layout = Layout(
120 mesh_shape=x_layout.mesh_shape,
121 alias_name=x_layout.alias_name,
122 rank_list=x_layout.rank_list
123 )
124 output_map = x_alias_map[:begin_norm_axis] + ("None",) * len(gamma_alias_map)
125 out_layout = output_layout(*output_map)
126 return ((x_layout, out_layout), None)
129class LayerNormDistributedOp(DistributedOp):
130 """Distributed implementation for torch.nn.functional.layer_norm."""
132 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
133 """
134 Preprocess arguments for layer_norm operator.
136 Args:
137 args (tuple): Positional arguments (input, normalized_shape, weight, bias, eps).
138 kwargs (dict): Keyword arguments (none expected for this functional API).
140 Returns:
141 tuple: (local_args, local_kwargs, cache_values)
142 """
143 args, kwargs = _normalize_layernorm_args(*args, **kwargs)
144 input_tensor, normalized_shape, weight, bias, eps = args
146 # Normalize normalized_shape: int → (int,), list → tuple
147 if isinstance(normalized_shape, int):
148 normalized_shape = (normalized_shape,)
149 elif isinstance(normalized_shape, list):
150 normalized_shape = tuple(normalized_shape)
152 local_args = [
153 input_tensor.to_local(),
154 normalized_shape,
155 weight.to_local() if weight is not None and hasattr(weight, 'to_local') else weight,
156 bias.to_local() if bias is not None and hasattr(bias, 'to_local') else bias,
157 eps,
158 ]
159 local_kwargs = {}
161 cache_values = [input_tensor.layout, normalized_shape]
162 return tuple(local_args), local_kwargs, cache_values
164 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
165 """
166 Infer output layout for layer_norm operator.
168 Rules:
169 1. Input must not have Partial status.
170 2. normalized_shape must be int, list, or tuple.
171 3. normalized_shape dimensions must be ≤ input ndim.
172 4. All dimensions in normalized_shape must be unsharded.
173 5. Output layout is identical to input layout.
175 Args:
176 cache_values (list): [input_layout, normalized_shape]
178 Returns:
179 tuple: ((output_layout,), None)
181 Raises:
182 ValueError: If any rule above is violated.
183 """
184 input_layout = cache_values[0]
185 if input_layout is None:
186 raise ValueError(f"{self.op_name} requires a valid input tensor layout.")
187 normalized_shape = cache_values[1]
188 # Check partial inputs
189 if not self._allow_partial_inputs:
190 self._check_partial_inputs([input_layout])
192 if normalized_shape is None:
193 raise ValueError(f"{self.op_name} requires normalized_shape.")
195 if not isinstance(normalized_shape, tuple):
196 raise ValueError(f"normalized_shape must be int, list, or tuple, got {type(normalized_shape)}")
198 in_alias_map = input_layout.alias_tensor_map
199 input_ndim = len(in_alias_map)
200 norm_ndim = len(normalized_shape)
202 if norm_ndim > input_ndim:
203 raise ValueError(
204 f"normalized_shape {normalized_shape} (dims={norm_ndim}) is larger than input ndim={input_ndim}."
205 )
207 # The last `norm_ndim` dimensions are going to be normalized
208 dims_to_normalize = list(range(input_ndim - norm_ndim, input_ndim))
210 # All normalized dims must be unsharded
211 for dim in dims_to_normalize:
212 alias_entry = in_alias_map[dim]
213 entries = alias_entry if isinstance(alias_entry, tuple) else (alias_entry,)
214 for name in entries:
215 if name == "None":
216 continue
217 raise ValueError(
218 f"Operation {self.op_name}: Cannot perform sharding on normalized dimension {dim}, "
219 f"but found sharding assignment: {in_alias_map[dim]}"
220 )
222 output_layout = Layout(
223 mesh_shape=input_layout.mesh_shape,
224 alias_name=input_layout.alias_name,
225 rank_list=input_layout.rank_list
226 )
227 output_layout = output_layout(*in_alias_map)
228 return ((output_layout,), None)