Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / mpipe / schedule.py: 69%

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"""The MPipe Transpose schedule (an Interleaved-1F1B variant). 

16 

17Layers the ``MPIPE_*`` transpose steps around the inherited Interleaved-1F1B 

18body order and registers their handlers via the generic custom-function 

19registry, so the core ``scheduler`` module carries no MPipe-specific code. 

20""" 

21from typing import Optional, TYPE_CHECKING 

22 

23from hyper_parallel.platform import get_platform 

24from hyper_parallel.platform.platform import PlatformType 

25from hyper_parallel.core.pipeline_parallel.scheduler import ( 

26 MetaStep, 

27 MetaStepType, 

28 ScheduleInterleaved1F1B, 

29) 

30from hyper_parallel.core.pipeline_parallel.mpipe.step_types import MpipeStepType 

31 

32if TYPE_CHECKING: 

33 from hyper_parallel.core.pipeline_parallel.utils import BatchDimSpec 

34 from hyper_parallel.dmodule.module import Module 

35 

36platform = get_platform() 

37 

38 

39class ScheduleMPipeTranspose(ScheduleInterleaved1F1B): 

40 """The MPipe Transpose schedule. 

41 

42 A variant of Interleaved 1F1B that shrinks the warmup pipeline bubble by 

43 **transposing** the forward of a model's first ``T`` layers (the 

44 *preprocess* block, which logically belongs to stage 0's first chunk). 

45 

46 Instead of stage 0 computing the preprocess forward serially for every 

47 micro-batch, the preprocess parameters are broadcast to all ``PP`` ranks 

48 and, for the first ``NT = min(PP, micro_batch_num)`` micro-batches, each 

49 rank ``i`` computes the preprocess forward of micro-batch ``i`` in parallel 

50 during what would otherwise be its warmup idle time. Each rank ``i > 0`` 

51 then ships to stage 0: 

52 

53 * the preprocess **output** (``MPIPE_FWD_SEND`` / ``MPIPE_FWD_RECV``) so 

54 stage 0 can run its body forward, and 

55 * the preprocess **input** (``MPIPE_GRAPH_SEND`` / ``MPIPE_GRAPH_RECV``) 

56 so the preprocess backward can be recomputed centrally on stage 0 

57 (gradients accumulate on stage 0 only). 

58 

59 The remaining ``micro_batch_num - NT`` micro-batches run the preprocess 

60 forward inline on stage 0 (graph-connected to the body, so their backward 

61 is automatic), exactly as ordinary Interleaved 1F1B. 

62 

63 The body model (stage 0 = the layers after the preprocess block, all other 

64 stages unchanged) is scheduled by the inherited Interleaved 1F1B logic; the 

65 preprocess steps are layered around it: a transpose-phase prefix per rank, 

66 plus inline preprocess forward / recompute backward steps on stage 0. 

67 

68 Args: 

69 stages (list[PipelineStage], PipelineStage): The body pipeline stages. 

70 Stage 0 must wrap only the layers **after** the preprocess block. 

71 micro_batch_num (int): The number of micro-batches. 

72 preprocess_module (Optional[Module]): The preprocess block (first ``T`` 

73 layers of stage 0). Following Option A, it must exist on **every** rank: on rank 0 

74 it holds the trained parameters; on other ranks it is a structural 

75 copy whose parameters are overwritten each step by the broadcast. 

76 num_transpose_layers (int): ``T`` — the number of preprocess layers, 

77 must be smaller than the layer count of stage 0's first chunk. 

78 ``0`` is allowed and means *only the data loading is transposed*: 

79 each rank loads its micro-batch and ships the raw input to stage 0, 

80 with no parameter broadcast, no preprocess compute, and no 

81 recompute backward. 

82 args_batch_dim (list, optional): See ``PipelineScheduleRuntime``. 

83 kwargs_batch_dim (dict, optional): See ``PipelineScheduleRuntime``. 

84 output_concat_dim (int, optional): See ``PipelineScheduleRuntime``. 

85 overlap_p2p (bool, optional): See ``ScheduleInterleaved1F1B``. 

86 Default ``False``. 

87 swap (bool, optional): Reserved for API compatibility. MPipe activation 

88 swap is not supported and ``True`` raises ``ValueError``. 

89 

90 Note: 

91 This class builds the schedule ordering and registers the ``MPIPE_*`` 

92 execution handlers; the handlers themselves live in the platform 

93 executors (see :class:`MPipeTransposeExecutorBase`). 

94 """ 

