Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / pipeline_parallel / mpipe / executor_base.py: 98%
111 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"""Platform-agnostic base for MPipe Transpose execution.
17The parameter broadcast, P2P transport, and step orchestration are identical
18across backends (they go through the ``platform`` abstraction), so they live
19here. Each backend subclass implements only the autograd-specific hooks:
20running the preprocess forward (detached vs graph-connected), marking a tensor
21as a grad-requiring leaf, and the recompute backward.
22"""
23from abc import ABC, abstractmethod
24from typing import TYPE_CHECKING
26from hyper_parallel.platform import get_platform
28if TYPE_CHECKING:
29 from hyper_parallel.core.pipeline_parallel.scheduler import MetaStep, PipelineContext
30 from hyper_parallel.core.pipeline_parallel.mpipe.schedule import ScheduleMPipeTranspose
32platform = get_platform()
35class MPipeTransposeExecutorBase(ABC):
36 """Backend-agnostic runtime for the ``MPIPE_*`` steps of MPipe Transpose.
38 Args:
39 schedule (ScheduleMPipeTranspose): The schedule instance, providing the
40 per-rank preprocess module, the body stages (for the PP group and
41 global-rank mapping), and the transposed micro-batch count.
43 Attributes:
44 nontransposed_connected (bool): Whether non-transposed micro-batches
45 can rely on the autograd graph flowing the body backward into the
46 preprocess automatically (torch). When ``False`` (mindspore), the
47 schedule also emits ``MPIPE_TRANSPOSE_BWD`` for them and the
48 preprocess forward is always detached.
49 """
51 nontransposed_connected = False
53 def __init__(self, schedule: "ScheduleMPipeTranspose") -> None:
54 """Bind the executor to its schedule and cache the per-rank transpose state."""
55 self._schedule = schedule
56 self._preprocess = schedule.preprocess_module
57 first_stage = schedule.stages[0]
58 self._device = first_stage.device
59 self._pp_group = first_stage.pp_group
60 self._this_rank = first_stage.stage_index % schedule.real_stage_num
61 self._num_transpose = schedule.num_transpose_micro_batches
62 # Trainable preprocess → mark shipped tensors grad-requiring (the body
63 # backward / recompute use them) and recompute the backward. Frozen or
64 # param-free (T=0 identity, a frozen visual tower) → ship the output
65 # as-is (e.g. integer input_ids must not be marked grad-requiring).
66 self._has_trainable = schedule.has_trainable_preprocess
67 # micro_index -> retained raw input args (recompute backward / shipping).
68 self._inputs = {}
69 # micro_index -> detached preprocess output buffer (stage 0 body input).
70 self._outputs = {}
72 def reset(self) -> None:
73 """Clear the per-step compute caches at the start of each schedule run."""
74 self._inputs.clear()
75 self._outputs.clear()
77 @staticmethod
78 def _send_meta(tensor, dst) -> None:
79 """Send a tensor's ``(shape, dtype)`` to ``dst``.
81 Exchanged every step (not cached): MPipe Transpose is used mostly with
82 dynamic shapes (variable sequence / image-token lengths), where a cached
83 shape would be wrong; the meta is tiny so the per-step cost is negligible
84 (``T = 0`` step time ~= 1F1B).
85 """
86 platform.send_object_list([tuple(tensor.shape), tensor.dtype], dst)
88 @staticmethod
89 def _recv_meta(src):
90 """Receive a tensor's ``(shape, dtype)`` from ``src`` (exchanged every step)."""
91 meta: list = [None, None]
92 platform.recv_object_list(meta, src)
93 return meta[0], meta[1]
95 def _global_rank(self, physical_rank: int) -> int:
96 """Global rank of physical pipeline rank ``physical_rank``."""
97 return platform.get_global_rank(self._pp_group, physical_rank)
99 @staticmethod
100 def _as_tuple(args):
101 """Normalize a micro-batch arg slot to a positional-args tuple."""
102 if isinstance(args, (list, tuple)):
103 return tuple(args)
104 return (args,)
106 @staticmethod
107 def _kwargs_for(ctx, micro):
108 """Keyword args for ``micro``'s preprocess forward (e.g. position_ids)."""
109 kwarg_mbs = getattr(ctx, "kwarg_mbs", None)
110 if not kwarg_mbs:
111 return {}
112 return kwarg_mbs[micro] or {}
114 def broadcast_params(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument
115 """Broadcast the preprocess parameters from stage 0 to all ranks.
117 Args:
118 step (MetaStep): The ``MPIPE_PARAM_BROADCAST`` schedule step (unused).
119 ctx (PipelineContext): The pipeline run context (unused).
120 """
121 src = self._global_rank(0)
122 for tensor in self._broadcast_tensors():
123 platform.broadcast(tensor, src, self._pp_group)
125 def transpose_forward(self, step: "MetaStep", ctx: "PipelineContext") -> None:
126 """Run the preprocess forward for ``step.micro_index``.
128 Non-transposed micro-batches stay graph-connected only when the backend
129 supports automatic backward into the preprocess; otherwise (and for all
130 transposed micro-batches) the output is detached and the input retained
131 for the recompute backward.
133 Args:
134 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
135 ctx (PipelineContext): The pipeline run context (``arg_mbs`` / ``kwarg_mbs``).
136 """
137 micro = step.micro_index
138 args = self._as_tuple(ctx.arg_mbs[micro])
139 kwargs = self._kwargs_for(ctx, micro)
140 if micro >= self._num_transpose and self.nontransposed_connected:
141 ctx.arg_mbs[micro] = [self._connected_forward(args, kwargs)]
142 return
143 self._inputs[micro] = args
144 out = self._detached_forward(args, kwargs)
145 if self._has_trainable:
146 self._mark_requires_grad(out)
147 self._outputs[micro] = out
148 if self._this_rank == 0:
149 ctx.arg_mbs[micro] = [out]
151 def fwd_send(self, step: "MetaStep", ctx: "PipelineContext") -> None:
152 """Send the preprocess output of ``step.micro_index`` to stage 0.
154 Args:
155 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
156 ctx (PipelineContext): The pipeline run context; the send handle is appended to it.
157 """
158 micro = step.micro_index
159 out = self._contiguous(self._outputs[micro])
160 dst = self._global_rank(0)
161 self._send_meta(out, dst)
162 # Deferred sends are drained at the end of ``run_microbatches`` via the
163 # schedule's ``_send_handles`` (each entry is a handle group).
164 ctx.schedule._send_handles.append([platform.isend(out, dst)]) # pylint: disable=protected-access
166 def fwd_recv(self, step: "MetaStep", ctx: "PipelineContext") -> None:
167 """Receive a transposed micro-batch's preprocess output into stage 0's input slot.
169 Args:
170 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
171 ctx (PipelineContext): The pipeline run context; the received buffer is placed in its ``arg_mbs``.
172 """
173 micro = step.micro_index
174 src = self._global_rank(micro)
175 shape, dtype = self._recv_meta(src)
176 buffer = platform.empty(shape, dtype=dtype, device=self._device)
177 platform.irecv(buffer, src).wait()
178 if self._has_trainable:
179 self._mark_requires_grad(buffer)
180 ctx.arg_mbs[micro] = [buffer]
182 def graph_send(self, step: "MetaStep", ctx: "PipelineContext") -> None:
183 """Send the preprocess input of ``step.micro_index`` to stage 0 (for recompute).
185 Args:
186 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
187 ctx (PipelineContext): The pipeline run context; send handles are appended to it.
188 """
189 micro = step.micro_index
190 dst = self._global_rank(0)
191 for tensor in self._inputs[micro]:
192 contiguous = self._contiguous(tensor)
193 self._send_meta(contiguous, dst)
194 ctx.schedule._send_handles.append([platform.isend(contiguous, dst)]) # pylint: disable=protected-access
196 def graph_recv(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument
197 """Receive a transposed micro-batch's preprocess input for the recompute backward.
199 Args:
200 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
201 ctx (PipelineContext): The pipeline run context (unused).
202 """
203 micro = step.micro_index
204 src = self._global_rank(micro)
205 # Stage 0 always owns transposed micro 0, so its retained input reveals
206 # how many input tensors each transposed micro-batch ships.
207 arity = len(self._inputs[0])
208 tensors = []
209 for _ in range(arity):
210 shape, dtype = self._recv_meta(src)
211 buffer = platform.empty(shape, dtype=dtype, device=self._device)
212 platform.irecv(buffer, src).wait()
213 tensors.append(buffer)
214 self._inputs[micro] = tuple(tensors)
216 def transpose_backward(self, step: "MetaStep", ctx: "PipelineContext") -> None:
217 """Recompute the preprocess forward on stage 0 and backprop the body's input grad.
219 Args:
220 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.
221 ctx (PipelineContext): The pipeline run context; the body's input grad is read from its ``arg_mbs``.
222 """
223 micro = step.micro_index
224 grad = ctx.arg_mbs[micro][0].grad
225 if grad is None:
226 return
227 self._recompute_backward(self._inputs[micro], self._kwargs_for(ctx, micro), grad)
229 @staticmethod
230 def _contiguous(tensor):
231 """Return a contiguous tensor suitable for P2P (overridden where needed)."""
232 return tensor
234 @abstractmethod
235 def _broadcast_tensors(self):
236 """Yield the preprocess tensors (params/buffers) to broadcast from stage 0."""
238 @abstractmethod
239 def _detached_forward(self, args, kwargs):
240 """Run the preprocess forward and return a detached output value.
242 The base marks it grad-requiring via :meth:`_mark_requires_grad` only for
243 a trainable preprocess; for a frozen / param-free one (frozen visual
244 tower, T=0 identity) the value is shipped as-is.
245 """
247 @abstractmethod
248 def _connected_forward(self, args, kwargs):
249 """Run the preprocess forward graph-connected to the body (backends that support it)."""
251 @abstractmethod
252 def _mark_requires_grad(self, tensor) -> None:
253 """Mark ``tensor`` as a grad-requiring leaf so the body backward deposits a grad on it."""
255 @abstractmethod
256 def _recompute_backward(self, inputs, kwargs, grad) -> None:
257 """Recompute ``preprocess(*inputs, **kwargs)`` and backprop ``grad``, accumulating preprocess grads."""