Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_norm.py: 76%
82 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 RmsNorm takes (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 if x_layout.mesh_shape != gamma_layout.mesh_shape:
95 raise ValueError(f"{self.op_name} inputs must have same mesh_shape")
96 x_alias_map = x_layout.alias_tensor_map
97 gamma_alias_map = gamma_layout.alias_tensor_map
98 if len(gamma_alias_map) > len(x_alias_map):
99 raise ValueError(
100 f"For {self.op_name}, gamma ndim {len(gamma_alias_map)} cannot exceed "
101 f"input ndim {len(x_alias_map)}."
102 )
103 begin_norm_axis = len(x_alias_map) - len(gamma_alias_map)
104 for alias_entry in x_alias_map[begin_norm_axis:]:
105 entries = alias_entry if isinstance(alias_entry, tuple) else (alias_entry,)
106 for name in entries:
107 if name == "None":
108 continue
109 axis_idx = x_layout.alias_name.index(name)
110 if x_layout.mesh_shape[axis_idx] > 1:
111 raise ValueError(f"{self.op_name} is disabled to support the splitting after "
112 f"begin_norm_axis {begin_norm_axis} for input 0.")
113 if x_alias_map[begin_norm_axis:] != gamma_alias_map:
114 raise ValueError(f"For {self.op_name}, input sharding from begin_norm_axis "
115 f"{begin_norm_axis}, {x_alias_map[begin_norm_axis:]}, should equal "
116 f"gamma sharding {gamma_alias_map}.")
117 output_layout = Layout(
118 mesh_shape=x_layout.mesh_shape,
119 alias_name=x_layout.alias_name,
120 rank_list=x_layout.rank_list
121 )
122 output_map = x_alias_map[:begin_norm_axis] + ("None",) * len(gamma_alias_map)
123 out_layout = output_layout(*output_map)
124 return ((x_layout, out_layout), None)
127class LayerNormDistributedOp(DistributedOp):
128 """Distributed implementation for torch.nn.functional.layer_norm."""
130 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
131 """
132 Preprocess arguments for layer_norm operator.
134 Args:
135 args (tuple): Positional arguments (input, normalized_shape, weight, bias, eps).
136 kwargs (dict): Keyword arguments (none expected for this functional API).
138 Returns:
139 tuple: (local_args, local_kwargs, cache_values)
140 """
141 args, kwargs = _normalize_layernorm_args(*args, **kwargs)
142 input_tensor, normalized_shape, weight, bias, eps = args
144 # Normalize normalized_shape: int → (int,), list → tuple
145 if isinstance(normalized_shape, int):
146 normalized_shape = (normalized_shape,)
147 elif isinstance(normalized_shape, list):
148 normalized_shape = tuple(normalized_shape)
150 local_args = [
151 input_tensor.to_local(),
152 normalized_shape,
153 weight.to_local() if weight is not None and hasattr(weight, 'to_local') else weight,
154 bias.to_local() if bias is not None and hasattr(bias, 'to_local') else bias,
155 eps,
156 ]
157 local_kwargs = {}
159 cache_values = [input_tensor.layout, normalized_shape]
160 return tuple(local_args), local_kwargs, cache_values
162 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
163 """
164 Infer output layout for layer_norm operator.
166 Rules:
167 1. Input must not have Partial status.
168 2. normalized_shape must be int, list, or tuple.
169 3. normalized_shape dimensions must be ≤ input ndim.
170 4. All dimensions in normalized_shape must be unsharded.
171 5. Output layout is identical to input layout.
173 Args:
174 cache_values (list): [input_layout, normalized_shape]
176 Returns:
177 tuple: ((output_layout,), None)
179 Raises:
180 ValueError: If any rule above is violated.
181 """
182 input_layout = cache_values[0]
183 if input_layout is None:
184 raise ValueError(f"{self.op_name} requires a valid input tensor layout.")
185 normalized_shape = cache_values[1]
186 # Check partial inputs
187 if not self._allow_partial_inputs:
188 self._check_partial_inputs([input_layout])
190 if normalized_shape is None:
191 raise ValueError(f"{self.op_name} requires normalized_shape.")
193 if not isinstance(normalized_shape, tuple):
194 raise ValueError(f"normalized_shape must be int, list, or tuple, got {type(normalized_shape)}")
196 in_alias_map = input_layout.alias_tensor_map
197 input_ndim = len(in_alias_map)
198 norm_ndim = len(normalized_shape)
200 if norm_ndim > input_ndim:
201 raise ValueError(
202 f"normalized_shape {normalized_shape} (dims={norm_ndim}) is larger than input ndim={input_ndim}."
203 )
205 # The last `norm_ndim` dimensions are going to be normalized
206 dims_to_normalize = list(range(input_ndim - norm_ndim, input_ndim))
208 # All normalized dims must be unsharded
209 for dim in dims_to_normalize:
210 alias_entry = in_alias_map[dim]
211 entries = alias_entry if isinstance(alias_entry, tuple) else (alias_entry,)
212 for name in entries:
213 if name == "None":
214 continue
215 raise ValueError(
216 f"Operation {self.op_name}: Cannot perform sharding on normalized dimension {dim}, "
217 f"but found sharding assignment: {in_alias_map[dim]}"
218 )
220 output_layout = Layout(
221 mesh_shape=input_layout.mesh_shape,
222 alias_name=input_layout.alias_name,
223 rank_list=input_layout.rank_list
224 )
225 output_layout = output_layout(*in_alias_map)
226 return ((output_layout,), None)