Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/core/pipeline

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

7# http://www.apache.org/licenses/LICENSE-2.0

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"""Torch execution and P2P transport for MPipe Transpose."""

16import logging

17import os

18from typing import TYPE_CHECKING

20import torch

21import torch.distributed as dist

23if TYPE_CHECKING:

24 from hyper_parallel.core.pipeline_parallel.scheduler import PipelineContext

25 from hyper_parallel.core.pipeline_parallel.mpipe.schedule import ScheduleMPipeTranspose

26 from hyper_parallel.core.pipeline_parallel.utils import MetaStep

28logger = logging.getLogger(__name__)

30# Stage 0's transposed-feature recvs are batched (all posted non-blocking, then

31# waited once) by default, so the NT-1 transfers pipeline and overlap stage 0's

32# own preprocess forward instead of being received-and-waited serially. Set

33# ``HYPER_PARALLEL_MPIPE_RECV_BATCH=0`` to fall back to the per-micro blocking path.

34_RECV_BATCH = os.environ.get("HYPER_PARALLEL_MPIPE_RECV_BATCH", "1").lower() not in ("", "0", "false", "no")

37class MPipeTransposeExecutor:

38 """Torch runtime for the ``MPIPE_*`` steps of MPipe Transpose.

40 Args:

41 schedule (ScheduleMPipeTranspose): The schedule instance, providing the

42 per-rank preprocess module, the body stages (for the PP group and

43 global-rank mapping), and the transposed micro-batch count.

44 """

46 def __init__(self, schedule: "ScheduleMPipeTranspose") -> None:

47 """Bind the executor to its schedule and cache the per-rank transpose state."""

48 self._schedule = schedule

49 self._preprocess = schedule.preprocess_module

50 first_stage = schedule.stages[0]

51 self._device = first_stage.device

52 self._mpipe_group = first_stage.pp_group

53 self._this_rank = first_stage.stage_index % schedule.real_stage_num

54 self._output_arity_for_comm = None

55 self._input_arity_for_comm = None

56 # Trainable: mark shipped tensors grad-requiring and run the stage-0

57 # backward. Frozen or param-free ships as-is (input_ids must not be marked).

58 self._has_trainable = schedule.has_trainable_preprocess

59 # micro_index -> retained raw input args (stage-0 backward / shipping).

60 self._inputs_for_explicit_forward = {}

61 # micro_index -> detached preprocess output buffer (stage 0 body input).

62 self._outputs_for_stage0 = {}

63 # Transposed micro-batches whose output stage 0 has already received (the

64 # batched recv posts all NT-1 at once, so later FWD_RECV steps no-op).

65 self._fwd_received = set()

66 # Owner-does-backward: the owner retains the tower graph, stage 0 ships

67 # dL/dfeatures back, and param-grads are SUM all-reduced across pp.

68 self._owner_backward = bool(getattr(schedule, "owner_backward", False))

69 # micro_index -> the feature tensors stage 0 placed; the body backward

70 # deposits dL/dfeatures on them for the stage-0 or owner backward.

71 self._keep_grad = {}

72 # Owner-backward: see reset() — lets GRAD_REDUCE reduce only this

73 # run's contribution.

74 self._grad_snapshot = None

75 # Owner-backward: micro_index -> the owner's graph-connected preprocess

76 # output tuple (same objects fwd_send ships detached).

77 self._outputs_for_bwd = {}

79 def reset(self) -> None:

80 """Clear the per-step compute caches at the start of each schedule run."""

81 self._output_arity_for_comm = None

82 self._input_arity_for_comm = None

83 # Every per-micro entry must have been consumed by the previous run.

84 for name, cache in (("inputs_for_explicit_forward", self._inputs_for_explicit_forward),

85 ("outputs_for_stage0", self._outputs_for_stage0),

86 ("outputs_for_bwd", self._outputs_for_bwd),

87 ("keep_grad", self._keep_grad)):

88 if cache:

89 raise RuntimeError(

90 f"MPipe executor cache '{name}' still holds micro-batches "

91 f"{sorted(cache)} from the previous run"

92 )

93 self._fwd_received.clear()

94 # Snapshot the already-reduced grads from prior runs so GRAD_REDUCE

