Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / shard / ops / parallel_matmul.py: 73%
256 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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 MatMul operator.
17"""
19from typing import Callable, Optional, Tuple
21from hyper_parallel.core.dtensor.layout import Layout
22from .parallel_ops import DistributedOp
25def _propagate_partial_from_inputs(out_layout, x_layout, w_layout):
26 """
27 Propagate Partial status from input layouts to the output layout for matmul-like operations.
29 For matmul ``y = x @ w``, the output should inherit Partial state from its inputs in addition
30 to any Partial state induced by the contracting dimension being sharded.
32 **Semantic rules for input Partial propagation:**
34 +-------------------+----------------------------+----------------------------------------------+
35 | x input | w / weight | output behavior |
36 +-------------------+----------------------------+----------------------------------------------+
37 | **Partial(d)** | **Replicate** (contracting) | **Propagate Partial(d)**. |
38 | | | Distributive law: |
39 | | | ``(x0 + x1) @ w = x0 @ w + x1 @ w``. |
40 | | | Output carries Partial(d). |
41 +-------------------+----------------------------+----------------------------------------------+
42 | Replicate | **Partial(d)** | **Propagate Partial(d)**. Symmetric to above.|
43 +-------------------+----------------------------+----------------------------------------------+
44 | **Partial(d1)** | **Partial(d2)**, d1 != d2 | **Propagate both**. |
45 | | | Each partial axis is independent; |
46 | | | ``(sum over d2)`` applied to x is legal. |
47 +-------------------+----------------------------+----------------------------------------------+
48 | **Partial(d)** | **Partial(d)** same axis | **Error**. Cross-terms ``x0 @ w1`` and |
49 | | same/different ops | ``x1 @ w0`` cannot be computed locally. |
50 +-------------------+----------------------------+----------------------------------------------+
51 | **Partial(d)** | **Shard(d)** on the same | **Error** naturally raised by |
52 | | device axis in the output | ``Layout.set_partial_by_dev_axis``: |
53 | | dimension map | "Partial dim must be replicate." |
54 +-------------------+----------------------------+----------------------------------------------+
56 Args:
57 out_layout (Layout): The partially-built output layout whose ``alias_tensor_map``
58 has already been set (via ``Layout.__call__``).
59 x_layout (Layout): Layout of the first input tensor (activations).
60 w_layout (Layout): Layout of the second input tensor (weight / matrix).
62 Raises:
63 ValueError: If both ``x_layout`` and ``w_layout`` have Partial on the same device
64 axis with different reduce operations (e.g. one is 'sum' and the other 'avg').
65 """
66 if x_layout is None or w_layout is None:
67 return
69 # Propagate x's partial status to output
70 for dev_idx, op in enumerate(x_layout.partial):
71 if op is not None:
72 out_layout.set_partial_by_dev_axis(
73 x_layout.alias_name[dev_idx], op
74 )
76 # Propagate w's partial status to output, checking for conflicts with x's partial
77 for dev_idx, op in enumerate(w_layout.partial):
78 if op is not None:
79 axis_alias = w_layout.alias_name[dev_idx]
80 existing = out_layout.get_partial_by_dev_id(axis_alias)
81 if existing is not None and existing != op:
82 raise ValueError(
83 f"Cannot propagate Partial from both input layouts: "
84 f"x has Partial({existing}) on axis '{axis_alias}' while "
85 f"w has Partial({op}) on the same axis. "
86 f"Partial on the same axis with different reduce ops for both inputs is invalid."
87 )
88 out_layout.set_partial_by_dev_axis(axis_alias, op)
91def _normalize_matmul_ext_args(x, w):
92 return (x, w), {}
95class MatMulExtDistributedOp(DistributedOp):
96 """Distributed implementation for MatMul operator."""
98 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
99 """
100 Preprocess arguments for MatMulExt operator.
102 Args:
103 args (tuple): Input arguments containing x and w tensors.
104 kwargs (dict): Keyword arguments (unused).
106 Returns:
107 tuple: (local_args, local_kwargs, cache_values) where local_args contains
108 local tensors for x and w; cache_values contains [x_layout, w_layout].
109 """
110 args, kwargs = _normalize_matmul_ext_args(*args, **kwargs)
111 x_tensor, w_tensor = args[0], args[1]
112 local_args = (x_tensor.to_local(), w_tensor.to_local())
113 local_kwargs = {}
114 cache_values = [x_tensor.layout, w_tensor.layout]
115 return local_args, local_kwargs, cache_values
117 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
118 """
119 Infer output layout for MatMul operator (output = x @ w).
121 Rules:
122 1. Inputs must share the same mesh_shape.
123 2. Contracting dimensions must have the same layout.
124 3. Output dimensions inherit layouts from non-contracting dimensions.
125 4. Input Partial status is propagated to the output.
127 Args:
128 cache_values (list): [x_layout, w_layout]
130 Returns:
131 tuple: ((output_layout,), None)
133 Raises:
134 ValueError: If any rule above is violated.
135 """
136 x_layout = cache_values[0]
137 w_layout = cache_values[1]
138 if not x_layout or not w_layout:
139 raise ValueError(
140 f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
141 )
142 x_mesh_shape = x_layout.mesh_shape
143 w_mesh_shape = w_layout.mesh_shape
144 if x_mesh_shape != w_mesh_shape:
145 raise ValueError(
146 f"For {self.op_name}, inputs must have same mesh_shape, "
147 f"but got x: {x_mesh_shape} and w: {w_mesh_shape}"
148 )
150 x_map = x_layout.alias_tensor_map
151 w_map = w_layout.alias_tensor_map
152 contract_dim = len(x_map) - 1
153 w_contract_dim = len(w_map) - 2
154 if x_map[contract_dim] != w_map[w_contract_dim]:
155 raise ValueError(
156 f"For {self.op_name}, contracting dimensions must have same layout, "
157 f"but got x: {x_map[contract_dim]} and w: {w_map[w_contract_dim]}"
158 )
160 output_dim = len(w_map) - 1
161 output_map = x_map[:-1] + (w_map[output_dim],)
163 output_layout = Layout(
164 mesh_shape=x_layout.mesh_shape,
165 alias_name=x_layout.alias_name,
166 rank_list=x_layout.rank_list
167 )
168 out_layout = output_layout(*output_map)
170 # Propagate Partial from inputs (e.g., x already has Partial from a prior matmul)
171 _propagate_partial_from_inputs(out_layout, x_layout, w_layout)
173 # Set partial status from contracting dimension sharding
174 if x_map[contract_dim] != "None":
175 if isinstance(x_map[contract_dim], tuple):
176 for axis in x_map[contract_dim]:
177 out_layout.set_partial_by_dev_axis(axis, 'sum')
178 else:
179 out_layout.set_partial_by_dev_axis(x_map[contract_dim], 'sum')
181 return ((out_layout,), None)
184def _normalize_matmul_args(x, w, transpose_a=False, transpose_b=False):
185 return (x, w, transpose_a, transpose_b), {}
188class MatMulDistributedOp(DistributedOp):
189 """Distributed implementation for MatMul operator."""
191 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
192 """
193 Preprocess arguments for MatMul operator.
195 Args:
196 args (tuple): Input arguments containing x, w tensors and optional transpose flags.
197 kwargs (dict): Keyword arguments (unused).
199 Returns:
200 tuple: (local_args, local_kwargs, cache_values) where local_args contains
201 local tensors for x and w; cache_values contains [x_layout, w_layout, transpose_a, transpose_b].
202 """
203 args, kwargs = _normalize_matmul_args(*args, **kwargs)
204 x_tensor, w_tensor, transpose_a, transpose_b = args
205 local_args = (x_tensor.to_local(), w_tensor.to_local(), transpose_a, transpose_b)
206 local_kwargs = {}
207 cache_values = [x_tensor.layout, w_tensor.layout, transpose_a, transpose_b]
208 return local_args, local_kwargs, cache_values
210 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
211 """
212 Infer output layout for MatMul operator (output = x @ w, with possible transpose).
214 Rules:
215 1. Inputs must share the same mesh_shape.
216 2. Contracting dimensions must have the same layout (adjusted by transpose flags).
217 3. Output dimensions inherit layouts from non-contracting dimensions.
218 4. Input Partial status is propagated to the output.
220 Args:
221 cache_values (list): [x_layout, w_layout, transpose_a, transpose_b]
223 Returns:
224 tuple: ((output_layout,), None)
226 Raises:
227 ValueError: If any rule above is violated.
228 """
229 x_layout = cache_values[0]
230 w_layout = cache_values[1]
231 transpose_a = cache_values[2]
232 transpose_b = cache_values[3]
234 if not x_layout or not w_layout:
235 raise ValueError(
236 f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
237 )
239 x_mesh_shape = x_layout.mesh_shape
240 w_mesh_shape = w_layout.mesh_shape
241 if x_mesh_shape != w_mesh_shape:
242 raise ValueError(
243 f"For {self.op_name}, inputs must have same mesh_shape, "
244 f"but got x: {x_mesh_shape} and w: {w_mesh_shape}"
245 )
247 x_map = x_layout.alias_tensor_map
248 w_map = w_layout.alias_tensor_map
250 # Determine contracting dimensions based on transpose flags
251 if transpose_a:
252 x_input_dim = len(x_map) - 1
253 x_contract_dim = len(x_map) - 2 # Second to last dimension
254 else:
255 x_input_dim = len(x_map) - 2
256 x_contract_dim = len(x_map) - 1 # Last dimension
258 if transpose_b:
259 w_output_dim = len(w_map) - 2
260 w_contract_dim = len(w_map) - 1 # Last dimension
261 else:
262 w_output_dim = len(w_map) - 1
263 w_contract_dim = len(w_map) - 2 # Second to last dimension
265 # Validate contracting dimensions
266 if x_map[x_contract_dim] != w_map[w_contract_dim]:
267 raise ValueError(
268 f"For {self.op_name}, contracting dimensions must have same layout, "
269 f"but got x: {x_map[x_contract_dim]} and w: {w_map[w_contract_dim]}"
270 )
272 # Create output layout
273 output_layout = Layout(
274 mesh_shape=x_layout.mesh_shape,
275 alias_name=x_layout.alias_name,
276 rank_list=x_layout.rank_list
277 )
278 output_map = list(x_map[:-2]) + [x_map[x_input_dim]] + [w_map[w_output_dim]]
279 out_layout = output_layout(*output_map)
281 # Propagate Partial from inputs (e.g., x already has Partial from a prior matmul)
282 _propagate_partial_from_inputs(out_layout, x_layout, w_layout)
284 # Set partial status
285 if x_map[x_contract_dim] != "None":
286 if isinstance(x_map[x_contract_dim], tuple):
287 for axis in x_map[x_contract_dim]:
288 out_layout.set_partial_by_dev_axis(axis, 'sum')
289 else:
290 out_layout.set_partial_by_dev_axis(x_map[x_contract_dim], 'sum')
292 return ((out_layout,), None)
295class BaseBatchMatMulDistributedOp(DistributedOp):
296 """Base class for BatchMatMul distributed implementations."""
298 def _merge_batch_entry(self, x_dims, w_dims):
299 """
300 Merge two batch tensor_map entries with broadcasting:
301 - none vs X -> X
302 - X vs none -> X
303 - X vs X (exact same after normalization) -> X
304 - otherwise -> conflict
305 """
306 if self._is_none_entry(x_dims) and self._is_none_entry(w_dims):
307 return "None"
308 if self._is_none_entry(x_dims):
309 return w_dims
310 if self._is_none_entry(w_dims):
311 return x_dims
312 if x_dims == w_dims:
313 return x_dims
314 raise ValueError(f"Incompatible batch sharding between inputs: {x_dims} vs {w_dims}")
316 def _is_none_entry(self, entry):
317 """An entry is 'none' (no sharding) if it is 'None' or tuple of all 'None'."""
318 if isinstance(entry, tuple):
319 return all(i == "None" for i in entry)
320 return entry == "None"
322 def _merge_batches(self, x_map, w_map):
323 """Right-align and merge batch dims from x_map and w_map."""
324 x_batch = list(x_map[:-2])
325 w_batch = list(w_map[:-2])
326 max_b = max(len(x_batch), len(w_batch))
327 x_batch = ["None"] * (max_b - len(x_batch)) + x_batch
328 w_batch = ["None"] * (max_b - len(w_batch)) + w_batch
329 merged_batch = []
330 for xb, wb in zip(x_batch, w_batch):
331 merged_batch.append(self._merge_batch_entry(xb, wb))
332 return merged_batch
334 def _build_output_layout(self, x_layout, w_layout, merged_batch, x_n, w_p, x_contract):
335 """Construct output layout from merged dims and set partial status if needed."""
336 output_map = tuple(merged_batch) + (x_n, w_p)
338 output_layout = Layout(
339 mesh_shape=x_layout.mesh_shape,
340 alias_name=x_layout.alias_name,
341 rank_list=x_layout.rank_list
342 )
343 output_layout = output_layout(*output_map)
345 # Propagate Partial from inputs
346 _propagate_partial_from_inputs(output_layout, x_layout, w_layout)
348 # Set partial status
349 if x_contract != "None":
350 if isinstance(x_contract, tuple):
351 for axis in x_contract:
352 output_layout.set_partial_by_dev_axis(axis, 'sum')
353 else:
354 output_layout.set_partial_by_dev_axis(x_contract, 'sum')
356 return output_layout
359def _normalize_batch_matmul_ext_args(x, w):
360 return (x, w), {}
363class BatchMatMulExtDistributedOp(BaseBatchMatMulDistributedOp):
364 """Distributed implementation for BatchMatMulExt operator."""
366 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
367 """
368 Preprocess arguments for BatchMatMulExt operator.
370 Args:
371 args (tuple): Input arguments containing x and w tensors.
372 kwargs (dict): Keyword arguments (unused).
374 Returns:
375 tuple: (local_args, local_kwargs, cache_values) where local_args contains
376 local tensors for x and w; cache_values contains [x_layout, w_layout].
377 """
378 args, kwargs = _normalize_batch_matmul_ext_args(*args, **kwargs)
379 x_tensor, w_tensor = args[0], args[1]
380 local_args = (x_tensor.to_local(), w_tensor.to_local())
381 local_kwargs = {}
382 cache_values = [x_tensor.layout, w_tensor.layout]
383 return local_args, local_kwargs, cache_values
385 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
386 """
387 Infer output layout for BatchMatMulExt operator (output = x @ w).
389 Inputs shape are x=[b, n, m] and w=[b, m, p].
391 Rules:
392 1. Inputs must share the same mesh_shape.
393 2. Contracting K dims must have identical layout: x[-1] == w[-2].
394 3. Batch dims are right-aligned with broadcast semantics.
395 4. Output batch dims = merged batch dims; N inherits x[-2], P inherits w[-1].
396 5. Input Partial status is propagated to the output.
398 Args:
399 cache_values (list): [x_layout, w_layout]
401 Returns:
402 tuple: ((output_layout,), None)
404 Raises:
405 ValueError: If any rule above is violated.
406 """
407 x_layout = cache_values[0]
408 w_layout = cache_values[1]
410 if not x_layout or not w_layout:
411 raise ValueError(
412 f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
413 )
415 if x_layout.mesh_shape != w_layout.mesh_shape:
416 raise ValueError(
417 f"For {self.op_name}, inputs must have same mesh_shape, "
418 f"but got x: {x_layout.mesh_shape} and w: {w_layout.mesh_shape}"
419 )
421 x_map = x_layout.alias_tensor_map
422 w_map = w_layout.alias_tensor_map
424 # contracting dims
425 x_contract = x_map[-1]
426 w_contract = w_map[-2]
427 if x_contract != w_contract:
428 raise ValueError(
429 f"For {self.op_name}, contracting (M) dim layouts must match, "
430 f"but got x: {x_contract} and w: {w_contract}"
431 )
433 merged_batch = self._merge_batches(x_map, w_map)
434 x_n = x_map[-2]
435 w_p = w_map[-1]
437 return ((self._build_output_layout(x_layout, w_layout, merged_batch, x_n, w_p, x_contract),), None)
440def _normalize_batch_matmul_args(x, w, transpose_a=False, transpose_b=False):
441 return (x, w, transpose_a, transpose_b), {}
444class BatchMatMulDistributedOp(BaseBatchMatMulDistributedOp):
445 """Distributed implementation for BatchMatMul operator."""
447 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
448 """
449 Preprocess arguments for BatchMatMul operator.
451 Args:
452 args (tuple): Input arguments containing x, w tensors and optional transpose flags.
453 kwargs (dict): Keyword arguments (unused).
455 Returns:
456 tuple: (local_args, local_kwargs, cache_values) where local_args contains
457 local tensors for x and w; cache_values contains
458 [x_layout, w_layout, transpose_a, transpose_b].
459 """
460 args, kwargs = _normalize_batch_matmul_args(*args, **kwargs)
461 x_tensor, w_tensor, transpose_a, transpose_b = args
462 local_args = (x_tensor.to_local(), w_tensor.to_local(), transpose_a, transpose_b)
463 local_kwargs = {}
464 cache_values = [x_tensor.layout, w_tensor.layout, transpose_a, transpose_b]
465 return local_args, local_kwargs, cache_values
467 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
468 """
469 Infer output layout for BatchMatMul operator (output = x @ w, with possible transpose).
471 Inputs shape are x=[b, n, m] and w=[b, m, p].
473 Rules:
474 1. Inputs must share the same mesh_shape.
475 2. Contracting K dims must have identical layout (adjusted by transpose flags).
476 3. Batch dims are right-aligned with broadcast semantics.
477 4. Output batch dims = merged batch dims; N/P dims inherit per transpose flags.
478 5. Input Partial status is propagated to the output.
480 Args:
481 cache_values (list): [x_layout, w_layout, transpose_a, transpose_b]
483 Returns:
484 tuple: ((output_layout,), None)
486 Raises:
487 ValueError: If any rule above is violated.
488 """
489 x_layout = cache_values[0]
490 w_layout = cache_values[1]
491 transpose_a = cache_values[2]
492 transpose_b = cache_values[3]
494 if not x_layout or not w_layout:
495 raise ValueError(
496 f"For {self.op_name}, x_layout: {x_layout}, w_layout: {w_layout}"
497 )
499 if x_layout.mesh_shape != w_layout.mesh_shape:
500 raise ValueError(
501 f"For {self.op_name}, inputs must have same mesh_shape, "
502 f"but got x: {x_layout.mesh_shape} and w: {w_layout.mesh_shape}"
503 )
505 x_map = x_layout.alias_tensor_map
506 w_map = w_layout.alias_tensor_map
508 # handle transpose
509 if transpose_a:
510 x_n = x_map[-1]
511 x_contract = x_map[-2]
512 else:
513 x_n = x_map[-2]
514 x_contract = x_map[-1]
516 if transpose_b:
517 w_contract = w_map[-1]
518 w_p = w_map[-2]
519 else:
520 w_contract = w_map[-2]
521 w_p = w_map[-1]
523 if x_contract != w_contract:
524 raise ValueError(
525 f"For {self.op_name}, contracting (M) dim layouts must match, "
526 f"but got x: {x_contract} and w: {w_contract}"
527 )
529 merged_batch = self._merge_batches(x_map, w_map)
531 return ((self._build_output_layout(x_layout, w_layout, merged_batch, x_n, w_p, x_contract),), None)
534def _normalize_linear_args(x, weight, bias=None):
535 return (x, weight, bias), {}
538class LinearDistributedOp(DistributedOp):
539 """Distributed implementation for Linear operator."""
541 def preprocess(self, args: tuple, kwargs: dict) -> tuple:
542 """
543 Preprocess arguments for Linear operator.
545 Args:
546 args (tuple): Input arguments containing x and weight tensors.
547 kwargs (dict): Keyword arguments, may contain bias.
549 Returns:
550 tuple: (local_args, local_kwargs, cache_values) where local_args contains
551 local tensors for x, weight, and bias; local_kwargs is empty; and
552 cache_values contains layouts and None-sentinel for absent bias.
553 """
554 args, kwargs = _normalize_linear_args(*args, **kwargs)
555 x_tensor, w_tensor, bias = args[0], args[1], args[2]
556 local_args = (
557 x_tensor.to_local(),
558 w_tensor.to_local(),
559 bias.to_local() if hasattr(bias, '_layout') else bias,
560 )
561 local_kwargs = {}
562 cache_values = [
563 x_tensor.layout,
564 w_tensor.layout,
565 bias.layout if hasattr(bias, '_layout') else None,
566 ]
567 return local_args, local_kwargs, cache_values
569 def infer_layout(self, cache_values: list) -> Tuple[tuple, None]:
570 """
571 Infer output layout for Linear operator (output = x @ weight.T + bias).
573 Rules:
574 1. x and weight must share the same mesh_shape.
575 2. weight must be 2D [out_features, in_features].
576 3. Contracting dimensions (in_features) must have the same layout.
577 4. Output batch dimensions inherit from x; output feature dim inherits from weight dim 0.
578 5. Partial state is set on the output when the contracting dimension is sharded.
580 Args:
581 cache_values (list): [x_layout, w_layout, bias_layout] where bias_layout may be None.
583 Returns:
584 tuple: ((out_layout,), None)
586 Raises:
587 ValueError: If cache_values length is not 3, layouts are invalid, mesh shapes differ,
588 weight is not 2D, contracting dims mismatch, or bias sharding is inconsistent.
589 """
590 if len(cache_values) != 3:
591 raise ValueError(
592 f"For {self.op_name}, cache_values length should be 3, but got {len(cache_values)}"
593 )
594 x_layout = cache_values[0]
595 w_layout = cache_values[1]
596 bias_layout = cache_values[2]
598 if not x_layout or not w_layout:
599 raise ValueError(f"x_layout : {x_layout}, w_layout : {w_layout}")
601 x_mesh_shape = x_layout.mesh_shape
602 w_mesh_shape = w_layout.mesh_shape
603 if x_mesh_shape != w_mesh_shape:
604 raise ValueError(
605 f"For {self.op_name}, x and weight must have the same mesh_shape, "
606 f"but got x: {x_mesh_shape} and weight: {w_mesh_shape}"
607 )
608 if bias_layout and bias_layout.mesh_shape != x_mesh_shape:
609 raise ValueError(
610 f"For {self.op_name}, bias and x must have the same mesh_shape, "
611 f"but got bias: {bias_layout.mesh_shape} and x: {x_mesh_shape}"
612 )
614 x_map = x_layout.alias_tensor_map
615 w_map = w_layout.alias_tensor_map
617 if len(w_map) != 2:
618 raise ValueError(
619 f"For {self.op_name}, weight should be 2D [out_features, in_features], "
620 f"but got {len(w_map)}D"
621 )
623 x_contract_dim = len(x_map) - 1
624 w_contract_dim = len(w_map) - 1
625 if x_map[x_contract_dim] != w_map[w_contract_dim]:
626 raise ValueError(
627 f"For {self.op_name}, contracting dimensions must have the same layout, "
628 f"but got x: {x_map[x_contract_dim]} and weight: {w_map[w_contract_dim]}"
629 )
631 output_dim = 0
632 output_map = x_map[:-1] + (w_map[output_dim],)
633 if bias_layout and bias_layout.alias_tensor_map[0] != w_map[output_dim]:
634 raise ValueError(
635 f"For {self.op_name}, bias output dim sharding must match weight output dim sharding, "
636 f"but got weight: {w_map[output_dim]} and bias: {bias_layout.alias_tensor_map[0]}"
637 )
639 output_layout = Layout(
640 mesh_shape=x_layout.mesh_shape,
641 alias_name=x_layout.alias_name,
642 rank_list=x_layout.rank_list,
643 )
644 out_layout = output_layout(*output_map)
646 # Propagate Partial from inputs (e.g., x already has Partial from a prior matmul)
647 _propagate_partial_from_inputs(out_layout, x_layout, w_layout)
649 # Set partial status when contracting dimension is sharded
650 if x_map[x_contract_dim] != "None":
651 if isinstance(x_map[x_contract_dim], tuple):
652 for axis in x_map[x_contract_dim]:
653 out_layout.set_partial_by_dev_axis(axis, 'sum')
654 else:
655 out_layout.set_partial_by_dev_axis(x_map[x_contract_dim], 'sum')
657 return ((out_layout,), None)
659 def get_expand_impl(self, func: Callable, infer_result: tuple,
660 cache_values: list) -> Optional[Callable]:
661 """
662 Return a custom expand implementation when bias scaling is needed.
664 When the contracting dimension is sharded each rank computes a partial sum
665 (x_shard @ w_shard.T + bias). After AllReduce the bias would accumulate
666 scaling_factor times. The returned closure pre-divides bias by scaling_factor
667 to keep the result numerically correct.
669 Args:
670 func: Original operator callable.
671 infer_result (tuple): ((out_layout,), None) from infer_layout.
672 cache_values (list): [x_layout, w_layout, bias_layout].
674 Returns:
675 callable | None: expand_impl closure when scaling is required, else None.
676 """
677 x_layout = cache_values[0]
678 bias_layout = cache_values[2]
679 x_map = x_layout.alias_tensor_map
680 x_contract_dim = len(x_map) - 1
682 # Guard: scaling only needed when contract dim is sharded AND bias is present
683 if x_map[x_contract_dim] == "None" or not bias_layout:
684 return None
686 output_layout = infer_result[0][0]
687 scaling_factor = 1
688 if isinstance(x_map[x_contract_dim], tuple):
689 for axis in x_map[x_contract_dim]:
690 scaling_factor *= output_layout.mesh.get_device_num_along_axis(axis)
691 else:
692 scaling_factor *= output_layout.mesh.get_device_num_along_axis(x_map[x_contract_dim])
694 def expand_impl(x: object, w: object, bias: object) -> object:
695 """Pre-scale bias to counteract the AllReduce accumulation over shards.
697 Args:
698 x (object): Local input activation tensor.
699 w (object): Local weight tensor.
700 bias (object): Local bias tensor to be pre-scaled.
702 Returns:
703 object: Result of the linear operation with pre-scaled bias.
704 """
705 return func(x, w, bias / scaling_factor)
707 return expand_impl