Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_reduce.py: 88%
189 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 Reduce operator.
17"""
19from copy import deepcopy
20from typing import Sequence, Union, Tuple, List
22from hyper_parallel.core.dtensor.layout import Layout
23from .parallel_ops import DistributedOp
26StrOrTuple = Union[str, Tuple["StrOrTuple", ...], List["StrOrTuple"]]
29def _normalize_reduce_args(input_tensor, dim=None, keepdim=False, dtype=None):
30 """Normalize reduce-family arguments to consistent positional form.
32 Handles torch.sum / torch.mean / torch.prod / torch.all, where *dim* and
33 *keepdim* are ordinary positional parameters. SumExt and MeanExt also pass
34 a trailing dtype slot positionally; normalize it here while keeping layout
35 inference based only on ``input_layout``, ``dim``, and ``keepdim``.
37 Args:
38 input_tensor: The input tensor (DTensor or Tensor).
39 dim: Dimension(s) to reduce. None means all dimensions, () means no reduction.
40 keepdim: Whether to retain reduced dimensions.
41 dtype: Optional output dtype.
43 Returns:
44 tuple: ``((input_tensor, dim, keepdim, dtype), {})`` — all-positional, empty kwargs.
45 """
46 return (input_tensor, dim, keepdim, dtype), {}
49class ReduceExtDistributedOpBase(DistributedOp):
50 """
51 Base class for distributed reduce operators.
53 Args:
54 op_name (str): Name of the operator to register.
55 partial_type (list): List of the operator for allreduce.
56 """
58 _TORCH_DTYPE_OP_NAMES = frozenset({"sum", "mean", "prod"})
59 _MS_POSITIONAL_DTYPE_OP_NAMES = frozenset({"SumExt", "MeanExt"})
61 def __init__(self, op_name, partial_type=None):
62 super().__init__(op_name)
63 if partial_type is None:
64 partial_type = ["sum"]
65 self.partial_type = partial_type
66 self._allow_partial_inputs = True
68 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
69 """
70 Preprocess arguments for reduce operators.
72 Normalizes ``(input, dim, keepdim)``, converts DTensor inputs to local
73 tensors, and builds ``cache_values`` for layout inference.
75 Args:
76 args (tuple): Positional arguments passed to the operator call.
77 kwargs (dict): Keyword arguments passed to the operator call.
79 Returns:
80 tuple: ``(local_args, local_kwargs, cache_values)`` where
81 ``local_args`` is ``(input_local, dim, keepdim)`` and
82 ``cache_values`` is ``[input_layout, dim, keepdim]``.
83 """
84 has_dtype_arg = len(args) >= 4 or "dtype" in kwargs
85 normalized_args, _ = _normalize_reduce_args(*args, **kwargs)
86 input_tensor, dim, keepdim, dtype = normalized_args
87 local_args = (input_tensor.to_local(), dim, keepdim)
88 local_kwargs = {}
89 if has_dtype_arg:
90 if self.op_name in self._TORCH_DTYPE_OP_NAMES:
91 local_kwargs["dtype"] = dtype
92 elif self.op_name in self._MS_POSITIONAL_DTYPE_OP_NAMES:
93 local_args += (dtype,)
94 else:
95 raise TypeError(
96 f"For {self.op_name}, the `dtype` argument is not supported."
97 )
98 cache_values = [input_tensor.layout, dim, keepdim]
99 return local_args, local_kwargs, cache_values
101 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
102 """
103 Infer output layout for reduce operators.
105 Rules:
106 1. Input must have a valid mesh_shape.
107 2. ``dim`` must be ``None``, ``int``, ``tuple[int]``, or ``list[int]`` —
108 never a ``Tensor``.
109 3. ``dim`` values must be valid axis indices for the input tensor.
110 4. Reduce axes are replaced with ``"None"`` (keepdim) or dropped;
111 Partial status is applied on sharded reduced axes.
113 Args:
114 cache_values (list): ``[input_layout, dim, keepdim]``.
116 Returns:
117 tuple: ``((output_layout,), None)``
119 Raises:
120 ValueError: If the input layout is missing or ``mesh_shape`` is ``None``.
121 TypeError: If ``dim`` is a ``Tensor``.
122 ValueError: If a ``dim`` axis index is out of range.
123 """
124 x_layout = cache_values[0]
125 dim = cache_values[1]
126 keepdim = cache_values[2]
128 if x_layout is None or x_layout.mesh_shape is None:
129 raise ValueError(
130 f"For {self.op_name}, input layout cannot be None."
131 )
133 # Check partial inputs
134 if not self._allow_partial_inputs:
135 self._check_partial_inputs([x_layout])
137 if dim is not None and not isinstance(dim, (int, tuple, list)):
138 raise TypeError(
139 f"For {self.op_name}, the `dim` argument should be `None`, `int`, "
140 f"`tuple[int]`, or `list[int]`, but got {type(dim).__name__}."
141 )
143 # Infer the output shape based on dim and keepdim
144 output_layout = self._infer_output_layout(x_layout, dim, keepdim)
146 return ((output_layout,), None)
148 def _infer_output_layout(self, x_layout, dim, keepdim):
149 """Infer output layout for reduce operator."""
150 # Case 1: dim is None — reduce all dimensions (scalar output or all-None when keepdim).
151 if dim is None:
152 return self._handle_all_axis_reduce(x_layout, keepdim)
154 # Case 2: dim is an empty tuple/list — reduce no dimensions, output layout equals input.
155 if isinstance(dim, (tuple, list)) and len(dim) == 0:
156 output_layout = Layout(
157 mesh_shape=x_layout.mesh_shape,
158 alias_name=x_layout.alias_name,
159 rank_list=x_layout.rank_list
160 )
161 return output_layout(*x_layout.alias_tensor_map)
163 # Case 3: dim is int, tuple, or list with at least one element.
164 output_layout = Layout(
165 mesh_shape=x_layout.mesh_shape,
166 alias_name=x_layout.alias_name,
167 rank_list=x_layout.rank_list
168 )
169 x_map = x_layout.alias_tensor_map
170 reduce_alias, x_map = self.replace_axis_with_none(dim, x_layout, keepdim)
171 output_layout = output_layout(*x_map)
172 self._apply_partial(output_layout, reduce_alias)
173 return output_layout
175 def _handle_all_axis_reduce(self, x_layout, keepdim):
176 """Handle the case where dim is None, meaning reduce all dimensions."""
177 layout = Layout(
178 mesh_shape=x_layout.mesh_shape,
179 alias_name=x_layout.alias_name,
180 rank_list=x_layout.rank_list
181 )
183 if not keepdim:
184 output_layout = layout()
185 else:
186 tensor_map = tuple(["None"] * len(x_layout.alias_tensor_map))
187 output_layout = layout(*tensor_map)
189 self._apply_partial(output_layout, x_layout.alias_tensor_map)
190 return output_layout
192 def replace_axis_with_none(self, dim, x_layout, keepdim):
193 """Replace or drop dimensions depending on keepdim."""
194 if not isinstance(dim, (tuple, list)):
195 dim = [dim]
196 else:
197 dim = list(dim)
199 rank = len(x_layout.alias_tensor_map)
200 for i, axis_id in enumerate(dim):
201 if axis_id < 0:
202 dim[i] = rank + axis_id
203 if not isinstance(axis_id, int) or dim[i] >= rank or dim[i] < 0:
204 raise ValueError(f"Invalid reduce axis index {axis_id} at position {i}.")
206 alias_tensor_map = x_layout.alias_tensor_map
207 reduce_alias = [alias_tensor_map[axis_id] for axis_id in dim if
208 alias_tensor_map[axis_id] is not None and alias_tensor_map[axis_id] != "None"]
209 reduce_alias = self._flatten_aliases(reduce_alias)
211 if keepdim:
212 return self._replace_keepdim(alias_tensor_map, reduce_alias)
213 return self._replace_dropdim(alias_tensor_map, reduce_alias, dim)
215 def _flatten_aliases(self, reduce_alias):
216 """Flatten reduce_alias into a list of atomic alias strings."""
217 flat = []
218 for alias in reduce_alias:
219 if isinstance(alias, (tuple, list)):
220 flat.extend(alias)
221 else:
222 flat.append(alias)
223 return flat
225 def _replace_keepdim(self, alias_tensor_map, reduce_alias):
226 """keepdim, replace reduce alias with 'None'."""
227 new_alias_map = []
228 for alias in alias_tensor_map:
229 if isinstance(alias, (tuple, list)):
230 new_alias = tuple("None" if item in reduce_alias else item for item in alias)
231 new_alias_map.append(new_alias)
232 else:
233 if alias in reduce_alias:
234 new_alias_map.append("None")
235 else:
236 new_alias_map.append(alias)
237 new_alias_map = self._compact_tensor_map(new_alias_map)
238 return reduce_alias, tuple(new_alias_map)
240 def _replace_dropdim(self, alias_tensor_map, reduce_alias, dim):
241 """Compress reduce dim."""
242 new_alias_map = []
243 for i, alias in enumerate(alias_tensor_map):
244 if i in dim:
245 continue
246 if isinstance(alias, (tuple, list)):
247 new_alias = tuple(item for item in alias if item not in reduce_alias)
248 if new_alias:
249 new_alias_map.append(new_alias)
250 else:
251 if alias in reduce_alias:
252 continue
253 new_alias_map.append(alias)
254 new_alias_map = self._compact_tensor_map(new_alias_map)
255 return reduce_alias, tuple(new_alias_map)
257 def _compact_tensor_map(self, alias_map: Sequence[StrOrTuple]) -> Tuple[StrOrTuple, ...]:
258 """Extend tensor map of 'None'."""
260 def _compress(elem: StrOrTuple) -> StrOrTuple:
261 if isinstance(elem, (list, tuple)):
262 compressed = tuple(_compress(e) for e in elem)
263 if len(compressed) == 1:
264 return compressed[0]
265 if all(x == 'None' for x in compressed):
266 return 'None'
267 return compressed
268 return elem
270 return tuple(_compress(elem) for elem in alias_map)
272 def _apply_partial(self, out_layout, alias):
273 """Apply all partial to given alias (string, tuple, list)."""
274 if alias == "None":
275 return
276 if isinstance(alias, (tuple, list)):
277 for elem in alias:
278 self._apply_partial(out_layout, elem)
279 else:
280 for ops in self.partial_type:
281 out_layout.set_partial_by_dev_axis(alias, ops)
284class SumExtDistributedOp(ReduceExtDistributedOpBase):
285 """Distributed implementation for SumExt operator."""
287 def __init__(self, op_name="SumExt"):
288 super().__init__(op_name, partial_type=["sum"])
291class MeanExtDistributedOp(ReduceExtDistributedOpBase):
292 """Distributed implementation for MeanExt operator."""
294 def __init__(self, op_name="MeanExt"):
295 super().__init__(op_name, partial_type=["avg"])
298class ReduceMaxDistributedOp(ReduceExtDistributedOpBase):
299 """Distributed implementation for ReduceMax operator."""
301 def __init__(self, op_name="ReduceMax"):
302 super().__init__(op_name, partial_type=["max"])
305class ProdExtDistributedOp(ReduceExtDistributedOpBase):
306 """
307 Distributed implementation for ProdExt operator (product of all elements or along a dim).
308 Compatible with torch.prod arguments.
309 """
311 def __init__(self, op_name="prod"):
312 super().__init__(op_name, partial_type=["prod"])
315class AllExtDistributedOp(ReduceExtDistributedOpBase):
316 """
317 Distributed implementation for All operator
318 Returns the cumulative sum of elements of input in the dimension dim.
319 """
321 def __init__(self, op_name="all"):
322 super().__init__(op_name, partial_type=["all"])
325class MaxDistributedOp(ReduceExtDistributedOpBase):
326 """
327 Distributed implementation for Pytorch style Max operator.
329 Supports three Pytorch behaviors:
330 1. torch.max(input) -> Global reduction (returns single Tensor)
331 2. torch.max(input, dim, keepdim=False) -> Dimension reduction (returns (values, indices))
332 3. torch.max(input, other) -> Element-wise max (returns single Tensor)
333 """
335 def __init__(self, op_name="max"):
336 super().__init__(op_name, partial_type=["max"])
338 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
339 """
340 Preprocess arguments for Max/Min operators.
342 Routes between element-wise mode (two DTensor inputs) and reduction
343 mode (input, dim, keepdim). All parameters are positional with empty
344 kwargs.
346 Args:
347 args (tuple): Positional arguments passed to the operator call.
348 kwargs (dict): Keyword arguments passed to the operator call.
350 Returns:
351 tuple: ``(local_args, local_kwargs, cache_values)``.
352 """
353 # Element-wise mode: torch.max(a, b) — second arg is a DTensor.
354 if len(args) >= 2 and hasattr(args[1], "_layout"):
355 input_tensor = args[0]
356 other_tensor = args[1]
357 local_args = (input_tensor.to_local(), other_tensor.to_local())
358 local_kwargs = {}
359 cache_values = [input_tensor.layout, other_tensor.layout]
360 return local_args, local_kwargs, cache_values
362 # Reduction mode: torch.max(input, dim, keepdim) / MaxDim(input, dim, keepdim).
363 has_dtype_arg = len(args) >= 4 or "dtype" in kwargs
364 args, kwargs = _normalize_reduce_args(*args, **kwargs)
365 input_tensor, dim, keepdim, _ = args
366 if has_dtype_arg:
367 raise TypeError(
368 f"For {self.op_name}, the `dtype` argument is not supported."
369 )
370 # torch.max(x) global reduction: only pass the tensor so the call
371 # remains torch.max(local_x), not torch.max(local_x, None, False).
372 if dim is None:
373 local_args = (input_tensor.to_local(),)
374 else:
375 local_args = (input_tensor.to_local(), dim, keepdim)
376 local_kwargs = {}
378 cache_values = [input_tensor.layout, dim, keepdim]
379 return local_args, local_kwargs, cache_values
381 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
382 """
383 Infer output layouts for Max/Min operators.
385 Rules:
386 1. Element-wise mode (two Layouts in *cache_values*): output layout
387 equals the first input layout.
388 2. Reduction mode: same rules as ``ReduceExtDistributedOpBase``.
389 3. When ``dim`` is ``None`` (global reduction), returns a single layout.
390 4. When ``dim`` is specified, returns ``(values_layout, indices_layout)``.
391 5. ``dim`` must be ``None``, ``int``, ``tuple[int]``, or ``list[int]``.
393 Args:
394 cache_values (list): ``[input_layout, dim, keepdim]`` for reduction,
395 or ``[input_layout, other_layout]`` for element-wise.
397 Returns:
398 tuple: ``((output_layout,), None)`` or ``((values_layout, indices_layout), None)``.
400 Raises:
401 ValueError: If the input layout is missing or ``mesh_shape`` is ``None``.
402 TypeError: If ``dim`` is not ``None``, ``int``, ``tuple``, or ``list``.
403 """
404 # Element-wise mode: two Layout objects in cache_values.
405 if len(cache_values) == 2 and hasattr(cache_values[1], "mesh_shape"):
406 # Check partial inputs
407 if not self._allow_partial_inputs:
408 self._check_partial_inputs(cache_values)
409 return ((deepcopy(cache_values[0]),), None)
411 x_layout = cache_values[0]
412 dim = cache_values[1]
413 keepdim = cache_values[2]
415 if x_layout is None or x_layout.mesh_shape is None:
416 raise ValueError(
417 f"For {self.op_name}, input layout cannot be None."
418 )
420 # Check partial inputs
421 if not self._allow_partial_inputs:
422 self._check_partial_inputs([x_layout])
424 if dim is not None and not isinstance(dim, (int, tuple, list)):
425 raise TypeError(
426 f"For {self.op_name}, the `dim` argument should be `None`, `int`, "
427 f"`tuple[int]`, or `list[int]`, but got {type(dim).__name__}."
428 )
430 values_layout = self._infer_output_layout(x_layout, dim, keepdim)
432 if dim is None:
433 # torch.max(input) -> Single Tensor
434 return ((values_layout,), None)
436 # torch.max(input, dim) -> (values, indices)
437 indices_layout = deepcopy(values_layout)
438 return ((values_layout, indices_layout), None)
441class MinDistributedOp(MaxDistributedOp):
442 """
443 Distributed implementation for Pytorch style Min operator.
445 Supports three Pytorch behaviors:
446 1. torch.min(input) -> Global reduction (returns single Tensor)
447 2. torch.min(input, dim, keepdim=False) -> Dimension reduction (returns (values, indices))
448 3. torch.min(input, other) -> Element-wise min (returns single Tensor)
449 """
451 def __init__(self, op_name="min"):
452 # Call the parent class (MaxDistributedOp) initialization
453 super().__init__(op_name=op_name)
454 # Override the partial_type to use "min" instead of "max" for the underlying communication
455 self.partial_type = ["min"]