95 # reduces only this run's added contribution.

96 if self._owner_backward:

97 self._grad_snapshot = self._snapshot_tower_grads()

99 def abort(self) -> None:

100 """Drop the per-micro caches after a run that raised.

101

102 Each entry is erased by the step that consumes it, so a step raising

103 mid-run strands whole micro-batches of device tensors -- including the

104 graph-connected tower output, which pins the autograd graph -- and the

105 next run's :meth:`reset` then refuses to start. The schedule calls this

106 on the error path only: on the normal path the caches must already be

107 empty, and ``reset`` still checks that.

108 """

109 self._inputs_for_explicit_forward.clear()

110 self._outputs_for_stage0.clear()

111 self._outputs_for_bwd.clear()

112 self._keep_grad.clear()

113 self._fwd_received.clear()

114 self._grad_snapshot = None

115

116 def _send_meta(self, tensor, dst) -> None:

117 """Send a tensor's ``(shape, dtype)`` to ``dst``.

118

119 Exchanged every step (not cached): MPipe Transpose is used mostly with

120 dynamic shapes (variable sequence / image-token lengths), where a cached

121 shape would be wrong; the meta is tiny so the per-step cost is negligible

122 (``T = 0`` step time ~= 1F1B).

123 """

124 dist.send_object_list([tuple(tensor.shape), tensor.dtype], dst, self._mpipe_group)

125

126 def _recv_meta(self, src):

127 """Receive a tensor's ``(shape, dtype)`` from ``src`` (exchanged every step)."""

128 meta: list = [None, None]

129 dist.recv_object_list(meta, src, self._mpipe_group)

130 return meta[0], meta[1]

131

132 def _output_arity(self):

133 """Number of output tensors to communicate (cached after first use)."""

134 if self._output_arity_for_comm is not None:

135 return self._output_arity_for_comm

136 # Fallback: derive the arity from the first retained micro.

137 if self._outputs_for_stage0.get(0):

138 self._output_arity_for_comm = len(self._outputs_for_stage0[0])

139 elif self._outputs_for_bwd.get(0):

140 self._output_arity_for_comm = len(self._outputs_for_bwd[0])

141 else:

142 raise ValueError("Cannot determine the arity, number of elements to be communicated")

143 return self._output_arity_for_comm

144

145 def _input_arity(self):

146 """Number of input tensors to communicate (cached after first use)."""

147 if self._input_arity_for_comm is not None:

148 return self._input_arity_for_comm

149 # Fallback: derive the arity from the first retained micro.

150 if self._inputs_for_explicit_forward.get(0):

151 self._input_arity_for_comm = len(self._inputs_for_explicit_forward[0])

152 else:

153 raise ValueError("Cannot determine the arity, number of elements to be communicated")

154 return self._input_arity_for_comm

155

156 def _global_rank(self, group_rank: int) -> int:

157 """Global rank of group pipeline rank ``group_rank``."""

158 return dist.get_global_rank(self._mpipe_group, group_rank)

159

160 @staticmethod

161 def _as_tuple(args):

162 """Normalize a micro-batch arg slot to a positional-args tuple."""

163 if isinstance(args, (list, tuple)):

164 return tuple(args)

165 return (args,)

166

167 @staticmethod

168 def _kwargs_for(ctx, micro):

169 """Keyword args for ``micro``'s preprocess forward (e.g. position_ids)."""

170 kwarg_mbs = getattr(ctx, "kwarg_mbs", None)

171 if not kwarg_mbs:

172 return {}

173 return kwarg_mbs[micro] or {}

174

