Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / mpipe / schedule.py: 99%
169 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""
21import logging
22from typing import Callable, Optional, TYPE_CHECKING
24from torch.nn import Module
26from hyper_parallel.core.pipeline_parallel.mpipe.executor import MPipeTransposeExecutor
28from hyper_parallel.core.pipeline_parallel.scheduler import ScheduleInterleaved1F1B
29from hyper_parallel.core.pipeline_parallel.mpipe.sampler import mpipe_owned_micros
30from hyper_parallel.core.pipeline_parallel.mpipe.step_types import MpipeStepType
31from hyper_parallel.core.pipeline_parallel.utils import MetaStep, MetaStepType
33if TYPE_CHECKING:
34 from hyper_parallel.core.pipeline_parallel.utils import BatchDimSpec
36logger = logging.getLogger(__name__)
38# (data, features, graph input) step types of the transpose-prefix transfers.
39_PREFIX_SEND_TYPES = (MetaStepType.DATA_SEND, MpipeStepType.MPIPE_FWD_SEND, MpipeStepType.MPIPE_GRAPH_SEND)
40_PREFIX_RECV_TYPES = (MetaStepType.DATA_RECV, MpipeStepType.MPIPE_FWD_RECV, MpipeStepType.MPIPE_GRAPH_RECV)
43def _prefix_owned_micros(rank, num_transpose_micro_batches, micro_batch_num, overflow_mode):
44 """Micros ``rank`` transposes in its prefix (under ``"min"`` only ``rank``
45 itself; rank 0's overflow micros are loaded inline in its body order)."""
46 nt = num_transpose_micro_batches
47 if rank >= nt:
48 return []
49 if overflow_mode == "min":
50 return [rank]
51 return [m for m in range(micro_batch_num) if m % nt == rank]
54def _prefix_transfer_steps(micros, step_types, is_dataload_only, ship_graph):
55 """Type-major transfer block: all data, then all features, then all graph inputs.
57 The sending ranks and rank 0 both build their block here, so the send and
58 receive orders match; otherwise the meta recv drifts and decodes garbage
59 shapes. Dataload-only ships the data alone; ``ship_graph`` adds the input
60 ship-back for the centralized stage-0 backward.
61 """
62 data_type, fwd_type, graph_type = step_types
63 steps = [MetaStep(m, data_type, -1) for m in micros]
64 if not is_dataload_only:
65 steps += [MetaStep(m, fwd_type, -1) for m in micros]
66 if ship_graph:
67 steps += [MetaStep(m, graph_type, -1) for m in micros]
68 return steps
71class ScheduleMPipeTranspose(ScheduleInterleaved1F1B):
72 """The MPipe Transpose schedule.
74 A variant of Interleaved 1F1B that shrinks the warmup pipeline bubble by
75 **transposing** the forward of a model's first ``T`` layers (the
76 *preprocess* block, which logically belongs to stage 0's first chunk).
78 Instead of stage 0 computing the preprocess forward serially for every
79 micro-batch, the preprocess parameters are broadcast to all ``PP`` ranks
80 and, for the first ``NT = min(PP, micro_batch_num)`` micro-batches, each
81 rank ``i`` computes the preprocess forward of micro-batch ``i`` in parallel
82 during what would otherwise be its warmup idle time. Each rank ``i > 0``
83 then ships to stage 0:
85 * the preprocess **output** (``MPIPE_FWD_SEND`` / ``MPIPE_FWD_RECV``) so
86 stage 0 can run its body forward, and
87 * the preprocess **input** (``MPIPE_GRAPH_SEND`` / ``MPIPE_GRAPH_RECV``)
88 so the preprocess backward can run centrally on stage 0
89 (gradients accumulate on stage 0 only).
91 The remaining ``micro_batch_num - NT`` micro-batches run the preprocess
92 forward inline on stage 0 (graph-connected to the body, so their backward
93 is automatic), exactly as ordinary Interleaved 1F1B.
95 The body model (stage 0 = the layers after the preprocess block, all other
96 stages unchanged) is scheduled by the inherited Interleaved 1F1B logic; the
97 preprocess steps are layered around it: a transpose-phase prefix per rank,
98 plus inline preprocess forward / stage-0 backward steps.
100 Args:
101 stages (list[PipelineStage], PipelineStage): The body pipeline stages,
102 given **after** stage 0's first chunk has been partitioned into the
103 preprocess block and the body. Stage 0 here must wrap only the
104 layers that remain *after* the preprocess block; the transposed
105 preprocess layers are passed separately as ``preprocess_module`` and
106 must not appear in ``stages``.
107 micro_batch_num (int): The number of micro-batches.
108 preprocess_module (Optional[Module]): The preprocess block (first ``T``
109 layers of stage 0). Following Option A, it must exist on **every** rank: on rank 0
110 it holds the trained parameters; on other ranks it is a structural
111 copy whose parameters are overwritten each step by the broadcast.
112 num_transpose_layers (int): ``T`` — the number of preprocess layers,
113 must be smaller than the layer count of stage 0's first chunk.
114 ``0`` is allowed and means *only the data loading is transposed*:
115 each rank loads its micro-batch and ships the raw input to stage 0,
116 with no parameter broadcast, no preprocess compute, and no
117 stage-0 backward.
118 args_batch_dim (list, optional): See ``PipelineScheduleRuntime``.
119 kwargs_batch_dim (dict, optional): See ``PipelineScheduleRuntime``.
120 output_concat_dim (int, optional): See ``PipelineScheduleRuntime``.
121 overlap_p2p (bool, optional): See ``ScheduleInterleaved1F1B``.
122 Default ``False``.
123 swap (bool, optional): Reserved for API compatibility. MPipe activation
124 swap is not supported and ``True`` raises ``ValueError``.
126 Note:
127 This class builds the schedule ordering and registers the ``MPIPE_*``
128 execution handlers implemented by :class:`MPipeTransposeExecutor`.
129 """
131 # Each rank runs one transposed preprocess forward, so every rank needs its
132 # own micro-batch input — drivers (the trainer) check this marker to feed all
133 # ranks rather than only the first stage.
134 requires_all_rank_input = True
136 def __init__(self,
137 stages: list,
138 micro_batch_num: int,
139 preprocess_module: "Optional[Module]",
140 num_transpose_layers: int,
141 num_visual_layers: Optional[int] = None,
142 args_batch_dim: "Optional[BatchDimSpec]" = None,
143 kwargs_batch_dim: "Optional[BatchDimSpec]" = None,
144 output_concat_dim: Optional[int] = None,
145 overlap_p2p: bool = False,
146 swap: bool = False,
147 output_consumer: Optional[Callable] = None,
148 owner_backward: bool = False,
149 less_memory: bool = False,
150 overflow_mode: str = "full") -> None:
151 """Build an interleaved-1F1B schedule that transposes the preprocess block.
153 Args:
154 stages (list): The local pipeline stages (as for :class:`ScheduleInterleaved1F1B`).
155 micro_batch_num (int): Number of micro-batches per optimizer step.
156 preprocess_module (Optional[Module]): The block transposed to every
157 rank — the first ``num_transpose_layers`` layers, a visual tower,
158 or a param-free identity for the dataload-only (``T = 0``) mode.
159 num_transpose_layers (int): ``T`` — the (informational) transposed-layer count.
160 num_visual_layers (int): total number of visual layers.
161 args_batch_dim (Optional[BatchDimSpec]): Positional-arg batch-dim spec (forwarded to the base).
162 kwargs_batch_dim (Optional[BatchDimSpec]): Keyword-arg batch-dim spec (forwarded to the base).
163 output_concat_dim (Optional[int]): Output concatenation dim (forwarded to the base).
164 overlap_p2p (bool): Whether to overlap P2P (forwarded to the base).
165 swap (bool): Reserved for API compatibility. Must be ``False``.
166 output_consumer: Optional ``(ctx, micro, out_tuple) -> None`` callback
167 routing a transposed micro-batch's preprocess output (``None`` =
168 default, the output becomes stage 0's body input).
169 less_memory (bool): Shallow-warmup interleaved body (forwarded to
170 the base) — see :class:`ScheduleInterleaved1F1B`. The transpose
171 prefix is unaffected.
172 owner_backward (bool): Opt-in owner-does-backward for a TRAINABLE
173 preprocess on torch. When effective, the owner rank retains the
174 forward graph, stage 0 ships the feature gradient back, and the
175 owner runs the tower backward in its cooldown (instead of the
176 centralized stage-0 backward); tower grads are SUM all-reduced
177 across the pp replicas. Ignored (stage-0 backward kept) for a frozen / param-free
178 preprocess.
179 overflow_mode (str): How to distribute the ``M > NT`` overflow micros.
180 ``"full"`` (default) — round-robin: each owner takes ``M/NT``
181 micros, ViT phase balanced. ``"min"`` — rank 0 absorbs the
182 overflow (saves inter-rank P2P at the cost of a longer rank-0
183 ViT phase).
184 """
185 if swap:
186 raise ValueError(
187 "MPipe Transpose with pipeline activation swap is not yet supported."
188 )
189 if not isinstance(num_transpose_layers, int) or num_transpose_layers < 0:
190 raise ValueError(
191 f"Argument 'num_transpose_layers' must be a non-negative int, "
192 f"but got {num_transpose_layers!r}."
193 )
194 if overflow_mode not in ("full", "min"):
195 raise ValueError(
196 f"overflow_mode must be 'full' or 'min', got {overflow_mode!r}"
197 )
198 self._overflow_mode = overflow_mode
199 # ``preprocess_module`` is the resolved block to transpose (first T text
200 # layers, the visual tower, or a param-free identity for dataload-only).
201 self._preprocess_module = preprocess_module
202 self._num_transpose_layers = num_transpose_layers
203 self._num_visual_layers = num_visual_layers
204 # ``None`` = the preprocess output is stage 0's body input (arg_mbs);
205 # VL passes a consumer that routes it into an injection kwarg instead.
206 self._output_consumer = output_consumer
207 # Trainable preprocess -> broadcast + stage-0 backward; frozen or
208 # param-free -> ship the output only (no broadcast, no recompute).
209 self._has_trainable_preprocess = self._module_has_trainable_params(preprocess_module)
210 self._owner_backward = bool(owner_backward) and self._has_trainable_preprocess
211 super().__init__(stages,
212 micro_batch_num,
213 args_batch_dim=args_batch_dim,
214 kwargs_batch_dim=kwargs_batch_dim,
215 output_concat_dim=output_concat_dim,
216 overlap_p2p=overlap_p2p,
217 overlap_b_f=False,
218 swap=swap,
219 less_memory=less_memory)
220 self._executor = None
221 self._setup_mpipe_execution()
223 @staticmethod
224 def _module_has_trainable_params(module) -> bool:
225 """Whether ``module`` has any trainable (grad-requiring) parameter."""
226 if module is None:
227 return False
228 return any(p.requires_grad for p in module.parameters())
230 @property
231 def preprocess_module(self) -> "Optional[Module]":
232 """The preprocess block transposed to every rank (present on every rank)."""
233 return self._preprocess_module
235 @property
236 def output_consumer(self) -> Optional[Callable]:
237 """Optional ``(ctx, micro, out_tuple)`` placement callback (None = body input)."""
238 return self._output_consumer
240 @property
241 def has_trainable_preprocess(self) -> bool:
242 """Whether the preprocess has trainable params (drives broadcast / stage-0 backward)."""
243 return self._has_trainable_preprocess
245 @property
246 def owner_backward(self) -> bool:
247 """Whether owner-does-backward is effective (trainable preprocess + torch)."""
248 return self._owner_backward
250 @property
251 def num_transpose_layers(self) -> int:
252 """The informational transposed-layer count ``T`` (``0`` = dataload only)."""
253 return self._num_transpose_layers
255 @property
256 def num_transpose_micro_batches(self) -> int:
257 """``NT = min(PP, micro_batch_num)`` — count of owner ranks per cycle."""
258 return min(self.real_stage_num, self.micro_batch_num)
260 def owned_micros(self, pp_rank: int) -> frozenset:
261 """Owner set for ``pp_rank`` under :attr:`overflow_mode` (see ``mpipe_owned_micros``)."""
262 return mpipe_owned_micros(self.real_stage_num, self.micro_batch_num,
263 pp_rank, mode=self._overflow_mode)
265 def owner_of(self, micro: int) -> int:
266 """Rank that owns micro ``micro`` under :attr:`overflow_mode`."""
267 nt = self.num_transpose_micro_batches
268 if self._overflow_mode == "min":
269 return 0 if micro >= nt else micro
270 return micro % nt
272 def _setup_mpipe_execution(self) -> None:
273 """Build the Torch executor and register the MPIPE_* handlers."""
274 self._executor = MPipeTransposeExecutor(self)
275 self._apply_backward_retain_flag()
276 handlers = {
277 MpipeStepType.MPIPE_PARAM_BROADCAST: self._executor.broadcast_params,
278 MpipeStepType.MPIPE_TRANSPOSE_FWD: self._executor.transpose_forward,
279 MpipeStepType.MPIPE_FWD_SEND: self._executor.fwd_send,
280 MpipeStepType.MPIPE_FWD_RECV: self._executor.fwd_recv,
281 MpipeStepType.MPIPE_GRAPH_SEND: self._executor.fwd_bwd_graph_input_send,
282 MpipeStepType.MPIPE_GRAPH_RECV: self._executor.fwd_bwd_graph_input_recv,
283 MpipeStepType.MPIPE_TRANSPOSE_BWD: self._executor.transpose_backward,
284 MpipeStepType.MPIPE_GRAD_SEND: self._executor.grad_send,
285 MpipeStepType.MPIPE_GRAD_RECV_WITH_BACKWARD: self._executor.grad_recv_with_backward,
286 MpipeStepType.MPIPE_GRAD_REDUCE: self._executor.reduce_tower_grads,
287 }
288 for step_type, handler in handlers.items():
289 self.register_custom_function(step_type, handler)
291 def _apply_backward_retain_flag(self) -> None:
292 """Flag the stages to retain their body-backward graph for a TRAINABLE tower.
294 A trainable transposed tower + FSDP shares one all-gather node across the
295 micro graphs that back-propagate into it (the connected micro-0 body graph
296 plus the deferred owner/stage-0 tower backwards). The first backward to
297 traverse that shared node frees it, so a later one raises "backward through
298 the graph a second time" (surfaces at ``vpp >= 2``). Retaining the body
299 graph keeps the shared node alive. A frozen tower (no shared trainable
300 node) and plain PP (no flag) are unaffected.
301 """
302 if self._has_trainable_preprocess:
303 for stage in self.stages:
304 stage.retain_backward_graph = True
306 def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:
307 """Reset the executor's per-step caches, then run the schedule.
309 A step raising mid-run leaves its micro-batches in those caches, so the
310 error path drops them (:meth:`MPipeTransposeExecutor.abort`) -- the
311 counterpart of the scheduler's in-flight P2P drain, which covers only
312 the comm handles. The normal path leaves the caches empty, which the
313 next run's reset still verifies.
315 Args:
316 arg_mbs (list): Per-micro-batch positional args.
317 kwarg_mbs (list): Per-micro-batch keyword args.
318 losses (list): Mutable list collecting per-step losses.
319 """
320 if self._executor is not None:
321 self._executor.reset()
322 try:
323 super().run_microbatches(arg_mbs, kwarg_mbs, losses)
324 except BaseException:
325 if self._executor is not None:
326 self._executor.abort()
327 raise
329 def construct_exec_order(self) -> None:
330 """Build the body Interleaved 1F1B order, then layer the preprocess
331 transpose phase and the centralized preprocess backward on top.
333 The parameter broadcast, the stage-0-backward input transport, and the
334 stage-0 backward are emitted only for a **trainable** preprocess; a
335 frozen (visual tower) or param-free (``T == 0``) one only transposes
336 the forward and ships its output.
338 When owner-does-backward is effective (trainable + torch + opt-in) the
339 stage-0-backward input transport and stage-0 backward are replaced by a
340 feature-gradient ship-back (``MPIPE_GRAD_SEND`` on stage 0,
341 ``MPIPE_GRAD_RECV_WITH_BACKWARD`` cooldown suffix on the owners) and a final
342 ``MPIPE_GRAD_REDUCE`` collective on every rank.
343 """
344 super().construct_exec_order()
345 body_order = self.exec_order
346 nt = self.num_transpose_micro_batches
347 micro_batch_num = self.micro_batch_num
348 has_trainable = self._has_trainable_preprocess
349 owner_backward = self._owner_backward
350 mode = self._overflow_mode
351 # Body DATA_LOAD steps are rescheduled into the mpipe prefix (and, under
352 # "min", inline onto rank 0's body for the overflow micros).
353 for i, order in body_order.items():
354 body_order[i] = [step for step in order
355 if step is not None and step.type != MetaStepType.DATA_LOAD]
356 body_order[0] = self._insert_rank0_preprocess_steps(
357 body_order[0], nt, has_trainable, mode, owner_backward=owner_backward,
358 num_transpose_layers=self._num_transpose_layers)
359 exec_order = {}
360 for rank in range(self.real_stage_num):
361 order = (
362 self._build_transpose_prefix(
363 rank, nt, micro_batch_num, has_trainable, mode,
364 owner_backward=owner_backward,
365 num_transpose_layers=self._num_transpose_layers)
366 + body_order[rank]
367 )
368 if owner_backward:
369 # Cooldown suffix: stage 0 ships dL/dfeatures, owners run the
370 # tower backward, then all ranks SUM-all-reduce the tower grads.
371 order = (
372 order
373 + self._build_owner_backward_suffix(rank, nt, micro_batch_num, mode)
374 + [MetaStep(None, MpipeStepType.MPIPE_GRAD_REDUCE, 0)]
375 )
376 exec_order[rank] = order
377 self.exec_order = exec_order
378 # frozen + all visual layers transposed, so stage 0 needs no pixels
379 if (not self._has_trainable_preprocess and
380 self._num_visual_layers is not None and
381 self._num_transpose_layers >= self._num_visual_layers):
382 self._data_keys = ("input_ids",) + tuple(
383 k for k in self.kwargs_batch_dim if k != "pixel_values"
384 )
385 else:
386 self._data_keys = ("input_ids",) + tuple(self.kwargs_batch_dim)
388 # Length-M so overflow micros and VPP chunks index in bounds; the
389 # prefix key -1 ships owner to rank 0, body stages chain via stages[0].
390 self._data_dst = {-1: [0 for _ in range(self.micro_batch_num)]}
391 self._data_src = {-1: [self.owner_of(m) for m in range(self.micro_batch_num)]}
392 for i in range(self._stage_num):
393 self._data_dst[i] = [self.stages[0].dst_stage for _ in range(self.micro_batch_num)]
394 self._data_src[i] = [self.stages[0].src_stage for _ in range(self.micro_batch_num)]
396 @staticmethod
397 def _build_transpose_prefix(rank, num_transpose_micro_batches, micro_batch_num,
398 has_trainable, overflow_mode,
399 owner_backward=False, num_transpose_layers=None):
400 """Build the transpose-phase prefix prepended to ``rank``'s body order.
402 Under ``overflow_mode="full"`` (round-robin) rank ``rank`` (``< NT``)
403 owns ``{m : m % NT == rank}`` — emits DATA_LOAD → MPIPE_TRANSPOSE_FWD
404 per owned micro, and (when not rank 0) ships each to rank 0. Rank 0
405 receives every non-owned micro. For a **trainable** preprocess every
406 rank also broadcasts its (trainable) parameters and ranks ``> 0``
407 additionally ship the input for the centralized stage-0 backward; for
408 a frozen (visual tower) preprocess only the transpose forward and its
409 output send/recv remain.
411 Under ``"min"`` each owner (including rank 0) emits exactly one
412 transposed micro (``{rank}``). Rank 0 receives only ``1..NT-1``; the
413 overflow micros ``NT..M-1`` are loaded inline on rank 0 by
414 :meth:`_insert_rank0_preprocess_steps`.
416 Dataload-only mode (``num_transpose_layers == 0``): every emission
417 of MPIPE_PARAM_BROADCAST, MPIPE_TRANSPOSE_FWD, MPIPE_FWD_SEND/RECV,
418 MPIPE_GRAPH_SEND/RECV is suppressed. Only DATA_LOAD + DATA_SEND/RECV
419 remain; stage 0 runs the visual tower itself during body forward.
420 Ownership (round-robin vs single-owner) still follows
421 ``overflow_mode`` so the dataload I/O sharding stays balanced.
423 Under owner-does-backward the input ship-back (``MPIPE_GRAPH_SEND`` /
424 ``MPIPE_GRAPH_RECV``) is dropped -- the owner keeps its forward graph and
425 receives only the feature gradient later (``MPIPE_GRAD_RECV_WITH_BACKWARD``).
426 """
427 nt = num_transpose_micro_batches
428 owned = _prefix_owned_micros(rank, nt, micro_batch_num, overflow_mode)
429 is_dataload_only = num_transpose_layers == 0
430 prefix = []
431 if has_trainable and not is_dataload_only:
432 prefix.append(MetaStep(rank, MpipeStepType.MPIPE_PARAM_BROADCAST, -1))
433 # DATA_LOAD + MPIPE_TRANSPOSE_FWD are local (no comm) — emit per-micro.
434 for m in owned:
435 prefix.append(MetaStep(m, MetaStepType.DATA_LOAD, -1))
436 if not is_dataload_only:
437 prefix.append(MetaStep(m, MpipeStepType.MPIPE_TRANSPOSE_FWD, -1))
438 ship_graph = has_trainable and not owner_backward
439 if rank != 0:
440 return prefix + _prefix_transfer_steps(owned, _PREFIX_SEND_TYPES, is_dataload_only, ship_graph)
441 # Non-owned = micros rank 0 didn't transpose locally ("min":
442 # overflow micros load inline, so they are not received).
443 if overflow_mode == "min":
444 non_owned = list(range(1, nt))
445 else:
446 non_owned = [m for m in range(micro_batch_num) if m not in owned]
447 return prefix + _prefix_transfer_steps(non_owned, _PREFIX_RECV_TYPES, is_dataload_only, ship_graph)
449 @staticmethod
450 def _build_owner_backward_suffix(rank, num_transpose_micro_batches, micro_batch_num, overflow_mode):
451 """Cooldown suffix (owner-backward), appended AFTER all body steps so the
452 feature-grad transport never interleaves with the body's P2P.
454 Stage 0 ships ``dL/dfeatures`` for every micro it does not own itself
455 (``MPIPE_GRAD_SEND``, one per micro whose owner != 0). Each owner rank
456 receives a gradient for **every** micro it owns (``MPIPE_GRAD_RECV_WITH_BACKWARD``,
457 one step per owned micro) and backprops its retained tower graph for each.
458 Ranks ``>= NT`` own no transposed micro-batch.
460 Under ``"min"`` an owner rank owns exactly one micro (``{rank}``), so this
461 degenerates to the original one-gradient-per-rank exchange. Under ``"full"``
462 (round-robin) with ``M > NT`` an owner rank can own several micros
463 (``{m : m % NT == rank}``); each gets its own SEND/RECV pair. Both sides
464 iterate ``range(micro_batch_num)`` in ascending order, so for any fixed
465 (rank 0, owner) pair the send order matches the receive order -- required
466 since HCCL serializes P2P per rank-pair on a single channel.
467 """
468 nt = num_transpose_micro_batches
470 def owner_of(micro: int) -> int:
471 """Rank that owns ``micro`` under the active overflow mode."""
472 if overflow_mode == "min":
473 return 0 if micro >= nt else micro
474 return micro % nt
476 if rank == 0:
477 return [MetaStep(micro, MpipeStepType.MPIPE_GRAD_SEND, 0)
478 for micro in range(micro_batch_num) if owner_of(micro) != 0]
479 if rank < nt:
480 return [MetaStep(micro, MpipeStepType.MPIPE_GRAD_RECV_WITH_BACKWARD, 0)
481 for micro in range(micro_batch_num) if owner_of(micro) == rank]
482 return []
484 @staticmethod
485 def _insert_rank0_preprocess_steps(order, num_transpose_micro_batches,
486 has_trainable, overflow_mode, owner_backward=False,
487 num_transpose_layers=None):
488 """Patch stage 0's body order for the overflow / recompute paths.
490 Under ``"full"``: every micro is transposed (somewhere). No inline
491 preprocess on rank 0; under trainable, MPIPE_TRANSPOSE_BWD fires after
492 every stage-0 BWD (the preprocess output is detached either way).
494 Under ``"min"``: rank 0 absorbs the overflow micros ``NT..M-1``. Before
495 each stage-0 FWD on an overflow micro we insert DATA_LOAD and (under
496 trainable or frozen) MPIPE_TRANSPOSE_FWD — running the preprocess
497 locally with a detached graph. For frozen this is just "rank 0
498 runs the ViT here so the body forward has the encoded features /
499 visual-injection kwarg"; for trainable the body BWD also needs an
500 MPIPE_TRANSPOSE_BWD stage-0 backward for the preprocess gradient. (The
501 previous design used a graph-connected forward to avoid the
502 stage-0 backward on overflow micros; that optimization isn't restored
503 here — `"min"` saves comm for slightly more rank-0 compute on the
504 trainable path.)
506 Under owner-does-backward (``owner_backward``) the centralized stage-0
507 backward is replaced: the body backward deposits ``dL/dfeatures`` on the
508 placed feature tensors, and a stage-0 cooldown SUFFIX
509 (:meth:`_build_owner_backward_suffix`) ships it back to the owners -- NOT
510 inline here, since the blocking grad meta-exchange would deadlock against
511 the body's P2P. So no per-BWD step is inserted for any micro; every
512 owned/inline forward on stage 0 already runs graph-connected
513 (``transpose_forward``'s owner-backward branch), so the body backward
514 flows into the tower automatically.
516 Dataload-only (``num_transpose_layers == 0``) suppresses the inline
517 MPIPE_TRANSPOSE_FWD, matching the contract documented on
518 :meth:`_build_transpose_prefix`: with an identity preprocess there is
519 nothing to run ahead of the body, and stage 0 consumes the raw loaded
520 data directly.
521 """
522 if not has_trainable and overflow_mode == "full":
523 return order
524 nt = num_transpose_micro_batches
525 is_dataload_only = num_transpose_layers == 0
526 patched = []
527 for step in order:
528 is_stage0_fwd = (
529 step is not None
530 and step.type == MetaStepType.FWD
531 and step.stage_index == 0
532 )
533 if (overflow_mode == "min" and is_stage0_fwd
534 and step.micro_index >= nt):
535 patched.append(MetaStep(step.micro_index, MetaStepType.DATA_LOAD, -1))
536 if not is_dataload_only:
537 # Populates the visual-injection kwarg via output_consumer;
538 # without it rank 0's body forward crashes.
539 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_FWD, -1))
540 patched.append(step)
541 is_stage0_bwd = (
542 step is not None
543 and step.type == MetaStepType.BWD
544 and step.stage_index == 0
545 )
546 if is_stage0_bwd and has_trainable and not owner_backward:
547 patched.append(MetaStep(step.micro_index, MpipeStepType.MPIPE_TRANSPOSE_BWD, -1))
548 return patched