Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_gather.py: 72%
239 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 Gather operator.
17"""
19from typing import Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from hyper_parallel.core.utils.communication import get_group_local_rank
23from .parallel_ops import DistributedOp
26def _normalize_index_select_args(input_tensor, dim, index):
27 return (input_tensor, dim, index), {}
30def _normalize_gatherd_args(input_tensor, dim, index, **kwargs):
31 """Normalize torch.gather / GatherD args, forwarding keyword-only params."""
32 return (input_tensor, dim, index), kwargs
35def _normalize_gathernd_args(input_tensor, indices):
36 return (input_tensor, indices), {}
39class IndexSelectDistributedOp(DistributedOp):
40 """Distributed implementation for Index Select operator."""
42 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
43 """
44 Preprocess arguments for IndexSelect operator.
46 Args:
47 args (tuple): Input arguments (input, dim, index).
48 kwargs (dict): Keyword arguments.
50 Returns:
51 tuple: (local_args, local_kwargs, cache_values)
52 """
53 args, _ = _normalize_index_select_args(*args, **kwargs)
54 input_tensor, dim, index = args[0], args[1], args[2]
55 local_args = (input_tensor.to_local(), dim, index.to_local())
56 local_kwargs = {}
57 cache_values = [input_tensor.layout, index.layout, dim]
58 return local_args, local_kwargs, cache_values
60 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
61 """
62 Infer output layouts for Index Select operations.
64 Rules:
65 1. Input and index must not have Partial status.
66 2. cache_values must contain [input_layout, index_layout, dim].
67 3. dim must be within the valid input rank range.
68 4. index must be one-dimensional.
69 5. Output replaces the selected input dimension with the index layout.
70 6. If the selected input dimension is sharded, output carries Partial('sum').
72 Args:
73 cache_values (list): [input_layout, index_layout, dim].
75 Returns:
76 tuple: ((output_layout,), None)
78 Raises:
79 ValueError: If input layouts are not compatible or have partial status.
80 """
81 if not self._allow_partial_inputs:
82 self._check_partial_inputs(cache_values[:2])
84 # Parse layout info
85 p_layout, i_layout, axis = cache_values[0], cache_values[1], cache_values[2]
87 p_tensor_map = p_layout.alias_tensor_map
88 i_tensor_map = i_layout.alias_tensor_map
90 # 1. Validate the axis range before any manipulation
91 if axis < -len(p_tensor_map) or axis >= len(p_tensor_map):
92 raise ValueError(
93 f"For {self.op_name}, dim value {axis} is out of valid range"
94 )
96 # 2. Convert negative axis to positive index to avoid Python slicing bugs
97 if axis < 0:
98 axis += len(p_tensor_map)
100 if len(i_tensor_map) != 1:
101 raise ValueError(
102 f"For {self.op_name}, index is not a one-dimensional Tensor"
103 )
105 # 3. Create output layout map
106 # We allow sharding on the `axis`. Since `index_select` replaces the `axis`
107 # dimension with the `index` dimension, if `axis` was sharded, that mesh
108 # dimension is removed from the output tensor map.
109 output_tensor_map = list(p_tensor_map[:axis]) + list(i_tensor_map) + list(p_tensor_map[axis + 1 :])
111 output_layout = Layout(
112 mesh_shape=p_layout.mesh_shape,
113 alias_name=p_layout.alias_name,
114 rank_list=p_layout.rank_list,
115 )
116 output_layout = output_layout(*output_tensor_map)
118 # 4. Implicit Communication via Partial Layout
119 # If the gather axis was sharded, the local output will only be a masked partial result.
120 # We set the output layout to Partial('sum') for that specific mesh dimension so the
121 # OpDispatcher handles the AllReduce automatically when this tensor is used later.
122 shard_mesh_dim_name = p_tensor_map[axis]
123 if shard_mesh_dim_name != "None":
124 # Handle possible multi-axis sharding tuple
125 if isinstance(shard_mesh_dim_name, tuple):
126 for dim_name in shard_mesh_dim_name:
127 if dim_name != "None":
128 output_layout.set_partial_by_dev_axis(dim_name, 'sum')
129 else:
130 output_layout.set_partial_by_dev_axis(shard_mesh_dim_name, 'sum')
132 return ((output_layout,), None)
134 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
135 """
136 Get the expanded execution implementation for Index Select.
137 """
138 p_layout = cache_values[0]
139 axis = cache_values[2]
140 if axis < 0:
141 axis += len(p_layout.alias_tensor_map)
143 shard_mesh_dim_name = p_layout.alias_tensor_map[axis]
145 # If the axis is NOT sharded, fallback to standard execution
146 if shard_mesh_dim_name == "None":
147 return None
149 # If the axis IS sharded, return a custom function with Masking ONLY.
150 # The explicit AllReduce is completely removed.
151 def expand_impl(input_tensor, dim, index, **kwargs):
152 mesh = p_layout.mesh
154 # Fetch the communication group for the sharded mesh dimension
155 if isinstance(shard_mesh_dim_name, tuple):
156 target_dim_name = next(d for d in shard_mesh_dim_name if d != "None")
157 else:
158 target_dim_name = shard_mesh_dim_name
160 comm_group_info = mesh.get_comm_group_by_axis(target_dim_name)
161 group = comm_group_info.group if hasattr(comm_group_info, 'group') else comm_group_info
163 # Get the rank of the current device within this specific communication group
164 group_rank = get_group_local_rank(group)
166 # Calculate global index boundaries for the local chunk
167 local_dim_size = input_tensor.shape[dim]
168 start_idx = group_rank * local_dim_size
169 end_idx = start_idx + local_dim_size
171 # 1. Compute mask: True for indices that belong to the current rank
172 mask = (index >= start_idx) & (index < end_idx)
174 # 2. Shift global indices to local indices
175 safe_index = index - start_idx
177 # Clamp safe_index to valid local ranges to prevent CUDA out-of-bounds
178 # errors during the local index_select (invalid ones will be masked out anyway).
179 safe_index = safe_index.clamp(min=0, max=local_dim_size - 1)
181 # 3. Perform local index_select using tensor's built-in method
182 local_out = input_tensor.index_select(dim, safe_index, **kwargs)
184 # 4. Mask out the invalid indices (set them to 0)
185 # Reshape the 1D mask to broadcast against the output shape
186 mask_shape = [1] * local_out.ndim
187 mask_shape[dim] = -1
188 mask_reshaped = mask.reshape(mask_shape).to(local_out.dtype)
190 local_out = local_out * mask_reshaped
192 # Return the partial local tensor directly. The framework's layout engine
193 # and OpDispatcher will trigger the AllReduce when this Partial tensor
194 # is redistributed to a non-partial layout.
195 return local_out
197 return expand_impl
200class GatherDDistributedOp(DistributedOp):
201 """Distributed implementation for GatherD operator.
203 GatherD gathers values along a specified axis from the input tensor using the index tensor.
205 Signature: GatherD(input, dim, index) -> output
207 Key constraints:
208 - Input and index must have the same number of dimensions
209 - Output inherits the sharding pattern of the input tensor
210 """
212 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
213 """
214 Preprocess arguments for GatherD operator.
216 Args:
217 args (tuple): Input arguments (input, dim, index).
218 kwargs (dict): Keyword arguments.
220 Returns:
221 tuple: (local_args, local_kwargs, cache_values)
222 """
223 args, kwargs = _normalize_gatherd_args(*args, **kwargs)
224 input_tensor, dim, index = args[0], args[1], args[2]
225 local_args = (input_tensor.to_local(), dim, index.to_local())
226 local_kwargs = kwargs
227 cache_values = [input_tensor.layout, index.layout, dim]
228 return local_args, local_kwargs, cache_values
230 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
231 """
232 Infer output layouts for GatherD operations.
234 Rules:
235 1. Input and index must not have Partial status.
236 2. cache_values must contain [input_layout, index_layout, dim].
237 3. Input and index must have the same rank.
238 4. dim must be within the valid input rank range.
239 5. Input and index must use the same sharding on non-dim axes.
240 6. Output inherits index layout and becomes Partial('sum') when dim is sharded.
242 Args:
243 cache_values (list): [input_layout, index_layout, dim].
245 Returns:
246 tuple: ((output_layout,), None)
248 Raises:
249 ValueError: If input layouts are not compatible or have partial status.
250 """
251 if not self._allow_partial_inputs:
252 self._check_partial_inputs(cache_values[:2])
254 input_layout, index_layout, dim = cache_values[0], cache_values[1], cache_values[2]
255 # Validate layouts exist
256 if input_layout is None or not hasattr(input_layout, "tensor_map"):
257 raise ValueError(f"For {self.op_name}, input layout cannot be None")
258 if index_layout is None or not hasattr(index_layout, "tensor_map"):
259 raise ValueError(f"For {self.op_name}, index layout cannot be None")
260 input_tensor_map = input_layout.alias_tensor_map
261 index_tensor_map = index_layout.alias_tensor_map
262 # Validate same rank
263 if len(input_tensor_map) != len(index_tensor_map):
264 raise ValueError(
265 f"For {self.op_name}, input and index must have the same number of dimensions. "
266 f"Got input rank={len(input_tensor_map)}, index rank={len(index_tensor_map)}"
267 )
268 # Validate dim is in valid range
269 rank = len(input_tensor_map)
270 if dim < -rank or dim >= rank:
271 raise ValueError(
272 f"For {self.op_name}, dim value {dim} is out of valid range [{-rank}, {rank-1}]"
273 )
274 # Normalize negative dim
275 if dim < 0:
276 dim = dim + rank
277 for axis, (input_axis_map, index_axis_map) in enumerate(zip(input_tensor_map, index_tensor_map)):
278 if axis == dim:
279 continue
280 if input_axis_map != index_axis_map:
281 raise ValueError(
282 f"For {self.op_name}, input and index must use the same sharding on non-dim axis {axis}. "
283 f"Got input tensor_map={input_tensor_map}, index tensor_map={index_tensor_map}, dim={dim}"
284 )
285 # Output inherits index layout
286 output_layout = Layout(
287 mesh_shape=index_layout.mesh_shape,
288 alias_name=index_layout.alias_name,
289 rank_list=index_layout.rank_list,
290 )
291 output_layout.set_tensor_map(index_layout.tensor_map)
292 dim_axis_name = input_tensor_map[dim]
293 if dim_axis_name != "None":
294 # pylint: disable=protected-access
295 # Inherit current partial state from index layout
296 output_layout._partial = list(index_layout.partial)
297 if isinstance(dim_axis_name, tuple):
298 for axis_name in dim_axis_name:
299 if axis_name != "None":
300 output_layout.set_partial_by_dev_axis(axis_name, 'sum')
301 else:
302 output_layout.set_partial_by_dev_axis(dim_axis_name, 'sum')
303 # pylint: disable=protected-access
304 # Rebuild readable alias tensor map
305 output_layout._alias_tensor_map = output_layout._build_readable_tensor_map()
306 # pylint: disable=protected-access
307 # Sync tensor_map to placement representation
308 output_layout.tensor_map_to_placement()
309 # Update compact string description
310 output_layout.update_compact_str()
311 return ((output_layout,), None)
313 def get_expand_impl(self, func, infer_result, cache_values): # pylint: disable=W0221
314 """
315 Returns the execution implementation wrapper for distributed GatherD.
317 When the dim axis is sharded, each rank gathers from its local slice of the input tensor.
318 The indices need to be adjusted to account for the local partition offset.
320 Args:
321 func: The original GatherD function to wrap
322 infer_result: The inferred output layouts and extra info
323 cache_values: [input_layout, index_layout, dim]
325 Returns:
326 Callable: Distributed implementation wrapper, or None if no sharding
327 """
328 input_layout = cache_values[0]
329 dim = cache_values[2]
330 if dim < 0:
331 dim += len(input_layout.tensor_map)
332 input_alias_map = input_layout.alias_tensor_map
333 # Check if dim axis is sharded (enhanced MP)
334 if input_alias_map[dim] == "None": # native sharding, no need for custom implementation
335 return None
337 dim_axis_name = input_alias_map[dim]
338 if isinstance(dim_axis_name, tuple):
339 dim_axis_name = next(axis for axis in dim_axis_name if axis != "None")
341 def distributed_gatherd_impl(*args, **kwargs):
342 """
343 Distributed GatherD implementation for sharded dim axis.
345 Each rank gathers from its local slice of input tensor.
346 Indices are adjusted by subtracting the local partition offset.
347 """
348 input_tensor = args[0]
349 index_tensor = args[2]
350 # Calculate local partition offset for the dim axis
351 mesh = input_layout.mesh
352 mesh_dim_idx = input_layout.alias_name.index(dim_axis_name)
353 # Get the coordinate of current rank along the mesh dimension
354 dim_coord = mesh.get_local_rank(mesh_dim_idx)
355 # Calculate the size of input tensor's dim dimension per partition
356 input_dim_size = input_tensor.shape[dim]
357 # Calculate the starting index of local partition
358 local_start_index = int(dim_coord * input_dim_size)
359 local_end_index = int(local_start_index + input_dim_size)
360 # Adjust indices: subtract local_start_index to map global indices to local range
361 # This is similar to how Embedding shifts indices for Row Parallelism
362 adjusted_index = index_tensor - local_start_index
363 # Create mask to identify out-of-bounds indices
364 # Indices outside [0, local_dim_size) belong to other partitions
365 mask = (index_tensor >= local_start_index) & (index_tensor < local_end_index)
366 # Cross-platform cast to matching int dtype
367 mask_int = mask.to(index_tensor.dtype)
368 # Zero out invalid indices to prevent out-of-bounds access
369 safe_index = adjusted_index * mask_int
370 # Replace original index tensor with adjusted index
371 new_args = list(args)
372 new_args[2] = safe_index
373 # Execute native GatherD with adjusted indices
374 output = func(*new_args, **kwargs)
375 # Zero out outputs corresponding to invalid indices
376 mask_int = mask_int.to(output.dtype)
377 output = output * mask_int
378 return output
379 return distributed_gatherd_impl
382class GatherNdDistributedOp(DistributedOp):
383 """Distributed implementation for GatherNd operator."""
385 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
386 """
387 Preprocess arguments for GatherNd operator.
389 Args:
390 args (tuple): Input arguments (input, indices).
391 kwargs (dict): Keyword arguments.
393 Returns:
394 tuple: (local_args, local_kwargs, cache_values)
395 """
396 args, _ = _normalize_gathernd_args(*args, **kwargs)
397 input_tensor, indices = args[0], args[1]
398 local_input = input_tensor.to_local() if hasattr(input_tensor, "_layout") else input_tensor
399 local_indices = indices.to_local()
400 local_args = (local_input, local_indices)
401 local_kwargs = {}
402 cache_values = [
403 input_tensor.layout if hasattr(input_tensor, "_layout") else None,
404 indices.layout,
405 input_tensor.shape,
406 indices.shape,
407 ]
408 return local_args, local_kwargs, cache_values
410 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
411 """
412 Infer output layout for GatherNd.
414 Rules:
415 1. Input and indices must not have Partial status.
416 2. cache_values must contain [input_layout_or_None, indices_layout, input_shape, indices_shape].
417 3. indices[-1] (K) must be replicated.
418 4. input indexed dims [0:K) must be replicated when input layout is provided.
419 5. Output inherits indices[:-1] sharding plus input trailing dims input[K:].
421 For GatherNd: out.shape = indices.shape[:-1] + input_x.shape[K:], where K = indices.shape[-1].
423 This implementation:
424 - Inherits sharding from indices[:-1].
425 - Allows sharding on input_x trailing dims input_x[K:].
426 - Requires input_x[:K] to be replicated ("None") if input_layout is provided.
427 - Requires indices[-1] (K dim) to be replicated ("None").
429 Output Layout:
430 output_tensor_map = indices_tensor_map[:-1] + input_tensor_map[K:]
431 If input_layout is None, input trailing dims are treated as replicated ("None").
433 Args:
434 cache_values (list): [input_layout_or_None, indices_layout, input_shape, indices_shape].
436 Returns:
437 tuple: ((output_layout,), None)
439 Raises:
440 ValueError: If input layouts, tensor maps, or shapes violate the rules above.
441 """
442 input_layout, indices_layout = self._parse_input_layouts(cache_values[:2])
443 if not self._allow_partial_inputs:
444 self._check_partial_inputs([input_layout, indices_layout])
446 input_shape, indices_shape = self._get_input_shapes(cache_values[2:])
447 k, trail_rank = self._get_k_and_trailing_rank(input_shape, indices_shape)
449 input_tensor_map, indices_tensor_map = self._validate_tensor_maps(
450 input_layout, indices_layout, k
451 )
453 # Output sharding: inherit indices[:-1] + input_x[K:].
454 if input_tensor_map is None:
455 output_tensor_map = tuple(indices_tensor_map[:-1]) + ("None",) * trail_rank
456 else:
457 output_tensor_map = tuple(indices_tensor_map[:-1]) + tuple(input_tensor_map[k:])
459 output_layout = Layout(
460 mesh_shape=indices_layout.mesh_shape,
461 alias_name=indices_layout.alias_name,
462 rank_list=indices_layout.rank_list,
463 )
465 if output_tensor_map:
466 output_layout = output_layout(*output_tensor_map)
467 else:
468 output_layout = output_layout("None")
470 return ((output_layout,), None)
472 def _parse_input_layouts(self, layouts):
473 """Parse and validate input layouts."""
474 if len(layouts) < 2:
475 raise ValueError(
476 f"For {self.op_name}, requires at least 2 input layouts, but got {len(layouts)}"
477 )
479 input_layout, indices_layout = layouts[0], layouts[1]
481 # Extra inputs are allowed only when they are non-tensor args (layout is None).
482 for extra_layout in layouts[2:]:
483 if extra_layout is not None:
484 raise ValueError(
485 f"For {self.op_name}, only supports 2 tensor inputs, but got extra tensor layout: "
486 f"{extra_layout}"
487 )
489 # For GatherNd: input_layout can be None (treated as fully replicated), but indices_layout must exist.
490 if indices_layout is None or not hasattr(indices_layout, "alias_tensor_map"):
491 raise ValueError(f"For {self.op_name}, indices layout cannot be None")
493 return input_layout, indices_layout
495 def _validate_tensor_maps(self, input_layout, indices_layout, k):
496 """Validate tensor maps constraints for GatherNd."""
497 indices_tensor_map = indices_layout.alias_tensor_map
499 # Validate: indices tensor_map must exist and last dimension cannot be split.
500 if not indices_tensor_map:
501 raise ValueError(f"For {self.op_name}, indices tensor_map cannot be empty")
503 last_axis = indices_tensor_map[-1]
504 if not self._is_none_axis(last_axis):
505 raise ValueError(
506 f"For {self.op_name}, the last dimension of indices cannot be split. "
507 f"Got indices[-1] = {last_axis}"
508 )
510 # Validate input only when layout is provided.
511 input_tensor_map = None
512 if input_layout is not None:
513 input_tensor_map = input_layout.alias_tensor_map
515 if k > len(input_tensor_map):
516 raise ValueError(
517 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank "
518 f"({len(input_tensor_map)})"
519 )
521 # Indexed dims [0:K) must be replicated.
522 for axis_name in input_tensor_map[:k]:
523 if not self._is_none_axis(axis_name):
524 raise ValueError(
525 f"For {self.op_name}, input_x cannot be split on indexed dims [0:{k}). "
526 f"These dims must be 'None', but got tensor_map: {input_tensor_map}"
527 )
529 return input_tensor_map, indices_tensor_map
531 def _get_input_shapes(self, shape_values):
532 """Get input and indices shapes from cache values."""
533 input_shapes = None
534 if shape_values and len(shape_values) == 2:
535 input_shapes = shape_values
537 if input_shapes is None:
538 raise ValueError(
539 f"For {self.op_name}, missing input_shapes in cache_values."
540 )
542 input_shape = input_shapes[0]
543 indices_shape = input_shapes[1]
544 if input_shape is None or indices_shape is None:
545 raise ValueError(f"For {self.op_name}, input_shapes contains None: {input_shapes}")
547 input_shape = self._normalize_shape(input_shape, "input")
548 indices_shape = self._normalize_shape(indices_shape, "indices")
550 if len(indices_shape) < 1:
551 raise ValueError(f"For {self.op_name}, indices shape invalid: {indices_shape}")
553 return input_shape, indices_shape
555 def _normalize_shape(self, shape, name):
556 """Normalize shape-like object to tuple of int."""
557 try:
558 norm = tuple(shape)
559 except TypeError as err:
560 raise ValueError(f"For {self.op_name}, {name} shape is not iterable: {shape}") from err
562 try:
563 norm = tuple(int(dim) for dim in norm)
564 except (TypeError, ValueError) as err:
565 raise ValueError(f"For {self.op_name}, {name} shape contains non-integer dims: {norm}") from err
567 return norm
569 def _get_k_and_trailing_rank(self, input_shape, indices_shape):
570 """Compute K and trailing rank = len(input_shape) - K, where K is indices_shape[-1]."""
571 k = indices_shape[-1]
572 try:
573 k = int(k)
574 except (TypeError, ValueError) as err:
575 raise ValueError(f"For {self.op_name}, indices last dim (K) is invalid: {k}") from err
577 if k <= 0:
578 raise ValueError(f"For {self.op_name}, indices last dim (K) must be positive, but got {k}")
580 trail_rank = len(input_shape) - k
581 if trail_rank < 0:
582 raise ValueError(
583 f"For {self.op_name}, indices last dim (K={k}) is larger than input rank ({len(input_shape)})"
584 )
586 return k, trail_rank
588 def _is_none_axis(self, axis_name):
589 """
590 Check if an axis name represents no sharding.
591 """
592 if axis_name == "None":
593 return True
595 if isinstance(axis_name, tuple):
596 return all(name == "None" for name in axis_name)
598 return False