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