Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_one_hot_ext.py: 68%
150 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 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 OneHotExt operator.
17"""
19from typing import Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from hyper_parallel.platform import get_platform
23from .parallel_ops import DistributedOp
25platform = get_platform()
28def _normalize_one_hot_ext_args(indices, num_classes, on_value, off_value, axis):
29 return (indices, num_classes, on_value, off_value, axis), {}
32class OneHotExtDistributedOp(DistributedOp):
33 """Distributed implementation for OneHotExt operator."""
35 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
36 """
37 Preprocess arguments for OneHotExt operator.
39 Args:
40 args (tuple): Input arguments (indices, num_classes, on_value, off_value, axis).
41 kwargs (dict): Keyword arguments (empty for this operator).
43 Returns:
44 tuple: (local_args, local_kwargs, cache_values)
45 """
46 args, kwargs = _normalize_one_hot_ext_args(*args, **kwargs)
47 indices, num_classes, on_value, off_value, axis = args
49 indices_local = indices.to_local()
50 on_value_local = on_value.to_local() if hasattr(on_value, '_layout') else on_value
51 off_value_local = off_value.to_local() if hasattr(off_value, '_layout') else off_value
53 on_value_layout = on_value.layout if hasattr(on_value, '_layout') else None
54 off_value_layout = off_value.layout if hasattr(off_value, '_layout') else None
56 local_args = (indices_local, num_classes, on_value_local, off_value_local, axis)
57 local_kwargs = {}
58 cache_values = [indices.layout, on_value_layout, off_value_layout, num_classes, axis]
59 return local_args, local_kwargs, cache_values
61 # pylint: disable=W0237
62 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
63 """
64 Infer output layout for OneHotExt.
66 Rules:
67 1. Indices must not have Partial status.
68 2. num_classes must be int >= -1.
69 3. axis must be in [-1, 1].
70 4. For multi-dimensional input (>1D), axis must be -1 and only dim0 may be sharded.
71 5. Non-indices inputs must be fully replicated.
72 6. Output layout inserts a replicated one-hot dimension at the specified axis.
74 Args:
75 cache_values (list): [indices_layout, on_value_layout, off_value_layout, num_classes, axis]
77 Returns:
78 tuple: ((output_layout,), None)
80 Raises:
81 ValueError: If any rule above is violated.
82 TypeError: If num_classes or axis has invalid type.
83 """
84 indices_layout = cache_values[0]
85 on_value_layout = cache_values[1]
86 off_value_layout = cache_values[2]
87 num_classes = cache_values[3]
88 axis = cache_values[4]
90 if indices_layout is None or indices_layout.mesh_shape is None:
91 raise ValueError(
92 f"For {self.op_name}, indices layout cannot be None."
93 )
95 if not self._allow_partial_inputs:
96 self._check_partial_inputs([indices_layout])
98 self._validate_num_classes(num_classes)
99 axis = self._validate_axis(axis)
101 in_tensor_map = indices_layout.tensor_map
102 if not in_tensor_map:
103 raise ValueError(
104 f"For {self.op_name}, indices tensor_map is empty."
105 )
107 self._validate_multi_dim_restriction(in_tensor_map, axis, indices_layout)
108 self._validate_inputs_layouts(
109 [indices_layout, on_value_layout, off_value_layout]
110 )
112 out_tensor_map = self._infer_output_tensor_map(in_tensor_map, axis)
113 out_layout = self._create_layout_from_tensor_map(indices_layout, out_tensor_map)
114 out_layout.tensor_map_to_placement()
116 return ((out_layout,), None)
118 # pylint: disable=W0237
119 def get_expand_impl(self, func, infer_result, cache_values):
120 """
121 Get expanded implementation for OneHotExt operator.
123 When indices are sharded and num_classes is -1 (auto-detect), returns a
124 closure that computes the global maximum index across all shards via
125 AllReduce(max) before calling the original operator.
127 Args:
128 func: Original operator callable.
129 infer_result: Result from infer_layout (unused).
130 cache_values (list): [indices_layout, on_value_layout, off_value_layout, num_classes, axis]
132 Returns:
133 Optional[callable]: Closure or None if no expansion is needed.
134 """
135 # pylint: disable=C0415
136 import mindspore as ms
137 from mindspore import ops, Tensor
139 indices_layout = cache_values[0]
140 if indices_layout is None or indices_layout.mesh_shape is None:
141 return None
143 sharded_axes = self._get_sharded_axes(indices_layout)
144 if not sharded_axes:
145 return None
147 original_op = func
148 reduce_max = ops.ReduceMax(keep_dims=False)
150 def expanded_one_hot(indices, num_classes, on_value, off_value, axis):
151 self._validate_num_classes(num_classes)
152 self._validate_indices_dtype(indices)
154 if num_classes != -1:
155 return original_op(indices, num_classes, on_value, off_value, axis)
157 local_max = reduce_max(indices, ())
158 if not isinstance(local_max, Tensor):
159 local_max = Tensor(local_max, ms.int64)
161 local_max_host = int(local_max.asnumpy())
162 if local_max_host > 2147483647:
163 raise ValueError(
164 f"For {self.op_name}, indices max value {local_max_host} "
165 f"exceeds int32 range."
166 )
168 zero_dim = local_max.ndim == 0
169 local_max_i32 = ops.cast(local_max, ms.int32)
171 if zero_dim:
172 local_max_i32 = ops.expand_dims(local_max_i32, 0)
174 global_max_i32 = local_max_i32
175 for axis_name in sharded_axes:
176 group = indices_layout.get_comm_group_by_axis(axis_name)
177 global_max_i32 = platform.differentiable_all_reduce(
178 global_max_i32, "max", group
179 )
181 if zero_dim:
182 global_max_i32 = ops.squeeze(global_max_i32, 0)
184 depth = int(global_max_i32.asnumpy()) + 1
185 return original_op(indices, depth, on_value, off_value, axis)
187 return expanded_one_hot
189 def _validate_num_classes(self, num_classes):
190 """Validate num_classes parameter."""
191 if not isinstance(num_classes, int):
192 raise TypeError(
193 f"For {self.op_name}, num_classes should be int, "
194 f"but got {type(num_classes).__name__}."
195 )
196 if num_classes < -1:
197 raise ValueError(
198 f"For {self.op_name}, num_classes should be >= -1, "
199 f"but got {num_classes}."
200 )
202 def _validate_indices_dtype(self, indices):
203 """Validate indices dtype."""
204 # pylint: disable=C0415
205 import mindspore as ms
207 if indices.dtype != ms.int64:
208 raise TypeError(
209 f"For {self.op_name}, indices dtype should be int64, "
210 f"but got {indices.dtype}."
211 )
213 def _get_sharded_axes(self, layout):
214 """Get all device axes that are used for sharding."""
215 sharded_axes = set()
217 if layout is None or layout.alias_tensor_map is None:
218 return []
220 for dim_alias in layout.alias_tensor_map:
221 if dim_alias == "None":
222 continue
224 if isinstance(dim_alias, tuple):
225 for axis_name in dim_alias:
226 if axis_name != "None":
227 sharded_axes.add(axis_name)
228 else:
229 sharded_axes.add(dim_alias)
231 return list(sharded_axes)
233 def _validate_axis(self, axis):
234 """Validate axis parameter."""
235 if not isinstance(axis, int):
236 raise TypeError(
237 f"For {self.op_name}, axis should be int, "
238 f"but got {type(axis).__name__}."
239 )
241 if axis > 1 or axis < -1:
242 raise ValueError(
243 f"For {self.op_name}, axis {axis} is out of range [-1, 1]."
244 )
246 return axis
248 def _validate_multi_dim_restriction(self, in_tensor_map, axis, indices_layout):
249 """Validate restriction for multi-dimensional inputs."""
250 in_rank = len(in_tensor_map)
251 if in_rank <= 1:
252 return
254 if axis != -1:
255 raise ValueError(
256 f"For {self.op_name}, when input dimension is > 1, axis should be -1, "
257 f"but got {axis}."
258 )
260 alias_map = indices_layout.alias_tensor_map
261 for i in range(1, len(alias_map)):
262 if alias_map[i] != "None":
263 raise ValueError(
264 f"For {self.op_name}, when input dimension is > 1, "
265 f"strategy should be data parallel, "
266 f"but dimension {i} is sharded on '{alias_map[i]}'."
267 )
269 def _validate_inputs_layouts(self, layouts):
270 """Validate that non-indices inputs are fully replicated."""
271 for layout in layouts[1:]:
272 if layout is None:
273 continue
274 alias_map = layout.alias_tensor_map
275 if alias_map and any(x != "None" for x in alias_map):
276 raise ValueError(
277 f"For {self.op_name}, non-indices inputs should be replicated, "
278 f"but got {alias_map}."
279 )
281 def _infer_output_tensor_map(self, in_tensor_map, axis):
282 """Infer output tensor map by inserting one-hot dimension at specified axis."""
283 in_rank = len(in_tensor_map)
285 if axis in (-1, in_rank):
286 insert_pos = in_rank
287 else:
288 insert_pos = axis
290 if insert_pos < 0 or insert_pos > in_rank:
291 raise ValueError(
292 f"For {self.op_name}, axis {axis} is out of range "
293 f"for input with rank {in_rank}."
294 )
296 out_tensor_map = list(in_tensor_map)
297 out_tensor_map.insert(insert_pos, -1)
298 return tuple(out_tensor_map)
300 def _create_layout_from_tensor_map(self, base_layout, out_tensor_map):
301 """Create output layout from tensor map."""
302 out_layout = Layout(
303 mesh_shape=base_layout.mesh_shape,
304 alias_name=base_layout.alias_name,
305 rank_list=base_layout.rank_list,
306 )
308 out_layout.set_tensor_map(out_tensor_map)
309 out_layout.set_alias_tensor_map(
310 self._tensor_map_to_alias_tensor_map(base_layout, out_tensor_map)
311 )
312 out_layout.update_compact_str()
313 return out_layout
315 def _tensor_map_to_alias_tensor_map(self, base_layout, tensor_map):
316 """Convert numeric tensor map to alias tensor map."""
317 alias_name = base_layout.alias_name
318 alias_tensor_map = []
320 for dim in tensor_map:
321 if dim == -1:
322 alias_tensor_map.append("None")
323 continue
325 if isinstance(dim, tuple):
326 names = tuple(
327 alias_name[len(alias_name) - 1 - d] for d in dim if d != -1
328 )
329 alias_tensor_map.append(names if names else "None")
330 continue
332 alias_tensor_map.append(alias_name[len(alias_name) - 1 - dim])
334 return tuple(alias_tensor_map)