Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / mindspore / pipeline_parallel / mpipe_transpose.py: 0%

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"""MindSpore autograd backend for the MPipe Transpose schedule. 

16 

17Only the autograd-specific hooks live here; the broadcast / P2P transport and 

18step orchestration are inherited from 

19:class:`~hyper_parallel.core.pipeline_parallel.mpipe.executor_base.MPipeTransposeExecutorBase`. 

20 

21Unlike torch, MindSpore's captured ``grad_fn`` is scoped to the body submodule's 

22own weights (see ``forward_and_gradfn`` / ``PipelineStage.forward_one_chunk``), 

23so the body backward only deposits the *input* gradient on the preprocess 

24output, never on the preprocess parameters. Therefore **every** micro-batch — 

25transposed and non-transposed — needs an explicit recompute backward, and the 

26preprocess forward is always detached. ``nontransposed_connected = False`` 

27signals the schedule to emit ``MPIPE_TRANSPOSE_BWD`` for non-transposed 

28micro-batches too. 

29 

30Note: 

31 This backend mirrors the established MindSpore pipeline-stage gradient 

32 pattern; it requires an Ascend/MindSpore environment for runtime validation. 

33""" 

34from hyper_parallel.core.pipeline_parallel.mpipe.executor_base import MPipeTransposeExecutorBase 

35from hyper_parallel.platform.mindspore.pipeline_parallel.backward import forward_and_gradfn 

36 

37 

38# This MindSpore backend is exercised by the MindSpore ST gate (msrun), not the 

39# torch/CPU coverage job (mindspore isn't importable there) — exclude from coverage. 

40class MPipeTransposeExecutor(MPipeTransposeExecutorBase): # pragma: no cover 

41 """MindSpore runtime for the ``MPIPE_*`` steps of MPipe Transpose.""" 

42 

43 nontransposed_connected = False 

44 

45 def _broadcast_tensors(self): 

46 # Only trainable params change after the optimizer step and need 

47 # re-syncing; frozen params are identical on every rank from init. 

48 for param in self._preprocess.get_parameters(): 

49 if param.requires_grad: 

50 yield param 

51 

52 def _detached_forward(self, args, kwargs): 

53 # A plain pynative Cell call builds no grad graph (it is naturally 

54 # detached). The base marks it grad-requiring (when T > 0) so the body's 

55 # grad_fn computes the input gradient on it. 

56 return self._preprocess(*args, **kwargs) 

57 

58 def _connected_forward(self, args, kwargs): 

59 # Unused (nontransposed_connected is False) but required by the base API. 

60 return self._preprocess(*args, **kwargs) 

61 

62 def _mark_requires_grad(self, tensor) -> None: 

63 tensor._requires_grad = True # pylint: disable=protected-access 

64 

65 def _recompute_backward(self, inputs, kwargs, grad) -> None: 

66 weights = tuple(self._preprocess.trainable_params()) 

67 _, grad_fn = forward_and_gradfn( 

68 self._preprocess, *inputs, weights=weights, grad_position=None, **kwargs) 

69 grad_fn(sens=grad)