95 

96 def __init__(self, 

97 stages: list, 

98 micro_batch_num: int, 

99 preprocess_module: "Optional[Module]", 

100 num_transpose_layers: int, 

101 args_batch_dim: "Optional[BatchDimSpec]" = None, 

102 kwargs_batch_dim: "Optional[BatchDimSpec]" = None, 

103 output_concat_dim: Optional[int] = None, 

104 overlap_p2p: bool = False, 

105 swap: bool = False) -> None: 

106 """Build an interleaved-1F1B schedule that transposes the preprocess block. 

107 

108 Args: 

109 stages (list): The local pipeline stages (as for :class:`ScheduleInterleaved1F1B`). 

110 micro_batch_num (int): Number of micro-batches per optimizer step. 

111 preprocess_module (Optional[Module]): The block transposed to every 

112 rank — the first ``num_transpose_layers`` layers, a visual tower, 

113 or a param-free identity for the dataload-only (``T = 0``) mode. 

114 num_transpose_layers (int): ``T`` — the (informational) transposed-layer count. 

115 args_batch_dim (Optional[BatchDimSpec]): Positional-arg batch-dim spec (forwarded to the base). 

116 kwargs_batch_dim (Optional[BatchDimSpec]): Keyword-arg batch-dim spec (forwarded to the base). 

117 output_concat_dim (Optional[int]): Output concatenation dim (forwarded to the base). 

118 overlap_p2p (bool): Whether to overlap P2P (forwarded to the base). 

119 swap (bool): Reserved for API compatibility. Must be ``False``. 

120 """ 

121 if swap: 

122 raise ValueError( 

123 "MPipe Transpose with pipeline activation swap is not yet supported." 

124 ) 

125 if not isinstance(num_transpose_layers, int) or num_transpose_layers < 0: 

126 raise ValueError( 

127 f"Argument 'num_transpose_layers' must be a non-negative int, " 

128 f"but got {num_transpose_layers!r}." 

129 ) 

130 # ``preprocess_module`` is the resolved block to transpose (first T text 

131 # layers, the visual tower, or a param-free identity for dataload-only). 

132 self._preprocess_module = preprocess_module 

133 self._num_transpose_layers = num_transpose_layers 

134 # Whether the preprocess has *trainable* params decides the path: 

135 # trainable -> broadcast (the trainable params only) + centralized 

136 # recompute backward; 

137 # frozen / param-free (T=0, a frozen visual tower) -> ship the output 

138 # only (no broadcast, no recompute). 

139 self._has_trainable_preprocess = self._module_has_trainable_params(preprocess_module) 

140 # MindSpore's grad_fn is scoped to the body submodule's weights, so a 

141 # *trainable* preprocess also needs an explicit recompute backward for the 

142 # non-transposed micro-batches; torch's autograd handles them via the 

143 # connected graph. 

144 self._explicit_nontransposed_backward = platform.platform_type == PlatformType.MINDSPORE 

145 super().__init__(stages, 

146 micro_batch_num, 

147 args_batch_dim=args_batch_dim, 

148 kwargs_batch_dim=kwargs_batch_dim, 

149 output_concat_dim=output_concat_dim, 

150 overlap_p2p=overlap_p2p, 

151 overlap_b_f=False, 

152 swap=swap) 

153 self._executor = None 

154 self._setup_mpipe_execution() 

155 

156 @staticmethod 

157 def _module_has_trainable_params(module) -> bool: 

158 """Whether ``module`` has any trainable (grad-requiring) parameter.""" 

159 if module is None: 

