Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / utils.py: 68%

1# Copyright 2025 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"""pipeline parallel utils""" 

16 

17 

18class BatchDimSpec: 

19 """ 

20 Specify the batch dimension of a Tensor. 

21 

22 Args: 

23 batch_dim (int): batch dimension. 

24 """ 

25 __slots__ = ("batch_dim",) 

26 

27 def __init__(self, batch_dim): 

28 if not isinstance(batch_dim, int): 

29 raise TypeError(f"batch_dim must be int, but got type {type(batch_dim)}.") 

30 self.batch_dim = batch_dim 

31 

32 def __repr__(self): 

33 return f"BatchDimSpec({self.batch_dim})" 

34 

35 def __str__(self): 

36 return f"BatchDim(dim={self.batch_dim})" 

37 

38 @staticmethod 

39 def from_tuple(batch_dims): 

40 """Create a tuple of BatchDimSpec from a tuple of batch dimensions.""" 

41 if not isinstance(batch_dims, tuple): 

42 raise TypeError(f"batch_dims must be tuple, but got type {type(batch_dims)}.") 

43 return tuple(BatchDimSpec(dim) for dim in batch_dims) 

44 

45 @staticmethod 

46 def from_dict(batch_dims): 

47 """Create a dict of BatchDimSpec from a dict mapping keys to batch dimensions.""" 

48 if not isinstance(batch_dims, dict): 

49 raise TypeError(f"batch_dims must be dict, but got type {type(batch_dims)}.") 

50 return {k: BatchDimSpec(v) for k, v in batch_dims.items()} 

51 

52 

53class _RecvInfo: 

54 """ 

55 Used for construct forward Receive operation and backward Send operation. 

56 

57 ``requires_grad`` mirrors the forward tensor's ``requires_grad`` so 

58 pipeline code can skip backward send/recv for tensors that have no 

59 gradient, keeping the bwd-send count consistent with the peer's 

60 bwd-recv count. 

61 """ 

62 

63 def __init__(self, global_rank, buffer=None, requires_grad: bool = True): 

64 self._global_rank = global_rank 

65 self._buffer = buffer 

66 self._requires_grad = bool(requires_grad) 

67 

68 @property 

69 def global_rank(self): 

70 """Return the global rank of the peer process.""" 

71 return self._global_rank 

72 

73 @property 

74 def buffer(self): 

75 """Return the receive/send buffer tensor.""" 

76 return self._buffer 

77 

78 @buffer.setter 

79 def buffer(self, val): 

80 """Set the receive/send buffer tensor.""" 

81 self._buffer = val 

82 

83 @property 

84 def requires_grad(self) -> bool: 

85 """Whether the corresponding forward tensor requires a gradient.""" 

86 return self._requires_grad 

87 

88 @requires_grad.setter 

89 def requires_grad(self, val: bool) -> None: 

90 """Set whether the corresponding forward tensor requires a gradient.""" 

91 self._requires_grad = bool(val)