Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_embedding.py: 90%
99 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"""Distributed implementation for the embedding operator."""
17from typing import Callable, Optional, Tuple
19from hyper_parallel.core.dtensor.layout import Layout
20from .parallel_ops import DistributedOp
23def _normalize_embedding_args(input_tensor, weight_tensor, padding_idx=None, max_norm=None,
24 norm_type=2.0, scale_grad_by_freq=False, sparse=False):
25 return (
26 input_tensor,
27 weight_tensor,
28 padding_idx,
29 max_norm,
30 norm_type,
31 scale_grad_by_freq,
32 sparse,
33 ), {}
36class EmbeddingDistributedOp(DistributedOp):
37 """
38 Distributed implementation for embedding operators.
39 Supports Column Parallelism (CP) and Row Parallelism (RP).
40 """
41 _MS_PRIMITIVE_OP_NAMES = frozenset({'Embedding'})
43 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
44 """
45 Preprocess arguments for Embedding operator.
47 Args:
48 args (tuple): Input arguments containing input and weight tensors.
49 kwargs (dict): Keyword arguments for embedding options.
51 Returns:
52 tuple: (local_args, local_kwargs, cache_values) where cache_values is
53 [input_layout, weight_layout].
54 """
55 args, _ = _normalize_embedding_args(*args, **kwargs)
56 input_tensor, weight_tensor = args[0], args[1]
57 if self.op_name in self._MS_PRIMITIVE_OP_NAMES:
58 local_args = (input_tensor.to_local(), weight_tensor.to_local()) + args[2:6]
59 else:
60 local_args = (input_tensor.to_local(), weight_tensor.to_local()) + args[2:]
61 local_kwargs = {}
62 cache_values = [input_tensor.layout, weight_tensor.layout]
63 return local_args, local_kwargs, cache_values
65 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]: # pylint: disable=W0221
66 """
67 Infer output layout for Embedding operator.
69 Rules:
70 1. Input and weight must not have Partial status.
71 2. Input and weight must share the same mesh_shape.
72 3. Weight must be 2D [vocab_size, embedding_dim].
73 4. Output shape is [*input_shape, embedding_dim], preserving input sharding
74 and inheriting weight embedding-dimension sharding.
75 5. If weight vocab dimension is sharded, output carries Partial('sum') on
76 that vocab sharding axis.
78 Args:
79 cache_values (list): [input_layout, weight_layout].
81 Returns:
82 tuple: ((output_layout,), None)
84 Raises:
85 ValueError: If layouts are missing, partial, incompatible, or weight is not 2D.
86 """
87 if len(cache_values) != 2:
88 raise ValueError(
89 f"For {self.op_name}, cache_values length should be 2, but got {len(cache_values)}"
90 )
92 input_layout = cache_values[0]
93 weight_layout = cache_values[1]
94 if not input_layout or not weight_layout:
95 raise ValueError(
96 f"For {self.op_name}, requires both input and weight layouts."
97 )
99 self._check_partial_inputs([input_layout, weight_layout])
101 if input_layout.mesh_shape != weight_layout.mesh_shape:
102 raise ValueError(
103 f"For {self.op_name}, input and weight must have the same mesh_shape, "
104 f"but got input: {input_layout.mesh_shape} and weight: {weight_layout.mesh_shape}"
105 )
107 weight_tensor_map = weight_layout.tensor_map
108 if len(weight_tensor_map) != 2:
109 raise ValueError(
110 f"For {self.op_name}, weight should be 2D [vocab_size, embedding_dim], "
111 f"but got {len(weight_tensor_map)}D"
112 )
114 # weight_tensor_map: [vocab_size_dim, embed_dim_dim]
115 w_shard_embed_axis = weight_tensor_map[1]
116 weight_alias_map = weight_layout.alias_tensor_map
117 w_shard_vocab_alias = weight_alias_map[0]
119 # Output shape is [*input_shape, embed_dim]
120 output_tensor_map = list(input_layout.tensor_map)
121 output_tensor_map.append(w_shard_embed_axis)
123 output_layout = Layout(
124 mesh_shape=input_layout.mesh_shape,
125 alias_name=input_layout.alias_name,
126 rank_list=input_layout.rank_list
127 )
128 output_layout.set_tensor_map(tuple(output_tensor_map))
130 # If vocab is sharded (Row Parallelism), output is in Partial Sum state
131 if w_shard_vocab_alias != "None":
132 # pylint: disable=protected-access
133 output_layout._partial = list(input_layout.partial)
134 output_layout.set_partial_by_dev_axis(w_shard_vocab_alias, 'sum')
135 # pylint: disable=protected-access
136 output_layout._alias_tensor_map = output_layout._build_readable_tensor_map()
137 # pylint: disable=protected-access
138 output_layout.tensor_map_to_placement()
139 output_layout.update_compact_str()
141 return ((output_layout,), None)
143 def _parse_params(self, args, kwargs):
144 """
145 Extracts padding_idx, max_norm, and scale_grad from args or kwargs.
146 F.embedding signature: (input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq...)
147 """
148 padding_idx = args[2] if len(args) > 2 else kwargs.get('padding_idx', None)
149 max_norm = args[3] if len(args) > 3 else kwargs.get('max_norm', None)
150 scale_grad = args[5] if len(args) > 5 else kwargs.get('scale_grad_by_freq', False)
151 return padding_idx, max_norm, scale_grad
153 def _handle_rp_input(self, input_tensor, weight_tensor, weight_layout, w_shard_vocab_axis,
154 new_args, kwargs, is_args_pad, padding_idx):
155 """
156 Processes Row Parallelism input: shifts indices, handles padding, and generates masks.
157 """
158 mesh = weight_layout.mesh
159 mesh_dim_idx = len(mesh.mesh_shape) - 1 - w_shard_vocab_axis
160 vocab_coord = mesh.get_local_rank(mesh_dim_idx)
162 vocab_size_per_partition = weight_tensor.shape[0]
163 vocab_start_index = int(vocab_coord * vocab_size_per_partition)
164 vocab_end_index = int(vocab_start_index + vocab_size_per_partition)
166 # Map global padding_idx to local rank range
167 if padding_idx is not None:
168 if vocab_start_index <= padding_idx < vocab_end_index:
169 mapped_padding_idx = int(padding_idx - vocab_start_index)
170 if is_args_pad:
171 new_args[2] = mapped_padding_idx
172 else:
173 kwargs['padding_idx'] = mapped_padding_idx
174 else:
175 if is_args_pad:
176 new_args[2] = None
177 else:
178 kwargs.pop('padding_idx', None)
180 # Calculate out-of-bounds mask
181 mask = (input_tensor >= vocab_start_index) & (input_tensor < vocab_end_index)
183 # Cross-platform cast to matching int dtype
184 mask_int = mask.to(input_tensor.dtype) if hasattr(mask, "to") else mask.astype(input_tensor.dtype)
186 # Shift global indices to local range using native scalar broadcast
187 local_input = input_tensor - vocab_start_index
189 # Zero out invalid indices mathematically instead of using .where() or clamp().
190 # This prevents NPU out-of-bounds memory access during the embedding lookup
191 # while keeping the code perfectly backend-neutral.
192 local_input = local_input * mask_int
194 return local_input, mask_int
196 def get_expand_impl(self, func: Optional[Callable], infer_result: tuple, # pylint: disable=W0221
197 cache_values: list) -> Optional[Callable]:
198 """
199 Returns the execution implementation wrapper.
200 Helper functions are used to keep Cyclomatic Complexity (CCN) low.
201 """
202 weight_layout = cache_values[1]
203 w_shard_vocab_axis = weight_layout.tensor_map[0]
204 weight_alias_map = weight_layout.alias_tensor_map
205 w_shard_vocab_alias = weight_alias_map[0]
206 w_shard_embed_alias = weight_alias_map[1]
208 # Use native implementation if no weight sharding is applied
209 if w_shard_vocab_alias == "None" and w_shard_embed_alias == "None":
210 return None
212 def distributed_embedding_impl(*args, **kwargs):
213 input_tensor, weight_tensor = args[0], args[1]
214 new_args, new_kwargs = list(args), kwargs.copy()
216 # 1. Parameter extraction and validation
217 padding_idx, max_norm, scale_grad = self._parse_params(args, kwargs)
219 # Check for max_norm with specific error messages for CP and RP
220 if max_norm is not None:
221 if w_shard_embed_alias != "None":
222 raise ValueError(
223 f"For {self.op_name}, Column-Parallel Embedding does not support `max_norm` parameter."
224 )
225 if w_shard_vocab_alias != "None":
226 raise ValueError(
227 f"For {self.op_name}, Row-Parallel Embedding does not support `max_norm` parameter."
228 )
230 # Check for scale_grad_by_freq with RP
231 if scale_grad and w_shard_vocab_alias != "None":
232 raise ValueError(
233 f"For {self.op_name}, Row-Parallel Embedding does not support `scale_grad_by_freq=True`."
234 )
236 # 2. Row Parallel Processing
237 input_mask_int = None
238 if w_shard_vocab_alias != "None":
239 is_args_pad = len(args) > 2
240 mapped_input, input_mask_int = self._handle_rp_input(
241 input_tensor, weight_tensor, weight_layout, w_shard_vocab_axis,
242 new_args, new_kwargs, is_args_pad, padding_idx
243 )
244 new_args[0] = mapped_input
246 # 3. Native Operator Execution
247 output = func(*new_args, **new_kwargs)
249 # 4. Erase invalid partial embeddings (Row-Parallel only)
250 if w_shard_vocab_alias != "None" and input_mask_int is not None:
251 expanded_mask = input_mask_int[..., None].to(output.dtype)
252 output = output * expanded_mask
254 return output
256 return distributed_embedding_impl