# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import logging import re from abc import ABC, abstractmethod from collections import Counter, defaultdict from enum import Enum from typing import ( TYPE_CHECKING, Any, Callable, NamedTuple, ) from paddle import nn from paddle.distributed.auto_parallel.pipelining.stage import PipelineStage if TYPE_CHECKING: from .stage import _PipelineStageBase import paddle import paddle.distributed as dist from paddle import profiler from .microbatch import ( TensorChunkSpec, _split_tensor, merge_chunks, split_args_kwargs_into_chunks, ) logger = logging.getLogger(__name__) class _ActType(Enum): FORWARD = 1 BACKWARD_INPUT = 2 BACKWARD_WEIGHT = 3 UNSHARD = 4 RESHARD = 5 SEND_F = 6 RECV_F = 7 SEND_B = 8 RECV_B = 9 FULL_BACKWARD = 10 def __str__(self): str_map = { _ActType.FORWARD: "F", _ActType.BACKWARD_INPUT: "I", _ActType.BACKWARD_WEIGHT: "W", _ActType.UNSHARD: "UNSHARD", _ActType.RESHARD: "RESHARD", _ActType.SEND_F: "SEND_F", _ActType.RECV_F: "RECV_F", _ActType.SEND_B: "SEND_B", _ActType.RECV_B: "RECV_B", _ActType.FULL_BACKWARD: "B", } return str_map[self] @staticmethod def from_str(action): if action == "F": return _ActType.FORWARD elif action == "I": return _ActType.BACKWARD_INPUT elif action == "W": return _ActType.BACKWARD_WEIGHT elif action == "UNSHARD": return _ActType.UNSHARD elif action == "RESHARD": return _ActType.RESHARD elif action == "SEND_F": return _ActType.SEND_F elif action == "RECV_F": return _ActType.RECV_F elif action == "SEND_B": return _ActType.SEND_B elif action == "RECV_B": return _ActType.RECV_B elif action == "B": return _ActType.FULL_BACKWARD else: raise RuntimeError(f"Invalid computation type {action}") FORWARD = _ActType.FORWARD BACKWARD_INPUT = _ActType.BACKWARD_INPUT BACKWARD_WEIGHT = _ActType.BACKWARD_WEIGHT UNSHARD = _ActType.UNSHARD RESHARD = _ActType.RESHARD SEND_F = _ActType.SEND_F RECV_F = _ActType.RECV_F SEND_B = _ActType.SEND_B RECV_B = _ActType.RECV_B FULL_BACKWARD = _ActType.FULL_BACKWARD # Convenience shorthand for compute actions only since they are used in 'simple schedule format' F = FORWARD I = BACKWARD_INPUT W = BACKWARD_WEIGHT B = FULL_BACKWARD # Helper to parse an action string like 1F0 into a tuple of (stage_index, computation_type, microbatch_index) _action_regex = re.compile( r"(\d+)(F|I|B|W|UNSHARD|RESHARD|SEND_F|RECV_F|SEND_B|RECV_B)(\d*)" ) class _Action(NamedTuple): stage_index: int computation_type: _ActType microbatch_index: int | None = None def __repr__(self): repr = str(self.stage_index) repr += str(self.computation_type) if self.microbatch_index is not None: repr += str(self.microbatch_index) return repr class _PipelineSchedule(ABC): def __init__( self, n_microbatches: int, loss_fn: Callable[..., paddle.Tensor] | None = None, args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None, kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None, output_merge_spec: dict[str, Any] | tuple[Any] | None = None, ): # From arguments self._n_microbatches = n_microbatches self._loss_fn = loss_fn # Chunking specification for positional inputs. (default: `None`) self._args_chunk_spec = args_chunk_spec # Chunking specification for keyword inputs. (default: `None`) self._kwargs_chunk_spec = kwargs_chunk_spec self._output_merge_spec = output_merge_spec """ # args_chunk_spec and kwargs_chunk_spec specify how to chunk inputs. # They are used to convert batch to microbatches in `step(x)`. See # `TensorChunkSpec` for helper methods for creating them. """ # Derived self._has_backward = self._loss_fn is not None # Holds the losses for each microbatch. self._internal_losses: list[paddle.Tensor] = [] logger.info("Using %s", self.__class__.__name__) def _maybe_compute_loss(self, stage, output, target_mbs, mb_index): if stage.is_last and self._has_backward: loss = self._compute_loss(output, target_mbs[mb_index]) # type: ignore[index] self._internal_losses.append(loss) def _maybe_get_loss(self, stage, mb_index): valid_index = 0 <= mb_index < len(self._internal_losses) if stage.is_last and self._has_backward and valid_index: return self._internal_losses[mb_index] elif len(self._internal_losses) != 0 and not valid_index: raise RuntimeError( f"Loss for microbatch {mb_index} is not available. " f"Available losses for microbatches: {self._internal_losses}" ) else: return None def _update_losses(self, stages, losses): """ Update the losses to those in the internal state """ # if stages not a list turn into a list if not isinstance(stages, list): stages = [stages] contains_last_stage = any(stage.is_last for stage in stages) # Return losses if there is a container passed in if contains_last_stage and losses is not None: if len(self._internal_losses) != self._n_microbatches: raise RuntimeError( f"Expecting {self._n_microbatches} losses but got {len(self._internal_losses)}" ) # Clean external container first losses.clear() # Copy internal losses to external container losses.extend(self._internal_losses) self._internal_losses.clear() @abstractmethod def _step_microbatches( self, arg_mbs: list | None = None, kwarg_mbs: list | None = None, target_mbs: list | None = None, losses: list | None = None, ): """ Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the schedule implementation. Args: microbatches: list of microbatch args. """ raise NotImplementedError @abstractmethod def step( self, *args, target=None, losses: list | None = None, return_output: bool = False, **kwargs, ): """ Run one iteration of the pipeline schedule with *whole-batch* input. Will chunk the input into microbatches automatically, and go through the microbatches according to the schedule implementation. args: positional arguments to the model (as in non-pipeline case). kwargs: keyword arguments to the model (as in non-pipeline case). target: target for the loss function. losses: a list to store the losses for each microbatch. """ raise NotImplementedError def _check_inputs( self, arg_mbs: list | None = None, kwarg_mbs: list | None = None, target_mbs: list | None = None, losses: list | None = None, ): """ Pre-process/check inputs """ def check_type_and_len(mbs, name: str): if not isinstance(mbs, list): raise TypeError(f"{name} must be a list but got a {type(mbs)}") if len(mbs) != self._n_microbatches: raise ValueError( f"Expecting {self._n_microbatches} {name} but got {len(mbs)}" ) if arg_mbs is not None: check_type_and_len(arg_mbs, "arg_mbs") else: arg_mbs = [()] * self._n_microbatches if kwarg_mbs is not None: check_type_and_len(kwarg_mbs, "kwarg_mbs") else: kwarg_mbs = [{}] * self._n_microbatches if target_mbs is not None: check_type_and_len(target_mbs, "target_mbs") if losses is not None: if not isinstance(losses, list): raise TypeError( f"losses must be a list but got a {type(losses)}" ) return arg_mbs, kwarg_mbs def _compute_loss(self, output, target): return self._loss_fn(output, target) # type: ignore[misc] def _split_inputs( self, args: tuple[Any, ...], kwargs: dict[str, Any] | None = None, ): """ Splits a full-batch input into chunks (i.e. microbatches) and returns the chunks """ if args or kwargs: args_split, kwargs_split = split_args_kwargs_into_chunks( args, kwargs, self._n_microbatches, self._args_chunk_spec, self._kwargs_chunk_spec, ) return args_split, kwargs_split else: # Empty inputs (e.g. when called on middle stages) # Return a list of empty tuples/dicts with matching length as chunks return [()] * self._n_microbatches, [{}] * self._n_microbatches def _merge_outputs(self, output_chunks: list[Any]) -> Any: """ Merge output chunks back to a batch state. If output_merge_spec is None, the utility will merge output chunks by dimension 0 (batch dim). """ return merge_chunks( output_chunks, self._output_merge_spec, ) class PipelineScheduleSingle(_PipelineSchedule): """ Base class for single-stage schedules. Implements the `step` method. Derived classes should implement `_step_microbatches`. """ def __init__( self, stage: _PipelineStageBase, n_microbatches: int, loss_fn: Callable | None = None, args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None, kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None, output_merge_spec: dict[str, Any] | tuple[Any] | None = None, ): # Init parent super().__init__( n_microbatches=n_microbatches, loss_fn=loss_fn, args_chunk_spec=args_chunk_spec, kwargs_chunk_spec=kwargs_chunk_spec, output_merge_spec=output_merge_spec, ) # Self attributes self._stage = stage self._num_stages = stage.num_stages # Set the same has_backward flag for stage object self._stage.has_backward = self._has_backward self._stage_initialized = False def _initialize_stage(self, args, kwargs, labels): if self._stage.is_first: next_stage_args = self._stage._prepare_forward_infra( self._n_microbatches, args, kwargs ) else: next_stage_args = self._stage._prepare_forward_infra( self._n_microbatches, (), kwargs ) loss = None if self._stage.is_last: loss = self._loss_fn(next_stage_args[0], labels) if self._has_backward: self._stage._prepare_backward_infra(self._n_microbatches, loss) self._stage_initialized = True def step( self, *args, target=None, losses: list | None = None, return_output: bool = False, **kwargs, ): """ Run one iteration of the pipeline schedule with *whole-batch* input. Will chunk the input into microbatches automatically, and go through the microbatches according to the schedule implementation. args: positional arguments to the model (as in non-pipeline case). kwargs: keyword arguments to the model (as in non-pipeline case). target: target for the loss function. losses: a list to store the losses for each microbatch. """ # Clean per iteration self._stage.clear_runtime_states() # Split inputs into microbatches args_split, kwargs_split = self._split_inputs(args, kwargs) # Split target into microbatches if target is not None: targets_split = list(_split_tensor(target, self._n_microbatches)) else: targets_split = None # Run microbatches self._step_microbatches(args_split, kwargs_split, targets_split, losses) # Return merged results per original format if return_output: if self._stage.is_last: return self._merge_outputs(self._stage.output_chunks) return None def _batch_p2p(p2p_ops: list[dist.P2POp], desc: str | None = None): """ Simple wrapper over batch_isend_irecv from paddle.distributed, which just adds a descriptive logger on top. """ if len(p2p_ops) == 0: return None desc_str = f"{desc}, " if desc else "" logger.info("batch_p2p %s%s", desc_str, p2p_ops) return dist.batch_isend_irecv(p2p_ops).pop() def _sorted_batch_p2p(p2p_ops: list[dist.P2POp], desc: str | None = None): """ Sorts the list of P2P ops by the peer rank, and then calls batch_isend_irecv. Return a dictionary of works by peer rank. This function helps us avoid hangs in case of skip connections. """ # Arrange p2p_ops by peer rank: # int is the peer rank; # list is the list of ops towards the peer ops_by_peer: dict[int, list[dist.P2POp]] = defaultdict(list) work_by_peer: dict[int, dist.Work] = {} if len(p2p_ops) == 0: return work_by_peer # Classify the ops by peer rank for op in p2p_ops: ops_by_peer[op.peer].append(op) # Call batch_isend_irecv per peer, in sorted order of the peers (to avoid hangs) for peer, ops in sorted(ops_by_peer.items()): work_by_peer[peer] = _batch_p2p(ops, desc=desc) return work_by_peer class ScheduleFThenB(PipelineScheduleSingle): """ The FThenB schedule. Will go through all the microbatches in a fill-drain manner. """ def _step_microbatches( self, arg_mbs: list | None = None, kwarg_mbs: list | None = None, target_mbs: list | None = None, losses: list | None = None, ): """ Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the FThenB schedule. Args: microbatches: list of microbatch args. """ arg_mbs, kwarg_mbs = self._check_inputs( arg_mbs, kwarg_mbs, target_mbs, losses ) if not self._stage_initialized: if target_mbs is not None: self._initialize_stage(arg_mbs[0], kwarg_mbs[0], target_mbs[0]) else: self._initialize_stage(arg_mbs[0], kwarg_mbs[0], None) # Delay send waits fwd_sends_to_wait: list[dist.Work] = [] # Run microbatches for i in range(self._n_microbatches): with profiler.RecordEvent(f"Forward {i}"): ops = self._stage.get_fwd_recv_ops(i) works = _sorted_batch_p2p(ops, desc="fwd_recv") for work in works.values(): work.wait() output = self._stage.forward_one_chunk( i, arg_mbs[i], kwarg_mbs[i] ) ops = self._stage.get_fwd_send_ops(i) works = _sorted_batch_p2p(ops, desc="fwd_send") fwd_sends_to_wait.extend(works.values()) logger.debug( "[%s] Forwarded microbatch %s", self._stage.stage_index, i ) self._maybe_compute_loss(self._stage, output, target_mbs, i) # Wait for all forward sends to finish # This should not have performance impact because by the time the first # backward arrives all the forward sends should have been finished. for work in fwd_sends_to_wait: work.wait() # No loss function, no need to run backward if not self._has_backward: return # Run backward # Delay send waits bwd_sends_to_wait: list[dist.Work] = [] for i in range(self._n_microbatches): with profiler.RecordEvent(f"Backward {i}"): ops = self._stage.get_bwd_recv_ops(i) works = _sorted_batch_p2p(ops, desc="bwd_recv") for work in works.values(): work.wait() loss = self._maybe_get_loss(self._stage, i) self._stage.backward_one_chunk( i, loss=loss, last_backward=i == self._n_microbatches - 1 ) ops = self._stage.get_bwd_send_ops(i) works = _sorted_batch_p2p(ops, desc="bwd_send") bwd_sends_to_wait.extend(works.values()) logger.debug( "[%s] Backwarded microbatch %s", self._stage.stage_index, i ) # Return losses if there is a container passed in self._update_losses(self._stage, losses) # Wait for all backward sends to finish for work in bwd_sends_to_wait: work.wait() # Synchronize the gradients of shared parameters. self._stage._sync_shared_param_grads() class PipelineChunk(nn.Layer): def __init__(self, layers=None, is_first=False, is_last=False): super().__init__() assert not (is_first and is_last), ( "Pipeline stage cannot be both first and last." ) self.layers = layers self.is_first = is_first self.is_last = is_last def forward(self, *args, **kwargs): if self.is_first: input_ids = kwargs.get("input_ids") attention_mask = kwargs.get("attention_mask") position_ids = kwargs.get("position_ids") outputs = (input_ids, attention_mask, position_ids) # decoder layers for idx, (decoder_layer) in enumerate(self.layers): outputs = decoder_layer(outputs) return outputs elif self.is_last: outputs = args # decoder layers for idx, (decoder_layer) in enumerate(self.layers): outputs = decoder_layer(outputs) if isinstance(outputs, tuple): outputs = outputs[0] else: outputs = args # decoder layers for idx, (decoder_layer) in enumerate(self.layers): outputs = decoder_layer(outputs) return outputs def _manual_model_split(model, stage_idx, group, mode, pp_degree): num_hidden_layers = model.config.num_hidden_layers virtual_pp_degree = model.config.virtual_pp_degree if mode == "VPP" else 1 chunk_size = num_hidden_layers // virtual_pp_degree // pp_degree chunk_num = virtual_pp_degree * pp_degree layer_lists = model.layers def _build_stage(model, stage_idx, group): new_model = None if stage_idx == 0: new_model = PipelineChunk( layer_lists[:chunk_size], is_first=True, is_last=False ) elif stage_idx == chunk_num - 1: new_model = PipelineChunk( layer_lists[ stage_idx * chunk_size : (stage_idx + 1) * chunk_size ], is_first=False, is_last=True, ) else: new_model = PipelineChunk( layer_lists[ stage_idx * chunk_size : (stage_idx + 1) * chunk_size ], is_first=False, is_last=False, ) stage = PipelineStage(new_model, stage_idx, chunk_num, group=group) return stage stages = [] for i in range(virtual_pp_degree): stage = _build_stage(model, stage_idx + i * pp_degree, group) stages.append(stage) return stages def get_pipeline_schedule(model, acc_steps, loss_fn, mode, pp_degree, group): assert mode in [ "VPP", "1F1B", "FThenB", ], ( f"Invalid pipeline schedule mode: {mode}, must be one of ['VPP', '1F1B', 'FThenB']" ) stages = _manual_model_split(model, group.rank, group, mode, pp_degree) if mode == "VPP": schedule = ScheduleVPP( stages, n_microbatches=acc_steps, loss_fn=loss_fn ) elif mode == "1F1B": schedule = Schedule1F1B( stages[0], n_microbatches=acc_steps, loss_fn=loss_fn ) else: schedule = ScheduleFThenB( stages[0], n_microbatches=acc_steps, loss_fn=loss_fn ) return schedule class Schedule1F1B(PipelineScheduleSingle): """ The 1F1B schedule. Will perform one forward and one backward on the microbatches in steady state. """ def _step_microbatches( self, arg_mbs: list | None = None, kwarg_mbs: list | None = None, target_mbs: list | None = None, losses: list | None = None, ): """ Run one iteration of the pipeline schedule with list of microbatches. Will go through all the microbatches according to the 1F1B schedule. Args: microbatches: list of microbatch args. """ arg_mbs, kwarg_mbs = self._check_inputs( arg_mbs, kwarg_mbs, target_mbs, losses ) if not self._stage_initialized: if target_mbs is not None: self._initialize_stage(arg_mbs[0], kwarg_mbs[0], target_mbs[0]) else: self._initialize_stage(arg_mbs[0], kwarg_mbs[0], None) # Last stage has 1 warmup, second-to-last 2 warmups, ... # first stage `num_stages` warmups warmup_chunks = min( self._n_microbatches, self._num_stages - self._stage.stage_index, ) # Chunk counters fwd_mb_index = 0 bwd_mb_index = 0 # Warmup phase send_work = None fwd_sends = [] for _ in range(warmup_chunks): # Receive activations fwd_recvs = self._stage.get_fwd_recv_ops(fwd_mb_index) if recv_work := _batch_p2p(fwd_recvs, desc="fwd_recv"): recv_work.wait() # Compute output = self._stage.forward_one_chunk( fwd_mb_index, arg_mbs[fwd_mb_index], kwarg_mbs[fwd_mb_index] ) # Clear previous chunk's forward sends (hopefully they have well # finished, otherwise, we are heavily communication bound, in which # case it doesn't create a lot of benefit to compute next chunk # eagerly either) if send_work: send_work.wait() # Send activations fwd_sends = self._stage.get_fwd_send_ops(fwd_mb_index) if fwd_mb_index != warmup_chunks - 1: # Safe to fire send_work = _batch_p2p(fwd_sends, desc="fwd_send") # otherwise: # The last forward send is left for fuse with first 1B in 1B1F below # Compute loss self._maybe_compute_loss( self._stage, output, target_mbs, fwd_mb_index ) fwd_mb_index += 1 # Now we should have send ops left over, to be fused with first 1B of 1B1F phase below. # 1B1F phase while True: # Don't worry, we have a break inside # We actually do 1B first as the `1B1F` name indicates, so prepare its recv ops bwd_recvs = self._stage.get_bwd_recv_ops(bwd_mb_index) # Now, we need to fire the fwd_sends and bwd_recvs together if fuse_work := _batch_p2p( fwd_sends + bwd_recvs, desc="fwd_send_bwd_recv" ): fuse_work.wait() # Backward one chunk loss = self._maybe_get_loss(self._stage, bwd_mb_index) self._stage.backward_one_chunk( bwd_mb_index, loss=loss, last_backward=bwd_mb_index == self._n_microbatches - 1, ) # Get the bwd send ops, but don't fire, to be fused with the 1F below bwd_sends = self._stage.get_bwd_send_ops(bwd_mb_index) bwd_mb_index += 1 if fwd_mb_index == self._n_microbatches: # We are done with 1B1F, so break with some left-over bwd_sends break # We prepare 1F of the `1B1F` fwd_recvs = self._stage.get_fwd_recv_ops(fwd_mb_index) # Fuse it with bwd_sends above if fuse_work := _batch_p2p( bwd_sends + fwd_recvs, desc="bwd_send_fwd_recv" ): fuse_work.wait() # Now do the fwd output = self._stage.forward_one_chunk( fwd_mb_index, arg_mbs[fwd_mb_index], kwarg_mbs[fwd_mb_index] ) # Compute loss self._maybe_compute_loss( self._stage, output, target_mbs, fwd_mb_index ) # Get the fwd send ops, but don't fire, leave it for the next iter (wrap-around) fwd_sends = self._stage.get_fwd_send_ops(fwd_mb_index) fwd_mb_index += 1 # Remember we still have some bwd_sends left over after the break? Now it is time to fire it send_work = _batch_p2p(bwd_sends, desc="bwd_send") # Cooldown while bwd_mb_index < self._n_microbatches: # prepare bwd recv ops bwd_recvs = self._stage.get_bwd_recv_ops(bwd_mb_index) if recv_work := _batch_p2p(bwd_recvs, desc="bwd_recv"): recv_work.wait() # Backward one chunk loss = self._maybe_get_loss(self._stage, bwd_mb_index) self._stage.backward_one_chunk( bwd_mb_index, loss=loss, last_backward=bwd_mb_index == self._n_microbatches - 1, ) # Clear previous chunk's backward sends (hopefully they have well finished) if send_work: send_work.wait() # Get the bwd send ops, fire it bwd_sends = self._stage.get_bwd_send_ops(bwd_mb_index) send_work = _batch_p2p(bwd_sends, desc="bwd_send") bwd_mb_index += 1 # Wait for the last backward send to finish if send_work: send_work.wait() # Return losses if there is a container passed in self._update_losses(self._stage, losses) # Synchronize the gradients of shared parameters. self._stage._sync_shared_param_grads() class PipelineScheduleMulti(_PipelineSchedule): """ Base class for multi-stage schedules. Implements the `step` method. """ def __init__( self, stages: list[_PipelineStageBase], n_microbatches: int, loss_fn: Callable | None = None, args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None, kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None, output_merge_spec: dict[str, Any] | tuple[Any] | None = None, stage_index_to_group_rank: dict[int, int] | None = None, use_full_backward: bool | None = None, ): # Init parent super().__init__( n_microbatches=n_microbatches, loss_fn=loss_fn, args_chunk_spec=args_chunk_spec, kwargs_chunk_spec=kwargs_chunk_spec, output_merge_spec=output_merge_spec, ) # Self attributes self._stages = stages self._num_stages = stages[0].num_stages self.pp_group_size = stages[0].group_size self.rank = stages[0].group_rank # Set the pipeline stage states if stage_index_to_group_rank is not None: for stage in self._stages: stage.stage_index_to_group_rank = stage_index_to_group_rank self.stage_index_to_group_rank = stages[0].stage_index_to_group_rank # Set the same has_backward flag for stage object for stage in self._stages: stage.has_backward = self._has_backward self._stages_initialized = False # avoid putting a reference to 'self' inside the lambda, it creates a ref cycle has_loss: bool = self._loss_fn is not None self._should_compute_loss = lambda stage: stage.is_last and has_loss # This will be set during init of derived schedules self.pipeline_order: dict[int, list[_Action | None]] = {} if use_full_backward is not None: logger.warning( "Deprecation warning: 'use_full_backward' is no longer supported. " "Simply stop passing it, and everything should still work fine." ) def _initialize_stages(self, args: tuple[Any, ...], kwargs, labels): # may be 'none' value (if this stage sends its output shapes to the next stage via P2P) # or real value (if this stage and next stage are on the same device) next_stage_args: tuple[Any, ...] = () for stage in self._stages: if stage.is_first: next_stage_args = stage._prepare_forward_infra( self._n_microbatches, args, kwargs ) else: next_stage_args = stage._prepare_forward_infra( self._n_microbatches, next_stage_args, kwargs ) loss = None last_stage = self._stages[-1] if last_stage.is_last: loss = self._loss_fn(next_stage_args[0], labels) if self._has_backward: for stage_reverse in reversed(self._stages): if stage_reverse.is_last: stage_reverse._prepare_backward_infra( self._n_microbatches, loss ) else: stage_reverse._prepare_backward_infra( self._n_microbatches, None ) self._stages_initialized = True def step( self, *args, target=None, losses: list | None = None, return_output: bool = False, **kwargs, ): """ Run one iteration of the pipeline schedule with *whole-batch* input. Will chunk the input into microbatches automatically, and go through the microbatches according to the schedule implementation. args: positional arguments to the model (as in non-pipeline case). kwargs: keyword arguments to the model (as in non-pipeline case). target: target for the loss function. losses: a list to store the losses for each microbatch. """ # Clean per iteration for stage in self._stages: stage.clear_runtime_states() # Split inputs into microbatches args_split, kwargs_split = self._split_inputs(args, kwargs) # Split target into microbatches if target is not None: targets_split = list(_split_tensor(target, self._n_microbatches)) else: targets_split = None # Run microbatches self._step_microbatches(args_split, kwargs_split, targets_split, losses) # Return merged results per original format if return_output: for stage in self._stages: if stage.is_last: return self._merge_outputs(stage.output_chunks) # Does not contain the last stage return None def _step_microbatches( self, arg_mbs: list | None = None, kwarg_mbs: list | None = None, target_mbs: list | None = None, losses: list | None = None, ): """ Operate on the microbatches for looped schedules (multiple stages on each rank). """ arg_mbs, kwarg_mbs = self._check_inputs( arg_mbs, kwarg_mbs, target_mbs, losses ) if not self._stages_initialized: if target_mbs is not None: self._initialize_stages(arg_mbs[0], kwarg_mbs[0], target_mbs[0]) else: self._initialize_stages(arg_mbs[0], kwarg_mbs[0], None) # Based on the plan in Step 1 created in __init__: # 2. Perform communication based on the pipeline_order stage_index_to_stage: dict[int, _PipelineStageBase] = { stage.stage_index: stage for stage in self._stages } # determine prev_rank and next_rank based on which ranks are next to # the stages in the pipeline_order all_prev_ranks: set[int] = set() all_next_ranks: set[int] = set() for stage_index in stage_index_to_stage.keys(): # TODO: assumption that stages only communicate from distances of +1/-1 (no skip connections) if stage_index > 0: all_prev_ranks.add( self.stage_index_to_group_rank[stage_index - 1] ) if stage_index < self._num_stages - 1: all_next_ranks.add( self.stage_index_to_group_rank[stage_index + 1] ) # count either full_backward or backward_weight together, to determine when to sync DP grads backward_counter: Counter[int] = Counter() for time_step, action in enumerate(self.pipeline_order[self.rank]): try: ops: list[dist.P2POp] = [] if action is not None: computation_type = action.computation_type mb_index = action.microbatch_index stage_index = action.stage_index assert mb_index is not None, ( "All currently supported action types require valid microbatch_index" ) if computation_type == _ActType.FORWARD: # perform forward computation stage = stage_index_to_stage[stage_index] output = stage.forward_one_chunk( mb_index, arg_mbs[mb_index], kwarg_mbs[mb_index] ) self._maybe_compute_loss( stage, output, target_mbs, mb_index ) ops.extend(stage.get_fwd_send_ops(mb_index)) elif computation_type == _ActType.FULL_BACKWARD: # perform backward computation stage = stage_index_to_stage[stage_index] loss = self._maybe_get_loss(stage, mb_index) backward_counter[stage_index] += 1 stage.backward_one_chunk( mb_index, loss=loss, full_backward=True, last_backward=backward_counter[stage_index] == self._n_microbatches, ) ops.extend(stage.get_bwd_send_ops(mb_index)) elif computation_type == _ActType.BACKWARD_INPUT: # perform backward computation stage = stage_index_to_stage[stage_index] loss = self._maybe_get_loss(stage, mb_index) stage.backward_one_chunk( mb_index, loss=loss, full_backward=False, last_backward=False, ) ops.extend(stage.get_bwd_send_ops(mb_index)) elif computation_type == _ActType.BACKWARD_WEIGHT: # perform weight update stage = stage_index_to_stage[stage_index] backward_counter[stage_index] += 1 stage.backward_weight_one_chunk( mb_index, last_backward=backward_counter[stage_index] == self._n_microbatches, ) else: raise ValueError( f"Unknown computation type {computation_type}" ) # Look at the neighboring ranks for this current timestep and determine whether # this current rank needs to do any recv communication for prev_rank in all_prev_ranks: prev_rank_ops = self.pipeline_order[prev_rank] prev_rank_action = None if time_step < len(prev_rank_ops): prev_rank_action = prev_rank_ops[time_step] if prev_rank_action is not None: computation_type = prev_rank_action.computation_type mb_index = prev_rank_action.microbatch_index stage_index = prev_rank_action.stage_index assert mb_index is not None, ( "All currently supported action types require valid microbatch_index" ) # Only handle sends for the forward from a previous rank if computation_type == _ActType.FORWARD: # If not the last stage, then receive fwd activations if stage_index + 1 in stage_index_to_stage: # TODO: We are assuming that stage will always receive from stage-1 # however that is not necessarily true of get_fwd_recv_ops stage = stage_index_to_stage[stage_index + 1] ops.extend(stage.get_fwd_recv_ops(mb_index)) elif computation_type in ( FULL_BACKWARD, BACKWARD_INPUT, BACKWARD_WEIGHT, ): # Previous rank doing backward has no influence for the current rank forward recv pass else: raise ValueError( f"Unknown computation type {computation_type}" ) for next_rank in all_next_ranks: next_rank_ops = self.pipeline_order[next_rank] next_rank_action = None if time_step < len(next_rank_ops): next_rank_action = next_rank_ops[time_step] if next_rank_action is not None: computation_type = next_rank_action.computation_type mb_index = next_rank_action.microbatch_index stage_index = next_rank_action.stage_index assert mb_index is not None, ( "All currently supported action types require valid microbatch_index" ) # Only handle receives for the backwards from a next rank if computation_type in (FORWARD, BACKWARD_WEIGHT): # Next rank doing forward or weight update has no influence for the current rank backward recv pass elif computation_type in ( BACKWARD_INPUT, FULL_BACKWARD, ): # If not the first stage, then receive bwd gradients if stage_index - 1 in stage_index_to_stage: # TODO: We are assuming that stage will always receive from stage+1 # however that is not necessarily true of get_bwd_recv_ops stage = stage_index_to_stage[stage_index - 1] ops.extend(stage.get_bwd_recv_ops(mb_index)) else: raise ValueError( f"Unknown computation type {computation_type}" ) # do the communication if ops: _batch_p2p(ops).wait() except Exception as e: logger.error( "[Rank %s] pipeline schedule %s caught the following exception \ at time_step %s when running action %s", self.rank, self.__class__.__name__, time_step, action, ) raise e # Return losses if there is a container passed in self._update_losses(self._stages, losses) # Synchronize the gradients of shared parameters. for stage in self._stages: stage._sync_shared_param_grads() def _get_1f1b_rank_ops( n_local_stages, pp_group_size, warmup_ops, fwd_bwd_ops, cooldown_ops, rank, forward_stage_index, backward_stage_index, num_1f1b_microbatches=0, enable_zero_bubble=False, ): # All stages start with handling microbatch 0 fwd_stage_mb_index: dict[int, int] = defaultdict(int) bwd_stage_mb_index: dict[int, int] = defaultdict(int) weight_stage_mb_index: dict[int, int] = defaultdict(int) # Store the list of operations used for that rank # Pre-padding, rank starts with no-ops based on the warmup. rank_ops: list[_Action | None] = [None for _ in range(rank)] # These are used to calculate the number of slots to fill with no-ops, to account for the delay in warmup # when we want to wait for the backward to trickle back up and start 1f1b to align all ranks. # Formula: # pre-padding + warmup_ops + post_warmup_ops = earliest time step of first backward # post_warmup_ops = [earliest time step of first backward] - (warmup_ops + pre-padding) # earliest time step of first backward = [local_stages * group_size + 2 * (group_size - 1 - rank)] # warmup_ops = calculated above post_warmup_ops = ( n_local_stages * pp_group_size + 2 * (pp_group_size - 1 - rank) ) - (warmup_ops + rank) if enable_zero_bubble: post_warmup_ops = pp_group_size - rank - 1 total_ops = warmup_ops + fwd_bwd_ops + cooldown_ops backward_op_ids = [] weight_op_count = 0 FULL_BACKWARD_OR_BACKWARD_INPUT = ( BACKWARD_INPUT if enable_zero_bubble else FULL_BACKWARD ) for op in range(total_ops): # Warmup phase if op < warmup_ops: fwd_stage_index = forward_stage_index(op) # This will assign the current microbatch index and update it as well fwd_stage_mb_index[fwd_stage_index] = ( mb_index := fwd_stage_mb_index[fwd_stage_index] ) + 1 rank_ops.append( _Action(fwd_stage_index, _ActType.FORWARD, mb_index) ) if op == warmup_ops - 1: # This is the last step in the warmup phase, so we need to wait for the backward to trickle back up rank_ops.extend([None] * post_warmup_ops) # 1F1B Phase (forward and backward) elif warmup_ops <= op < warmup_ops + fwd_bwd_ops: fwd_stage_index = forward_stage_index(op) fwd_stage_mb_index[fwd_stage_index] = ( fwd_mb_index := fwd_stage_mb_index[fwd_stage_index] ) + 1 rank_ops.append( _Action(fwd_stage_index, _ActType.FORWARD, fwd_mb_index) ) bwd_stage_index = backward_stage_index(op) bwd_stage_mb_index[bwd_stage_index] = ( bwd_mb_index := bwd_stage_mb_index[bwd_stage_index] ) + 1 rank_ops.append( _Action( bwd_stage_index, FULL_BACKWARD_OR_BACKWARD_INPUT, bwd_mb_index, ) ) backward_op_ids.append(op) if enable_zero_bubble and op - warmup_ops >= num_1f1b_microbatches: weight_stage_index = backward_stage_index( backward_op_ids[weight_op_count] ) weight_stage_mb_index[weight_stage_index] = ( weight_mb_index := weight_stage_mb_index[weight_stage_index] ) + 1 rank_ops.append( _Action( weight_stage_index, _ActType.BACKWARD_WEIGHT, weight_mb_index, ) ) weight_op_count += 1 # Cooldown phase else: # During cooldown phase, we need steps to align with 1f1b happening in other ranks # TODO: we don't need to always append, after all 1f1b are finished we can stop appending None if not enable_zero_bubble: rank_ops.append(None) bwd_stage_index = backward_stage_index(op) bwd_stage_mb_index[bwd_stage_index] = ( bwd_mb_index := bwd_stage_mb_index[bwd_stage_index] ) + 1 rank_ops.append( _Action( bwd_stage_index, FULL_BACKWARD_OR_BACKWARD_INPUT, bwd_mb_index, ) ) backward_op_ids.append(op) if enable_zero_bubble and op - warmup_ops >= num_1f1b_microbatches: weight_stage_index = backward_stage_index( backward_op_ids[weight_op_count] ) weight_stage_mb_index[weight_stage_index] = ( weight_mb_index := weight_stage_mb_index[weight_stage_index] ) + 1 rank_ops.append( _Action( weight_stage_index, _ActType.BACKWARD_WEIGHT, weight_mb_index, ) ) weight_op_count += 1 while enable_zero_bubble and weight_op_count < len(backward_op_ids): weight_stage_index = backward_stage_index( backward_op_ids[weight_op_count] ) weight_stage_mb_index[weight_stage_index] = ( weight_mb_index := weight_stage_mb_index[weight_stage_index] ) + 1 rank_ops.append( _Action( weight_stage_index, _ActType.BACKWARD_WEIGHT, weight_mb_index ) ) weight_op_count += 1 return rank_ops class ScheduleVPP(PipelineScheduleMulti): """ The VPP schedule. See https://arxiv.org/pdf/2104.04473 for details. Will perform one forward and one backward on the microbatches in steady state and supports multiple stages per rank. When microbatches are ready for multiple local stages, VPP prioritizes the earlier microbatch (also called "depth first"). This schedule is mostly similar to the original paper. It differs by being relaxing the requirement of num_microbatch % pp_size == 0. Using the flex_pp schedule, we will have num_rounds = max(1, n_microbatches // pp_group_size) and it works as long as n_microbatches % num_rounds is 0. As a few examples, support 1. pp_group_size = 4, n_microbatches = 10. We will have num_rounds = 2 and n_microbatches % 2 is 0. 2. pp_group_size = 4, n_microbatches = 3. We will have num_rounds = 1 and n_microbatches % 1 is 0. """ def __init__( self, stages: list[_PipelineStageBase], n_microbatches: int, loss_fn: Callable | None = None, args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None, kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None, output_merge_spec: dict[str, Any] | tuple[Any] | None = None, ): self.pp_group_size = stages[0].group_size super().__init__( stages=stages, n_microbatches=n_microbatches, loss_fn=loss_fn, args_chunk_spec=args_chunk_spec, kwargs_chunk_spec=kwargs_chunk_spec, output_merge_spec=output_merge_spec, ) self.n_local_stages = len(stages) self.rank = stages[0].group_rank self.number_of_rounds = max(1, n_microbatches // self.pp_group_size) self.microbatches_per_round = n_microbatches // self.number_of_rounds if n_microbatches % self.number_of_rounds != 0: raise ValueError( "VPP requires the number of microbatches to be a " f"multiple of the number of rounds ({self.number_of_rounds}), " f"but got {n_microbatches}." ) # 1. Create the pipeline_order (all ranks do this calculation) # This will be used to keep track of the current state of the entire pipeline # pipeline_order[rank] = [Action(computation_type, microbatch_index, stage_index), ...] self.pipeline_order: dict[int, list[_Action | None]] = {} for rank in range(self.pp_group_size): rank_ops = self._calculate_single_rank_operations(rank) self.pipeline_order[rank] = rank_ops def _calculate_single_rank_operations(self, rank) -> list[_Action | None]: def get_rank_warmup_ops(rank): # Warms up operations for last stage warmups_ops_last_stage = ( self.n_local_stages - 1 ) * self.microbatches_per_round # Increment warmup operations by 2 for each hop away from the last stage multiply_factor = 2 warmup_ops = warmups_ops_last_stage + multiply_factor * ( (self.pp_group_size - 1) - rank ) # We cannot have more warmup operations than there are number of microbatches, so cap it there return min(warmup_ops, self._n_microbatches * self.n_local_stages) warmup_ops = get_rank_warmup_ops(rank) microbatch_ops = self.n_local_stages * self._n_microbatches # fwd_bwd_ops should encompass the remaining forwards fwd_bwd_ops = microbatch_ops - warmup_ops # cooldown_ops should encompass the remaining backwards cooldown_ops = microbatch_ops - fwd_bwd_ops # total ops encompass both forward and backward ops total_ops = warmup_ops + fwd_bwd_ops + cooldown_ops # warmup_ops + fwd_bwd_ops * 2 + cooldown_ops == microbatch_ops * 2 logger.debug( "rank %s, warmup_ops %s, 1f1b %s, cooldown_ops %s total_ops %s", rank, warmup_ops, fwd_bwd_ops, cooldown_ops, total_ops, ) # Calculates the stage index based on step and pp_group_size def forward_stage_index(step): # Get the local index from 0 to n_local_stages-1 local_index = ( step // self.microbatches_per_round ) % self.n_local_stages return (local_index * self.pp_group_size) + rank def backward_stage_index(step): local_index = ( self.n_local_stages - 1 - ((step - warmup_ops) // self.microbatches_per_round) % self.n_local_stages ) return (local_index * self.pp_group_size) + rank return _get_1f1b_rank_ops( self.n_local_stages, self.pp_group_size, warmup_ops, fwd_bwd_ops, cooldown_ops, rank, forward_stage_index, backward_stage_index, )