175 def broadcast_params(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument

176 """Broadcast the preprocess parameters from stage 0 to all ranks.

177

178 Args:

179 step (MetaStep): The ``MPIPE_PARAM_BROADCAST`` schedule step (unused).

180 ctx (PipelineContext): The pipeline run context (unused).

181 """

182 src = self._global_rank(0)

183 for tensor in self._broadcast_tensors():

184 dist.broadcast(tensor, src, self._mpipe_group)

185

186 def transpose_forward(self, step: "MetaStep", ctx: "PipelineContext") -> None:

187 """Run the preprocess forward for ``step.micro_index``.

188

189 The output is currently detached and the input retained for the stage-0

190 recompute backward -- every micro takes this uniform path. (The old

191 graph-connected inline forward for non-transposed / overflow micros was

192 dropped in the round-robin change; the in-flight owner-does-backward work

193 for a non-stage-0 backward may reintroduce a connected forward on the

194 owner.)

195

196 Args:

197 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

198 ctx (PipelineContext): The pipeline run context (``arg_mbs`` / ``kwarg_mbs``).

199 """

200 micro = step.micro_index

201 args = self._as_tuple(ctx.arg_mbs[micro])

202 kwargs = self._kwargs_for(ctx, micro)

203 self._input_arity_for_comm = len(args)

204 if self._owner_backward:

205 # Stage 0's own micros backprop the connected tower during the body

206 # backward, so only the other ranks retain a graph for the owner backward.

207 out = self._as_tuple(self._connected_forward(args, kwargs))

208 if self._this_rank == 0:

209 self._output_arity_for_comm = len(out)

210 self._transfer_out_for_fwd(ctx, micro, out)

211 else:

212 self._outputs_for_stage0[micro] = out # shipped detached by fwd_send

213 self._outputs_for_bwd[micro] = out # retained graph for the owner backward

214 else:

215 if self._has_trainable:

216 self._inputs_for_explicit_forward[micro] = args

217 # The output may be one tensor or a tuple (VL: image_embeds + DeepStack

218 # levels); normalize so transport and placement are uniform.

219 out = self._as_tuple(self._detached_forward(args, kwargs))

220 if self._has_trainable:

221 for tensor in out:

222 self._mark_requires_grad(tensor)

223 if self._this_rank == 0:

224 self._output_arity_for_comm = len(out)

225 self._save_grad_for_bwd(micro, out)

226 self._transfer_out_for_fwd(ctx, micro, out)

227 else:

228 self._outputs_for_stage0[micro] = out

229 # Stages built under meta-init only learn their real device here.

230 if out and getattr(self._device, "type", None) in (None, "meta"):

231 self._device = out[0].device

232

233 def _save_grad_for_bwd(self, micro, out) -> None:

234 """Record a transposed micro-batch's placed feature tensors for the backward.

235

236 The stage-0 backward reads ``dL/dfeatures`` off these objects after the

237 body backward deposits them. Routing-agnostic: it holds the same tensors

238 whether they were handed to the forward as ``arg_mbs`` (a text body input)

239 or as a kwarg (a VL ``mpipe_visual`` payload).

240 """

241 if self._has_trainable:

242 self._keep_grad[micro] = tuple(out)

243

244 def _transfer_out_for_fwd(self, ctx, micro, out) -> None:

245 """Deliver a transposed micro-batch's preprocess output to stage 0's forward.

246

247 Default: the output is stage 0's body input (placed into ``arg_mbs``). A

248 model may override this via the schedule's ``output_consumer`` -- e.g. VL

249 routes the visual payload into an injection kwarg the stage forward reads,

250 since the visual output is injected mid-forward, not as the body input. The

251 stage forward consumes these slots on the later ``FWD`` step, via

252 ``forward_one_chunk(micro, arg_mbs[micro], kwarg_mbs[micro])``.

253

254 Args:

255 ctx: The pipeline run context.

256 micro: The transposed micro-batch index.

257 out: The preprocess output as a tuple of tensors.

258 """

259 consumer = self._schedule.output_consumer

260 if consumer is not None:

261 consumer(ctx, micro, out)

262 else:

263 ctx.arg_mbs[micro] = list(out)

264

265 def fwd_send(self, step: "MetaStep", ctx: "PipelineContext") -> None:

266 """Send the preprocess output of ``step.micro_index`` to stage 0.

267

268 Args:

269 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

270 ctx (PipelineContext): The pipeline run context; the send handles are appended to it.

271 """

272 micro = step.micro_index

273 dst = self._global_rank(0)

274 # Under owner-backward the output is graph-connected, so ship a detached

275 # wire copy rather than the graph tensor (keeps the owner's graph local).

276 for tensor in self._outputs_for_stage0[micro]:

277 wire = self._detach_for_wire(tensor) if self._owner_backward else self._contiguous(tensor)

278 self._send_meta(wire, dst)

279 ctx.schedule._send_handles.append([dist.isend(wire, dst, self._mpipe_group)]) # pylint: disable=protected-access

280 # What is in flight is the wire copy, not this cache entry, so the

281 # source can go now; the handles are drained by ``run_microbatches``.

282 del self._outputs_for_stage0[micro]

283

284 def fwd_recv(self, step: "MetaStep", ctx: "PipelineContext") -> None:

285 """Receive transposed micro-batches' preprocess outputs and hand them to stage 0.

286

287 Batched path (default): the first ``MPIPE_FWD_RECV`` step posts the recvs

288 for every transposed micro-batch non-blocking, then waits once, so the

289 ``NT - 1`` transfers pipeline and overlap stage 0's own preprocess forward

290 (whose kernels are still in flight on the compute stream) instead of being

291 received-and-waited serially. Per-peer P2P order and tags are unchanged

292 (only the waits are deferred), so collective ordering is unperturbed; the

293 later ``MPIPE_FWD_RECV`` steps then find their micro already received. Set

294 ``HYPER_PARALLEL_MPIPE_RECV_BATCH=0`` for the per-micro blocking path.

295

296 Args:

297 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

298 ctx (PipelineContext): The pipeline run context; received buffers

299 are handed to the forward via ``_transfer_out_for_fwd`` and

300 recorded via ``_save_grad_for_bwd``.

301 """

302 if not _RECV_BATCH:

303 self._fwd_recv_one(step.micro_index, ctx)

304 return

305 if step.micro_index in self._fwd_received:

306 return

307 self._fwd_recv_all(ctx)

308

309 def _recv_device(self):

310 """Real (non-``meta``) device for receive buffers.

311

312 ``self._device`` is captured from ``first_stage.device`` at construction.

313 Under deferred / meta init (FSDP2 builds the stage on ``meta`` and only

314 materializes it later) that capture is ``meta``; allocating a recv buffer

315 there yields a Meta tensor that ``c10d::recv_`` rejects on the wire

316 (``NotImplementedError: ... Meta tensor``). Resolve lazily from a tensor

317 stage 0 has already materialized this step -- its own micro-0 preprocess

318 output (``_outputs_for_stage0[0]``) or retained input (``_inputs_for_explicit_forward[0]``)

319 or grad (``_keep_grad[0]``) -- cache it, and fall back to the captured device.

320 """

321 if getattr(self._device, "type", None) != "meta":

322 return self._device

323 for materialized in (self._outputs_for_stage0.get(0),

324 self._inputs_for_explicit_forward.get(0),

325 self._keep_grad.get(0)):

326 if materialized:

327 self._device = materialized[0].device

328 return self._device

329 return self._device

330

331 def _recv_payload(self, micro):

332 """Post the recv of ``micro``'s preprocess output; return ``(buffers, handles)``.

333

334 Exchanges each tensor's meta, allocates its buffer and issues a

335 non-blocking ``irecv`` (the caller decides when to wait). Stage 0 always

336 owns transposed micro 0, so its own output reveals the per-micro arity

337 (1 for a text body input, >1 for a VL tuple payload).

338 """

339 src = self._global_rank(self._schedule.owner_of(micro))

340 arity = self._output_arity()

341 buffers = []

342 handles = []

343 for _ in range(arity):

344 shape, dtype = self._recv_meta(src)

345 buffer = torch.empty(shape, dtype=dtype, device=self._recv_device())

346 if self._has_trainable:

347 self._mark_requires_grad(buffer)

348 handles.append(dist.irecv(buffer, src, self._mpipe_group))

349 buffers.append(buffer)

350 return buffers, handles

351

352 def _fwd_recv_one(self, micro: int, ctx: "PipelineContext") -> None:

353 """Receive a single transposed micro-batch's output, waiting inline (fallback)."""

354 buffers, handles = self._recv_payload(micro)

355 for handle in handles:

356 handle.wait()

357 out = tuple(buffers)

358 self._save_grad_for_bwd(micro, out)

359 self._transfer_out_for_fwd(ctx, micro, out)

360 self._fwd_received.add(micro)

361

362 def _fwd_recv_all(self, ctx: "PipelineContext") -> None:

363 """Post every outstanding transposed recv, then wait once so the transfers

364 overlap each other and stage 0's own preprocess forward."""

365 owned = self._schedule.owned_micros(self._this_rank)

366 pending = [(micro, *self._recv_payload(micro))

367 for micro in range(self._schedule.micro_batch_num)

368 if micro not in owned and micro not in self._fwd_received]

369 for micro, buffers, handles in pending:

370 for handle in handles:

371 handle.wait()

372 out = tuple(buffers)

373 self._save_grad_for_bwd(micro, out)

374 self._transfer_out_for_fwd(ctx, micro, out)

375 self._fwd_received.add(micro)

376

377 def fwd_bwd_graph_input_send(self, step: "MetaStep", ctx: "PipelineContext") -> None:

378 """Send the preprocess input of ``step.micro_index`` to stage 0 (for the stage-0 backward).

379

380 Args:

381 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

382 ctx (PipelineContext): The pipeline run context; send handles are appended to it.

383 """

384 micro = step.micro_index

385 dst = self._global_rank(0)

386 for tensor in self._inputs_for_explicit_forward[micro]:

387 contiguous = self._contiguous(tensor)

388 self._send_meta(contiguous, dst)

389 ctx.schedule._send_handles.append([dist.isend(contiguous, dst, self._mpipe_group)]) # pylint: disable=protected-access

390 # What is in flight is the wire copy, not this cache entry, so the

391 # source can go now; the handles are drained by ``run_microbatches``.

392 del self._inputs_for_explicit_forward[micro]

393

394 def fwd_bwd_graph_input_recv(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument

395 """Receive a transposed micro-batch's preprocess input for the stage-0 backward.

396

397 Args:

398 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

399 ctx (PipelineContext): The pipeline run context (unused).

400 """

401 micro = step.micro_index

402 src = self._global_rank(self._schedule.owner_of(micro))

403 # Stage 0 always owns transposed micro 0, so its retained input reveals

404 # how many input tensors each transposed micro-batch ships.

405 arity = self._input_arity()

406 tensors = []

407 for _ in range(arity):

408 shape, dtype = self._recv_meta(src)

409 buffer = torch.empty(shape, dtype=dtype, device=self._recv_device())

410 dist.irecv(buffer, src, self._mpipe_group).wait()

411 tensors.append(buffer)

412 self._inputs_for_explicit_forward[micro] = tuple(tensors)

413

414 def grad_send(self, step: "MetaStep", ctx: "PipelineContext") -> None:

415 """Owner-backward: ship ``dL/dfeatures`` for transposed ``micro`` back to its owner.

416

417 The body backward has deposited the gradient on the feature tensors stage

418 0 placed for ``micro`` (``self._keep_grad[micro]``); each is shipped (absent

419 per-tensor grads become zeros so the owner's backward is well-formed). The

420 destination is ``self._schedule.owner_of(micro)`` -- under round-robin

421 overflow (``"full"``) that is not ``micro`` itself once ``micro >= NT``.

422 Stage 0's own micro(s) already backpropagated via the connected tower

423 during the body backward, so those are a no-op (never actually scheduled;

424 the guard is defensive).

425

426 Args:

427 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

428 ctx (PipelineContext): The pipeline run context; send handles are appended to it.

429 """

430 micro = step.micro_index

431 dst = self._global_rank(self._schedule.owner_of(micro))

432 if dst == self._global_rank(0):

433 # Nothing to ship, but the entry still has to go: no later step

434 # erases it, so leaving it would strand the feature tensors and

435 # make the next run's reset() refuse to start.

436 self._keep_grad.pop(micro, None)

437 return

438 for tensor in self._keep_grad[micro]:

439 grad = getattr(tensor, "grad", None)

440 wire = self._detach_for_wire(grad) if grad is not None else self._zeros_like(tensor)

441 self._send_meta(wire, dst)

442 ctx.schedule._send_handles.append([dist.isend(wire, dst, self._mpipe_group)]) # pylint: disable=protected-access

443 # What is in flight is the wire copy, not this cache entry, so the

444 # source can go now; the handles are drained by ``run_microbatches``.

445 del self._keep_grad[micro]

446

447 def grad_recv_with_backward(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument

448 """Owner-backward: receive ``dL/dfeatures`` and backprop the retained tower graph.

449

450 The graph is ``self._outputs_for_bwd[micro]`` (this owner's connected

451 preprocess output); the received grads are its cotangents.

452

453 Args:

454 step (MetaStep): The schedule step; ``step.micro_index`` selects the owned micro-batch.

455 ctx (PipelineContext): The pipeline run context (unused).

456 """

457 micro = step.micro_index

458 src = self._global_rank(0)

459 retained = self._outputs_for_bwd[micro]

460 device = retained[0].device

461 grads = []

462 handles = []

463 # Post every recv, then wait once: the payloads overlap in flight

464 # (same per-peer posting order as the sends, so matching is unchanged).

465 for _ in range(len(retained)):

466 shape, dtype = self._recv_meta(src)

467 buffer = torch.empty(shape, dtype=dtype, device=device)

468 handles.append(dist.irecv(buffer, src, self._mpipe_group))

469 grads.append(buffer)

470 for handle in handles:

471 handle.wait()

472 self._owner_transpose_backward(retained, grads)

473 del self._outputs_for_bwd[micro]

474

475 def reduce_tower_grads(self, step: "MetaStep", ctx: "PipelineContext") -> None: # pylint: disable=unused-argument

476 """Owner-backward: SUM all-reduce the tower param-grads across the pp replicas.

477

478 A collective -- emitted on every pp rank; ranks that produced no grad for a

479 param contribute zeros so the all-reduce stays symmetric. Only this run's

480 contribution is reduced (``grad - snapshot``), so prior gradient-

481 accumulation passes (already reduced) are not re-reduced.

482

483 Args:

484 step (MetaStep): The ``MPIPE_GRAD_REDUCE`` step (unused).

485 ctx (PipelineContext): The pipeline run context (unused).

486 """

487 self._reduce_grads(self._mpipe_group, self._grad_snapshot)

488 # Dead once the delta is reduced: drop the clone rather than hold it

489 # across the optimizer step; the next reset() takes a fresh one.

490 self._grad_snapshot = None

491

492 def transpose_backward(self, step: "MetaStep", ctx: "PipelineContext") -> None:

493 """Recompute the preprocess forward on stage 0 and backprop dL/dfeatures.

494

495 The body backward deposits the feature gradient on the tensors stage 0

496 placed for ``micro`` (``self._keep_grad[micro]`` -- the same objects whether

497 they were routed into ``arg_mbs`` (a text body input) or into a kwarg (the

498 VL ``mpipe_visual`` payload)). Recompute the preprocess forward from the

499 retained input and backprop those grads into it. Reading ``_keep_grad`` (not

500 ``arg_mbs[micro][0]``) is what makes the stage-0 backward correct for a

501 trainable VL tower, whose features live in the kwarg.

502 Assume that ``self._keep_grad[micro]`` was previously fill before calling the BWD

503

504 Args:

505 step (MetaStep): The schedule step; ``step.micro_index`` selects the micro-batch.

506 ctx (PipelineContext): The pipeline run context.

507 """

508 micro = step.micro_index

509 grads = [getattr(tensor, "grad", None) for tensor in self._keep_grad[micro]]

510 if all(grad is None for grad in grads):

511 self._inputs_for_explicit_forward.pop(micro, None)

512 del self._keep_grad[micro]

513 return

514 self._explicit_forward_before_backward(

515 self._inputs_for_explicit_forward[micro], self._kwargs_for(ctx, micro), grads)

516 del self._inputs_for_explicit_forward[micro]

517 del self._keep_grad[micro]

518

519 @staticmethod

520 def _local(tensor):

521 # An FSDP-sharded tower param/grad is a DTensor; operate on its local

522 # shard so pp collectives skip the DTensor dispatcher. Plain: pass through.

523 to_local = getattr(tensor, "to_local", None)

524 return to_local() if callable(to_local) else tensor

525

526 def _broadcast_tensors(self):

527 """Yield each trainable tower param's local shard, resharded first.

528

529 Frozen params and buffers stay identical from init, so only the

530 trainable ones need re-syncing after the optimizer step.

531 """

532 # Reshard first: under owner-backward stage 0 holds the tower unsharded

533 # while the replicas are sharded, so the sizes would disagree and hang.

534 for module in self._preprocess.modules():

535 reshard = getattr(module, "reshard", None)

536 if callable(reshard):

537 reshard()

538 for param in self._preprocess.parameters():

539 if param.requires_grad:

540 yield self._local(param.data)

541

542 def _detached_forward(self, args, kwargs):

543 with torch.no_grad():

544 out = self._preprocess(*args, **kwargs)

545 # The preprocess output may be a single tensor (text body input) or a

546 # tuple (e.g. a VL visual payload: image_embeds + DeepStack levels).

547 if isinstance(out, (tuple, list)):

548 return tuple(t.detach() for t in out)

549 return out.detach()

550

551 def _connected_forward(self, args, kwargs):

552 return self._preprocess(*args, **kwargs)

553

554 @staticmethod

555 def _mark_requires_grad(tensor) -> None:

556 tensor.requires_grad_(True)

557

558 def _explicit_forward_before_backward(self, inputs, kwargs, grads) -> None:

559 out = self._preprocess(*inputs, **kwargs)

560 out = out if isinstance(out, (tuple, list)) else (out,)

561 # Backprop only the outputs that received a gradient; an absent

562 # dL/dfeature is a zero contribution.

563 pairs = [(out_i, g) for out_i, g in zip(out, grads) if g is not None]

564 if not pairs:

565 return

566 outs, grad_tensors = zip(*pairs)

567 torch.autograd.backward(outs, grad_tensors=grad_tensors)

568

569 @staticmethod

570 def _contiguous(tensor):

571 return tensor.contiguous()

572

573 # --- owner-does-backward hooks (opt-in, trainable tower) -----------------

574

575 @staticmethod

576 def _detach_for_wire(tensor):

577 # Ship a detached, contiguous copy so the owner keeps the autograd graph.

578 return tensor.detach().contiguous()

579

580 @staticmethod

581 def _zeros_like(tensor):

582 # Contiguous zero grad for a feature tensor that received none.

583 return torch.zeros_like(tensor).contiguous()

584

585 @staticmethod

586 def _owner_transpose_backward(retained_out, grads) -> None:

587 # Backprop dL/dfeatures through the retained connected tower graph on this

588 # rank's replica; skip non-grad-requiring outputs (else autograd raises).

589 pairs = [(out_i, g) for out_i, g in zip(retained_out, grads) if out_i.requires_grad]

590 if pairs:

591 outs, grad_tensors = zip(*pairs)

592 torch.autograd.backward(outs, grad_tensors=grad_tensors)

593 else:

594 logger.debug("[mpipe] owner backward skipped: no retained output requires grad.")

595

596 def _snapshot_tower_grads(self):

597 # Clone each trainable tower grad's LOCAL shard (None where absent) so

598 # _reduce_grads reduces only this run's contribution.

599 return [None if p.grad is None else self._local(p.grad).detach().clone()

600 for p in self._preprocess.parameters() if p.requires_grad]

601

602 def _reduce_grads(self, group, snapshot) -> None:

603 """SUM-reduce each trainable tower param's this-run contribution

604 (grad - snapshot) over the pp replicas on the LOCAL shard, so an

605 FSDP-sharded DTensor grad never meets a raw c10d collective. FSDP owns the

606 orthogonal dp reduce (mean over dp_shard/dp_replicate); this owns only the

607 pp sum. Reducing the delta -- not the total -- keeps prior grad-accumulation

608 passes from being re-reduced; the write-back is in place so the DTensor grad

609 wrapper the optimizer reads is preserved (reassigning param.grad drops it).

610 """

611 params = [p for p in self._preprocess.parameters() if p.requires_grad]

612 if snapshot is None: # first run (no prior accumulation) -> reduce full grad

613 snapshot = [None] * len(params)

614 for param, snap in zip(params, snapshot):

615 if param.grad is None:

616 param.grad = torch.zeros_like(param)

617 g_local = self._local(param.grad)

618 delta = g_local if snap is None else (g_local - snap)

619 reduced = delta.contiguous()

620 dist.all_reduce(reduced, group=group)

621 g_local.copy_(reduced if snap is None else (snap + reduced))

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

288 statements