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

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"""Torch 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 

21Torch's ``autograd.backward`` traverses the whole connected graph to all leaves, 

22so non-transposed micro-batches can run a graph-connected preprocess forward and 

23have the body backward flow into the preprocess parameters automatically (no 

24recompute step) — hence ``nontransposed_connected = True``. 

25""" 

26import torch 

27 

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

29 

30 

31class MPipeTransposeExecutor(MPipeTransposeExecutorBase): 

32 """Torch runtime for the ``MPIPE_*`` steps of MPipe Transpose.""" 

33 

34 nontransposed_connected = True 

35 

36 def _broadcast_tensors(self): 

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

38 # re-syncing; frozen params and constant buffers are identical on every 

39 # rank from init, so broadcasting them would be wasted bandwidth. 

40 for param in self._preprocess.parameters(): 

41 if param.requires_grad: 

42 yield param.data 

43 

44 def _detached_forward(self, args, kwargs): 

45 with torch.no_grad(): 

46 out = self._preprocess(*args, **kwargs) 

47 return out.detach() 

48 

49 def _connected_forward(self, args, kwargs): 

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

51 

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

53 tensor.requires_grad_(True) 

54 

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

56 out = self._preprocess(*inputs, **kwargs) 

57 torch.autograd.backward(out, grad_tensors=grad) 

58 

59 @staticmethod 

60 def _contiguous(tensor): 

61 return tensor.contiguous()