Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / tensor_parallel / mc2.py: 43%
130 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"""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
49__all__ = [
50 "get_hcomm_info",
51 "AllGatherMatmulFunction",
52 "MatmulReduceScatterFunction",
53 "MC2Linear",
54]
57def get_hcomm_info(group: Any) -> str:
58 """Resolve Ascend HCCL communicator name for a torch ProcessGroup.
60 Args:
61 group (Any): ``torch.distributed.ProcessGroup`` for the TP mesh axis.
63 Returns:
64 HCCL communicator handle name expected by ``torch_npu`` MC2 kernels.
65 """
66 rank = torch.distributed.get_rank(group)
67 if torch.__version__ > "2.0":
68 global_rank = torch.distributed.get_global_rank(group, rank)
69 # torch.distributed ProcessGroup exposes HCCL via a private backend API.
70 return group._get_backend(torch.device("npu")).get_hccl_comm_name( # pylint: disable=protected-access
71 global_rank
72 )
73 return group.get_hccl_comm_name(rank)
76def _require_torch_npu():
77 try:
78 import torch_npu # pylint: disable=import-outside-toplevel
79 except ImportError as exc:
80 raise RuntimeError(
81 "MC2 fused kernels require torch_npu "
82 "(npu_all_gather_base_mm / npu_mm_reduce_scatter_base)."
83 ) from exc
84 return torch_npu
87def _normalize_sequence_dim(sequence_dim: int, ndim_leading: int) -> int:
88 """Resolve a possibly-negative sequence dim against leading-rank count."""
89 seq_dim = sequence_dim
90 if seq_dim < 0:
91 seq_dim += ndim_leading
92 if seq_dim < 0 or seq_dim >= ndim_leading:
93 raise RuntimeError(
94 f"MC2Linear sequence_dim={sequence_dim} is out of range "
95 f"for leading rank {ndim_leading}."
96 )
97 return seq_dim
100def _move_dim_to_front(tensor: torch.Tensor, dim: int) -> torch.Tensor:
101 """Permute ``dim`` to axis 0 so fused AG/RS on flattened dim-0 is SP-correct."""
102 if dim == 0:
103 return tensor
104 order = (dim,) + tuple(i for i in range(tensor.dim()) if i != dim)
105 return tensor.permute(*order).contiguous()
108def _move_front_to_dim(tensor: torch.Tensor, dim: int) -> torch.Tensor:
109 """Inverse of :func:`_move_dim_to_front`."""
110 if dim == 0:
111 return tensor
112 # front -> dim: [1..dim] + [0] + [dim+1..]
113 order = tuple(range(1, dim + 1)) + (0,) + tuple(range(dim + 1, tensor.dim()))
114 return tensor.permute(*order).contiguous()
117class AllGatherMatmulFunction(torch.autograd.Function):
118 """Column-parallel fused all-gather + matmul with custom backward.
120 Forward (local view): ``out = AllGather_m(x) @ w^T``.
121 """
123 @staticmethod
124 def forward(ctx, x, w, group, world_size, bias): # pylint: disable=arguments-differ
125 """Run fused all-gather + matmul and stash tensors for backward."""
126 torch_npu = _require_torch_npu()
127 hcom = get_hcomm_info(group)
128 # x2 = w.T -> physical (k, n_local); kernel computes AG(x) @ x2
129 out, gathered = torch_npu.npu_all_gather_base_mm(
130 x,
131 w.t(),
132 hcom,
133 world_size,
134 bias=None,
135 gather_index=0,
136 gather_output=True,
137 )
138 if bias is not None:
139 out = out + bias
140 ctx.save_for_backward(gathered, w)
141 ctx.group = group
142 ctx.world_size = world_size
143 ctx.hcom = hcom
144 ctx.has_bias = bias is not None
145 return out
147 @staticmethod
148 def backward(ctx, grad_out): # pylint: disable=arguments-differ
149 """Gradient: dx via fused matmul_reduce_scatter, dw via gathered-input matmul.
151 Matches MindFormers: ``dX = matmul_reduce_scatter(dY, W)``. The fused
152 kernel contracts over ``n_local`` (``W.shape[0]``), which must be in
153 ``[256, 65535)``.
154 """
155 torch_npu = _require_torch_npu()
156 gathered, w = ctx.saved_tensors
157 # input (m_full, n_local) @ x2 (n_local, k) -> RS_m -> (m_local, k)
158 grad_x = torch_npu.npu_mm_reduce_scatter_base(
159 grad_out.contiguous(),
160 w,
161 ctx.hcom,
162 ctx.world_size,
163 reduce_op="sum",
164 bias=None,
165 )
166 grad_w = grad_out.t().matmul(gathered)
167 grad_bias = grad_out.sum(dim=0) if ctx.has_bias else None
168 return grad_x, grad_w, None, None, grad_bias
171class MatmulReduceScatterFunction(torch.autograd.Function):
172 """Row-parallel fused matmul + reduce-scatter with custom backward.
174 Forward (local view): ``out = ReduceScatter_m(x @ w^T)``.
175 """
177 @staticmethod
178 def forward(ctx, x, w, group, world_size, bias): # pylint: disable=arguments-differ
179 """Run fused matmul + reduce-scatter and stash tensors for backward."""
180 torch_npu = _require_torch_npu()
181 hcom = get_hcomm_info(group)
182 out = torch_npu.npu_mm_reduce_scatter_base(
183 x,
184 w.t(),
185 hcom,
186 world_size,
187 reduce_op="sum",
188 bias=None,
189 )
190 if bias is not None:
191 out = out + bias
192 ctx.save_for_backward(x, w)
193 ctx.group = group
194 ctx.world_size = world_size
195 ctx.hcom = hcom
196 ctx.has_bias = bias is not None
197 return out
199 @staticmethod
200 def backward(ctx, grad_out): # pylint: disable=arguments-differ
201 """Gradient: dx via fused all-gather+matmul, dw via gathered-grad matmul."""
202 torch_npu = _require_torch_npu()
203 x, w = ctx.saved_tensors
204 # AG(dY) @ W : pass W (n, k) so the kernel uses transposed-x2 semantics.
205 grad_x, grad_out_full = torch_npu.npu_all_gather_base_mm(
206 grad_out,
207 w,
208 ctx.hcom,
209 ctx.world_size,
210 bias=None,
211 gather_index=0,
212 gather_output=True,
213 )
214 grad_w = grad_out_full.t().matmul(x)
215 grad_bias = grad_out_full.sum(dim=0) if ctx.has_bias else None
216 return grad_x, grad_w, None, None, grad_bias
219class MC2Linear(nn.Linear):
220 """``nn.Linear`` that uses fused matmul + TP communication kernels."""
222 def configure_mc2(
223 self,
224 mode: str,
225 group: Any,
226 world_size: int,
227 sequence_dim: int = 0,
228 ) -> None:
229 """Configure the fused collective used by this layer.
231 Args:
232 mode (str): ``\"all_gather\"`` (column) or ``\"reduce_scatter\"`` (row).
233 group (Any): TP process group.
234 world_size (int): Size of ``group``.
235 sequence_dim (int, optional): Tensor dim that carries sequence sharding
236 under SP. Default: ``0``.
237 """
238 if mode not in ("all_gather", "reduce_scatter"):
239 raise ValueError(
240 "For MC2Linear.configure_mc2, mode should be 'all_gather' or "
241 f"'reduce_scatter', but got {mode}."
242 )
243 self.mc2_mode = mode
244 self.mc2_group = group
245 self.mc2_world_size = world_size
246 self.mc2_sequence_dim = sequence_dim
248 @classmethod
249 def from_linear(cls, linear: nn.Linear) -> "MC2Linear":
250 """Convert a Linear in place while preserving parameters and module state."""
251 if not isinstance(linear, nn.Linear):
252 raise TypeError(
253 f"MC2Linear can only replace nn.Linear, but got {type(linear).__name__}."
254 )
255 linear.__class__ = cls
256 return linear
258 def _mc2_forward(self, input_: DTensor, weight: DTensor) -> DTensor:
259 """Run the configured fused kernel on local tensors.
261 Ascend MC2 kernels all-gather / reduce-scatter the **flattened dim-0**
262 of a 2-D activation. That is only layout-correct when the SP sequence
263 axis is the outermost leading dim. For ``sequence_dim != 0`` (e.g.
264 batch-first ``[B, S, H]`` with ``Shard(1)``), move the sequence dim to
265 front before flatten, then restore after the fused op.
266 """
267 input_local = input_.to_local()
268 weight_local = weight.to_local()
269 leading_global = tuple(int(s) for s in input_.shape[:-1])
270 seq_dim = _normalize_sequence_dim(self.mc2_sequence_dim, len(leading_global))
272 # [..., S_local_or_full, ..., H] -> [S_*, *other_leading, H] -> 2-D
273 x_seq_first = _move_dim_to_front(input_local, seq_dim)
274 other_leading = tuple(x_seq_first.shape[1:-1])
275 input_2d = x_seq_first.reshape(-1, x_seq_first.shape[-1])
277 bias_local = None
278 if self.bias is not None:
279 bias = self.bias
280 bias_local = bias.to_local() if isinstance(bias, DTensor) else bias
282 if self.mc2_mode == "all_gather":
283 output_2d = AllGatherMatmulFunction.apply(
284 input_2d,
285 weight_local,
286 self.mc2_group,
287 self.mc2_world_size,
288 bias_local,
289 )
290 # AG expands the sequence dim; other leading dims stay local sizes.
291 seq_global = leading_global[seq_dim]
292 output = output_2d.reshape(seq_global, *other_leading, output_2d.shape[-1])
293 output = _move_front_to_dim(output, seq_dim)
294 return DTensor.from_local(output, input_.device_mesh, (Shard(-1),))
296 seq_global = leading_global[seq_dim]
297 if seq_global % self.mc2_world_size != 0:
298 raise RuntimeError(
299 f"MC2Linear reduce_scatter requires sequence dim {seq_dim} "
300 f"(size {seq_global}) divisible by world_size "
301 f"{self.mc2_world_size}."
302 )
303 output_2d = MatmulReduceScatterFunction.apply(
304 input_2d,
305 weight_local,
306 self.mc2_group,
307 self.mc2_world_size,
308 bias_local,
309 )
310 seq_local = seq_global // self.mc2_world_size
311 output = output_2d.reshape(seq_local, *other_leading, output_2d.shape[-1])
312 output = _move_front_to_dim(output, seq_dim)
313 return DTensor.from_local(output, input_.device_mesh, (Shard(seq_dim),))
315 def forward(self, input_: torch.Tensor, weight: Optional[torch.Tensor] = None) -> torch.Tensor:
316 """Forward using MC2 for DTensor inputs and ``nn.Linear`` otherwise."""
317 if not isinstance(input_, DTensor):
318 return super().forward(input_)
319 if not hasattr(self, "mc2_mode"):
320 raise RuntimeError(
321 "MC2Linear must be configured by an MC2 parallel style before use."
322 )
324 if weight is None:
325 weight = self.weight
326 if not isinstance(weight, DTensor):
327 raise TypeError(
328 "MC2Linear expects a DTensor weight after tensor-parallel sharding, "
329 f"but got {type(weight).__name__}."
330 )
331 return self._mc2_forward(input_, weight)