Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / mpipe / schedule.py: 68%
91 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
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).
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
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
32if TYPE_CHECKING:
33 from hyper_parallel.core.pipeline_parallel.utils import BatchDimSpec
34 from hyper_parallel.dmodule.module import Module
36platform = get_platform()
39class ScheduleMPipeTranspose(ScheduleInterleaved1F1B):
40 """The MPipe Transpose schedule.
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).
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:
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).
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.
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.
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): Whether to inject activation-swap steps.
88 Default ``False``.
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 """
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.
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): Whether to enable activation swapping (forwarded to the base).
120 """
121 if not isinstance(num_transpose_layers, int) or num_transpose_layers < 0:
122 raise ValueError(
123 f"Argument 'num_transpose_layers' must be a non-negative int, "
124 f"but got {num_transpose_layers!r}."
125 )
126 # ``preprocess_module`` is the resolved block to transpose (first T text
127 # layers, the visual tower, or a param-free identity for dataload-only).
128 self._preprocess_module = preprocess_module
129 self._num_transpose_layers = num_transpose_layers
130 # Whether the preprocess has *trainable* params decides the path:
131 # trainable -> broadcast (the trainable params only) + centralized
132 # recompute backward;
133 # frozen / param-free (T=0, a frozen visual tower) -> ship the output
134 # only (no broadcast, no recompute).
135 self._has_trainable_preprocess = self._module_has_trainable_params(preprocess_module)
136 # MindSpore's grad_fn is scoped to the body submodule's weights, so a
137 # *trainable* preprocess also needs an explicit recompute backward for the
138 # non-transposed micro-batches; torch's autograd handles them via the
139 # connected graph.
140 self._explicit_nontransposed_backward = platform.platform_type == PlatformType.MINDSPORE
141 super().__init__(stages,
142 micro_batch_num,
143 args_batch_dim=args_batch_dim,
144 kwargs_batch_dim=kwargs_batch_dim,
145 output_concat_dim=output_concat_dim,
146 overlap_p2p=overlap_p2p,
147 overlap_b_f=False,
148 swap=swap)
149 self._executor = None
150 self._setup_mpipe_execution()
152 @staticmethod
153 def _module_has_trainable_params(module) -> bool:
154 """Whether ``module`` has any trainable (grad-requiring) parameter."""
155 if module is None:
156 return False
157 if platform.platform_type == PlatformType.PYTORCH:
158 return any(p.requires_grad for p in module.parameters())
159 if platform.platform_type == PlatformType.MINDSPORE:
160 return any(p.requires_grad for p in module.get_parameters())
161 raise NotImplementedError(
162 f"MPipe Transpose is not implemented for platform {platform.platform_type}."
163 )
165 @property
166 def preprocess_module(self) -> "Optional[Module]":
167 """The preprocess block transposed to every rank (present on every rank)."""
168 return self._preprocess_module
170 @property
171 def has_trainable_preprocess(self) -> bool:
172 """Whether the preprocess has trainable params (drives broadcast/recompute)."""
173 return self._has_trainable_preprocess
175 @property
176 def num_transpose_layers(self) -> int:
177 """``T`` — informational transposed-layer count (``0`` = dataload only)."""
178 return self._num_transpose_layers
180 @property
181 def num_transpose_micro_batches(self) -> int:
182 """``NT = min(PP, micro_batch_num)`` — the count of transposed micro-batches."""
183 return min(self.real_stage_num, self.micro_batch_num)
185 def _setup_mpipe_execution(self) -> None:
186 """Build the platform execution backend and register the MPIPE_* handlers."""
187 # Lazy import: the backend executor pulls in torch/mindspore and is
188 # resolved only when a schedule is actually constructed.
189 if platform.platform_type == PlatformType.PYTORCH:
190 from hyper_parallel.platform.torch.pipeline_parallel.mpipe_transpose import ( # pylint: disable=C0415
191 MPipeTransposeExecutor,
192 )
193 elif platform.platform_type == PlatformType.MINDSPORE:
194 from hyper_parallel.platform.mindspore.pipeline_parallel.mpipe_transpose import ( # pylint: disable=C0415
195 MPipeTransposeExecutor,
196 )
197 else:
198 raise NotImplementedError(
199 f"MPipe Transpose execution is not implemented for platform {platform.platform_type}."
200 )
201 self._executor = MPipeTransposeExecutor(self)
202 handlers = {
203 MpipeStepType.MPIPE_PARAM_BROADCAST: self._executor.broadcast_params,
204 MpipeStepType.MPIPE_TRANSPOSE_FWD: self._executor.transpose_forward,
205 MpipeStepType.MPIPE_FWD_SEND: self._executor.fwd_send,
206 MpipeStepType.MPIPE_FWD_RECV: self._executor.fwd_recv,
207 MpipeStepType.MPIPE_GRAPH_SEND: self._executor.graph_send,
208 MpipeStepType.MPIPE_GRAPH_RECV: self._executor.graph_recv,
209 MpipeStepType.MPIPE_TRANSPOSE_BWD: self._executor.transpose_backward,
210 }
211 for step_type, handler in handlers.items():
212 self.register_custom_function(step_type, handler)
214 def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
215 """Reset the executor's per-step caches, then run the schedule.
217 Args:
218 arg_mbs (list): Per-micro-batch positional args.
219 kwarg_mbs (list): Per-micro-batch keyword args.
220 losses (list): Mutable list collecting per-step losses.
221 """
222 if self._executor is not None:
223 self._executor.reset()
224 super().run_microbatches(arg_mbs, kwarg_mbs, losses)
226 def construct_exec_order(self) -> None:
227 """Build the body Interleaved 1F1B order, then layer the preprocess
228 transpose phase and the centralized preprocess backward on top.
230 The parameter broadcast, the recompute-input transport, and the
231 recompute backward are emitted only for a **trainable** preprocess; a
232 frozen or param-free one (``T == 0``, a frozen visual tower) only
233 transposes the forward and ships its output.
234 """
235 super().construct_exec_order()
236 body_order = self.exec_order
237 num_transpose = self.num_transpose_micro_batches
238 has_trainable = self._has_trainable_preprocess
239 body_order[0] = self._insert_rank0_preprocess_steps(
240 body_order[0], num_transpose, has_trainable,
241 backward_all=self._explicit_nontransposed_backward)
242 self.exec_order = {
243 rank: self._build_transpose_prefix(rank, num_transpose, has_trainable) + body_order[rank]
244 for rank in range(self.real_stage_num)
245 }
247 @staticmethod
248 def _build_transpose_prefix(rank, num_transpose, has_trainable):
249 """Build the transpose-phase prefix prepended to ``rank``'s body order.
251 A rank that owns a transposed micro-batch (``rank < num_transpose``)
252 computes its preprocess forward and ranks ``> 0`` ship the output to
253 stage 0. For a **trainable** preprocess every rank also broadcasts its
254 (trainable) parameters and ranks ``> 0`` additionally ship the input for
255 the centralized recompute backward; for a frozen / param-free preprocess
256 only the transpose forward and its output send/recv remain.
257 """
258 prefix = []
259 if has_trainable:
260 prefix.append(MetaStep(None, MpipeStepType.MPIPE_PARAM_BROADCAST, 0))
261 if rank < num_transpose:
262 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_TRANSPOSE_FWD, 0))
263 if rank != 0:
264 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_FWD_SEND, 0))
265 if has_trainable:
266 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_GRAPH_SEND, 0))
267 if rank == 0:
268 for micro_index in range(1, num_transpose):
269 prefix.append(MetaStep(micro_index, MpipeStepType.MPIPE_FWD_RECV, 0))
270 if has_trainable:
271 for micro_index in range(1, num_transpose):
272 prefix.append(MetaStep(micro_index, MpipeStepType.MPIPE_GRAPH_RECV, 0))
273 return prefix
275 @staticmethod
276 def _insert_rank0_preprocess_steps(order, num_transpose, has_trainable, backward_all=False):
277 """Patch stage 0's (rank 0) body order with preprocess fwd/bwd steps.
279 For a **trainable** preprocess: before each ``FWD(stage 0, micro >=
280 num_transpose)`` an inline ``MPIPE_TRANSPOSE_FWD`` runs the preprocess
281 forward for that non-transposed micro-batch, and an
282 ``MPIPE_TRANSPOSE_BWD`` is inserted after each ``BWD(stage 0, micro)``
283 needing a centralized recompute backward: the transposed micro-batches
284 (``micro < num_transpose``) always, and — when ``backward_all`` is set
285 (MindSpore, whose body backward does not flow into the preprocess) — the
286 non-transposed ones too. On torch (``backward_all`` False) non-transposed
287 micro-batches backprop into the preprocess via the connected graph.
289 For a frozen / param-free preprocess (no trainable params) there is no
290 recompute backward, so the body order is returned unchanged (transposed
291 outputs are placed by ``MPIPE_FWD_RECV``; non-transposed micro-batches
292 are handled by stage 0 directly). VL's frozen-visual injection is wired
293 per-model rather than through this text-style body-input path.
294 """
295 if not has_trainable:
296 return order
297 patched = []
298 for step in order:
299 is_stage0_fwd = (
300 step is not None
301 and step.type == MetaStepType.FWD
302 and step.stage_index == 0
303 )
304 if is_stage0_fwd and step.micro_index >= num_transpose:
305 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_FWD, 0))
306 patched.append(step)
307 is_stage0_bwd = (
308 step is not None
309 and step.type == MetaStepType.BWD
310 and step.stage_index == 0
311 )
312 if is_stage0_bwd and (step.micro_index < num_transpose or backward_all):
313 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_BWD, 0))
314 return patched