160 return False 

161 if platform.platform_type == PlatformType.PYTORCH: 

162 return any(p.requires_grad for p in module.parameters()) 

163 if platform.platform_type == PlatformType.MINDSPORE: 

164 return any(p.requires_grad for p in module.get_parameters()) 

165 raise NotImplementedError( 

166 f"MPipe Transpose is not implemented for platform {platform.platform_type}." 

167 ) 

168 

169 @property 

170 def preprocess_module(self) -> "Optional[Module]": 

171 """The preprocess block transposed to every rank (present on every rank).""" 

172 return self._preprocess_module 

173 

174 @property 

175 def has_trainable_preprocess(self) -> bool: 

176 """Whether the preprocess has trainable params (drives broadcast/recompute).""" 

177 return self._has_trainable_preprocess 

178 

179 @property 

180 def num_transpose_layers(self) -> int: 

181 """``T`` — informational transposed-layer count (``0`` = dataload only).""" 

182 return self._num_transpose_layers 

183 

184 @property 

185 def num_transpose_micro_batches(self) -> int: 

186 """``NT = min(PP, micro_batch_num)`` — the count of transposed micro-batches.""" 

187 return min(self.real_stage_num, self.micro_batch_num) 

188 

189 def _setup_mpipe_execution(self) -> None: 

190 """Build the platform execution backend and register the MPIPE_* handlers.""" 

191 # Lazy import: the backend executor pulls in torch/mindspore and is 

192 # resolved only when a schedule is actually constructed. 

193 if platform.platform_type == PlatformType.PYTORCH: 

194 from hyper_parallel.platform.torch.pipeline_parallel.mpipe_transpose import ( # pylint: disable=C0415 

195 MPipeTransposeExecutor, 

196 ) 

197 elif platform.platform_type == PlatformType.MINDSPORE: 

198 from hyper_parallel.platform.mindspore.pipeline_parallel.mpipe_transpose import ( # pylint: disable=C0415 

199 MPipeTransposeExecutor, 

200 ) 

201 else: 

202 raise NotImplementedError( 

203 f"MPipe Transpose execution is not implemented for platform {platform.platform_type}." 

204 ) 

205 self._executor = MPipeTransposeExecutor(self) 

206 handlers = { 

207 MpipeStepType.MPIPE_PARAM_BROADCAST: self._executor.broadcast_params, 

208 MpipeStepType.MPIPE_TRANSPOSE_FWD: self._executor.transpose_forward, 

209 MpipeStepType.MPIPE_FWD_SEND: self._executor.fwd_send, 

210 MpipeStepType.MPIPE_FWD_RECV: self._executor.fwd_recv, 

211 MpipeStepType.MPIPE_GRAPH_SEND: self._executor.graph_send, 

212 MpipeStepType.MPIPE_GRAPH_RECV: self._executor.graph_recv, 

213 MpipeStepType.MPIPE_TRANSPOSE_BWD: self._executor.transpose_backward, 

214 } 

215 for step_type, handler in handlers.items(): 

216 self.register_custom_function(step_type, handler) 

217 

218 def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None: 

219 """Reset the executor's per-step caches, then run the schedule. 

220 

221 Args: 

222 arg_mbs (list): Per-micro-batch positional args. 

223 kwarg_mbs (list): Per-micro-batch keyword args. 

224 losses (list): Mutable list collecting per-step losses. 

225 """ 

226 if self._executor is not None: 

227 self._executor.reset() 

228 super().run_microbatches(arg_mbs, kwarg_mbs, losses) 

229 

230 def construct_exec_order(self) -> None: 

231 """Build the body Interleaved 1F1B order, then layer the preprocess 

232 transpose phase and the centralized preprocess backward on top. 

233 

234 The parameter broadcast, the recompute-input transport, and the 

235 recompute backward are emitted only for a **trainable** preprocess; a 

236 frozen or param-free one (``T == 0``, a frozen visual tower) only 

237 transposes the forward and ships its output. 

238 """ 

239 super().construct_exec_order() 

