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

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

16from enum import Enum, auto 

17 

18 

19class MetaStepType(Enum): 

20 """Specify the enumeration type for MetaStep.""" 

21 DATA_LOAD = auto() 

22 DATA_SEND = auto() 

23 DATA_RECV = auto() 

24 FWD = auto() 

25 BWD = auto() 

26 BWD_INPUT = auto() 

27 BWD_WEIGHT = auto() 

28 FWD_RECV = auto() 

29 FWD_SEND = auto() 

30 BWD_RECV = auto() 

31 BWD_SEND = auto() 

32 # Composite P2P: a contiguous run of FWD_SEND/FWD_RECV/BWD_SEND/BWD_RECV 

33 # coalesced by ``coalesce_p2p`` into one step whose ``sub_steps`` the runtime 

34 # groups by peer and issues as ``batch_isend_irecv`` (same-peer send+recv -> 

35 # duplex). Produced under ``p2p_transport="batch"`` or ``"multi_stream"``. 

36 BATCH_SEND_RECV = auto() 

37 OVERLAP_F_B = auto() 

38 OVERLAP_B_F = auto() 

39 FSDP_UNSHARD = auto() 

40 FSDP_RESHARD = auto() 

41 FSDP_REDUCE_GRAD = auto() 

42 SWAP_LAUNCH_OFFLOAD = auto() 

43 SWAP_WAIT_OFFLOAD = auto() 

44 SWAP_LAUNCH_LOAD = auto() 

45 SWAP_WAIT_LOAD = auto() 

46 

47 

48class MetaStep: 

49 """ 

50 Meta step of PipelineSchedule. 

51 An execution list composed of MetaStep can be constructed 

52 and fed into the PipelineSchedule for execution. 

53 

54 Args: 

55 micro_index (int | None): The index of micro-batch. ``None`` for 

56 composite types (``OVERLAP_F_B`` / ``OVERLAP_B_F``) whose real 

57 micro index lives in each ``sub_steps`` entry. 

58 type (MetaStepType): Specify the type of current step. 

59 stage_index (int | None): Stage index of current step. ``None`` 

60 for composite types; use ``sub_steps`` to get each direction's 

61 stage. 

62 sub_steps (tuple[MetaStep, MetaStep] | None): For composite types 

63 only: ``(fwd, bwd)`` for ``OVERLAP_F_B``, ``(bwd, fwd)`` for 

64 ``OVERLAP_B_F``. 

65 boundary_p2p (tuple[MetaStep, ...] | None): For ``OVERLAP_B_F`` under 

66 the ``"boundary"`` P2P transport only: P2P steps to issue at the 

67 fwd/bwd boundary inside the overlap (the forward's ``FWD_SEND`` 

68 plus the next slot's recvs), hoisted out of the following gap by 

69 ``attach_fwd_boundary_p2p``. Issued via 

70 :meth:`PipelineScheduleRuntime.exec_boundary_p2p`. 

71 """ 

72 def __init__(self, micro_index, meta_type, stage_index, sub_steps=None, 

73 boundary_p2p=None): 

74 self._type = meta_type 

75 self._micro_index = micro_index 

76 self._stage_index = stage_index 

77 self._sub_steps = sub_steps 

78 self._boundary_p2p = boundary_p2p 

79 

80 @property 

81 def micro_index(self): 

82 """Return the micro-batch index of this step.""" 

83 return self._micro_index 

84 

85 @property 

86 def stage_index(self): 

87 """Return the stage index of this step.""" 

88 return self._stage_index 

89 

90 @property 

91 def type(self): 

92 """Return the MetaStepType of this step.""" 

93 return self._type 

94 

95 @property 

96 def sub_steps(self): 

97 """Return this step's sub-steps. 

98 

99 ``(fwd, bwd)`` for OVERLAP_F_B, ``(bwd, fwd)`` for OVERLAP_B_F, or 

100 ``None`` for every other type. 

101 """ 

102 return self._sub_steps 

103 

104 @property 

105 def boundary_p2p(self): 

106 """P2P steps to issue at the overlap's fwd/bwd boundary, or ``None``.""" 

107 return self._boundary_p2p 

108 

109 def __eq__(self, value): 

110 if not isinstance(value, MetaStep): 

111 return NotImplemented 

112 return (self.type == value.type 

113 and self.micro_index == value.micro_index 

114 and self.stage_index == value.stage_index 

115 and self.sub_steps == value.sub_steps) 

116 

117 def __ne__(self, value): 

118 if not isinstance(value, MetaStep): 

119 return NotImplemented 

120 return not self.__eq__(value) 

121 

122 def __hash__(self): 

123 return hash((self.type, self.micro_index, self.stage_index)) 

124 

125 def __str__(self): 

126 if self.sub_steps: 

127 sub = ", ".join(str(s) for s in self.sub_steps) 

128 return (f"MetaStep(type={self.type}, micro_index={self.micro_index}, " 

129 f"stage_index={self.stage_index}, sub_steps=[{sub}])") 

130 return f"MetaStep(type={self.type}, micro_index={self.micro_index}, stage_index={self.stage_index})" 

131 

132 def __repr__(self): 

133 return self.__str__() 

134 

135 @staticmethod 

136 def from_str(step_str): 

137 """Parse a MetaStep from its string representation.""" 

138 

139 

140class BatchDimSpec: 

141 """ 

142 Specify the batch dimension of a Tensor. 

143 

144 Args: 

145 batch_dim (int): batch dimension. 

146 """ 

147 __slots__ = ("batch_dim",) 

148 

149 def __init__(self, batch_dim): 

150 if not isinstance(batch_dim, int): 

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

152 self.batch_dim = batch_dim 

153 

154 def __repr__(self): 

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

156 

157 def __str__(self): 

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

159 

160 @staticmethod 

161 def from_tuple(batch_dims): 

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

163 if not isinstance(batch_dims, tuple): 

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

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

166 

167 @staticmethod 

168 def from_dict(batch_dims): 

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

170 if not isinstance(batch_dims, dict): 

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

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

173 

174 

175class _RecvInfo: 

176 """ 

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

178 

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

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

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

182 bwd-recv count. 

183 """ 

184 

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

186 self._global_rank = global_rank 

187 self._buffer = buffer 

188 self._requires_grad = bool(requires_grad) 

189 

190 @property 

191 def global_rank(self): 

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

193 return self._global_rank 

194 

195 @property 

196 def buffer(self): 

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

198 return self._buffer 

199 

200 @buffer.setter 

201 def buffer(self, val): 

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

203 self._buffer = val 

204 

205 @property 

206 def requires_grad(self) -> bool: 

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

208 return self._requires_grad 

209 

210 @requires_grad.setter 

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

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

213 self._requires_grad = bool(val)