Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / pipeline_swap.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"""Pipeline-parallel activation swap scheduling helpers.""" 

16 

17from collections import defaultdict 

18from enum import IntEnum 

19 

20from hyper_parallel.core.activation_checkpoint.swap import SwapManager 

21 

22MIN_SWAP_GAP = 4 

23 

24 

25class _BeforeActionPriority(IntEnum): 

26 SET_GROUP = 10 

27 WAIT_LOAD = 20 

28 LAUNCH_LOAD = 30 

29 

30 

31class _AfterActionPriority(IntEnum): 

32 WAIT_OFFLOAD = 10 

33 LAUNCH_OFFLOAD = 20 

34 

35 

36def pp_swap_group_name(stage_index: int, micro_index: int) -> str: 

37 """Return the swap group name for a pipeline chunk.""" 

38 return f"pp_swap_s{stage_index}_m{micro_index}" 

39 

40 

41def _is_compute_step(step) -> bool: 

42 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415 

43 

44 return step is not None and step.type in (MetaStepType.FWD, MetaStepType.BWD) 

45 

46 

47def _is_comm_step(step) -> bool: 

48 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415 

49 

50 return step is not None and step.type in ( 

51 MetaStepType.FWD_RECV, 

52 MetaStepType.FWD_SEND, 

53 MetaStepType.BWD_RECV, 

54 MetaStepType.BWD_SEND, 

55 ) 

56 

57 

58def _is_composite_compute_step(step) -> bool: 

59 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415 

60 

61 return ( 

62 step is not None 

63 and step.type in (MetaStepType.OVERLAP_F_B, MetaStepType.OVERLAP_B_F) 

64 and step.sub_steps 

65 ) 

66 

67 

68class _ComputeLeaf: 

69 """A real FWD/BWD leaf and the top-level container that owns it.""" 

70 

71 __slots__ = ("step", "container_index", "compute_index") 

72 

73 def __init__(self, step, container_index, compute_index): 

74 self.step = step 

75 self.container_index = container_index 

76 self.compute_index = compute_index 

77 

78 

79def _iter_compute_leaf_steps(step): 

80 """Yield real FWD/BWD steps, expanding composite containers.""" 

81 if _is_compute_step(step): 

82 yield step 

83 return 

84 if _is_composite_compute_step(step): 

85 for sub_step in step.sub_steps: 

86 if _is_compute_step(sub_step): 

87 yield sub_step 

88 

89 

90def _collect_compute_leaves(order): 

91 """Collect compute leaves while counting each composite as one slot.""" 

92 leaves = [] 

93 container_by_compute_index = {} 

94 compute_index = 0 

95 for container_index, step in enumerate(order): 

96 leaf_steps = list(_iter_compute_leaf_steps(step)) 

97 if not leaf_steps: 

98 continue 

99 container_by_compute_index[compute_index] = container_index 

100 for leaf_step in leaf_steps: 

101 leaves.append(_ComputeLeaf(leaf_step, container_index, compute_index)) 

102 compute_index += 1 

103 return leaves, container_by_compute_index 

104 

105 

106def _append_after(after_steps, index, priority, step): 

107 after_steps[index].append((priority, step)) 

108 

109 

110def _append_before(before_steps, index, priority, step): 

111 before_steps[index].append((priority, step)) 

112 

113 

114def _iter_steps_by_priority(priority_steps): 

115 """Yield steps from high priority to low priority.""" 

116 for _, step in sorted(priority_steps, key=lambda item: item[0], reverse=True): 

117 yield step 

118 

119 

120def _comm_block_anchor(order, index): 

121 """Return the last immediately following communication step.""" 

122 anchor = index 

123 for next_index in range(index + 1, len(order)): 

124 next_step = order[next_index] 

125 if _is_comm_step(next_step): 

126 anchor = next_index 

127 continue 

128 break 

129 return anchor 

130 

131 

132def _post_compute_anchor(order, index, leaf_step=None): 

133 """Return the safe index after which post-compute swap steps may run.""" 

134 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStepType # pylint: disable=C0415 

135 

136 step = leaf_step if leaf_step is not None else order[index] 

137 fallback_anchor = _comm_block_anchor(order, index) 

138 if step.type == MetaStepType.FWD: 

139 send_type = MetaStepType.FWD_SEND 

140 elif step.type == MetaStepType.BWD: 

141 send_type = MetaStepType.BWD_SEND 

142 else: 

143 return fallback_anchor 

144 

145 for next_index in range(index + 1, fallback_anchor + 1): 

146 next_step = order[next_index] 

147 next_valid = (next_step is not None and next_step.type == send_type 

148 and next_step.stage_index == step.stage_index and next_step.micro_index == step.micro_index) 

149 if next_valid: 

150 return next_index 

151 return fallback_anchor 

152 

153 

154def inject_pipeline_swap_steps(order): 

155 """Inject SWAP_* steps into one rank's pipeline order. 

156 

157 The injected order preserves the required lifecycle: 

158 SET_GROUP -> FWD -> LAUNCH_OFFLOAD -> WAIT_OFFLOAD -> 

159 LAUNCH_LOAD -> WAIT_LOAD -> BWD. 

160 """ 

161 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStep, MetaStepType # pylint: disable=C0415 

162 

163 fwd_index = {} 

164 bwd_index = {} 

165 compute_leaves, container_by_compute_index = _collect_compute_leaves(order) 

166 for leaf in compute_leaves: 

167 step = leaf.step 

168 key = (step.stage_index, step.micro_index) 

