# Source code for tensorplay.distributed.pipelining.schedules

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/distributed/pipelining/schedules.html

```
"""Microbatch pipeline schedules."""

import csv
import logging
import re
from collections import Counter, defaultdict
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Literal, Protocol, cast

import tensorplay as tp

from .. import distributed_core as dist
from ..fsdp._fully_shard import FSDPModule, UnshardHandle
from ._utils import InferenceMode, generate_rank_to_stage_mapping, generate_stage_to_rank_mapping
from .microbatch import TensorChunkSpec, merge_chunks, split_args_kwargs_into_chunks, _split_tensor

logger = logging.getLogger(__name__)

__all__ = [
    "get_schedule_class",
    "PipelineScheduleSingle",
    "PipelineScheduleMulti",
    "Schedule1F1B",
    "ScheduleGPipe",
    "ScheduleInterleaved1F1B",
    "ScheduleLoopedBFS",
    "ScheduleInterleavedZeroBubble",
    "ScheduleZBVZeroBubble",
    "ScheduleDualPipeV",
]

class _ComputationType(str, Enum):
    FORWARD = "F"
    BACKWARD_INPUT = "I"
    BACKWARD_WEIGHT = "W"
    UNSHARD = "UNSHARD"
    RESHARD = "RESHARD"
    SEND_F = "SEND_F"
    RECV_F = "RECV_F"
    SEND_B = "SEND_B"
    RECV_B = "RECV_B"
    FULL_BACKWARD = "B"
    OVERLAP_F_B = "OVERLAP_F_B"
    REDUCE_GRAD = "REDUCE_GRAD"

    @staticmethod
    def from_str(action: str) -> "_ComputationType":
        return _ComputationType(action)

class _Action(tuple):
    __slots__ = ()

    def __new__(cls, stage_index: int, computation_type: _ComputationType, microbatch_index: int | None = None, sub_actions: tuple["_Action", ...] | None = None):
        return tuple.__new__(cls, (stage_index, computation_type, microbatch_index, sub_actions))

    @property
    def stage_index(self) -> int:
        return self[0]

    @property
    def computation_type(self) -> _ComputationType:
        return self[1]

    @property
    def microbatch_index(self) -> int | None:
        return self[2]

    @property
    def sub_actions(self) -> tuple["_Action", ...] | None:
        return self[3]

    @property
    def is_compute_op(self) -> bool:
        return self.computation_type in {_ComputationType.FORWARD, _ComputationType.BACKWARD_INPUT, _ComputationType.BACKWARD_WEIGHT, _ComputationType.FULL_BACKWARD, _ComputationType.OVERLAP_F_B}

    def __repr__(self) -> str:
        if self.sub_actions is not None:
            return f"({';'.join(map(repr, self.sub_actions))}){self.computation_type.value}"
        return f"{self.stage_index}{self.computation_type.value}{'' if self.microbatch_index is None else self.microbatch_index}"

    def __str__(self) -> str:
        return self.__repr__()

    @staticmethod
    def from_str(action_string: str) -> "_Action | None":
        action_string = action_string.strip()
        if not action_string:
            return None
        if action_string.startswith("(") and ")" in action_string:
            end = action_string.index(")")
            sub = tuple(item for item in (_Action.from_str(part) for part in action_string[1:end].split(";")) if item is not None)
            return _Action(-1, _ComputationType.from_str(action_string[end + 1:]), None, sub)
        match = re.fullmatch(r"(\d+)(F|I|B|W|UNSHARD|RESHARD|REDUCE_GRAD|SEND_F|RECV_F|SEND_B|RECV_B)(\d*)", action_string)
        if match is None:
            raise ValueError(f"invalid pipeline action: {action_string}")
        stage, kind, microbatch = match.groups()
        return _Action(int(stage), _ComputationType(kind), int(microbatch) if microbatch else None)

def _get_profiler_function_name(action: _Action) -> str:
    return f"TP:{action}"

def _format_pipeline_order(pipeline_order: dict[int, list[_Action | None]], error_step_number: int | None = None) -> str:
    steps = max((len(actions) for actions in pipeline_order.values()), default=0)
    rows = ["step " + " ".join(f"rank {rank}" for rank in sorted(pipeline_order))]
    for index in range(steps):
        values = [str(pipeline_order.get(rank, [None] * steps)[index] or "") for rank in sorted(pipeline_order)]
        suffix = " <error>" if index == error_step_number else ""
        rows.append(f"{index}: " + " ".join(values) + suffix)
    return "\n".join(rows)

class _PipelineSchedule:
    def __init__(self, n_microbatches: int, loss_fn: Any = None, args_chunk_spec: Any = None, kwargs_chunk_spec: Any = None, output_merge_spec: Any = None, scale_grads: bool = True) -> None:
        if n_microbatches <= 0:
            raise ValueError("n_microbatches must be positive")
        self._n_microbatches = int(n_microbatches)
        self._loss_fn = loss_fn
        self._args_chunk_spec = args_chunk_spec
        self._kwargs_chunk_spec = kwargs_chunk_spec
        self._output_merge_spec = output_merge_spec
        self._scale_grads = scale_grads
        self._has_backward = loss_fn is not None
        self._stages: list[Any] = []

    def _maybe_compute_loss(self, stage: Any, output: Any, target_mbs: Any, mb_index: int, loss_kwargs: dict[str, Any] | None) -> Any:
        if not getattr(stage, "is_last", False) or self._loss_fn is None or target_mbs is None:
            return None
        loss = self._loss_fn(output, target_mbs[mb_index], **(loss_kwargs or {}))
        self._losses.append(loss)
        return loss

    def _maybe_get_loss(self, stage: Any, mb_index: int) -> Any:
        if not getattr(stage, "is_last", False):
            return None
        return self._losses[mb_index] if 0 <= mb_index < len(self._losses) else None

    def _update_losses(self, stages: Any, losses: list[Any] | None) -> None:
        stage_list = stages if isinstance(stages, (list, tuple)) else [stages]
        if losses is not None and any(getattr(stage, "is_last", False) for stage in stage_list):
            if len(self._losses) != self._n_microbatches:
                raise RuntimeError(
                    f"expected {self._n_microbatches} losses, got {len(self._losses)}"
                )
            losses.clear()
            losses.extend(self._losses)
        self._losses.clear()

    def _warmup_p2p(self, stages: Any, has_backward: bool, p2p_done: Any) -> None:
        del p2p_done
        pipeline_stages = [
            stage for stage in stages if hasattr(stage, "_user_meta")
        ]
        if len(pipeline_stages) != len(stages):
            if not dist.is_initialized():
                return
            operations = [
                operation
                for stage in stages
                for operation in stage._get_init_p2p_neighbors_ops()
            ]
            _wait_batch_p2p(_batch_p2p(operations))
            return
        if not dist.is_initialized():
            for stage in pipeline_stages:
                stage._inference_mode = (
                    InferenceMode.DYNAMIC
                    if InferenceMode.needs_dynamic(stage._user_meta, has_backward)
                    else InferenceMode.STATIC
                )
            return
        has_cross_rank = any(
            (not stage.is_first and not stage._is_same_rank(stage.stage_index - 1))
            or (not stage.is_last and not stage._is_same_rank(stage.stage_index + 1))
            for stage in pipeline_stages
        )
        if has_cross_rank and any(
            dist.get_backend(stage.group) == "fake" for stage in pipeline_stages
        ):
            for stage in pipeline_stages:
                if InferenceMode.needs_dynamic(stage._user_meta, has_backward):
                    raise RuntimeError(
                        f"Stage {stage.stage_index} requires dynamic metadata with a fake process group"
                    )
                stage._inference_mode = InferenceMode.STATIC
            return
        accumulated = None
        for stage in pipeline_stages:
            accumulated = stage._warmup_forward_vote(
                has_backward,
                received_acc=accumulated,
            )
        result = accumulated
        for stage in reversed(pipeline_stages):
            result = stage._warmup_backward_result(received_result=result)
            stage._inference_mode = (
                InferenceMode.STATIC if int(result.item()) == 1 else InferenceMode.DYNAMIC
            )

    def _initialize_pp_stages(
        self,
        stages: list[Any],
        args: Any,
        kwargs: Any,
        target: Any,
        fwd_initialized: Any,
        bwd_initialized: Any,
        loss_kwargs: Any,
    ) -> tuple[bool, bool]:
        if fwd_initialized and self._has_backward != bwd_initialized:
            fwd_initialized = False
            bwd_initialized = False
        if not fwd_initialized:
            self._warmup_p2p(stages, self._has_backward, fwd_initialized)
            for stage in stages:
                stage.has_backward = self._has_backward
                backup = getattr(stage, "_pre_metadata_inference_backup", None)
                if callable(backup):
                    backup()
            try:
                next_stage_args = None
                for stage in stages:
                    stage_args = args if stage.is_first else next_stage_args
                    next_stage_args = stage._prepare_forward_infra(
                        self._n_microbatches,
                        stage_args,
                        kwargs,
                        self._has_backward,
                    )
                    fwd_initialized = True
                if self._has_backward and not bwd_initialized:
                    previous_grad_meta = None
                    for stage in reversed(stages):
                        previous_grad_meta = stage._prepare_backward_infra(
                            self._n_microbatches,
                            loss_fn=self._loss_fn,
                            target=target,
                            received_grad_meta=previous_grad_meta,
                            loss_kwargs=loss_kwargs,
                        )
                    bwd_initialized = True
            finally:
                for stage in stages:
                    cleanup = getattr(stage, "_post_metadata_inference_cleanup", None)
                    if callable(cleanup):
                        cleanup()
        elif self._has_backward and not bwd_initialized:
            previous_grad_meta = None
            for stage in reversed(stages):
                previous_grad_meta = stage._prepare_backward_infra(
                    self._n_microbatches,
                    loss_fn=self._loss_fn,
                    target=target,
                    received_grad_meta=previous_grad_meta,
                    loss_kwargs=loss_kwargs,
                )
            bwd_initialized = True
        return fwd_initialized, bwd_initialized

    def _step_microbatches(self, arg_mbs: list[tuple[Any, ...]], kwarg_mbs: list[dict[str, Any]], target_mbs: list[Any] | None, losses: list[Any] | None, return_outputs: bool = True, loss_kwargs: dict[str, Any] | None = None) -> Any:
        raise NotImplementedError

    def step(self, *args: Any, target: Any = None, losses: list[Any] | None = None, return_outputs: bool = True, loss_kwargs: dict[str, Any] | None = None, arg_mbs: Any = None, kwarg_mbs: Any = None, target_mbs: Any = None, **kwargs: Any) -> Any:
        arg_mbs, kwarg_mbs, target_mbs = self._get_microbatch_inputs(args, kwargs, target, arg_mbs, kwarg_mbs, target_mbs)
        self._initialize_for_step(args, kwargs, arg_mbs, kwarg_mbs, target, loss_kwargs)
        self._losses = []
        return self._step_microbatches(arg_mbs or [], kwarg_mbs or [], target_mbs, losses, return_outputs, loss_kwargs)

    def _initialize_for_step(
        self,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
        arg_mbs: list[Any],
        kwarg_mbs: list[Any],
        target: Any,
        loss_kwargs: Any,
    ) -> None:
        init_args = tuple(arg_mbs[0]) if arg_mbs else args
        init_kwargs = dict(kwarg_mbs[0]) if kwarg_mbs else kwargs
        if isinstance(self, PipelineScheduleSingle):
            self._initialize_stage(init_args, init_kwargs, target, loss_kwargs)
        elif self._stages:
            self._initialize_stages(init_args, init_kwargs, target, loss_kwargs)

    def eval(self, *args: Any, target: Any = None, losses: list[Any] | None = None, arg_mbs: Any = None, kwarg_mbs: Any = None, target_mbs: Any = None, **kwargs: Any) -> Any:
        old = self._has_backward
        self._has_backward = False
        try:
            return self.step(*args, target=target, losses=losses, arg_mbs=arg_mbs, kwarg_mbs=kwarg_mbs, target_mbs=target_mbs, **kwargs)
        finally:
            self._has_backward = old

    def _check_inputs(
        self,
        arg_mbs: Any = None,
        kwarg_mbs: Any = None,
        target_mbs: Any = None,
        losses: Any = None,
    ) -> tuple[list[Any], list[Any]]:
        def check_type_and_len(value: Any, name: str) -> None:
            if not isinstance(value, list):
                raise TypeError(f"{name} must be a list but got a {type(value)}")
            if len(value) != self._n_microbatches:
                raise ValueError(
                    f"Expecting {self._n_microbatches} {name} but got {len(value)}"
                )

        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 and 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: Any, target: Any, loss_kwargs: dict[str, Any] | None = None) -> Any:
        return self._loss_fn(output, target, **(loss_kwargs or {}))

    def _split_inputs(
        self, args: tuple[Any, ...], kwargs: dict[str, Any]
    ) -> tuple[list[tuple[Any, ...]], list[dict[str, Any]]]:
        if args or kwargs:
            return split_args_kwargs_into_chunks(
                args,
                kwargs,
                self._n_microbatches,
                self._args_chunk_spec,
                self._kwargs_chunk_spec,
            )
        return [()] * self._n_microbatches, [{}] * self._n_microbatches

    def _get_microbatch_inputs(
        self,
        args: tuple[Any, ...],
        kwargs: dict[str, Any],
        target: Any,
        arg_mbs: Any,
        kwarg_mbs: Any,
        target_mbs: Any,
    ) -> tuple[list[Any], list[Any], list[Any] | None]:
        pre_split = any(
            value is not None for value in (arg_mbs, kwarg_mbs, target_mbs)
        )
        if not pre_split:
            args_split, kwargs_split = self._split_inputs(args, kwargs)
            target_split = (
                list(_split_tensor(target, TensorChunkSpec(0), self._n_microbatches))
                if target is not None
                else None
            )
            return args_split, kwargs_split, target_split

        if args:
            raise ValueError(
                "When using pre-split inputs, pass pre-split positional inputs "
                "through arg_mbs=... instead of positional args."
            )
        if kwargs:
            names = ", ".join(sorted(kwargs))
            raise ValueError(
                f"Unexpected keyword arguments with pre-split inputs: {names}. "
                "Pass pre-split keyword inputs through kwarg_mbs=."
            )
        if target is not None:
            raise ValueError(
                "When using pre-split inputs, pass pre-split targets through "
                "target_mbs=... instead of target=."
            )

        arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs)
        for mb_index, (arg_mb, kwarg_mb) in enumerate(
            zip(arg_mbs, kwarg_mbs, strict=True)
        ):
            if not isinstance(arg_mb, tuple):
                raise TypeError(
                    "arg_mbs must be a list of tuples, but "
                    f"arg_mbs[{mb_index}] is a {type(arg_mb)}"
                )
            if not isinstance(kwarg_mb, dict):
                raise TypeError(
                    "kwarg_mbs must be a list of dicts, but "
                    f"kwarg_mbs[{mb_index}] is a {type(kwarg_mb)}"
                )
        return arg_mbs, kwarg_mbs, target_mbs

    def _merge_outputs(self, output_chunks: list[Any]) -> Any:
        if self._output_merge_spec is None:
            output_spec = _default_merge_spec(output_chunks[0])
        else:
            output_spec = self._output_merge_spec
        return merge_chunks(output_chunks, output_spec)

[docs]
class PipelineScheduleSingle(_PipelineSchedule):
    def __init__(self, stage: Any, n_microbatches: int, loss_fn: Any = None, args_chunk_spec: Any = None, kwargs_chunk_spec: Any = None, output_merge_spec: Any = None, scale_grads: bool = True) -> None:
        super().__init__(n_microbatches, loss_fn, args_chunk_spec, kwargs_chunk_spec, output_merge_spec, scale_grads)
        self._stage = stage
        self._stages = [stage]
        self._num_stages = int(stage.num_stages)
        self._stage_forward_initialized = False
        self._stage_backward_initialized = False
        self.pipeline_order = self._get_pipeline_order()
        self._stage.has_backward = self._has_backward

    def _initialize_stage(self, args: Any, kwargs: Any, target: Any = None, loss_kwargs: Any = None) -> None:
        (
            self._stage_forward_initialized,
            self._stage_backward_initialized,
        ) = self._initialize_pp_stages(
            [self._stage],
            args,
            kwargs,
            target,
            self._stage_forward_initialized,
            self._stage_backward_initialized,
            loss_kwargs,
        )

    def step(
        self,
        *args: Any,
        target: Any = None,
        losses: list[Any] | None = None,
        return_outputs: bool = True,
        loss_kwargs: dict[str, Any] | None = None,
        arg_mbs: Any = None,
        kwarg_mbs: Any = None,
        target_mbs: Any = None,
        **kwargs: Any,
    ) -> Any:
        if self._has_backward and not tp.is_grad_enabled():
            raise RuntimeError(
                "step() requires gradients to be enabled for backward computation"
            )
        self._stage.has_backward = self._has_backward
        self._stage.clear_runtime_states()
        args_split, kwargs_split, targets_split = self._get_microbatch_inputs(
            args,
            kwargs,
            target,
            arg_mbs,
            kwarg_mbs,
            target_mbs,
        )
        self._losses = []
        self._initialize_stage(
            tuple(args_split[0]) if args_split else args,
            dict(kwargs_split[0]) if kwargs_split else kwargs,
            targets_split[0] if targets_split else None,
            loss_kwargs,
        )
        self._step_microbatches(
            args_split,
            kwargs_split,
            targets_split,
            losses,
            return_outputs,
            loss_kwargs=loss_kwargs,
        )
        if self._stage.is_last and return_outputs and self._stage.output_chunks:
            return self._merge_outputs(self._stage.output_chunks)
        return None

    def _step_microbatches(self, arg_mbs: list[tuple[Any, ...]], kwarg_mbs: list[dict[str, Any]], target_mbs: list[Any] | None, losses: list[Any] | None, return_outputs: bool = True, loss_kwargs: dict[str, Any] | None = None) -> Any:
        self._stage.clear_runtime_states()
        outputs = []
        forward_sends = []
        for index, (args, kwargs) in enumerate(zip(arg_mbs, kwarg_mbs)):
            for work in _run_p2p(self._stage.get_fwd_recv_ops(index)):
                work.wait()
            output = self._stage.forward_one_chunk(
                index,
                args,
                kwargs,
                save_forward_output=return_outputs,
            )
            outputs.append(output)
            forward_sends.extend(_run_p2p(self._stage.get_fwd_send_ops(index)))
        for work in forward_sends:
            work.wait()

        if self._has_backward:
            backward_sends = []
            for index in range(len(outputs)):
                for work in _run_p2p(self._stage.get_bwd_recv_ops(index)):
                    work.wait()
                loss = self._maybe_compute_loss(self._stage, outputs[index], target_mbs, index, loss_kwargs)
                if self._stage.is_last:
                    if loss is not None:
                        self._stage.backward_one_chunk(index, loss=loss)
                else:
                    self._stage.backward_one_chunk(
                        index,
                        last_backward=index == len(outputs) - 1,
                    )
                backward_sends.extend(_run_p2p(self._stage.get_bwd_send_ops(index)))
            for work in backward_sends:
                work.wait()
            if self._scale_grads:
                self._stage.scale_grads(self._n_microbatches)
        if losses is not None and self._stage.is_last:
            losses.extend(self._losses)
        if not return_outputs or not self._stage.is_last:
            return None
        return self._merge_outputs(outputs)

    def _get_pipeline_order(self) -> dict[int, list[_Action]]:
        actions: list[_Action | None] = [
            _Action(self._stage.stage_index, _ComputationType.FORWARD, index)
            for index in range(self._n_microbatches)
        ]
        if self._has_backward:
            actions.extend(
                _Action(self._stage.stage_index, _ComputationType.FULL_BACKWARD, index)
                for index in range(self._n_microbatches)
            )
            actions = _add_reduce_grad(actions, self._n_microbatches)
        return {int(self._stage.group_rank): actions}

def _batch_p2p(operations: list[Any], desc: str | None = None) -> list[Any]:
    del desc
    if not operations:
        return []
    operations_by_group: dict[str, list[Any]] = defaultdict(list)
    for operation in operations:
        group = operation.group
        group_name = getattr(group, "group_name", None)
        operations_by_group[str(group_name if group_name is not None else group)].append(
            operation
        )
    if len(operations_by_group) > 1:
        works: list[Any] = []
        for _, group_operations in sorted(operations_by_group.items()):
            works.extend(_batch_p2p(group_operations))
        return works

    operation_types = {operation.op for operation in operations}
    if operation_types == {dist.isend}:
        return [
            work
            for work in (
                operation.op(
                    operation.tensor,
                    group=operation.group,
                    tag=operation.tag,
                    group_dst=operation.group_peer,
                )
                for operation in operations
            )
            if work is not None
        ]
    if operation_types == {dist.irecv}:
        return [
            work
            for work in (
                operation.op(
                    operation.tensor,
                    group=operation.group,
                    tag=operation.tag,
                    group_src=operation.group_peer,
                )
                for operation in operations
            )
            if work is not None
        ]
    return dist.batch_isend_irecv(operations)

def _sorted_batch_p2p(
    operations: list[Any], desc: str | None = None
) -> dict[int, list[Any]]:
    del desc
    operations_by_peer: dict[int, list[Any]] = defaultdict(list)
    works_by_peer: dict[int, list[Any]] = {}
    for operation in operations:
        operations_by_peer[int(operation.peer)].append(operation)
    for peer, peer_operations in sorted(operations_by_peer.items()):
        works_by_peer[peer] = _batch_p2p(peer_operations)
    return works_by_peer

def _wait_batch_p2p(works: list[Any]) -> None:
    for work in works:
        work.wait()

def _run_p2p(operations: list[Any]) -> list[Any]:
    return [
        work
        for peer_works in _sorted_batch_p2p(operations).values()
        for work in peer_works
    ]

def _normalize_stage_args(value: Any) -> tuple[Any, ...]:
    if isinstance(value, tuple):
        return value
    if isinstance(value, list):
        return tuple(value)
    return (value,)

class _ScheduleForwardOnly(PipelineScheduleSingle):
    def _step_microbatches(self, *args: Any, **kwargs: Any) -> Any:
        arg_mbs = kwargs.pop("arg_mbs", args[0] if args else None)
        kwarg_mbs = kwargs.pop("kwarg_mbs", args[1] if len(args) > 1 else None)
        target_mbs = kwargs.pop("target_mbs", args[2] if len(args) > 2 else None)
        losses = kwargs.pop("losses", args[3] if len(args) > 3 else None)
        return_outputs = kwargs.pop(
            "return_outputs", args[4] if len(args) > 4 else True
        )
        if target_mbs is not None or losses is not None:
            raise RuntimeError("forward-only schedule does not support loss computation")
        arg_mbs, kwarg_mbs = self._check_inputs(
            arg_mbs, kwarg_mbs, target_mbs, losses
        )
        self._initialize_stage(arg_mbs[0], kwarg_mbs[0])
        self._stage.clear_runtime_states()
        send_works: list[Any] = []
        for index in range(self._n_microbatches):
            _wait_batch_p2p(
                _batch_p2p(
                    self._stage.get_fwd_recv_ops(index), desc="fwd_recv"
                )
            )
            self._stage.forward_one_chunk(
                index,
                arg_mbs[index],
                kwarg_mbs[index],
                save_forward_output=return_outputs,
            )
            send_works.extend(
                _batch_p2p(
                    self._stage.get_fwd_send_ops(index), desc="fwd_send"
                )
            )
        _wait_batch_p2p(send_works)
        if not return_outputs or not self._stage.is_last:
            return None
        return self._merge_outputs(self._stage.output_chunks)

[docs]
class ScheduleGPipe(PipelineScheduleSingle):
    """Execute all forward microbatches before draining their backwards."""

    def _step_microbatches(
        self,
        arg_mbs: list[tuple[Any, ...]],
        kwarg_mbs: list[dict[str, Any]],
        target_mbs: list[Any] | None,
        losses: list[Any] | None,
        return_outputs: bool = True,
        loss_kwargs: dict[str, Any] | None = None,
    ) -> Any:
        arg_mbs, kwarg_mbs = self._check_inputs(
            arg_mbs, kwarg_mbs, target_mbs, losses
        )
        outputs: list[Any] = []
        forward_sends: list[Any] = []
        for index in range(self._n_microbatches):
            _wait_batch_p2p(
                [
                    work
                    for works in _sorted_batch_p2p(
                        self._stage.get_fwd_recv_ops(index), desc="fwd_recv"
                    ).values()
                    for work in works
                ]
            )
            output = self._stage.forward_one_chunk(
                index,
                arg_mbs[index],
                kwarg_mbs[index],
                save_forward_output=return_outputs,
            )
            outputs.append(output)
            forward_sends.extend(
                work
                for works in _sorted_batch_p2p(
                    self._stage.get_fwd_send_ops(index), desc="fwd_send"
                ).values()
                for work in works
            )
            self._maybe_compute_loss(
                self._stage, output, target_mbs, index, loss_kwargs
            )

        _wait_batch_p2p(forward_sends)

        backward_sends: list[Any] = []
        if self._has_backward:
            for index in range(self._n_microbatches):
                _wait_batch_p2p(
                    [
                        work
                        for works in _sorted_batch_p2p(
                            self._stage.get_bwd_recv_ops(index), desc="bwd_recv"
                        ).values()
                        for work in works
                    ]
                )
                self._stage.backward_one_chunk(
                    index,
                    loss=self._maybe_get_loss(self._stage, index),
                    last_backward=index == self._n_microbatches - 1,
                )
                backward_sends.extend(
                    work
                    for works in _sorted_batch_p2p(
                        self._stage.get_bwd_send_ops(index), desc="bwd_send"
                    ).values()
                    for work in works
                )
            _wait_batch_p2p(backward_sends)
            self._stage.perform_reduce_grad(
                self._n_microbatches if self._scale_grads else 1
            )

        self._update_losses(self._stage, losses)
        if not return_outputs or not self._stage.is_last:
            return None
        return self._merge_outputs(outputs)

    def _get_pipeline_order(self) -> dict[int, list[_Action | None]]:
        group_size = int(self._stage.group_size)
        if group_size != self._num_stages:
            raise ValueError("GPipe requires one stage per pipeline rank")
        pipeline_order: dict[int, list[_Action | None]] = {}
        for rank in range(group_size):
            actions: list[_Action | None] = [None] * rank
            actions.extend(
                _Action(rank, _ComputationType.FORWARD, microbatch)
                for microbatch in range(self._n_microbatches)
            )
            if self._has_backward:
                actions.extend(
                    [None] * (3 * (group_size - 1 - rank))
                )
                actions.extend(
                    _Action(rank, _ComputationType.FULL_BACKWARD, microbatch)
                    for microbatch in range(self._n_microbatches)
                )
                pipeline_order[rank] = _add_reduce_grad(
                    actions, self._n_microbatches
                )
            else:
                pipeline_order[rank] = actions
        return pipeline_order

[docs]
class Schedule1F1B(PipelineScheduleSingle):
    """Overlap forward and backward microbatches after warmup."""

    def __init__(self, stage: Any, n_microbatches: int, loss_fn: Any = None, args_chunk_spec: Any = None, kwargs_chunk_spec: Any = None, output_merge_spec: Any = None, scale_grads: bool = True) -> None:
        super().__init__(
            stage,
            n_microbatches,
            loss_fn,
            args_chunk_spec,
            kwargs_chunk_spec,
            output_merge_spec,
            scale_grads,
        )
        if self._has_backward and n_microbatches < self._num_stages:
            raise ValueError("1F1B requires at least one microbatch per stage")

    def _step_microbatches(
        self,
        arg_mbs: list[tuple[Any, ...]],
        kwarg_mbs: list[dict[str, Any]],
        target_mbs: list[Any] | None,
        losses: list[Any] | None,
        return_outputs: bool = True,
        loss_kwargs: dict[str, Any] | None = None,
    ) -> Any:
        arg_mbs, kwarg_mbs = self._check_inputs(
            arg_mbs, kwarg_mbs, target_mbs, losses
        )
        warmup = min(self._n_microbatches, self._num_stages - self._stage.stage_index)
        outputs: list[Any] = []
        fwd_mb_index = 0
        bwd_mb_index = 0
        send_work: list[Any] = []
        fwd_sends: list[Any] = []

        for _ in range(warmup):
            _wait_batch_p2p(
                _batch_p2p(
                    self._stage.get_fwd_recv_ops(fwd_mb_index), desc="fwd_recv"
                )
            )
            output = self._stage.forward_one_chunk(
                fwd_mb_index,
                arg_mbs[fwd_mb_index],
                kwarg_mbs[fwd_mb_index],
                save_forward_output=return_outputs,
            )
            outputs.append(output)
            _wait_batch_p2p(send_work)
            fwd_sends = self._stage.get_fwd_send_ops(fwd_mb_index)
            if not self._has_backward or fwd_mb_index != warmup - 1:
                send_work = _batch_p2p(fwd_sends, desc="fwd_send")
            self._maybe_compute_loss(
                self._stage, output, target_mbs, fwd_mb_index, loss_kwargs
            )
            fwd_mb_index += 1

        if not self._has_backward:
            for fwd_mb_index in range(fwd_mb_index, self._n_microbatches):
                _wait_batch_p2p(
                    _batch_p2p(
                        self._stage.get_fwd_recv_ops(fwd_mb_index),
                        desc="fwd_recv",
                    )
                )
                output = self._stage.forward_one_chunk(
                    fwd_mb_index,
                    arg_mbs[fwd_mb_index],
                    kwarg_mbs[fwd_mb_index],
                    save_forward_output=return_outputs,
                )
                outputs.append(output)
                _wait_batch_p2p(send_work)
                send_work = _batch_p2p(
                    self._stage.get_fwd_send_ops(fwd_mb_index), desc="fwd_send"
                )
                self._maybe_compute_loss(
                    self._stage, output, target_mbs, fwd_mb_index, loss_kwargs
                )
            _wait_batch_p2p(send_work)
        else:
            while True:
                _wait_batch_p2p(
                    _batch_p2p(
                        fwd_sends + self._stage.get_bwd_recv_ops(bwd_mb_index),
                        desc="fwd_send_bwd_recv",
                    )
                )
                self._stage.backward_one_chunk(
                    bwd_mb_index,
                    loss=self._maybe_get_loss(self._stage, bwd_mb_index),
                    last_backward=bwd_mb_index == self._n_microbatches - 1,
                )
                bwd_sends = self._stage.get_bwd_send_ops(bwd_mb_index)
                bwd_mb_index += 1
                if fwd_mb_index == self._n_microbatches:
                    break

                _wait_batch_p2p(
                    _batch_p2p(
                        bwd_sends + self._stage.get_fwd_recv_ops(fwd_mb_index),
                        desc="bwd_send_fwd_recv",
                    )
                )
                output = self._stage.forward_one_chunk(
                    fwd_mb_index,
                    arg_mbs[fwd_mb_index],
                    kwarg_mbs[fwd_mb_index],
                    save_forward_output=return_outputs,
                )
                outputs.append(output)
                self._maybe_compute_loss(
                    self._stage, output, target_mbs, fwd_mb_index, loss_kwargs
                )
                fwd_sends = self._stage.get_fwd_send_ops(fwd_mb_index)
                fwd_mb_index += 1

            send_work = _batch_p2p(bwd_sends, desc="bwd_send")
            while bwd_mb_index < self._n_microbatches:
                _wait_batch_p2p(
                    _batch_p2p(
                        self._stage.get_bwd_recv_ops(bwd_mb_index),
                        desc="bwd_recv",
                    )
                )
                self._stage.backward_one_chunk(
                    bwd_mb_index,
                    loss=self._maybe_get_loss(self._stage, bwd_mb_index),
                    last_backward=bwd_mb_index == self._n_microbatches - 1,
                )
                _wait_batch_p2p(send_work)
                send_work = _batch_p2p(
                    self._stage.get_bwd_send_ops(bwd_mb_index), desc="bwd_send"
                )
                bwd_mb_index += 1
            _wait_batch_p2p(send_work)

            self._stage.perform_reduce_grad(
                self._n_microbatches if self._scale_grads else 1
            )

        self._update_losses(self._stage, losses)
        if not return_outputs or not self._stage.is_last:
            return None
        return self._merge_outputs(outputs)

    def _get_pipeline_order(self) -> dict[int, list[_Action | None]]:
        group_size = int(self._stage.group_size)
        if group_size != self._num_stages:
            raise ValueError("1F1B requires one stage per pipeline rank")
        pipeline_order: dict[int, list[_Action | None]] = {}
        for rank in range(group_size):
            actions: list[_Action | None] = [None] * rank
            warmup = min(self._n_microbatches, group_size - 1 - rank)
            actions.extend(
                _Action(rank, _ComputationType.FORWARD, microbatch)
                for microbatch in range(warmup)
            )
            actions.extend([None] * (2 * (group_size - 1 - rank)))
            next_forward = warmup
            next_backward = 0
            while next_forward < self._n_microbatches:
                actions.append(
                    _Action(rank, _ComputationType.FORWARD, next_forward)
                )
                next_forward += 1
                actions.append(
                    _Action(rank, _ComputationType.FULL_BACKWARD, next_backward)
                )
                next_backward += 1
            while next_backward < self._n_microbatches:
                if rank != group_size - 1:
                    actions.append(None)
                actions.append(
                    _Action(rank, _ComputationType.FULL_BACKWARD, next_backward)
                )
                next_backward += 1
            pipeline_order[rank] = _add_reduce_grad(
                actions, self._n_microbatches
            )
        return pipeline_order

[docs]
class PipelineScheduleMulti(_PipelineSchedule):
    def __init__(self, stages: list[Any], n_microbatches: int, loss_fn: Any = None, args_chunk_spec: Any = None, kwargs_chunk_spec: Any = None, output_merge_spec: Any = None, use_full_backward: bool | None = None, scale_grads: bool = True, backward_requires_autograd: bool = True) -> None:
        if not stages:
            raise ValueError("at least one pipeline stage is required")
        super().__init__(n_microbatches, loss_fn, args_chunk_spec, kwargs_chunk_spec, output_merge_spec, scale_grads)
        self._stages = list(stages)
        self.use_full_backward = use_full_backward
        self.backward_requires_autograd = backward_requires_autograd
        self._backward_requires_autograd = backward_requires_autograd
        self._num_stages = int(stages[0].num_stages)
        self.pp_group_size = int(stages[0].group_size)
        if self._num_stages <= 0 or self.pp_group_size <= 0:
            raise ValueError("pipeline dimensions must be positive")
        if any(int(stage.num_stages) != self._num_stages for stage in stages):
            raise ValueError("all pipeline stages must use the same stage count")
        if len({int(stage.stage_index) for stage in stages}) != len(stages):
            raise ValueError("a pipeline stage cannot be listed more than once")
        self.rank = int(stages[0].group_rank)
        self.stage_index_to_group_rank = {
            index: index % self.pp_group_size for index in range(self._num_stages)
        }
        for stage in self._stages:
            stage.stage_index_to_group_rank = dict(self.stage_index_to_group_rank)
        self._stages_forward_initialized = False
        self._stages_backward_initialized = False
        self.pipeline_order: dict[int, list[_Action | None]] = {}
        if use_full_backward is not None:
            logger.warning(
                "use_full_backward is no longer supported; omit it from the schedule"
            )

    def _initialize_stages(self, args: Any, kwargs: Any, target: Any = None, loss_kwargs: Any = None) -> None:
        reinit_for_mode_switch = self._stages_forward_initialized and (
            self._has_backward != self._stages_backward_initialized
        )
        forward_initialized_before = self._stages_forward_initialized
        (
            self._stages_forward_initialized,
            self._stages_backward_initialized,
        ) = self._initialize_pp_stages(
            self._stages,
            args,
            kwargs,
            target,
            self._stages_forward_initialized,
            self._stages_backward_initialized,
            loss_kwargs,
        )
        if self._stages_forward_initialized and (
            not forward_initialized_before or reinit_for_mode_switch
        ):
            self._validate_adjacent_stage_communication()

    def step(
        self,
        *args: Any,
        target: Any = None,
        losses: list[Any] | None = None,
        return_outputs: bool = True,
        loss_kwargs: dict[str, Any] | None = None,
        arg_mbs: Any = None,
        kwarg_mbs: Any = None,
        target_mbs: Any = None,
        **kwargs: Any,
    ) -> Any:
        if (
            self._has_backward
            and self._backward_requires_autograd
            and not tp.is_grad_enabled()
        ):
            raise RuntimeError(
                "step() requires gradients to be enabled for backward computation"
            )
        for stage in self._stages:
            stage.has_backward = self._has_backward
            stage.clear_runtime_states()
        args_split, kwargs_split, targets_split = self._get_microbatch_inputs(
            args,
            kwargs,
            target,
            arg_mbs,
            kwarg_mbs,
            target_mbs,
        )
        self._losses = []
        self._initialize_stages(
            tuple(args_split[0]) if args_split else args,
            dict(kwargs_split[0]) if kwargs_split else kwargs,
            targets_split[0] if targets_split else None,
            loss_kwargs,
        )
        self._step_microbatches(
            args_split,
            kwargs_split,
            targets_split,
            losses,
            return_outputs,
            loss_kwargs=loss_kwargs,
        )
        if return_outputs:
            for stage in self._stages:
                if stage.is_last and stage.output_chunks:
                    return self._merge_outputs(stage.output_chunks)
        return None

    def _step_microbatches(self, arg_mbs: list[tuple[Any, ...]], kwarg_mbs: list[dict[str, Any]], target_mbs: list[Any] | None, losses: list[Any] | None, return_outputs: bool = True, loss_kwargs: dict[str, Any] | None = None) -> Any:
        if self._stages_are_local():
            return self._step_local_stages(
                arg_mbs,
                kwarg_mbs,
                target_mbs,
                losses,
                return_outputs,
                loss_kwargs,
            )
        return self._step_distributed_stages(
            arg_mbs,
            kwarg_mbs,
            target_mbs,
            losses,
            return_outputs,
            loss_kwargs,
        )

    def _stages_are_local(self) -> bool:
        if not self._stages:
            return True
        if not dist.is_initialized():
            return True
        if len(self._stages) == 1:
            stage = self._stages[0]
            return stage.num_stages <= 1 or stage.group_size <= 1
        first_rank = self._stages[0].group_rank
        return all(stage.group_rank == first_rank for stage in self._stages)

    def _step_local_stages(
        self,
        arg_mbs: list[tuple[Any, ...]],
        kwarg_mbs: list[dict[str, Any]],
        target_mbs: list[Any] | None,
        losses: list[Any] | None,
        return_outputs: bool,
        loss_kwargs: dict[str, Any] | None,
    ) -> Any:
        for stage in self._stages:
            stage.clear_runtime_states()
        outputs = []
        for index, (args, kwargs) in enumerate(zip(arg_mbs, kwarg_mbs)):
            value = self._stages[0].forward_one_chunk(index, args, kwargs)
            for stage in self._stages[1:]:
                stage.set_local_fwd_input(value, index)
                value = stage.forward_one_chunk(
                    index, _normalize_stage_args(value), {}
                )
            outputs.append(value)
            loss = self._maybe_compute_loss(self._stages[-1], value, target_mbs, index, loss_kwargs)
        if self._has_backward:
            for index in reversed(range(len(outputs))):
                loss = self._maybe_get_loss(self._stages[-1], index)
                next_grad = self._stages[-1].backward_one_chunk(index, loss=loss)
                for stage_index in range(len(self._stages) - 2, -1, -1):
                    stage = self._stages[stage_index]
                    stage.set_local_bwd_input(next_grad, index)
                    next_grad = stage.backward_one_chunk(index)
        if self._has_backward and self._scale_grads:
            for stage in self._stages:
                stage.scale_grads(self._n_microbatches)
        if losses is not None:
            losses.extend(self._losses)
        return self._merge_outputs(outputs) if return_outputs else None

    def _step_distributed_stages(
        self,
        arg_mbs: list[tuple[Any, ...]],
        kwarg_mbs: list[dict[str, Any]],
        target_mbs: list[Any] | None,
        losses: list[Any] | None,
        return_outputs: bool,
        loss_kwargs: dict[str, Any] | None,
    ) -> Any:
        for stage in self._stages:
            stage.clear_runtime_states()
        outputs: list[Any] = []
        send_works: list[Any] = []
        for index, (args, kwargs) in enumerate(zip(arg_mbs, kwarg_mbs)):
            for stage in self._stages:
                for work in _run_p2p(stage.get_fwd_recv_ops(index)):
                    work.wait()
                output = stage.forward_one_chunk(index, args, kwargs)
                if stage.is_last:
                    outputs.append(output)
                    loss = self._maybe_compute_loss(
                        stage, output, target_mbs, index, loss_kwargs
                    )
                send_works.extend(_run_p2p(stage.get_fwd_send_ops(index)))
        for work in send_works:
            work.wait()
        if self._has_backward:
            for index in reversed(range(self._n_microbatches)):
                for stage in reversed(self._stages):
                    for work in _run_p2p(stage.get_bwd_recv_ops(index)):
                        work.wait()
                    loss = self._maybe_get_loss(stage, index)
                    stage.backward_one_chunk(index, loss=loss)
                    for work in _run_p2p(stage.get_bwd_send_ops(index)):
                        work.wait()
            if self._scale_grads:
                for stage in self._stages:
                    stage.scale_grads(self._n_microbatches)
        if losses is not None:
            losses.extend(self._losses)
        if not return_outputs or not outputs:
            return None
        return self._merge_outputs(outputs)

    def _validate_adjacent_stage_communication(self) -> None:
        def check_stage_indices(
            stage_index: int,
            direction: str,
            actual: set[int],
            expected: set[int],
        ) -> None:
            non_adjacent = actual - expected
            if non_adjacent:
                raise RuntimeError(
                    f"stage {stage_index} has non-adjacent {direction} stages "
                    f"{sorted(non_adjacent)}; expected only {sorted(expected)}"
                )

        for stage in self._stages:
            stage_index = stage.stage_index
            forward_sources = {
                int(getattr(info, "source"))
                for info in stage.args_recv_info.get(0, ())
                if getattr(info, "source", None) is not None
            }
            expected_sources = set() if stage.is_first else {stage_index - 1}
            check_stage_indices(
                stage_index,
                "forward receive",
                forward_sources,
                expected_sources,
            )
            forward_destinations = {
                int(destination)
                for destinations in stage.act_send_info.values()
                for destination in destinations
                if destination is not None
            }
            expected_destinations = set() if stage.is_last else {stage_index + 1}
            check_stage_indices(
                stage_index,
                "forward send",
                forward_destinations,
                expected_destinations,
            )

    def _validate_and_set_stage_mapping(self, actions: Any) -> None:
        self.stage_index_to_group_rank = _validate_schedule(
            actions,
            self.pp_group_size,
            self._num_stages,
            self._n_microbatches,
        )
        for stage in self._stages:
            stage.stage_index_to_group_rank = dict(self.stage_index_to_group_rank)

    def _dump_csv(self, filename: str) -> None:
        with open(filename, "w", newline="", encoding="utf-8") as stream:
            writer = csv.writer(stream)
            for rank in sorted(self.pipeline_order):
                writer.writerow(self.pipeline_order[rank])

    def _load_csv(
        self,
        filename: str,
        format: Literal["compute_only", "compute_comms"] = "compute_only",
    ) -> dict[int, list[_Action | None]]:
        if format != "compute_only":
            raise AssertionError(f"format must be compute_only, got {format}")
        with open(filename, newline="", encoding="utf-8") as stream:
            actions = {
                rank: [_Action.from_str(value) for value in row]
                for rank, row in enumerate(csv.reader(stream))
            }
        self.pipeline_order = actions
        self._validate_and_set_stage_mapping(actions)
        return actions

    def _get_pipeline_order(self) -> dict[int, list[_Action | None]]:
        owned: dict[int, list[int]] = {rank: [] for rank in range(self.pp_group_size)}
        for stage_index in range(self._num_stages):
            rank = self.stage_index_to_group_rank[stage_index]
            owned.setdefault(rank, []).append(stage_index)

        pipeline_order: dict[int, list[_Action | None]] = {}
        for rank in range(self.pp_group_size):
            actions: list[_Action | None] = [None] * rank
            for stage_index in owned.get(rank, ()):
                actions.extend(
                    _Action(stage_index, _ComputationType.FORWARD, microbatch)
                    for microbatch in range(self._n_microbatches)
                )
            if self._has_backward:
                actions.extend([None] * (2 * max(0, self.pp_group_size - 1 - rank)))
                for stage_index in reversed(owned.get(rank, ())):
                    actions.extend(
                        _Action(stage_index, _ComputationType.FULL_BACKWARD, microbatch)
                        for microbatch in reversed(range(self._n_microbatches))
                    )
                actions = _add_reduce_grad(actions, self._n_microbatches)
            pipeline_order[rank] = actions
        return pipeline_order

@dataclass
class _PipelineContext:
    schedule_ref: _PipelineSchedule
    arg_mbs: list[tuple[Any, ...]] | None = None
    kwarg_mbs: list[dict[str, Any]] | None = None
    target_mbs: list[Any] | None = None
    losses: list[Any] | None = None

class _CustomFunctionProtocol(Protocol):
    def __call__(self, action: _Action, ctx: _PipelineContext) -> None: ...

class _PipelineScheduleRuntime(PipelineScheduleMulti):
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self._defer_pp_recv = bool(kwargs.pop("defer_pp_recv", False))
        max_active_stages = kwargs.pop("max_active_stages", 3)
        self._max_active_stages = 3 if max_active_stages is None else int(max_active_stages)
        super().__init__(*args, **kwargs)
        self._comp_type_to_function_map: dict[_ComputationType, Callable[..., Any]] = {}
        self.backward_counter: Counter[int] = Counter()
        self.bwd_recv_ops: dict[tuple[int, int], list[Any]] = {}
        self.fwd_recv_ops: dict[tuple[int, int], list[Any]] = {}
        self.unshard_ops: dict[int, list[UnshardHandle]] = defaultdict(list)
        self.unsharded_stages: set[int] = set()
        self.pipeline_order_with_comms: dict[int, list[_Action | None]] | None = None

    def register_custom_function(
        self,
        computation_type: _ComputationType,
        custom_function: _CustomFunctionProtocol,
    ) -> None:
        supported = {
            _ComputationType.FORWARD,
            _ComputationType.FULL_BACKWARD,
            _ComputationType.BACKWARD_INPUT,
            _ComputationType.BACKWARD_WEIGHT,
            _ComputationType.OVERLAP_F_B,
            _ComputationType.UNSHARD,
            _ComputationType.RESHARD,
            _ComputationType.REDUCE_GRAD,
        }
        if computation_type not in supported:
            raise ValueError(f"invalid computation type {computation_type}")
        if computation_type in self._comp_type_to_function_map:
            logger.warning(
                "computation type %s is already registered; replacing it",
                computation_type,
            )
        self._comp_type_to_function_map[computation_type] = custom_function

    def _prepare_schedule_with_comms(
        self,
        actions: dict[int, list[_Action | None]],
        format: Literal["compute_only", "compute_comms"] = "compute_only",
    ) -> None:
        super()._validate_and_set_stage_mapping(actions)
        if format == "compute_comms":
            lowered: dict[int, list[_Action | None]] = {}
            for rank, rank_actions in actions.items():
                if any(action is None for action in rank_actions):
                    raise AssertionError("communication schedules cannot contain empty actions")
                lowered[rank] = list(rank_actions)
            self.pipeline_order_with_comms = lowered
            return
        if format != "compute_only":
            raise NotImplementedError(f"{format=} is not implemented")
        for rank, rank_actions in actions.items():
            for index, action in enumerate(rank_actions):
                if action is not None and not action.is_compute_op:
                    raise ValueError(
                        f"expected compute-only action at rank {rank}, position {index}: {action}"
                    )
        lowered = {
            rank: _add_unshard_reshard(
                rank_actions, max_active_stages=self._max_active_stages
            )
            for rank, rank_actions in actions.items()
        }
        lowered = {
            rank: _add_reduce_grad(rank_actions, self._n_microbatches)
            for rank, rank_actions in lowered.items()
        }
        lowered = _add_send_recv(
            lowered,
            stage_to_rank=lambda stage: self.stage_index_to_group_rank[stage],
            num_stages=self._num_stages,
        )
        if self._defer_pp_recv:
            lowered = _defer_recv_ops(
                lowered,
                stage_to_rank=lambda stage: self.stage_index_to_group_rank[stage],
            )
        self.pipeline_order_with_comms = lowered

    def _load_csv(
        self,
        filename: str,
        format: Literal["compute_only", "compute_comms"] = "compute_only",
    ) -> None:
        if format == "compute_only":
            actions = super()._load_csv(filename)
            self.pipeline_order = actions
            self._prepare_schedule_with_comms(actions)
            return
        if format != "compute_comms":
            raise NotImplementedError(f"{format=} is not implemented")
        with open(filename, newline="", encoding="utf-8") as stream:
            actions = {
                rank: [_Action.from_str(value) for value in row]
                for rank, row in enumerate(csv.reader(stream))
            }
        self._prepare_schedule_with_comms(actions, format=format)

    def _dump_csv(
        self,
        filename: str,
        format: Literal["compute_only", "compute_comms"] = "compute_comms",
    ) -> None:
        if format == "compute_only":
            actions = self.pipeline_order
        elif format == "compute_comms":
            actions = self.pipeline_order_with_comms
        else:
            raise NotImplementedError(f"{format=} is not implemented")
        with open(filename, "w", newline="", encoding="utf-8") as stream:
            writer = csv.writer(stream)
            for rank in sorted(actions):
                writer.writerow(actions[rank])

    def _simulate(self) -> Any:
        return _simulate_comms_compute(
            self.pipeline_order_with_comms,
            lambda stage: self.stage_index_to_group_rank[stage],
            self._num_stages,
        )

    def _assert_unsharded(self, stage: Any) -> None:
        if not isinstance(stage.submod, FSDPModule):
            return
        stage_index = stage.stage_index
        if stage_index in self.unshard_ops:
            for handle in self.unshard_ops[stage_index]:
                handle.wait()
            del self.unshard_ops[stage_index]
            self.unsharded_stages.add(stage_index)
        if stage_index not in self.unsharded_stages:
            raise AssertionError(f"attempted to compute on sharded stage {stage_index}")

    def _step_microbatches(
        self,
        arg_mbs: list[tuple[Any, ...]] | None = None,
        kwarg_mbs: list[dict[str, Any]] | None = None,
        target_mbs: list[Any] | None = None,
        losses: list[Any] | None = None,
        return_outputs: bool = True,
        loss_kwargs: dict[str, Any] | None = None,
    ) -> None:
        arg_mbs, kwarg_mbs = self._check_inputs(arg_mbs, kwarg_mbs, target_mbs, losses)
        first_target = target_mbs[0] if target_mbs is not None else None
        self._initialize_stages(arg_mbs[0], kwarg_mbs[0], first_target, loss_kwargs)
        stage_index_to_stage = {
            stage.stage_index: stage for stage in self._stages
        }
        if self.pipeline_order_with_comms is None:
            raise AssertionError(
                "must prepare a schedule with communication actions before execution"
            )
        self.fwd_recv_ops.clear()
        self.bwd_recv_ops.clear()
        self.unshard_ops.clear()
        self.unsharded_stages.clear()
        send_ops: list[list[Any]] = []

        def perform_action(action: _Action) -> None:
            computation_type = action.computation_type
            microbatch = action.microbatch_index
            if microbatch is None and computation_type not in {
                _ComputationType.UNSHARD,
                _ComputationType.RESHARD,
                _ComputationType.REDUCE_GRAD,
            }:
                raise AssertionError(f"{action=} is missing a microbatch index")
            stage = stage_index_to_stage[action.stage_index]
            stage_index = action.stage_index
            stage_uses_fsdp = isinstance(stage.submod, FSDPModule)
            next_local = stage_index + 1 in stage_index_to_stage
            previous_local = stage_index - 1 in stage_index_to_stage
            mb_index = -1 if microbatch is None else microbatch

            if computation_type is _ComputationType.SEND_F:
                send_ops.append(_batch_p2p(stage.get_fwd_send_ops(mb_index)))
            elif computation_type is _ComputationType.SEND_B:
                send_ops.append(_batch_p2p(stage.get_bwd_send_ops(mb_index)))
            elif computation_type is _ComputationType.RECV_F:
                key = (stage_index, mb_index)
                if key in self.fwd_recv_ops:
                    raise AssertionError(f"forward receive repeated for {key}")
                self.fwd_recv_ops[key] = _batch_p2p(stage.get_fwd_recv_ops(mb_index))
            elif computation_type is _ComputationType.RECV_B:
                key = (stage_index, mb_index)
                if key in self.bwd_recv_ops:
                    raise AssertionError(f"backward receive repeated for {key}")
                self.bwd_recv_ops[key] = _batch_p2p(stage.get_bwd_recv_ops(mb_index))
            elif computation_type is _ComputationType.UNSHARD:
                if stage_uses_fsdp:
                    if stage_index in self.unsharded_stages or stage_index in self.unshard_ops:
                        raise AssertionError(f"unsharding stage {stage_index} twice")
                    for submodule in stage.submod.modules():
                        if isinstance(submodule, FSDPModule):
                            handle = cast(UnshardHandle, submodule.unshard(async_op=True))
                            self.unshard_ops[stage_index].append(handle)
            elif computation_type is _ComputationType.RESHARD:
                if stage_uses_fsdp:
                    if stage_index not in self.unsharded_stages:
                        raise AssertionError(f"resharding stage {stage_index} without unsharding")
                    if stage_index in self.unshard_ops:
                        raise AssertionError(f"resharding stage {stage_index} before unshard completion")
                    for submodule in stage.submod.modules():
                        if isinstance(submodule, FSDPModule):
                            submodule.reshard()
                    self.unsharded_stages.remove(stage_index)
            elif computation_type is _ComputationType.FORWARD:
                self._assert_unsharded(stage)
                if not stage.is_first and not previous_local:
                    key = (stage_index, mb_index)
                    if key not in self.fwd_recv_ops:
                        raise AssertionError(f"forward action {action} has no receive")
                    _wait_batch_p2p(self.fwd_recv_ops.pop(key))
                output = stage.forward_one_chunk(
                    mb_index,
                    arg_mbs[mb_index],
                    kwarg_mbs[mb_index],
                    save_forward_output=return_outputs,
                )
                self._maybe_compute_loss(stage, output, target_mbs, mb_index, loss_kwargs)
                if next_local:
                    stage_index_to_stage[stage_index + 1].set_local_fwd_input(output, mb_index)
            elif computation_type is _ComputationType.FULL_BACKWARD:
                self._assert_unsharded(stage)
                if not stage.is_last and not next_local:
                    key = (stage_index, mb_index)
                    if key not in self.bwd_recv_ops:
                        raise AssertionError(f"backward action {action} has no receive")
                    _wait_batch_p2p(self.bwd_recv_ops.pop(key))
                self.backward_counter[stage_index] += 1
                last_backward = self.backward_counter[stage_index] == self._n_microbatches
                stage.backward_one_chunk(
                    mb_index,
                    loss=self._maybe_get_loss(stage, mb_index),
                    full_backward=True,
                    last_backward=last_backward,
                )
                if previous_local:
                    stage_index_to_stage[stage_index - 1].set_local_bwd_input(
                        stage.get_local_bwd_output(mb_index), mb_index
                    )
            elif computation_type is _ComputationType.BACKWARD_INPUT:
                self._assert_unsharded(stage)
                if not stage.is_last and not next_local:
                    key = (stage_index, mb_index)
                    if key not in self.bwd_recv_ops:
                        raise AssertionError(f"backward action {action} has no receive")
                    _wait_batch_p2p(self.bwd_recv_ops.pop(key))
                stage.backward_one_chunk(
                    mb_index,
                    loss=self._maybe_get_loss(stage, mb_index),
                    full_backward=False,
                    last_backward=False,
                )
                if previous_local:
                    stage_index_to_stage[stage_index - 1].set_local_bwd_input(
                        stage.get_local_bwd_output(mb_index), mb_index
                    )
            elif computation_type is _ComputationType.BACKWARD_WEIGHT:
                self._assert_unsharded(stage)
                self.backward_counter[stage_index] += 1
                last_backward = self.backward_counter[stage_index] == self._n_microbatches
                stage.backward_weight_one_chunk(mb_index, last_backward=last_backward)
            elif computation_type is _ComputationType.REDUCE_GRAD:
                scale = self._n_microbatches if self._scale_grads else 1
                stage.perform_reduce_grad(scale)
            else:
                raise ValueError(f"unknown or unsupported action {action}")

        self.backward_counter.clear()
        for time_step, action in enumerate(self.pipeline_order_with_comms[self.rank]):
            if action is None:
                continue
            try:
                custom = self._comp_type_to_function_map.get(action.computation_type)
                if custom is not None:
                    custom(
                        action,
                        _PipelineContext(self, arg_mbs, kwarg_mbs, target_mbs, losses),
                    )
                elif action.computation_type is _ComputationType.OVERLAP_F_B:
                    if action.sub_actions is None:
                        raise AssertionError("overlap action must contain sub-actions")
                    for sub_action in action.sub_actions:
                        perform_action(sub_action)
                else:
                    perform_action(action)
            except Exception:
                logger.error(
                    "pipeline runtime failed at step %d on action %s",
                    time_step,
                    action,
                )
                logger.error(
                    "%s",
                    _format_pipeline_order(
                        self.pipeline_order_with_comms,
                        error_step_number=time_step,
                    ),
                )
                raise
        while send_ops:
            _wait_batch_p2p(send_ops.pop())
        if self.unshard_ops:
            raise AssertionError("unused unshard operations")
        self._update_losses(self._stages, losses)

[docs]
class ScheduleLoopedBFS(_PipelineScheduleRuntime):
    def __init__(self, stages: list[Any], n_microbatches: int, loss_fn: Any = None, output_merge_spec: Any = None, scale_grads: bool = True, backward_requires_autograd: bool = True, defer_pp_recv: bool = False, max_active_stages: int | None = None) -> None:
        super().__init__(
            stages,
            n_microbatches,
            loss_fn,
            output_merge_spec=output_merge_spec,
            scale_grads=scale_grads,
            backward_requires_autograd=backward_requires_autograd,
            defer_pp_recv=defer_pp_recv,
            max_active_stages=max_active_stages,
        )
        self.defer_pp_recv = self._defer_pp_recv
        self.max_active_stages = self._max_active_stages
        self.pipeline_order = {
            rank: self._calculate_single_rank_operations(rank)
            for rank in range(self.pp_group_size)
        }
        self._prepare_schedule_with_comms(self.pipeline_order)

    def _calculate_single_rank_operations(self, rank: int) -> list[_Action | None]:
        local_stage_count = len(self._stages)
        stage_indices = range(
            rank,
            self.pp_group_size * local_stage_count,
            self.pp_group_size,
        )
        rank_actions: list[_Action | None] = [None for _ in range(rank)]
        for stage_index in stage_indices:
            rank_actions.extend(
                _Action(stage_index, _ComputationType.FORWARD, microbatch)
                for microbatch in range(self._n_microbatches)
            )
        rank_actions.extend(
            [None] * (2 * (self.pp_group_size - 1 - rank))
        )
        for stage_index in reversed(stage_indices):
            rank_actions.extend(
                _Action(
                    stage_index,
                    _ComputationType.FULL_BACKWARD,
                    microbatch,
                )
                for microbatch in reversed(range(self._n_microbatches))
            )
        return rank_actions

[docs]
class ScheduleInterleaved1F1B(_PipelineScheduleRuntime):
    def __init__(
        self,
        stages: list[Any],
        n_microbatches: int,
        loss_fn: Callable[..., Any] | None = None,
        args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None,
        kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None,
        output_merge_spec: Any = None,
        scale_grads: bool = True,
        backward_requires_autograd: bool = True,
        defer_pp_recv: bool = False,
        max_active_stages: int = 3,
    ) -> 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,
            scale_grads=scale_grads,
            backward_requires_autograd=backward_requires_autograd,
            defer_pp_recv=defer_pp_recv,
            max_active_stages=max_active_stages,
        )
        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(
                "Interleaved 1F1B requires the microbatch count to be a multiple "
                f"of the round count ({self.number_of_rounds}), got {n_microbatches}"
            )
        self.pipeline_order = {
            rank: self._calculate_single_rank_operations(rank)
            for rank in range(self.pp_group_size)
        }
        self._prepare_schedule_with_comms(self.pipeline_order)

    def _calculate_single_rank_operations(self, rank: int) -> list[_Action | None]:
        warmup_ops = _get_warmup_ops(
            rank,
            self.n_local_stages,
            self.microbatches_per_round,
            self.pp_group_size,
            self._n_microbatches,
            multiply_factor=2,
        )
        microbatch_ops = self.n_local_stages * self._n_microbatches
        forward_backward_ops = microbatch_ops - warmup_ops
        cooldown_ops = microbatch_ops - forward_backward_ops

        def forward_stage_index(step: int) -> int:
            local_index = (step // self.microbatches_per_round) % self.n_local_stages
            return local_index * self.pp_group_size + rank

        def backward_stage_index(step: int) -> int:
            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

        logger.debug(
            "rank %s: warmup=%s steady=%s cooldown=%s",
            rank,
            warmup_ops,
            forward_backward_ops,
            cooldown_ops,
        )
        return _get_1f1b_rank_ops(
            self.n_local_stages,
            self.pp_group_size,
            warmup_ops,
            forward_backward_ops,
            cooldown_ops,
            rank,
            forward_stage_index,
            backward_stage_index,
        )

[docs]
class ScheduleInterleavedZeroBubble(_PipelineScheduleRuntime):
    def __init__(
        self,
        stages: list[Any],
        n_microbatches: int,
        loss_fn: Callable[..., Any] | None = None,
        args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None,
        kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None,
        output_merge_spec: Any = None,
        scale_grads: bool = True,
        backward_requires_autograd: bool = True,
        defer_pp_recv: bool = False,
        max_active_stages: int = 3,
    ) -> None:
        _check_torch_compile_compatibility(stages, self.__class__.__name__)
        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,
            scale_grads=scale_grads,
            backward_requires_autograd=backward_requires_autograd,
            defer_pp_recv=defer_pp_recv,
            max_active_stages=max_active_stages,
        )
        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(
                "Zero bubble requires the microbatch count to be a multiple of "
                f"the round count ({self.number_of_rounds}), got {n_microbatches}"
            )
        self.pipeline_order = {
            rank: self._calculate_single_rank_operations(rank)
            for rank in range(self.pp_group_size)
        }
        self.pipeline_order = self._add_bubbles_to_actions(
            self.n_local_stages * self.pp_group_size
        )
        self._prepare_schedule_with_comms(self.pipeline_order)

    def _calculate_single_rank_operations(self, rank: int) -> list[_Action | None]:
        warmup_ops = _get_warmup_ops(
            rank,
            self.n_local_stages,
            self.microbatches_per_round,
            self.pp_group_size,
            self._n_microbatches,
            multiply_factor=1,
        )
        microbatch_ops = self.n_local_stages * self._n_microbatches
        forward_backward_ops = microbatch_ops - warmup_ops
        cooldown_ops = microbatch_ops - forward_backward_ops

        def forward_stage_index(step: int) -> int:
            local_index = (step // self.microbatches_per_round) % self.n_local_stages
            return local_index * self.pp_group_size + rank

        def backward_stage_index(step: int) -> int:
            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,
            forward_backward_ops,
            cooldown_ops,
            rank,
            forward_stage_index,
            backward_stage_index,
            rank,
            enable_zero_bubble=True,
        )

    def _add_bubbles_to_actions(
        self, num_stages_global: int
    ) -> dict[int, list[_Action | None]]:
        actions = self.pipeline_order

        def need_bubble(
            stage: int,
            operation: _ComputationType,
            microbatch: int | None,
            seen_ops: set[tuple[int, _ComputationType, int]],
        ) -> bool:
            if operation is _ComputationType.FORWARD:
                return stage != 0 and (
                    stage - 1,
                    operation,
                    microbatch,
                ) not in seen_ops
            if operation is _ComputationType.FULL_BACKWARD:
                if stage == num_stages_global - 1:
                    return (
                        stage,
                        _ComputationType.FORWARD,
                        microbatch,
                    ) not in seen_ops
                return (
                    stage + 1,
                    operation,
                    microbatch,
                ) not in seen_ops
            return False

        seen_ops: set[tuple[int, _ComputationType, int]] = set()
        result: dict[int, list[_Action | None]] = {
            rank: [] for rank in range(self.pp_group_size)
        }
        next_pointer = {rank: 0 for rank in range(self.pp_group_size)}
        bubbles_added = {rank: 0 for rank in range(self.pp_group_size)}
        total_bubbles_added = 0

        while True:
            should_stop = True
            temporary_seen: set[tuple[int, _ComputationType, int]] = set()
            for rank in range(self.pp_group_size):
                timestamp = next_pointer[rank]
                if timestamp >= len(actions[rank]):
                    continue
                should_stop = False
                action = actions[rank][timestamp]
                if action is None:
                    result[rank].append(None)
                    next_pointer[rank] += 1
                    continue
                stage_index, operation, microbatch, _ = action
                if not need_bubble(
                    stage_index,
                    operation,
                    microbatch,
                    seen_ops,
                ):
                    result[rank].append(action)
                    if microbatch is not None:
                        temporary_seen.add((stage_index, operation, microbatch))
                    next_pointer[rank] += 1
                else:
                    result[rank].append(None)
                    bubbles_added[rank] += 1
            seen_ops.update(temporary_seen)
            if should_stop:
                break
        if total_bubbles_added > 0:
            logger.warning(
                "non-zero bubbles added: total=%s by-rank=%s",
                total_bubbles_added,
                bubbles_added,
            )
        return result

[docs]
class ScheduleZBVZeroBubble(_PipelineScheduleRuntime):
    def __init__(
        self,
        stages: list[Any],
        n_microbatches: int,
        loss_fn: Callable[..., Any] | None = None,
        args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None,
        kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None,
        output_merge_spec: Any = None,
        scale_grads: bool = True,
        backward_requires_autograd: bool = True,
        defer_pp_recv: bool = False,
        max_active_stages: int = 3,
    ) -> None:
        _check_torch_compile_compatibility(stages, self.__class__.__name__)
        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,
            scale_grads=scale_grads,
            backward_requires_autograd=backward_requires_autograd,
            defer_pp_recv=defer_pp_recv,
            max_active_stages=max_active_stages,
        )
        self.stage_index_to_group_rank = generate_stage_to_rank_mapping(
            self.pp_group_size,
            self._num_stages,
            style="v",
        )
        for stage in self._stages:
            stage.stage_index_to_group_rank = self.stage_index_to_group_rank
        self.n_local_stages = len(stages)
        if self.n_local_stages != 2:
            raise ValueError(
                "ZBV requires exactly two stages per rank, "
                f"got {self.n_local_stages}"
            )
        self.rank = stages[0].group_rank
        self.num_stages = stages[0].num_stages
        self.pipeline_order = {
            rank: self._calculate_single_rank_operations(rank)
            for rank in range(self.pp_group_size)
        }
        self._prepare_schedule_with_comms(self.pipeline_order)

    def _calculate_single_rank_operations(self, rank: int) -> list[_Action | None]:
        microbatch_count = max(2 * self.pp_group_size - 1, self._n_microbatches)
        rank_actions: list[_Action | None] = [None for _ in range(rank)]
        forward_chunk0 = 0
        forward_chunk1 = 0
        backward_chunk0 = 0
        backward_chunk1 = 0

        warmup_first = 2 * (self.pp_group_size - rank) - 1
        stage_chunk0 = rank
        stage_chunk1 = self.num_stages - 1 - rank
        for _ in range(warmup_first):
            rank_actions.append(
                _Action(
                    stage_chunk0,
                    _ComputationType.FORWARD,
                    forward_chunk0,
                )
            )
            forward_chunk0 += 1

        warmup_second = rank
        for _ in range(warmup_second):
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.FORWARD, forward_chunk1)
            )
            forward_chunk1 += 1
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.FORWARD, forward_chunk0)
            )
            forward_chunk0 += 1

        warmup_third = self.pp_group_size - rank
        for _ in range(warmup_third):
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.FORWARD, forward_chunk1)
            )
            forward_chunk1 += 1
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_INPUT, backward_chunk1)
            )
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_WEIGHT, backward_chunk1)
            )
            backward_chunk1 += 1

        while forward_chunk1 < forward_chunk0 or forward_chunk0 < microbatch_count:
            if forward_chunk0 < microbatch_count:
                rank_actions.append(
                    _Action(stage_chunk0, _ComputationType.FORWARD, forward_chunk0)
                )
                forward_chunk0 += 1
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_INPUT, backward_chunk0)
            )
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_WEIGHT, backward_chunk0)
            )
            backward_chunk0 += 1
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.FORWARD, forward_chunk1)
            )
            forward_chunk1 += 1
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_INPUT, backward_chunk1)
            )
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_WEIGHT, backward_chunk1)
            )
            backward_chunk1 += 1

        weight_chunk0 = backward_chunk0
        weight_chunk1 = backward_chunk1
        for _ in range(rank):
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_INPUT, backward_chunk0)
            )
            backward_chunk0 += 1
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_INPUT, backward_chunk1)
            )
            backward_chunk1 += 1

        for _ in range(self.pp_group_size - rank):
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_INPUT, backward_chunk0)
            )
            backward_chunk0 += 1
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_WEIGHT, weight_chunk0)
            )
            weight_chunk0 += 1

        while weight_chunk1 < backward_chunk1:
            rank_actions.append(
                _Action(stage_chunk1, _ComputationType.BACKWARD_WEIGHT, weight_chunk1)
            )
            weight_chunk1 += 1
        while weight_chunk0 < backward_chunk0:
            rank_actions.append(
                _Action(stage_chunk0, _ComputationType.BACKWARD_WEIGHT, weight_chunk0)
            )
            weight_chunk0 += 1

        if not (weight_chunk0 == backward_chunk0 == forward_chunk0):
            raise AssertionError(
                "stage chunk 0 action counts do not match"
            )
        if not (weight_chunk1 == backward_chunk1 == forward_chunk1):
            raise AssertionError(
                "stage chunk 1 action counts do not match"
            )
        return [
            action
            if action is not None
            and action.microbatch_index is not None
            and action.microbatch_index < self._n_microbatches
            else None
            for action in rank_actions
        ]

[docs]
class ScheduleDualPipeV(_PipelineScheduleRuntime):
    """Run the bidirectional local stage schedule."""

    def __init__(
        self,
        stages: list[Any],
        n_microbatches: int,
        loss_fn: Callable[..., Any] | None = None,
        args_chunk_spec: tuple[TensorChunkSpec, ...] | None = None,
        kwargs_chunk_spec: dict[str, TensorChunkSpec] | None = None,
        output_merge_spec: Any = None,
        scale_grads: bool = True,
        backward_requires_autograd: bool = True,
        defer_pp_recv: bool = False,
        max_active_stages: int = 3,
    ) -> None:
        _check_torch_compile_compatibility(stages, self.__class__.__name__)
        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,
            scale_grads=scale_grads,
            backward_requires_autograd=backward_requires_autograd,
            defer_pp_recv=defer_pp_recv,
            max_active_stages=max_active_stages,
        )
        self.stage_index_to_group_rank = generate_stage_to_rank_mapping(
            self.pp_group_size,
            self._num_stages,
            style="v",
        )
        for stage in self._stages:
            stage.stage_index_to_group_rank = self.stage_index_to_group_rank
        self.n_local_stages = len(stages)
        if self.n_local_stages != 2:
            raise ValueError(
                "ZBV requires exactly 2 stages per rank, but got "
                f"{self.n_local_stages}."
            )
        if n_microbatches < self._num_stages:
            raise ValueError(
                "DualPipeV requires at least as many microbatches as stages, but got "
                f"{n_microbatches} microbatches and {self._num_stages} stages."
            )
        self.rank = stages[0].group_rank
        self.num_stages = stages[0].num_stages
        self.pipeline_order = {
            rank: self._calculate_single_rank_operations(rank)
            for rank in range(self.pp_group_size)
        }
        self._prepare_schedule_with_comms(self.pipeline_order)

    def _calculate_single_rank_operations(self, rank: int) -> list[_Action | None]:
        actions: list[_Action | None] = []
        counters: dict[tuple[int, _ComputationType], int] = {}
        weight_queue: list[tuple[int, int]] = []

        num_ranks = self.pp_group_size
        num_chunks = self._n_microbatches
        rank_to_stages = generate_rank_to_stage_mapping(
            num_ranks, num_ranks * 2, style="v"
        )
        stage0_index, stage1_index = rank_to_stages[rank]

        def increment_backward_counts(stage_index: int) -> None:
            input_key = (stage_index, _ComputationType.BACKWARD_INPUT)
            weight_key = (stage_index, _ComputationType.BACKWARD_WEIGHT)
            counters[input_key] = counters.get(input_key, 0) + 1
            counters[weight_key] = counters.get(weight_key, 0) + 1

        def add_overlap_f_b(
            forward_stage: int,
            backward_stage: int,
        ) -> None:
            forward_key = (forward_stage, _ComputationType.FORWARD)
            backward_key = (backward_stage, _ComputationType.BACKWARD_INPUT)
            forward_mb = counters.get(forward_key, 0)
            backward_mb = counters.get(backward_key, 0)
            sub_actions = (
                _Action(forward_stage, _ComputationType.FORWARD, forward_mb),
                _Action(backward_stage, _ComputationType.FULL_BACKWARD, backward_mb),
            )
            actions.append(
                _Action(-1, _ComputationType.OVERLAP_F_B, None, sub_actions)
            )
            counters[forward_key] = forward_mb + 1
            increment_backward_counts(backward_stage)

        def add_action(
            stage_index: int,
            computation_type: _ComputationType,
        ) -> None:
            key = (
                (stage_index, computation_type)
                if computation_type != _ComputationType.FULL_BACKWARD
                else (stage_index, _ComputationType.BACKWARD_INPUT)
            )
            mb_index = counters.get(key, 0)
            actions.append(_Action(stage_index, computation_type, mb_index))
            if computation_type == _ComputationType.FULL_BACKWARD:
                increment_backward_counts(stage_index)
            else:
                if computation_type == _ComputationType.BACKWARD_INPUT:
                    weight_queue.append((stage_index, mb_index))
                counters[key] = mb_index + 1

        def add_weight_action_if_pending() -> None:
            if not weight_queue:
                return
            actual_stage_index, weight_mb_index = weight_queue.pop(0)
            actions.append(
                _Action(
                    actual_stage_index,
                    _ComputationType.BACKWARD_WEIGHT,
                    weight_mb_index,
                )
            )
            weight_key = (actual_stage_index, _ComputationType.BACKWARD_WEIGHT)
            counters[weight_key] = counters.get(weight_key, 0) + 1

        step_1 = (num_ranks - rank - 1) * 2
        for _ in range(step_1):
            add_action(stage0_index, _ComputationType.FORWARD)

        step_2 = rank + 1
        for _ in range(step_2):
            add_action(stage0_index, _ComputationType.FORWARD)
            add_action(stage1_index, _ComputationType.FORWARD)

        step_3 = num_ranks - rank - 1
        for _ in range(step_3):
            add_action(stage1_index, _ComputationType.BACKWARD_INPUT)
            add_weight_action_if_pending()
            add_action(stage1_index, _ComputationType.FORWARD)

        step_4 = num_chunks - num_ranks * 2 + rank + 1
        for index in range(step_4):
            if index == 0 and rank == num_ranks - 1:
                add_action(stage0_index, _ComputationType.FORWARD)
                add_action(stage1_index, _ComputationType.FULL_BACKWARD)
            else:
                add_overlap_f_b(stage0_index, stage1_index)
            add_overlap_f_b(stage1_index, stage0_index)

        step_5 = num_ranks - rank - 1
        for _ in range(step_5):
            add_action(stage1_index, _ComputationType.FULL_BACKWARD)
            add_overlap_f_b(stage1_index, stage0_index)

        step_6 = rank + 1
        enable_zb = False
        for index in range(step_6):
            if index == step_6 // 2 and rank % 2 == 1:
                enable_zb = True
            comp_type = (
                _ComputationType.BACKWARD_INPUT
                if enable_zb
                else _ComputationType.FULL_BACKWARD
            )
            add_action(stage1_index, comp_type)
            if index == step_6 // 2 and rank % 2 == 0:
                enable_zb = True
            comp_type = (
                _ComputationType.BACKWARD_INPUT
                if enable_zb
                else _ComputationType.FULL_BACKWARD
            )
            add_action(stage0_index, comp_type)

        step_7 = num_ranks - rank - 1
        for _ in range(step_7):
            add_weight_action_if_pending()
            comp_type = (
                _ComputationType.BACKWARD_INPUT
                if enable_zb
                else _ComputationType.FULL_BACKWARD
            )
            add_action(stage0_index, comp_type)

        step_8 = rank + 1
        for _ in range(step_8):
            add_weight_action_if_pending()
        return actions

def _requires_reduce_grad(action_type: _ComputationType) -> bool:
    return action_type in {
        _ComputationType.BACKWARD_WEIGHT,
        _ComputationType.FULL_BACKWARD,
    }

def _add_reduce_grad(
    actions: list[_Action | None], n_microbatches: int
) -> list[_Action | None]:
    if n_microbatches <= 0:
        raise ValueError("n_microbatches must be positive")
    result: list[_Action | None] = []
    counts: dict[int, int] = defaultdict(int)
    for action in actions:
        if action is None:
            result.append(None)
            continue
        result.append(action)
        leaves = action.sub_actions or (action,)
        for leaf in leaves:
            if not _requires_reduce_grad(leaf.computation_type):
                continue
            counts[leaf.stage_index] += 1
            if counts[leaf.stage_index] == n_microbatches:
                result.append(
                    _Action(leaf.stage_index, _ComputationType.REDUCE_GRAD, None)
                )
    return result

def _split_backward_pipeline_order(
    num_stages: int,
    pp_group_size: int,
    n_microbatches: int,
    stage_to_rank: dict[int, int],
) -> dict[int, list[_Action | None]]:
    owned: dict[int, list[int]] = {rank: [] for rank in range(pp_group_size)}
    for stage_index in range(num_stages):
        owned[int(stage_to_rank[stage_index])].append(stage_index)
    result: dict[int, list[_Action | None]] = {}
    for rank in range(pp_group_size):
        actions: list[_Action | None] = [None] * rank
        for stage_index in owned[rank]:
            actions.extend(
                _Action(stage_index, _ComputationType.FORWARD, microbatch)
                for microbatch in range(n_microbatches)
            )
        for stage_index in reversed(owned[rank]):
            actions.extend(
                _Action(stage_index, _ComputationType.BACKWARD_INPUT, microbatch)
                for microbatch in reversed(range(n_microbatches))
            )
            actions.extend(
                _Action(stage_index, _ComputationType.BACKWARD_WEIGHT, microbatch)
                for microbatch in reversed(range(n_microbatches))
            )
        result[rank] = _add_reduce_grad(actions, n_microbatches)
    return result

def _add_unshard_reshard(
    compute_actions: list[_Action | None], max_active_stages: int = 3
) -> list[_Action]:
    if max_active_stages <= 0:
        raise ValueError("max_active_stages must be positive")
    active: set[int] = set()
    result: list[_Action] = []

    def next_stage_indices(
        count: int, actions: list[_Action | None]
    ) -> list[int]:
        seen: set[int] = set()
        stages: list[int] = []
        for action in actions:
            if action is None:
                continue
            leaves = action.sub_actions or (action,)
            for leaf in leaves:
                if leaf.stage_index not in seen:
                    seen.add(leaf.stage_index)
                    stages.append(leaf.stage_index)
            if len(stages) >= count:
                break
        return stages

    actions = list(compute_actions)
    for index, action in enumerate(actions):
        if action is None:
            continue
        upcoming = next_stage_indices(max_active_stages, actions[index:])
        for stage_index in [stage for stage in active if stage not in upcoming]:
            active.remove(stage_index)
            result.append(_Action(stage_index, _ComputationType.RESHARD))
        for stage_index in upcoming:
            if stage_index not in active:
                active.add(stage_index)
                result.append(_Action(stage_index, _ComputationType.UNSHARD))
        result.append(action)
    for stage_index in list(active):
        result.append(_Action(stage_index, _ComputationType.RESHARD))
    return result

def _merge_bw(compute_actions: list[_Action | None]) -> list[_Action]:
    pending = list(compute_actions)
    result: list[_Action] = []
    while pending:
        action = pending.pop(0)
        if action is None:
            continue
        while pending and pending[0] is None:
            pending.pop(0)
        if (
            action.computation_type is _ComputationType.BACKWARD_INPUT
            and pending
            and pending[0] is not None
            and pending[0].computation_type is _ComputationType.BACKWARD_WEIGHT
            and action.stage_index == pending[0].stage_index
            and action.microbatch_index == pending[0].microbatch_index
        ):
            result.append(
                _Action(
                    action.stage_index,
                    _ComputationType.FULL_BACKWARD,
                    action.microbatch_index,
                )
            )
            pending.pop(0)
        else:
            result.append(action)
    return result

def _add_send_recv(
    compute_actions: dict[int, list[_Action | None]],
    stage_to_rank: Any,
    num_stages: int,
) -> dict[int, list[_Action | None]]:
    if callable(stage_to_rank):
        rank_of = stage_to_rank
    else:
        rank_of = lambda stage: stage_to_rank[int(stage)]
    remaining = {
        int(rank): list(actions) for rank, actions in compute_actions.items()
    }
    result: dict[int, list[_Action | None]] = {rank: [] for rank in remaining}
    previous: dict[int, set[_Action]] = {rank: set() for rank in remaining}

    def leaves(action: _Action) -> tuple[_Action, ...]:
        return action.sub_actions or (action,)

    def communicates(action: _Action) -> bool:
        if action.computation_type is _ComputationType.FORWARD:
            return (
                action.stage_index < num_stages - 1
                and rank_of(action.stage_index) != rank_of(action.stage_index + 1)
            )
        if action.computation_type in {
            _ComputationType.BACKWARD_INPUT,
            _ComputationType.FULL_BACKWARD,
        }:
            return (
                action.stage_index > 0
                and rank_of(action.stage_index) != rank_of(action.stage_index - 1)
            )
        return False

    def ready(action: _Action, done: set[_Action]) -> bool:
        if (
            action.computation_type is _ComputationType.FORWARD
            and action.stage_index > 0
        ):
            return (
                _Action(
                    action.stage_index,
                    _ComputationType.RECV_F,
                    action.microbatch_index,
                )
                in done
                or _Action(
                    action.stage_index - 1,
                    _ComputationType.FORWARD,
                    action.microbatch_index,
                )
                in done
            )
        if (
            action.computation_type
            in {_ComputationType.BACKWARD_INPUT, _ComputationType.FULL_BACKWARD}
            and action.stage_index < num_stages - 1
        ):
            return (
                _Action(
                    action.stage_index,
                    _ComputationType.RECV_B,
                    action.microbatch_index,
                )
                in done
                or _Action(
                    action.stage_index + 1,
                    _ComputationType.BACKWARD_INPUT,
                    action.microbatch_index,
                )
                in done
                or _Action(
                    action.stage_index + 1,
                    _ComputationType.FULL_BACKWARD,
                    action.microbatch_index,
                )
                in done
            )
        return True

    while remaining:
        progress = False
        for rank in sorted(tuple(remaining)):
            actions = remaining[rank]
            if not actions:
                del remaining[rank]
                continue
            action = actions[0]
            if action is None:
                result[rank].append(None)
                actions.pop(0)
                progress = True
                if not actions:
                    del remaining[rank]
                continue
            action_leaves = leaves(action)
            if not all(ready(leaf, previous[rank]) for leaf in action_leaves):
                continue
            result[rank].append(action)
            for leaf in action_leaves:
                previous[rank].add(leaf)
                if not communicates(leaf):
                    continue
                is_forward = leaf.computation_type is _ComputationType.FORWARD
                send_kind = (
                    _ComputationType.SEND_F
                    if is_forward
                    else _ComputationType.SEND_B
                )
                recv_kind = (
                    _ComputationType.RECV_F
                    if is_forward
                    else _ComputationType.RECV_B
                )
                peer_stage = leaf.stage_index + 1 if is_forward else leaf.stage_index - 1
                send = _Action(leaf.stage_index, send_kind, leaf.microbatch_index)
                recv = _Action(peer_stage, recv_kind, leaf.microbatch_index)
                result[rank].append(send)
                previous[rank].add(send)
                peer_rank = int(rank_of(peer_stage))
                if peer_rank not in result:
                    raise ValueError(f"stage mapping points to unknown rank {peer_rank}")
                result[peer_rank].append(recv)
                previous[peer_rank].add(recv)
            actions.pop(0)
            progress = True
            if not actions:
                del remaining[rank]
        if not progress:
            raise ValueError("malformed pipeline schedule")
    return result

def _defer_recv_ops(
    actions: list[_Action | None] | dict[int, list[_Action | None]],
    stage_to_rank: Any,
) -> list[_Action | None] | dict[int, list[_Action | None]]:
    rank_of = stage_to_rank if callable(stage_to_rank) else lambda stage: stage_to_rank[int(stage)]
    was_list = isinstance(actions, list)
    by_rank = {0: list(actions)} if was_list else {
        int(rank): list(rank_actions) for rank, rank_actions in actions.items()
    }
    result: dict[int, list[_Action | None]] = {}
    recv_types = {_ComputationType.RECV_F, _ComputationType.RECV_B}
    send_types = {_ComputationType.SEND_F, _ComputationType.SEND_B}

    def recv_peer(action: _Action) -> int:
        peer_stage = action.stage_index - 1 if action.computation_type is _ComputationType.RECV_F else action.stage_index + 1
        return int(rank_of(peer_stage))

    def send_peer(action: _Action) -> int:
        peer_stage = action.stage_index + 1 if action.computation_type is _ComputationType.SEND_F else action.stage_index - 1
        return int(rank_of(peer_stage))

    for rank, rank_actions in by_rank.items():
        deferred: dict[tuple[int, _ComputationType, int | None], _Action] = {}
        output: list[_Action | None] = []
        for action in rank_actions:
            if action is None:
                output.append(None)
                continue
            if action.computation_type in recv_types:
                key = (action.stage_index, action.computation_type, action.microbatch_index)
                deferred[key] = action
                continue
            if action.computation_type in send_types:
                peer = send_peer(action)
                if rank < peer:
                    for key in tuple(deferred):
                        if recv_peer(deferred[key]) == peer:
                            output.append(deferred.pop(key))
            leaves = action.sub_actions or (action,)
            for leaf in leaves:
                if leaf.computation_type is _ComputationType.FORWARD:
                    key = (leaf.stage_index, _ComputationType.RECV_F, leaf.microbatch_index)
                elif leaf.computation_type in {
                    _ComputationType.FULL_BACKWARD,
                    _ComputationType.BACKWARD_INPUT,
                }:
                    key = (leaf.stage_index, _ComputationType.RECV_B, leaf.microbatch_index)
                else:
                    continue
                if key in deferred:
                    output.append(deferred.pop(key))
            output.append(action)
        if deferred:
            raise ValueError("every receive action must have a consuming compute action")
        result[rank] = output
    return result[0] if was_list else result

def _validate_schedule(
    actions: Any,
    pp_group_size: int,
    num_stages: int,
    num_microbatches: int,
) -> dict[int, int]:
    if pp_group_size <= 0 or num_stages <= 0 or num_microbatches <= 0:
        raise ValueError("pipeline dimensions must be positive")
    if not isinstance(actions, dict) or len(actions) != pp_group_size:
        raise ValueError("schedule must provide one action list per rank")
    if set(actions) != set(range(pp_group_size)):
        raise ValueError("schedule ranks must be contiguous")

    stage_actions: dict[int, dict[_ComputationType, set[int]]] = {
        stage: {
            _ComputationType.FORWARD: set(),
            _ComputationType.BACKWARD_INPUT: set(),
            _ComputationType.BACKWARD_WEIGHT: set(),
            _ComputationType.FULL_BACKWARD: set(),
        }
        for stage in range(num_stages)
    }
    stage_index_to_rank: dict[int, int] = {}
    seen: set[tuple[int, _ComputationType, int | None]] = set()

    compute_types = {
        _ComputationType.FORWARD,
        _ComputationType.FULL_BACKWARD,
        _ComputationType.BACKWARD_INPUT,
        _ComputationType.BACKWARD_WEIGHT,
    }
    communication_types = {
        _ComputationType.SEND_F,
        _ComputationType.RECV_F,
        _ComputationType.SEND_B,
        _ComputationType.RECV_B,
    }

    def process_action(action: _Action, rank: int, step: int) -> None:
        if action.sub_actions is not None:
            if action.computation_type is not _ComputationType.OVERLAP_F_B:
                raise ValueError("only overlap actions may contain sub-actions")
            if not action.sub_actions:
                raise ValueError("an overlap action must contain sub-actions")
            for sub_action in action.sub_actions:
                if not isinstance(sub_action, _Action):
                    raise TypeError("sub-actions must be _Action instances")
                process_action(sub_action, rank, step)
            return

        stage = action.stage_index
        kind = action.computation_type
        microbatch = action.microbatch_index
        if not 0 <= stage < num_stages:
            raise ValueError("action stage is outside the pipeline")
        if kind not in compute_types | communication_types | {
            _ComputationType.UNSHARD,
            _ComputationType.RESHARD,
            _ComputationType.REDUCE_GRAD,
        }:
            raise ValueError(f"unsupported pipeline action {kind!r}")
        if kind in compute_types | communication_types:
            if microbatch is None or not 0 <= microbatch < num_microbatches:
                raise ValueError("action microbatch is outside the schedule")
        elif microbatch is not None:
            raise ValueError("non-compute actions cannot carry a microbatch")

        previous_rank = stage_index_to_rank.get(stage)
        if previous_rank is not None and previous_rank != rank:
            raise ValueError(
                f"stage {stage} is assigned to ranks {previous_rank} and {rank}"
            )
        stage_index_to_rank[stage] = rank
        if kind not in compute_types:
            return

        key = (stage, kind, microbatch)
        if key in seen:
            raise ValueError("a compute action occurs more than once")
        seen.add(key)
        if kind is _ComputationType.FORWARD:
            stage_actions[stage][kind].add(microbatch)
            return
        if kind is _ComputationType.FULL_BACKWARD:
            if microbatch not in stage_actions[stage][_ComputationType.FORWARD]:
                raise ValueError("backward ran before its forward")
            stage_actions[stage][kind].add(microbatch)
            return
        if kind is _ComputationType.BACKWARD_INPUT:
            if microbatch not in stage_actions[stage][_ComputationType.FORWARD]:
                raise ValueError("backward input ran before its forward")
            stage_actions[stage][kind].add(microbatch)
            return
        if microbatch not in stage_actions[stage][_ComputationType.BACKWARD_INPUT]:
            raise ValueError("backward weight ran before its input backward")
        stage_actions[stage][kind].add(microbatch)

    for rank, rank_actions in actions.items():
        if not isinstance(rank_actions, list):
            raise TypeError(f"actions for rank {rank} must be a list")
        for step, action in enumerate(rank_actions):
            if action is None:
                continue
            if not isinstance(action, _Action):
                raise TypeError("schedule entries must be _Action instances")
            process_action(action, rank, step)
    for stage in range(num_stages):
        counts = stage_actions[stage]
        if len(counts[_ComputationType.FORWARD]) != num_microbatches:
            raise ValueError("schedule is missing a forward action")
        if len(counts[_ComputationType.BACKWARD_INPUT]) != len(
            counts[_ComputationType.BACKWARD_WEIGHT]
        ):
            raise ValueError("input and weight backward counts must match")
        if len(counts[_ComputationType.FULL_BACKWARD]) + len(
            counts[_ComputationType.BACKWARD_INPUT]
        ) != num_microbatches:
            raise ValueError("schedule is missing a backward action")
    if len(stage_index_to_rank) != num_stages:
        raise ValueError("schedule does not assign every pipeline stage")
    return stage_index_to_rank

def _get_1f1b_rank_ops(
    n_local_stages: int,
    pp_group_size: int,
    warmup_ops: int,
    fwd_bwd_ops: int,
    cooldown_ops: int,
    rank: int,
    forward_stage_index: Any,
    backward_stage_index: Any,
    num_1f1b_microbatches: int = 0,
    enable_zero_bubble: bool = False,
) -> list[_Action | None]:
    if min(n_local_stages, pp_group_size) <= 0:
        raise ValueError("pipeline dimensions must be positive")
    if rank < 0 or rank >= pp_group_size:
        raise ValueError("rank is outside the pipeline group")
    if min(warmup_ops, fwd_bwd_ops, cooldown_ops) < 0:
        raise ValueError("operation counts must be non-negative")
    forward_counts: dict[int, int] = defaultdict(int)
    backward_counts: dict[int, int] = defaultdict(int)
    weight_counts: dict[int, int] = defaultdict(int)
    result: list[_Action | None] = [None] * rank
    backward_ids: list[int] = []
    total_ops = warmup_ops + fwd_bwd_ops + cooldown_ops
    post_warmup = (
        n_local_stages * pp_group_size
        + 2 * (pp_group_size - 1 - rank)
        - warmup_ops
        - rank
    )
    if enable_zero_bubble:
        post_warmup = pp_group_size - rank - 1

    for operation in range(total_ops):
        if operation < warmup_ops:
            stage = int(forward_stage_index(operation))
            microbatch = forward_counts[stage]
            forward_counts[stage] += 1
            result.append(_Action(stage, _ComputationType.FORWARD, microbatch))
            if operation == warmup_ops - 1:
                result.extend([None] * max(0, post_warmup))
            continue

        if operation < warmup_ops + fwd_bwd_ops:
            stage = int(forward_stage_index(operation))
            microbatch = forward_counts[stage]
            forward_counts[stage] += 1
            result.append(_Action(stage, _ComputationType.FORWARD, microbatch))
            backward_stage = int(backward_stage_index(operation))
            microbatch = backward_counts[backward_stage]
            backward_counts[backward_stage] += 1
            backward_kind = (
                _ComputationType.BACKWARD_INPUT
                if enable_zero_bubble
                else _ComputationType.FULL_BACKWARD
            )
            result.append(
                _Action(backward_stage, backward_kind, microbatch)
            )
            backward_ids.append(operation)
            if (
                enable_zero_bubble
                and operation - warmup_ops >= num_1f1b_microbatches
            ):
                weight_index = sum(weight_counts.values())
                weight_stage = int(backward_stage_index(backward_ids[weight_index]))
                microbatch = weight_counts[weight_stage]
                weight_counts[weight_stage] += 1
                result.append(
                    _Action(
                        weight_stage,
                        _ComputationType.BACKWARD_WEIGHT,
                        microbatch,
                    )
                )
            continue

        if not enable_zero_bubble:
            result.append(None)
        backward_stage = int(backward_stage_index(operation))
        microbatch = backward_counts[backward_stage]
        backward_counts[backward_stage] += 1
        backward_kind = (
            _ComputationType.BACKWARD_INPUT
            if enable_zero_bubble
            else _ComputationType.FULL_BACKWARD
        )
        result.append(_Action(backward_stage, backward_kind, microbatch))
        backward_ids.append(operation)
        if (
            enable_zero_bubble
            and operation - warmup_ops >= num_1f1b_microbatches
        ):
            weight_index = sum(weight_counts.values())
            weight_stage = int(backward_stage_index(backward_ids[weight_index]))
            microbatch = weight_counts[weight_stage]
            weight_counts[weight_stage] += 1
            result.append(
                _Action(
                    weight_stage,
                    _ComputationType.BACKWARD_WEIGHT,
                    microbatch,
                )
            )

    while enable_zero_bubble and sum(weight_counts.values()) < len(backward_ids):
        weight_index = sum(weight_counts.values())
        weight_stage = int(backward_stage_index(backward_ids[weight_index]))
        microbatch = weight_counts[weight_stage]
        weight_counts[weight_stage] += 1
        result.append(
            _Action(weight_stage, _ComputationType.BACKWARD_WEIGHT, microbatch)
        )
    return result

def _get_warmup_ops(
    rank: int,
    n_local_stages: int,
    microbatches_per_round: int,
    pp_group_size: int,
    n_microbatches: int,
    multiply_factor: int = 2,
) -> int:
    warmups_last_stage = (n_local_stages - 1) * microbatches_per_round
    warmup_ops = warmups_last_stage + multiply_factor * (pp_group_size - 1 - rank)
    return min(warmup_ops, n_microbatches * n_local_stages)

def get_schedule_class(schedule_name: str) -> type[_PipelineSchedule]:
    mapping = {
        "GPipe": ScheduleGPipe,
        "1F1B": Schedule1F1B,
        "Interleaved1F1B": ScheduleInterleaved1F1B,
        "LoopedBFS": ScheduleLoopedBFS,
        "InterleavedZeroBubble": ScheduleInterleavedZeroBubble,
        "ZBVZeroBubble": ScheduleZBVZeroBubble,
        "DualPipeV": ScheduleDualPipeV,
        "PipelineScheduleSingle": PipelineScheduleSingle,
        "PipelineScheduleMulti": PipelineScheduleMulti,
    }
    if not isinstance(schedule_name, str):
        raise TypeError("schedule name must be a string")
    normalized = {name.lower(): cls for name, cls in mapping.items()}
    try:
        return normalized[schedule_name.lower()]
    except KeyError as exc:
        raise ValueError(f"unknown pipeline schedule {schedule_name!r}") from exc

def _simulate_comms_compute(pipeline_order: Any, stage_to_rank: Any, num_stages: int) -> Any:
    if not isinstance(pipeline_order, dict):
        raise TypeError("pipeline_order must be a rank-to-actions mapping")
    if callable(stage_to_rank):
        rank_of = stage_to_rank
    else:
        rank_of = lambda stage: stage_to_rank[int(stage)]
    pending = {
        int(rank): [action for action in actions if action is not None]
        for rank, actions in pipeline_order.items()
    }
    schedule: dict[int, list[_Action | None]] = {
        rank: [] for rank in sorted(pending)
    }
    completed: dict[int, set[_Action]] = {
        rank: set() for rank in pending
    }

    def leaves(action: _Action) -> tuple[_Action, ...]:
        return action.sub_actions or (action,)

    def ready_leaf(action: _Action, owner: int) -> bool:
        if action.stage_index < 0 or action.stage_index >= num_stages:
            raise ValueError("action stage is outside the pipeline")
        owner = int(rank_of(action.stage_index)) if action.stage_index >= 0 else owner
        done = completed[owner]
        kind = action.computation_type
        stage = action.stage_index
        microbatch = action.microbatch_index
        if kind is _ComputationType.FORWARD:
            if stage == 0:
                return True
            return (
                _Action(stage, _ComputationType.RECV_F, microbatch) in done
                or _Action(stage - 1, _ComputationType.FORWARD, microbatch) in done
            )
        if kind in {
            _ComputationType.BACKWARD_INPUT,
            _ComputationType.FULL_BACKWARD,
        }:
            if stage == num_stages - 1:
                return True
            return (
                _Action(stage, _ComputationType.RECV_B, microbatch) in done
                or _Action(stage + 1, _ComputationType.BACKWARD_INPUT, microbatch) in done
                or _Action(stage + 1, _ComputationType.FULL_BACKWARD, microbatch) in done
            )
        if kind is _ComputationType.SEND_F:
            return _Action(stage, _ComputationType.FORWARD, microbatch) in done
        if kind is _ComputationType.RECV_F:
            peer = stage - 1
            return _Action(peer, _ComputationType.SEND_F, microbatch) in completed[int(rank_of(peer))]
        if kind is _ComputationType.SEND_B:
            return (
                _Action(stage, _ComputationType.BACKWARD_INPUT, microbatch) in done
                or _Action(stage, _ComputationType.FULL_BACKWARD, microbatch) in done
            )
        if kind is _ComputationType.RECV_B:
            peer = stage + 1
            return _Action(peer, _ComputationType.SEND_B, microbatch) in completed[int(rank_of(peer))]
        if kind is _ComputationType.BACKWARD_WEIGHT:
            return True
        if kind in {
            _ComputationType.UNSHARD,
            _ComputationType.RESHARD,
            _ComputationType.REDUCE_GRAD,
        }:
            return True
        raise ValueError(f"unsupported pipeline action {kind!r}")

    def ready(action: _Action, owner: int) -> bool:
        return all(ready_leaf(leaf, owner) for leaf in leaves(action))

    def mark_completed(action: _Action, owner: int) -> None:
        completed[owner].add(action)
        for leaf in leaves(action):
            leaf_owner = int(rank_of(leaf.stage_index))
            completed[leaf_owner].add(leaf)

    while pending:
        progress = False
        for rank in sorted(tuple(pending)):
            actions = pending[rank]
            if not actions:
                del pending[rank]
                continue
            action = actions[0]
            if not ready(action, rank):
                schedule[rank].append(None)
                continue
            schedule[rank].append(action)
            mark_completed(action, rank)
            actions.pop(0)
            progress = True
            if not actions:
                del pending[rank]

        for rank in sorted(tuple(pending)):
            if not schedule[rank] or schedule[rank][-1] is not None:
                continue
            action = pending[rank][0]
            if ready(action, rank):
                schedule[rank][-1] = action
                mark_completed(action, rank)
                pending[rank].pop(0)
                if not pending[rank]:
                    del pending[rank]
                progress = True
        if not progress:
            raise ValueError("pipeline schedule cannot make progress")
    return schedule

def _dump_chrometrace(schedule: Any, filename: str) -> None:
    events: list[dict[str, Any]] = []
    for rank in sorted(schedule):
        for timestep, action in enumerate(schedule[rank]):
            if action is None:
                continue
            events.append(
                {
                    "name": str(action),
                    "cat": (
                        "computation"
                        if action.computation_type
                        in {
                            _ComputationType.FORWARD,
                            _ComputationType.FULL_BACKWARD,
                            _ComputationType.BACKWARD_WEIGHT,
                        }
                        else "communication"
                    ),
                    "ph": "X",
                    "pid": rank,
                    "tid": rank,
                    "ts": timestep,
                    "dur": 1,
                }
            )
    import json

    with open(filename, "w", encoding="utf-8") as stream:
        json.dump({"traceEvents": events}, stream)

def _check_torch_compile_compatibility(stages: Any, schedule_name: str) -> None:
    del stages, schedule_name

def _default_merge_spec(value: Any) -> Any:
    if isinstance(value, tuple):
        return tuple(_default_merge_spec(item) for item in value)
    if isinstance(value, list):
        return [_default_merge_spec(item) for item in value]
    if isinstance(value, dict):
        return {key: _default_merge_spec(item) for key, item in value.items()}
    return TensorChunkSpec(0) if hasattr(value, "shape") else None
```