240 body_order = self.exec_order 

241 num_transpose = self.num_transpose_micro_batches 

242 has_trainable = self._has_trainable_preprocess 

243 body_order[0] = self._insert_rank0_preprocess_steps( 

244 body_order[0], num_transpose, has_trainable, 

245 backward_all=self._explicit_nontransposed_backward) 

246 self.exec_order = { 

247 rank: self._build_transpose_prefix(rank, num_transpose, has_trainable) + body_order[rank] 

248 for rank in range(self.real_stage_num) 

249 } 

250 

251 @staticmethod 

252 def _build_transpose_prefix(rank, num_transpose, has_trainable): 

253 """Build the transpose-phase prefix prepended to ``rank``'s body order. 

254 

255 A rank that owns a transposed micro-batch (``rank < num_transpose``) 

256 computes its preprocess forward and ranks ``> 0`` ship the output to 

257 stage 0. For a **trainable** preprocess every rank also broadcasts its 

258 (trainable) parameters and ranks ``> 0`` additionally ship the input for 

259 the centralized recompute backward; for a frozen / param-free preprocess 

260 only the transpose forward and its output send/recv remain. 

261 """ 

262 prefix = [] 

263 if has_trainable: 

264 prefix.append(MetaStep(None, MpipeStepType.MPIPE_PARAM_BROADCAST, 0)) 

265 if rank < num_transpose: 

266 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_TRANSPOSE_FWD, 0)) 

267 if rank != 0: 

268 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_FWD_SEND, 0)) 

269 if has_trainable: 

270 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_GRAPH_SEND, 0)) 

271 if rank == 0: 

272 for micro_index in range(1, num_transpose): 

273 prefix.append(MetaStep(micro_index, MpipeStepType.MPIPE_FWD_RECV, 0)) 

274 if has_trainable: 

275 for micro_index in range(1, num_transpose): 

276 prefix.append(MetaStep(micro_index, MpipeStepType.MPIPE_GRAPH_RECV, 0)) 

277 return prefix 

278 

279 @staticmethod 

280 def _insert_rank0_preprocess_steps(order, num_transpose, has_trainable, backward_all=False): 

281 """Patch stage 0's (rank 0) body order with preprocess fwd/bwd steps. 

282 

283 For a **trainable** preprocess: before each ``FWD(stage 0, micro >= 

284 num_transpose)`` an inline ``MPIPE_TRANSPOSE_FWD`` runs the preprocess 

285 forward for that non-transposed micro-batch, and an 

286 ``MPIPE_TRANSPOSE_BWD`` is inserted after each ``BWD(stage 0, micro)`` 

287 needing a centralized recompute backward: the transposed micro-batches 

288 (``micro < num_transpose``) always, and — when ``backward_all`` is set 

289 (MindSpore, whose body backward does not flow into the preprocess) — the 

290 non-transposed ones too. On torch (``backward_all`` False) non-transposed 

291 micro-batches backprop into the preprocess via the connected graph. 

292 

293 For a frozen / param-free preprocess (no trainable params) there is no 

294 recompute backward, so the body order is returned unchanged (transposed 

295 outputs are placed by ``MPIPE_FWD_RECV``; non-transposed micro-batches 

296 are handled by stage 0 directly). VL's frozen-visual injection is wired 

297 per-model rather than through this text-style body-input path. 

298 """ 

299 if not has_trainable: 

300 return order 

301 patched = [] 

302 for step in order: 

303 is_stage0_fwd = ( 

304 step is not None 

305 and step.type == MetaStepType.FWD 

306 and step.stage_index == 0 

307 ) 

308 if is_stage0_fwd and step.micro_index >= num_transpose: 

309 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_FWD, 0)) 

310 patched.append(step) 

311 is_stage0_bwd = ( 

312 step is not None 

313 and step.type == MetaStepType.BWD 

314 and step.stage_index == 0 

315 ) 

316 if is_stage0_bwd and (step.micro_index < num_transpose or backward_all): 

317 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_BWD, 0)) 

318 return patched