wehub-resource-sync
1384 lines
52 KiB
Python
1384 lines
52 KiB
Python
# 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,
|
|
NamedTuple,
|
|
)
|
|
|
|
from paddle import nn
|
|
from paddle.distributed.auto_parallel.pipelining.stage import PipelineStage
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
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,
|
|
)
|