Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / mc2.py: 44%
132 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +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"""MC2 (matmul + communication) fused linear primitives for tensor parallelism.
17PyTorch / Ascend exposes fused kernels via ``torch_npu.npu_all_gather_base_mm`` and
18``torch_npu.npu_mm_reduce_scatter_base``. These kernels do not ship a complete
19autograd formula for every training path, so this module wraps each fused forward
20with a custom backward so they can be used inside a trainable TP + SP model.
22The two kernels are mathematical duals (aligned with MindFormers MC2):
24* ``all_gather_matmul`` forward : ``Y = AllGather_m(X) @ W^T``
25 backward: ``dX = matmul_reduce_scatter(dY, W)``
26 ``dW = dY^T @ AllGather_m(X)``
28* ``matmul_reduce_scatter`` forward : ``Y = ReduceScatter_m(X @ W^T)``
29 backward: ``dX, dY_full = all_gather_matmul(dY, W)``
30 ``dW = AllGather_m(dY)^T @ X``
32Both forward and backward ``dX`` paths use fused Ascend MC2 kernels. Callers must
33satisfy kernel constraints (notably contraction dim ``k ∈ [256, 65535)``); for
34column-parallel backward that ``k`` is local ``out_features`` (``n_local``).
36All autograd functions operate on **local** (non-DTensor) 2-D tensors.
37``W`` uses the native ``nn.Linear`` layout ``(out, in)``.
38"""
39from __future__ import annotations
41from typing import Any, Optional
43import torch
44from torch import nn
46from hyper_parallel.core.dtensor.dtensor import DTensor
47from hyper_parallel.core.dtensor.placement_types import Shard
48from hyper_parallel.platform import get_platform
50platform = get_platform()
52__all__ = [
53 "get_hcomm_info",
54 "AllGatherMatmulFunction",
55 "MatmulReduceScatterFunction",
56 "MC2Linear",
57]
60def get_hcomm_info(group: Any) -> str:
61 """Resolve Ascend HCCL communicator name for a torch ProcessGroup.
63 Args:
64 group (Any): ``torch.distributed.ProcessGroup`` for the TP mesh axis.
66 Returns:
67 HCCL communicator handle name expected by ``torch_npu`` MC2 kernels.
68 """
69 rank = torch.distributed.get_rank(group)
70 if torch.__version__ > "2.0":
71 global_rank = torch.distributed.get_global_rank(group, rank)
72 # torch.distributed ProcessGroup exposes HCCL via a private backend API.
73 return group._get_backend(torch.device("npu")).get_hccl_comm_name( # pylint: disable=protected-access
74 global_rank
75 )
76 return group.get_hccl_comm_name(rank)
79def _require_torch_npu():
80 try:
81 import torch_npu # pylint: disable=import-outside-toplevel
82 except ImportError as exc:
83 raise RuntimeError(
84 "MC2 fused kernels require torch_npu "
85 "(npu_all_gather_base_mm / npu_mm_reduce_scatter_base)."
86 ) from exc
87 return torch_npu
90def _normalize_sequence_dim(sequence_dim: int, ndim_leading: int) -> int:
91 """Resolve a possibly-negative sequence dim against leading-rank count."""
92 seq_dim = sequence_dim
93 if seq_dim < 0:
94 seq_dim += ndim_leading
95 if seq_dim < 0 or seq_dim >= ndim_leading:
96 raise RuntimeError(
97 f"MC2Linear sequence_dim={sequence_dim} is out of range "
98 f"for leading rank {ndim_leading}."
99 )
100 return seq_dim
103def _move_dim_to_front(tensor: torch.Tensor, dim: int) -> torch.Tensor:
104 """Permute ``dim`` to axis 0 so fused AG/RS on flattened dim-0 is SP-correct."""
105 if dim == 0:
106 return tensor
107 order = (dim,) + tuple(i for i in range(tensor.dim()) if i != dim)
108 return tensor.permute(*order).contiguous()
111def _move_front_to_dim(tensor: torch.Tensor, dim: int) -> torch.Tensor:
112 """Inverse of :func:`_move_dim_to_front`."""
113 if dim == 0:
114 return tensor
115 # front -> dim: [1..dim] + [0] + [dim+1..]
116 order = tuple(range(1, dim + 1)) + (0,) + tuple(range(dim + 1, tensor.dim()))
117 return tensor.permute(*order).contiguous()
120class AllGatherMatmulFunction(platform.Function):
121 """Column-parallel fused all-gather + matmul with custom backward.
123 Forward (local view): ``out = AllGather_m(x) @ w^T``.
124 """
126 @staticmethod
127 def forward(ctx, x, w, group, world_size, bias): # pylint: disable=arguments-differ
128 """Run fused all-gather + matmul and stash tensors for backward."""
129 torch_npu = _require_torch_npu()
130 hcom = get_hcomm_info(group)
131 # x2 = w.T -> physical (k, n_local); kernel computes AG(x) @ x2
132 out, gathered = torch_npu.npu_all_gather_base_mm(
133 x,
134 w.t(),
135 hcom,
136 world_size,
137 bias=None,
138 gather_index=0,
139 gather_output=True,
140 )
141 if bias is not None:
142 out = out + bias
143 ctx.save_for_backward(gathered, w)
144 ctx.group = group
145 ctx.world_size = world_size
146 ctx.hcom = hcom
147 ctx.has_bias = bias is not None
148 return out
150 @staticmethod
151 def backward(ctx, grad_out): # pylint: disable=arguments-differ
152 """Gradient: dx via fused matmul_reduce_scatter, dw via gathered-input matmul.
154 Matches MindFormers: ``dX = matmul_reduce_scatter(dY, W)``. The fused
155 kernel contracts over ``n_local`` (``W.shape[0]``), which must be in
156 ``[256, 65535)``.
157 """
158 torch_npu = _require_torch_npu()
159 gathered, w = ctx.saved_tensors
160 # input (m_full, n_local) @ x2 (n_local, k) -> RS_m -> (m_local, k)
161 grad_x = torch_npu.npu_mm_reduce_scatter_base(
162 grad_out.contiguous(),
163 w,
164 ctx.hcom,
165 ctx.world_size,
166 reduce_op="sum",
167 bias=None,
168 )
169 grad_w = grad_out.t().matmul(gathered)
170 grad_bias = grad_out.sum(dim=0) if ctx.has_bias else None
171 return grad_x, grad_w, None, None, grad_bias
174class MatmulReduceScatterFunction(platform.Function):
175 """Row-parallel fused matmul + reduce-scatter with custom backward.
177 Forward (local view): ``out = ReduceScatter_m(x @ w^T)``.
178 """
180 @staticmethod
181 def forward(ctx, x, w, group, world_size, bias): # pylint: disable=arguments-differ
182 """Run fused matmul + reduce-scatter and stash tensors for backward."""
183 torch_npu = _require_torch_npu()
184 hcom = get_hcomm_info(group)
185 out = torch_npu.npu_mm_reduce_scatter_base(
186 x,
187 w.t(),
188 hcom,
189 world_size,
190 reduce_op="sum",
191 bias=None,
192 )
193 if bias is not None:
194 out = out + bias
195 ctx.save_for_backward(x, w)
196 ctx.group = group
197 ctx.world_size = world_size
198 ctx.hcom = hcom
199 ctx.has_bias = bias is not None
200 return out
202 @staticmethod
203 def backward(ctx, grad_out): # pylint: disable=arguments-differ
204 """Gradient: dx via fused all-gather+matmul, dw via gathered-grad matmul."""
205 torch_npu = _require_torch_npu()
206 x, w = ctx.saved_tensors
207 # AG(dY) @ W : pass W (n, k) so the kernel uses transposed-x2 semantics.
208 grad_x, grad_out_full = torch_npu.npu_all_gather_base_mm(
209 grad_out,
210 w,
211 ctx.hcom,
212 ctx.world_size,
213 bias=None,
214 gather_index=0,
215 gather_output=True,
216 )
217 grad_w = grad_out_full.t().matmul(x)
218 grad_bias = grad_out_full.sum(dim=0) if ctx.has_bias else None
219 return grad_x, grad_w, None, None, grad_bias
222class MC2Linear(nn.Linear):
223 """``nn.Linear`` that uses fused matmul + TP communication kernels."""
225 def configure_mc2(
226 self,
227 mode: str,
228 group: Any,
229 world_size: int,
230 sequence_dim: int = 0,
231 ) -> None:
232 """Configure the fused collective used by this layer.
234 Args:
235 mode (str): ``\"all_gather\"`` (column) or ``\"reduce_scatter\"`` (row).
236 group (Any): TP process group.
237 world_size (int): Size of ``group``.
238 sequence_dim (int, optional): Tensor dim that carries sequence sharding
239 under SP. Default: ``0``.
240 """
241 if mode not in ("all_gather", "reduce_scatter"):
242 raise ValueError(
243 "For MC2Linear.configure_mc2, mode should be 'all_gather' or "
244 f"'reduce_scatter', but got {mode}."
245 )
246 self.mc2_mode = mode
247 self.mc2_group = group
248 self.mc2_world_size = world_size
249 self.mc2_sequence_dim = sequence_dim
251 @classmethod
252 def from_linear(cls, linear: nn.Linear) -> "MC2Linear":
253 """Convert a Linear in place while preserving parameters and module state."""
254 if not isinstance(linear, nn.Linear):
255 raise TypeError(
256 f"MC2Linear can only replace nn.Linear, but got {type(linear).__name__}."
257 )
258 linear.__class__ = cls
259 return linear
261 def _mc2_forward(self, input_: DTensor, weight: DTensor) -> DTensor:
262 """Run the configured fused kernel on local tensors.
264 Ascend MC2 kernels all-gather / reduce-scatter the **flattened dim-0**
265 of a 2-D activation. That is only layout-correct when the SP sequence
266 axis is the outermost leading dim. For ``sequence_dim != 0`` (e.g.
267 batch-first ``[B, S, H]`` with ``Shard(1)``), move the sequence dim to
268 front before flatten, then restore after the fused op.
269 """
270 input_local = input_.to_local()
271 weight_local = weight.to_local()
272 leading_global = tuple(int(s) for s in input_.shape[:-1])
273 seq_dim = _normalize_sequence_dim(self.mc2_sequence_dim, len(leading_global))
275 # [..., S_local_or_full, ..., H] -> [S_*, *other_leading, H] -> 2-D
276 x_seq_first = _move_dim_to_front(input_local, seq_dim)
277 other_leading = tuple(x_seq_first.shape[1:-1])
278 input_2d = x_seq_first.reshape(-1, x_seq_first.shape[-1])
280 bias_local = None
281 if self.bias is not None:
282 bias = self.bias
283 bias_local = bias.to_local() if isinstance(bias, DTensor) else bias
285 if self.mc2_mode == "all_gather":
286 output_2d = AllGatherMatmulFunction.apply(
287 input_2d,
288 weight_local,
289 self.mc2_group,
290 self.mc2_world_size,
291 bias_local,
292 )
293 # AG expands the sequence dim; other leading dims stay local sizes.
294 seq_global = leading_global[seq_dim]
295 output = output_2d.reshape(seq_global, *other_leading, output_2d.shape[-1])
296 output = _move_front_to_dim(output, seq_dim)
297 return DTensor.from_local(output, input_.device_mesh, (Shard(-1),))
299 seq_global = leading_global[seq_dim]
300 if seq_global % self.mc2_world_size != 0:
301 raise RuntimeError(
302 f"MC2Linear reduce_scatter requires sequence dim {seq_dim} "
303 f"(size {seq_global}) divisible by world_size "
304 f"{self.mc2_world_size}."
305 )
306 output_2d = MatmulReduceScatterFunction.apply(
307 input_2d,
308 weight_local,
309 self.mc2_group,
310 self.mc2_world_size,
311 bias_local,
312 )
313 seq_local = seq_global // self.mc2_world_size
314 output = output_2d.reshape(seq_local, *other_leading, output_2d.shape[-1])
315 output = _move_front_to_dim(output, seq_dim)
316 return DTensor.from_local(output, input_.device_mesh, (Shard(seq_dim),))
318 def forward(self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
319 """Forward using MC2 for DTensor inputs and ``nn.Linear`` otherwise."""
320 if not isinstance(input_, DTensor):
321 return super().forward(input_)
322 if not hasattr(self, "mc2_mode"):
323 raise RuntimeError(
324 "MC2Linear must be configured by an MC2 parallel style before use."
325 )
327 if weight is None:
328 weight = self.weight
329 if not isinstance(weight, DTensor):
330 raise TypeError(
331 "MC2Linear expects a DTensor weight after tensor-parallel sharding, "
332 f"but got {type(weight).__name__}."
333 )
334 return self._mc2_forward(input_, weight)