169 if step.type == MetaStepType.FWD: 

170 fwd_index[key] = leaf 

171 elif step.type == MetaStepType.BWD: 

172 bwd_index[key] = leaf 

173 

174 before_steps = defaultdict(list) 

175 after_steps = defaultdict(list) 

176 for key, fwd_leaf in fwd_index.items(): 

177 bwd_leaf = bwd_index.get(key) 

178 if bwd_leaf is None: 

179 continue 

180 if bwd_leaf.compute_index - fwd_leaf.compute_index < MIN_SWAP_GAP: 

181 continue 

182 compute_between = [ 

183 container_by_compute_index[index] 

184 for index in range(fwd_leaf.compute_index + 1, bwd_leaf.compute_index) 

185 ] 

186 if not compute_between: 

187 continue 

188 stage_index, micro_index = key 

189 

190 _append_before( 

191 before_steps, fwd_leaf.container_index, _BeforeActionPriority.SET_GROUP, 

192 MetaStep(micro_index, MetaStepType.SWAP_SET_GROUP, stage_index), 

193 ) 

194 fwd_anchor = _post_compute_anchor(order, fwd_leaf.container_index, fwd_leaf.step) 

195 first_between_anchor = _post_compute_anchor(order, compute_between[0]) 

196 

197 _append_after( 

198 after_steps, fwd_anchor, _AfterActionPriority.LAUNCH_OFFLOAD, 

199 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_OFFLOAD, stage_index), 

200 ) 

201 

202 _append_after( 

203 after_steps, first_between_anchor, _AfterActionPriority.WAIT_OFFLOAD, 

204 MetaStep(micro_index, MetaStepType.SWAP_WAIT_OFFLOAD, stage_index), 

205 ) 

206 

207 _append_before( 

208 before_steps, compute_between[-1], _BeforeActionPriority.LAUNCH_LOAD, 

209 MetaStep(micro_index, MetaStepType.SWAP_LAUNCH_LOAD, stage_index), 

210 ) 

211 _append_before( 

212 before_steps, bwd_leaf.container_index, _BeforeActionPriority.WAIT_LOAD, 

213 MetaStep(micro_index, MetaStepType.SWAP_WAIT_LOAD, stage_index), 

214 ) 

215 

216 injected = [] 

217 for index, step in enumerate(order): 

218 injected.extend(_iter_steps_by_priority(before_steps[index])) 

219 injected.append(step) 

220 injected.extend(_iter_steps_by_priority(after_steps[index])) 

221 return injected 

222 

223 

224def _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs) -> None: 

225 """Keep pipeline-owned boundary tensors alive on device. 

226 

227 Swap offload clears the device storage of saved tensors after D2H copy. 

228 If a saved tensor aliases a pipeline boundary tensor, clearing it would 

229 also invalidate the object still held by the pipeline runtime. The alias 

230 protection below marks those saved tensors as keep-on-device. 

231 """ 

232 stage = schedule._stage_dict[step.stage_index] # pylint: disable=protected-access 

233 group_name = pp_swap_group_name(step.stage_index, step.micro_index) 

234 manager = SwapManager() 

235 

236 if stage.is_first_stage: 

237 # First-stage inputs come from split_microbatches(), outside the 

238 # wrapped stage. They are not stage outputs, but they can alias 

239 # tensors saved by the first layer and must not have their storage 

240 # resized by the swap group. 

241 manager.protect_alias_tensors( 

242 group_name, 

243 (arg_mbs[step.micro_index], kwarg_mbs[step.micro_index]), 

244 ) 

245 

246 if stage.is_last_stage: 

247 # Last-stage outputs are consumed by the schedule as losses / sens 

248 # roots, so they must stay device-valid until backward has used them. 

249 outputs = stage.fwd_outputs_cache.get(step.micro_index) 

250 if outputs is None: 

251 outputs = stage.last_stage_outputs 

252 if outputs is not None: 

253 manager.protect_alias_tensors(group_name, outputs) 

254 

255 

256def swap_set_group(step) -> None: 

257 """Set the active SwapManager group for the next pipeline forward chunk.""" 

258 group_name = pp_swap_group_name(step.stage_index, step.micro_index) 

259 manager = SwapManager() 

260 manager.ensure_group(group_name) 

261 manager.set_current_group_name(group_name) 

262 

263 

264def swap_launch_offload(step, schedule, arg_mbs, kwarg_mbs) -> None: 

265 """Launch D2H for a pipeline swap group.""" 

266 group_name = pp_swap_group_name(step.stage_index, step.micro_index) 

267 manager = SwapManager() 

268 _protect_pipeline_owned_tensors(step, schedule, arg_mbs, kwarg_mbs) 

269 manager.launch_offload(group_name) 

270 # Only the immediately preceding forward belongs to this swap group. 

271 manager.set_current_group_name("") 

272 

273 

274def swap_wait_offload(step) -> None: 

275 """Wait for a pipeline swap group's D2H and release device storage.""" 

276 SwapManager().wait_offload(pp_swap_group_name(step.stage_index, step.micro_index)) 

277 

278 

279def swap_launch_load(step) -> None: 

280 """Launch H2D for a pipeline swap group.""" 

281 SwapManager().launch_load(pp_swap_group_name(step.stage_index, step.micro_index)) 

282 

283 

284def swap_wait_load(step) -> None: 

285 """Wait for a pipeline swap group's H2D before backward uses activations.""" 

286 SwapManager().wait_load(pp_swap_group_name(step.stage_index, step.micro_index))