Source code for tensorplay.distributed.pipelining.stage

"""Pipeline stage execution and metadata management."""

from abc import ABC
from dataclasses import dataclass
import operator
from typing import Any, Callable

import tensorplay as tp
from .. import config as dist_config
from .. import distributed_core as dist

from ._backward import (
    _autograd_grad_for_inputs,
    stage_backward,
    stage_backward_input,
    stage_backward_weight,
)
from ._utils import (
    _MeshCache,
    PipeliningMetadataError,
    _StageBackwardMeta,
    _StageForwardMeta,
    _StageMeta,
    _DTensorMeta,
    _TensorMeta,
    _derive_grad_metas,
    flatten_args,
    _make_tensor_from_meta,
    InferenceMode,
    extract_tensor_meta,
    extract_tensor_metas,
    to_local_if_dtensor,
    validate_static_arg_grad_correspondence,
    validate_tensors_metadata,
)
from ..tensor import DTensor

__all__ = ["PipelineStage", "build_stage"]


def _normalize_model_output_as_tuple(output: Any) -> tuple[Any, ...]:
    if isinstance(output, list):
        return tuple(output)
    return output if isinstance(output, tuple) else (output,)


@dataclass
class _RecvInfo:
    input_name: str
    source: int | None
    buffer: Any
    tensor_meta: Any
    is_root_arg: bool = False

    def __init__(
        self,
        input_name: str,
        source: int | None,
        buffer: Any,
        tensor_meta: Any,
        is_root_arg: bool = False,
    ) -> None:
        self.input_name = input_name
        self.source = source
        self.buffer = buffer
        self.tensor_meta = tensor_meta
        self.is_root_arg = is_root_arg

    def __repr__(self) -> str:
        if self.is_root_arg:
            return f"_RecvInfo(input={self.input_name}, root_arg=True)"
        meta_type = type(self.tensor_meta).__name__ if self.tensor_meta else "None"
        buffer_shape = self.buffer.size() if self.buffer is not None else "None"
        return f"_RecvInfo(input={self.input_name}, source={self.source}, shape={buffer_shape}, meta={meta_type})"


def _build_p2p_direction_groups(group: Any) -> tuple[Any, Any]:
    if not dist.is_initialized():
        return group, group
    parent = group if group is not None else dist._get_default_group()
    if parent.size() <= 1:
        return group, group
    cache = getattr(_build_p2p_direction_groups, "_cache", None)
    if cache is None:
        cache = _build_p2p_direction_groups._cache = {}
    key = id(parent)
    cached = cache.get(key)
    if cached is not None and cached[0] is parent:
        return cached[1], cached[2]
    split_ranks = [list(range(parent.size()))]
    downstream = dist.split_group(
        parent_pg=parent,
        split_ranks=split_ranks,
        group_desc="pipeline_downstream",
    )
    upstream = dist.split_group(
        parent_pg=parent,
        split_ranks=split_ranks,
        group_desc="pipeline_upstream",
    )
    if downstream is dist.GroupMember.NON_GROUP_MEMBER or upstream is dist.GroupMember.NON_GROUP_MEMBER:
        raise RuntimeError("pipeline direction groups must contain the current rank")
    cache[key] = (parent, downstream, upstream)
    return downstream, upstream


class _PipelineStageBase(ABC):
    def __init__(self, submodule: Any, stage_index: int, num_stages: int, device: Any = None, group: Any = None, dw_builder: Callable[[], Callable[..., None]] | None = None) -> None:
        if stage_index < 0 or stage_index >= num_stages:
            raise ValueError("stage_index is outside the pipeline")
        self.submod = submodule
        self.stage_index = stage_index
        self.num_stages = num_stages
        self.device = device
        self.group = group
        self.dw_builder = dw_builder
        self.p2p_per_direction = bool(dist_config.pipeline_per_direction_p2p)
        if self.p2p_per_direction:
            self._downstream_group, self._upstream_group = _build_p2p_direction_groups(group)
        else:
            self._downstream_group = group
            self._upstream_group = group
        try:
            self.group_rank = int(dist.get_rank(group)) if dist.is_initialized() else stage_index
            self.group_size = int(dist.get_world_size(group)) if dist.is_initialized() else num_stages
        except (RuntimeError, ValueError):
            self.group_rank = stage_index
            self.group_size = num_stages
        if self.group_size > num_stages:
            raise ValueError("pipeline group cannot contain more ranks than stages")
        self.stage_index_to_group_rank = {
            index: index % self.group_size for index in range(num_stages)
        }
        self._has_backward = False
        self.fwd_cache: dict[int, tuple[Any, tuple[Any, ...]]] = {}
        self.bwd_cache: dict[int, Any] = {}
        self.output_chunks: list[Any] = []
        self.args_recv_info: dict[int, tuple[_RecvInfo, ...]] = {}
        self.act_send_info: dict[int, list[Any]] = {}
        self.grad_recv_info: dict[int, tuple[_RecvInfo, ...]] = {}
        self.grad_send_info: list[Any] | None = None
        self.chunks: int | None = None
        self._stage_meta = _StageMeta()
        self._mesh_cache = _MeshCache()
        self._input_chunks: dict[int, tuple[Any, ...]] = {}
        self._forward_inputs: dict[int, tuple[Any, ...]] = {}
        self.backward_state: dict[int, tuple[Any, Any, Any, Any]] = {}
        self.dw_runner: dict[int, Callable[[], Any]] = {}

    @property
    def has_backward(self) -> bool:
        return self._has_backward

    @has_backward.setter
    def has_backward(self, value: bool) -> None:
        self._has_backward = bool(value)

    @property
    def is_first(self) -> bool:
        return self.stage_index == 0

    @property
    def is_last(self) -> bool:
        return self.stage_index == self.num_stages - 1

    def _validate_stage_tensors(self, desc: str, expected: tuple[Any, ...] | None, actual: tuple[Any, ...]) -> None:
        if expected is None:
            raise PipeliningMetadataError(f"{desc}: metadata is unavailable")
        validate_tensors_metadata(desc, expected, actual)

    def _check_chunk_id(self, chunk_id: int) -> None:
        if self.chunks is None or chunk_id < 0 or chunk_id >= self.chunks:
            raise RuntimeError("chunk id is outside the configured range")

    def _create_grad_send_info(self, args_recv_info: tuple[_RecvInfo, ...]) -> list[Any]:
        return [item.source if isinstance(item, _RecvInfo) else None for item in args_recv_info]

    def _prepare_forward_infra(self, num_microbatches: int, args: Any, kwargs: Any, has_backward: bool) -> Any:
        self.chunks = num_microbatches
        self.has_backward = has_backward
        self._stage_meta.forward.input_metas = tuple(meta for meta in (extract_tensor_meta(value) for value in args) if meta is not None)
        self.args_recv_info = {index: tuple(_RecvInfo(str(pos), None, None, extract_tensor_meta(value), True) for pos, value in enumerate(args)) for index in range(num_microbatches)}

    def _prepare_backward_infra(
        self,
        num_microbatches: int,
        loss_fn: Any = None,
        target: Any = None,
        received_grad_meta: Any = None,
        loss_kwargs: Any = None,
    ) -> None:
        del loss_fn, target, loss_kwargs
        self.chunks = num_microbatches
        self.has_backward = True
        self._stage_meta.backward.output_grad_metas = tuple(received_grad_meta or ())
        self.grad_recv_info = {
            index: self._create_grad_recv_info(self.act_send_info)
            for index in range(num_microbatches)
        }
        self.grad_send_info = self._create_grad_send_info(
            self.args_recv_info.get(0, ())
        )

    def _setup_backward_recv_info(self, num_microbatches: int) -> None:
        self.chunks = num_microbatches
        self.grad_recv_info = {
            index: self._create_grad_recv_info(self.act_send_info)
            for index in range(num_microbatches)
        }

    def _create_grad_recv_info(self, act_send_info: Any) -> tuple[_RecvInfo, ...]:
        del act_send_info
        return ()

    def _resolve_peer_global_rank(self, stage_idx: int) -> int:
        peer_group_rank = self.stage_index_to_group_rank[int(stage_idx)]
        if self.group is None:
            return int(peer_group_rank)
        return int(dist.get_global_rank(self.group, peer_group_rank))

    def _get_recv_ops(self, recv_infos: Any, group: Any) -> list[Any]:
        if not dist.is_initialized():
            return []
        process_group = self.group if group is None else group
        operations = []
        for info in recv_infos:
            if not isinstance(info, _RecvInfo) or info.source is None or info.buffer is None:
                continue
            peer_group_rank = self.stage_index_to_group_rank[int(info.source)]
            peer = (
                peer_group_rank
                if process_group is None
                else dist.get_global_rank(process_group, peer_group_rank)
            )
            operations.append(dist.P2POp(dist.irecv, info.buffer, peer, process_group))
        return operations

    def set_local_fwd_input(self, prev_stage_outputs: Any, mb_index: int) -> None:
        values = _normalize_model_output_as_tuple(prev_stage_outputs)
        recv_infos = self.args_recv_info[mb_index]
        if len(recv_infos) != len(values):
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: local forward input count does not match "
                f"the receive metadata ({len(values)} != {len(recv_infos)})"
            )
        if self.is_first:
            raise AssertionError("local forward input is only valid for a non-first stage")
        for info, value in zip(recv_infos, values, strict=True):
            if info.is_root_arg:
                raise AssertionError("local forward input cannot replace a root argument")
            local_value = to_local_if_dtensor(value)
            if isinstance(local_value, tp.Tensor):
                local_value = local_value.detach()
                if (
                    info.tensor_meta is not None
                    and info.tensor_meta.requires_grad
                    and (local_value.is_floating_point() or local_value.is_complex())
                ):
                    local_value.requires_grad_(True)
            info.buffer = local_value
        self._input_chunks[mb_index] = tuple(info.buffer for info in recv_infos)

    def get_local_bwd_output(self, mb_index: int) -> Any:
        if not self.has_backward:
            raise AssertionError("cannot get a backward output without backward enabled")
        if self.is_first:
            raise AssertionError("the first stage has no local backward output")
        self._check_chunk_id(mb_index)
        return self.bwd_cache.pop(mb_index)

    def set_local_bwd_input(self, next_stage_bwd_outputs: Any, mb_index: int) -> None:
        values = next_stage_bwd_outputs
        if not isinstance(values, tuple):
            raise AssertionError(f"expected a tuple of gradients, got {type(values)}")
        if not self.has_backward:
            raise AssertionError("cannot set a backward input without backward enabled")
        if self.is_last:
            raise AssertionError("the last stage has no local backward input")
        recv_infos = self.grad_recv_info[mb_index]
        if len(recv_infos) != len(values):
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: local backward input count does not match "
                f"the receive metadata ({len(values)} != {len(recv_infos)})"
            )
        for info, value in zip(recv_infos, values, strict=True):
            if value is None:
                if info.buffer is not None:
                    info.buffer.zero_()
                continue
            if info.is_root_arg:
                raise AssertionError("local backward input cannot target a root argument")
            info.buffer = to_local_if_dtensor(value)

    def get_fwd_recv_ops(self, fwd_chunk_id: int) -> list[Any]:
        self._check_chunk_id(fwd_chunk_id)
        return self._get_recv_ops(
            self.args_recv_info.get(fwd_chunk_id, ()), self._downstream_group
        )

    def get_bwd_recv_ops(self, bwd_chunk_id: int) -> list[Any]:
        self._check_chunk_id(bwd_chunk_id)
        if not self.has_backward or self.is_last:
            return []
        return self._get_recv_ops(
            self.grad_recv_info.get(bwd_chunk_id, ()), self._upstream_group
        )

    def get_fwd_send_ops(self, fwd_chunk_id: int) -> list[Any]:
        self._check_chunk_id(fwd_chunk_id)
        output = self.fwd_cache[fwd_chunk_id][0]
        values = _normalize_model_output_as_tuple(output)
        operations = []
        for index, value in enumerate(values):
            for destination in self.act_send_info.get(index, ()):
                if destination is None:
                    continue
                value = to_local_if_dtensor(value, detach=True)
                if not isinstance(value, tp.Tensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: activation {index} is not a tensor"
                    )
                peer_group_rank = self.stage_index_to_group_rank[int(destination)]
                peer = (
                    peer_group_rank
                    if self._downstream_group is None
                    else dist.get_global_rank(self._downstream_group, peer_group_rank)
                )
                operations.append(
                    dist.P2POp(dist.isend, value, peer, self._downstream_group)
                )
        return operations

    def _get_grad_send_meta(self, input_idx: int) -> Any:
        input_grads = self._stage_meta.input_grads
        if input_grads is not None and input_idx < len(input_grads):
            return input_grads[input_idx]
        inputs = self._stage_meta.inputs
        if inputs is not None and input_idx < len(inputs):
            meta = inputs[input_idx]
            if meta is not None:
                return _derive_grad_metas((meta,))[0]
        raise PipeliningMetadataError(
            f"Stage {self.stage_index}: backward produced a gradient for input "
            f"{input_idx}, but no gradient metadata is available"
        )

    def get_bwd_send_ops(self, bwd_chunk_id: int) -> list[Any]:
        self._check_chunk_id(bwd_chunk_id)
        if not self.has_backward or self.is_first:
            return []
        if self.grad_send_info is None:
            self.grad_send_info = self._create_grad_send_info(
                self.args_recv_info.get(bwd_chunk_id, ())
            )
        gradients = self.bwd_cache.pop(bwd_chunk_id, ())
        operations = []
        for index, (gradient, destination) in enumerate(
            zip(gradients or (), self.grad_send_info, strict=True)
        ):
            if destination is None:
                if gradient is not None:
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: input {index} has a gradient but "
                        "no previous stage receives it"
                    )
                continue
            grad_meta = self._get_grad_send_meta(index)
            if grad_meta is None:
                if gradient is not None:
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: input {index} produced a gradient "
                        "without gradient metadata"
                    )
                continue
            if gradient is None:
                send_tensor = _make_tensor_from_meta(grad_meta, self.device).zero_()
            else:
                send_tensor = to_local_if_dtensor(gradient, detach=True)
                if not isinstance(send_tensor, tp.Tensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: input {index} gradient is not a tensor"
                    )
            peer_group_rank = self.stage_index_to_group_rank[int(destination)]
            peer = (
                peer_group_rank
                if self._upstream_group is None
                else dist.get_global_rank(self._upstream_group, peer_group_rank)
            )
            operations.append(
                dist.P2POp(dist.isend, send_tensor, peer, self._upstream_group)
            )
        return operations

    def clear_runtime_states(self) -> None:
        self.fwd_cache.clear()
        self.bwd_cache.clear()
        self.output_chunks.clear()
        self._input_chunks.clear()
        self._forward_inputs.clear()
        self.backward_state.clear()
        self.dw_runner.clear()
        for recv_infos in self.args_recv_info.values():
            for info in recv_infos:
                if not info.is_root_arg and isinstance(info.buffer, tp.Tensor):
                    info.buffer.grad = None

    def _map_tensor_from_recv_info(self, recv_infos: Any) -> tuple[Any, ...]:
        values = []
        for item in recv_infos:
            if item.is_root_arg:
                raise PipeliningMetadataError("root arguments are not received tensors")
            values.append(item.buffer)
        return tuple(values)

    def _retrieve_recv_activations(self, fwd_chunk_id: int) -> tuple[Any, ...]:
        recv_infos = self.args_recv_info.get(fwd_chunk_id, ())
        values = []
        for index, info in enumerate(recv_infos):
            if info.is_root_arg:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: root input cannot be received"
                )
            if info.buffer is None or info.tensor_meta is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: activation {index} has no receive buffer or metadata"
                )
            effective_requires_grad = bool(
                info.tensor_meta.requires_grad
                and self.has_backward
                and tp.is_grad_enabled()
            )
            if isinstance(info.tensor_meta, _DTensorMeta):
                local = info.buffer
                if not isinstance(local, tp.Tensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: DTensor activation buffer is not a tensor"
                    )
                local = local.detach()
                if effective_requires_grad and (
                    local.is_floating_point() or local.is_complex()
                ):
                    local.requires_grad_(True)
                mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key)
                values.append(
                    DTensor.from_local(
                        local,
                        device_mesh=mesh,
                        placements=info.tensor_meta.placements,
                        shape=info.tensor_meta.global_shape,
                        stride=info.tensor_meta.global_stride,
                        run_check=False,
                    )
                )
            else:
                value = info.buffer
                if not isinstance(value, tp.Tensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: activation {index} is not a tensor"
                    )
                value.requires_grad_(
                    effective_requires_grad
                    and (value.is_floating_point() or value.is_complex())
                )
                values.append(value)
        return tuple(values)

    def _retrieve_recv_grads(self, bwd_chunk_id: int) -> tuple[Any, ...]:
        recv_infos = self.grad_recv_info.get(bwd_chunk_id, ())
        values = []
        for index, info in enumerate(recv_infos):
            if info.is_root_arg:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: root input cannot receive a gradient"
                )
            if info.buffer is None:
                if info.tensor_meta is not None:
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: gradient {index} has metadata but no buffer"
                    )
                values.append(None)
                continue
            if info.tensor_meta is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: gradient {index} has a buffer but no metadata"
                )
            if isinstance(info.tensor_meta, _DTensorMeta):
                mesh = self._mesh_cache.get_mesh(info.tensor_meta.mesh_cache_key)
                values.append(
                    DTensor.from_local(
                        info.buffer,
                        device_mesh=mesh,
                        placements=info.tensor_meta.placements,
                        shape=info.tensor_meta.global_shape,
                        stride=info.tensor_meta.global_stride,
                        run_check=False,
                    )
                )
            else:
                values.append(info.buffer)
        return tuple(values)

    def forward_maybe_with_nosync(self, *args: Any, **kwargs: Any) -> Any:
        from ...nn.parallel.distributed import DistributedDataParallel

        if isinstance(self.submod, DistributedDataParallel):
            with self.submod.no_sync():
                return self.submod(*args, **kwargs)
        return self.submod(*args, **kwargs)

    def scale_grads(self, grad_scale_factor: float) -> None:
        for param in self.submod.parameters():
            if getattr(param, "grad", None) is not None:
                param.grad.div_(grad_scale_factor)

    def backward_maybe_with_nosync(self, backward_type: Any, bwd_kwargs: dict[str, Any], last_backward: bool = False) -> Any:
        del last_backward

        fsdp_flags = (
            ("set_is_last_backward", False),
            ("set_reshard_after_backward", False),
            ("set_requires_gradient_sync", False),
        )
        for method_name, value in fsdp_flags:
            method = getattr(self.submod, method_name, None)
            if callable(method):
                method(value)
        if backward_type == "full":
            return stage_backward(
                bwd_kwargs["stage_output"],
                bwd_kwargs["output_grads"],
                bwd_kwargs["input_values"],
            ), None
        if backward_type == "input":
            return stage_backward_input(
                bwd_kwargs["stage_output"],
                bwd_kwargs["output_grads"],
                bwd_kwargs["input_values"],
                self.submod.parameters(),
            )
        if backward_type == "weight":
            return stage_backward_weight(
                self.submod.parameters(),
                bwd_kwargs["param_groups"] or [],
            ), None
        raise RuntimeError(f"unknown backward type {backward_type!r}")

    def forward_one_chunk(self, fwd_chunk_id: int, args: tuple[Any, ...], kwargs: dict[str, Any], save_forward_output: bool = True) -> Any:
        self._check_chunk_id(fwd_chunk_id)
        composite_args = args if self.is_first else self._retrieve_recv_activations(fwd_chunk_id)
        output = self.forward_maybe_with_nosync(*composite_args, **kwargs)
        self._forward_inputs[fwd_chunk_id] = tuple(
            value
            for value in flatten_args(composite_args)
            if isinstance(value, tp.Tensor) or value is not None
        ) + tuple(
            value
            for value in flatten_args(kwargs)
            if isinstance(value, tp.Tensor) or value is not None
        )
        output_tuple = _normalize_model_output_as_tuple(output)
        self.fwd_cache[fwd_chunk_id] = (output, output_tuple)
        if save_forward_output:
            while len(self.output_chunks) <= fwd_chunk_id:
                self.output_chunks.append(None)
            self.output_chunks[fwd_chunk_id] = output
        self._stage_meta.forward.output_metas = tuple(meta for meta in (extract_tensor_meta(value) for value in output_tuple) if meta is not None)
        return output

    def backward_one_chunk(self, bwd_chunk_id: int, loss: Any = None, full_backward: bool = True, last_backward: bool = False) -> Any:
        if not self.has_backward:
            return None
        self._check_chunk_id(bwd_chunk_id)
        output, output_values = self.fwd_cache.pop(bwd_chunk_id)
        if self.is_last:
            stage_output = output if loss is None else loss
            output_grads = None
        else:
            stage_output = output_values
            output_grads = self._retrieve_recv_grads(bwd_chunk_id)
        input_values = self._forward_inputs.pop(bwd_chunk_id, ())
        bwd_kwargs = {
            "stage_output": stage_output,
            "output_grads": output_grads,
            "input_values": input_values,
        }
        grads_input: tuple[Any, ...] = ()
        if self.dw_builder is not None:
            grads_input, _ = self.backward_maybe_with_nosync(
                "full", bwd_kwargs, last_backward=last_backward
            )
            if full_backward:
                self.dw_builder()()
            else:
                self.dw_runner[bwd_chunk_id] = self.dw_builder()
        elif full_backward:
            grads_input, _ = self.backward_maybe_with_nosync(
                "full", bwd_kwargs, last_backward=last_backward
            )
        else:
            param_groups = None
            if not self.is_first:
                grads_input, param_groups = self.backward_maybe_with_nosync(
                    "input", bwd_kwargs, last_backward=last_backward
                )
            self.backward_state[bwd_chunk_id] = (
                input_values,
                param_groups,
                stage_output,
                output_grads,
            )
            self.dw_runner[bwd_chunk_id] = lambda: None
        num_forward_inputs = len(self._stage_meta.inputs or ())
        self.bwd_cache[bwd_chunk_id] = tuple(grads_input[:num_forward_inputs])
        return self.bwd_cache[bwd_chunk_id]

    def backward_weight_one_chunk(self, bwd_chunk_id: int, last_backward: bool = False) -> Any:
        if not self.has_backward:
            return None
        runner = self.dw_runner.pop(bwd_chunk_id, None)
        if runner is None:
            raise AssertionError(
                f"backward weight requested for chunk {bwd_chunk_id} without input backward"
            )
        if self.dw_builder is not None:
            return runner()
        input_values, param_groups, stage_output, output_grads = self.backward_state.pop(
            bwd_chunk_id
        )
        if self.is_first:
            self.backward_maybe_with_nosync(
                "full",
                {
                    "stage_output": stage_output,
                    "output_grads": output_grads,
                    "input_values": input_values,
                },
                last_backward=last_backward,
            )
        else:
            self.backward_maybe_with_nosync(
                "weight",
                {"param_groups": param_groups},
                last_backward=last_backward,
            )
        return None

    def _get_init_p2p_neighbors_ops(self) -> list[Any]:
        operations: list[Any] = []
        next_stage_peer_rank = self.stage_index_to_group_rank.get(
            self.stage_index + 1
        )
        previous_stage_peer_rank = self.stage_index_to_group_rank.get(
            self.stage_index - 1
        )
        downstream_recv_tensor = tp.zeros(
            1, device=self.device, dtype=tp.float32
        )
        upstream_recv_tensor = tp.zeros(
            1, device=self.device, dtype=tp.float32
        )
        send_tensor = tp.tensor(
            self.stage_index, device=self.device, dtype=tp.float32
        )
        if not self.is_first:
            operations.append(
                dist.P2POp(
                    dist.irecv,
                    downstream_recv_tensor,
                    group_peer=previous_stage_peer_rank,
                    group=self._downstream_group,
                )
            )
        if not self.is_last:
            operations.append(
                dist.P2POp(
                    dist.isend,
                    send_tensor,
                    group_peer=next_stage_peer_rank,
                    group=self._downstream_group,
                )
            )
        if not self.is_first:
            operations.append(
                dist.P2POp(
                    dist.isend,
                    send_tensor,
                    group_peer=previous_stage_peer_rank,
                    group=self._upstream_group,
                )
            )
        if not self.is_last:
            operations.append(
                dist.P2POp(
                    dist.irecv,
                    upstream_recv_tensor,
                    group_peer=next_stage_peer_rank,
                    group=self._upstream_group,
                )
            )
        return operations

    def perform_reduce_grad(self, grad_scale_factor: float) -> None:
        state_getter = getattr(self.submod, "_get_fsdp_state", None)
        if not callable(state_getter):
            state_getter = getattr(self.submod, "_get_replicate_state", None)
        if callable(state_getter):
            for method_name, value in (
                ("set_is_last_backward", True),
                ("set_reshard_after_backward", True),
                ("set_requires_gradient_sync", True),
            ):
                method = getattr(self.submod, method_name, None)
                if callable(method):
                    method(value)
            state = state_getter()
            state_context = getattr(state, "_state_ctx", None)
            states = (
                getattr(state_context, "all_states", None)
                or getattr(state_context, "states", None)
                or [state]
            )
            for state_item in states:
                groups_getter = getattr(state_item, "_all_param_groups", None)
                if callable(groups_getter):
                    for param_group in groups_getter():
                        param_group.post_backward()
            callback = getattr(state, "_root_post_backward_final_callback", None)
            if callable(callback):
                callback()
        self.scale_grads(grad_scale_factor)


class _PipelineStage(_PipelineStageBase):
    def __init__(self, stage_module: Any, stage_index: int, pipe_info: Any, device: Any = None, group: Any = None) -> None:
        super().__init__(stage_module, stage_index, pipe_info.num_stages, device, group)
        self.pipe_info = pipe_info
        graph_owner = getattr(pipe_info, "graph", None)
        self.graph = getattr(graph_owner, "graph", graph_owner)
        submod_nodes = [
            node
            for node in getattr(self.graph, "nodes", ())
            if getattr(node, "op", None) == "call_module"
        ]
        if len(submod_nodes) != self.num_stages:
            raise PipeliningMetadataError(
                f"Number of submodules in pipe graph {len(submod_nodes)} does not match "
                f"number of stages {self.num_stages}"
            )
        self.node = submod_nodes[stage_index]
        self.name = self.node.name
        self.submod_to_stage_index = {
            getattr(node, "name", ""): index
            for index, node in enumerate(submod_nodes)
        }
        self._move_submod_to_device()

    def _move_submod_to_device(self) -> None:
        parameters = getattr(self.submod, "parameters", None)
        if callable(parameters) and any(
            bool(getattr(parameter, "is_meta", False))
            for parameter in parameters()
        ):
            return
        if self.device is not None and hasattr(self.submod, "to"):
            self.submod.to(self.device)

    def get_stage_index_of_submod(self, submod_name: str) -> int:
        try:
            return self.submod_to_stage_index[submod_name]
        except KeyError as exc:
            raise PipeliningMetadataError(
                f"stage {submod_name!r} is not present"
            ) from exc

    def _tensor_from_meta(self, meta: Any, value: Any = None) -> Any:
        if isinstance(value, tp.Tensor):
            result = value.detach().clone()
        elif meta is not None and hasattr(meta, "to_tensor"):
            result = meta.to_tensor(self.device)
        elif meta is not None and hasattr(meta, "shape"):
            result = tp.empty(tuple(meta.shape), dtype=meta.dtype, device=self.device)
        else:
            result = value
        if isinstance(result, tp.Tensor) and self.has_backward:
            if result.is_floating_point() or result.is_complex():
                result.requires_grad_(True)
        return result

    def _create_act_recv_info(self) -> tuple[_RecvInfo, ...]:
        if self.node is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: graph stage node is unavailable"
            )
        stage_graph = getattr(self.submod, "graph", None)
        placeholders = [
            node
            for node in getattr(stage_graph, "nodes", ())
            if getattr(node, "op", None) == "placeholder"
        ]
        outer_args = tuple(getattr(self.node, "args", ()))
        result: list[_RecvInfo] = []
        if len(placeholders) != len(outer_args):
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: graph placeholder and dependency counts differ"
            )
        for placeholder, arg_node in zip(placeholders, outer_args, strict=True):
            meta_value = getattr(placeholder, "meta", {}).get("val")
            if meta_value is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: placeholder metadata is unavailable"
                )
            if isinstance(meta_value, DTensor):
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: distributed tensor metadata is unsupported for graph stages"
                )
            if getattr(arg_node, "op", None) == "placeholder":
                result.append(
                    _RecvInfo(
                        f"root_input_{getattr(placeholder, 'name', 'input')}",
                        None,
                        None,
                        _TensorMeta.from_tensor(meta_value),
                        True,
                    )
                )
                continue
            while getattr(arg_node, "target", None) is operator.getitem:
                arg_node = arg_node.args[0]
            if getattr(arg_node, "op", None) != "call_module":
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: expected a stage dependency"
                )
            source = self.get_stage_index_of_submod(getattr(arg_node, "name", ""))
            meta = _TensorMeta(
                shape=tuple(meta_value.shape),
                stride=tuple(meta_value.stride()),
                dtype=meta_value.dtype,
                requires_grad=bool(
                    self.has_backward
                    and (
                        meta_value.is_floating_point()
                        or meta_value.is_complex()
                    )
                ),
            )
            result.append(
                _RecvInfo(
                    getattr(arg_node, "name", getattr(placeholder, "name", "input")),
                    source,
                    _make_tensor_from_meta(meta, self.device),
                    meta,
                )
            )
        return tuple(result)

    def _prepare_forward_infra(self, num_microbatches: int, args: Any, kwargs: Any, has_backward: bool) -> Any:
        del kwargs
        self.chunks = int(num_microbatches)
        self.has_backward = bool(has_backward)
        for index in range(self.chunks):
            self.args_recv_info[index] = self._create_act_recv_info()
        recv_infos = self.args_recv_info[0]
        if self.is_first:
            if not isinstance(args, tuple):
                raise AssertionError("first stage requires real tensor args")
            self._stage_meta.inputs = tuple(
                info.tensor_meta for info in recv_infos[: len(args)]
            )
        else:
            self._stage_meta.inputs = tuple(
                info.tensor_meta for info in recv_infos if not info.is_root_arg
            )
        self.act_send_info = self._create_act_send_info()

    def _prepare_backward_infra(
        self,
        num_microbatches: int,
        loss_fn: Any = None,
        target: Any = None,
        received_grad_meta: Any = None,
        loss_kwargs: Any = None,
    ) -> None:
        del loss_fn, target, received_grad_meta, loss_kwargs
        if self._stage_meta.inputs is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: inputs metadata required for backward inference."
            )
        self._stage_meta.input_grads = _derive_grad_metas(self._stage_meta.inputs)
        self._setup_backward_recv_info(num_microbatches)
        return None

    def find_dst_rank(self, user: Any) -> int:
        if getattr(user, "op", None) != "call_module":
            return None
        return self.get_stage_index_of_submod(getattr(user, "name", ""))

    def _create_act_send_info(self) -> dict[int, list[int]]:
        if self.node is None:
            return {0: [self.stage_index + 1]} if not self.is_last else {0: []}
        result: dict[int, list[int]] = {}
        for user in getattr(self.node, "users", ()):
            if getattr(user, "target", None) is operator.getitem:
                output_index = int(user.args[1])
                destinations = result.setdefault(output_index, [])
                for child in getattr(user, "users", ()):
                    destination = self.find_dst_rank(child)
                    if destination is not None and destination not in destinations:
                        destinations.append(destination)
            else:
                destination = self.find_dst_rank(user)
                if destination is not None:
                    destinations = result.setdefault(0, [])
                    if destination not in destinations:
                        destinations.append(destination)
        output_node = self._get_output_node()
        if output_node is not None:
            values = output_node.args[0] if getattr(output_node, "args", ()) else ()

            def flatten_graph_values(value: Any) -> list[Any]:
                if isinstance(value, (tuple, list)):
                    result_values: list[Any] = []
                    for item in value:
                        result_values.extend(flatten_graph_values(item))
                    return result_values
                if isinstance(value, dict):
                    result_values = []
                    for item in value.values():
                        result_values.extend(flatten_graph_values(item))
                    return result_values
                return [value]

            output_metas: list[_TensorMeta] = []
            for index, value in enumerate(flatten_graph_values(values)):
                example_value = getattr(value, "meta", {}).get("val")
                if example_value is None:
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: output metadata is unavailable at index {index}"
                    )
                if isinstance(example_value, DTensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: distributed tensor metadata is unsupported for graph stages"
                    )
                if not isinstance(example_value, tp.Tensor):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: output {index} is not a tensor"
                    )
                output_metas.append(
                    _TensorMeta(
                        shape=tuple(example_value.shape),
                        stride=tuple(example_value.stride()),
                        dtype=example_value.dtype,
                        requires_grad=bool(
                            self.has_backward
                            and (
                                example_value.is_floating_point()
                                or example_value.is_complex()
                            )
                        ),
                    )
                )
            self._stage_meta.outputs = tuple(output_metas)
        return result

    def _create_grad_recv_info(self, act_send_info: Any) -> tuple[_RecvInfo, ...]:
        if self._stage_meta.outputs is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: outputs metadata required for grad recv info."
            )

        outputs_meta = self._stage_meta.outputs
        output_grads_metas: list[Any] = []
        grad_recv_infos: list[_RecvInfo] = []
        for out_idx, out_meta in enumerate(outputs_meta):
            dst_list = act_send_info.get(out_idx, [])
            grad_src = dst_list[0] if dst_list else self.stage_index + 1
            if not dst_list or not out_meta.requires_grad:
                output_grads_metas.append(None)
                grad_recv_infos.append(
                    _RecvInfo(
                        f"recv_grad_for_{self.stage_index}_none_{out_idx}",
                        grad_src,
                        None,
                        None,
                    )
                )
                continue
            grad_meta = _TensorMeta(
                shape=out_meta.shape,
                stride=out_meta.stride,
                dtype=out_meta.dtype,
                requires_grad=False,
            )
            output_grads_metas.append(grad_meta)
            if len(dst_list) != 1:
                raise PipeliningMetadataError(
                    "Backward of skip connections not supported yet"
                )
            grad_recv_infos.append(
                _RecvInfo(
                    f"recv_grad_for_{self.stage_index}_from_{grad_src}",
                    grad_src,
                    _make_tensor_from_meta(grad_meta, self.device),
                    grad_meta,
                )
            )
        self._stage_meta.output_grads = tuple(output_grads_metas)
        if self._stage_meta.inputs is not None:
            self._stage_meta.input_grads = _derive_grad_metas(self._stage_meta.inputs)
        return tuple(grad_recv_infos)

    def backward_one_chunk(self, bwd_chunk_id: int, loss: Any = None, full_backward: bool = True, last_backward: bool = False) -> Any:
        self._check_chunk_id(bwd_chunk_id)
        return super().backward_one_chunk(
            bwd_chunk_id,
            loss=loss,
            full_backward=full_backward,
            last_backward=last_backward,
        )

    def _get_output_node(self) -> Any:
        for graph in (getattr(self.submod, "graph", None), self.graph):
            output_node = next(
                (
                    node
                    for node in getattr(graph, "nodes", ())
                    if getattr(node, "op", None) == "output"
                ),
                None,
            )
            if output_node is not None:
                return output_node
        return None



[docs]
def build_stage(stage_module: Any, stage_index: int, pipe_info: Any, device: Any = None, group: Any = None) -> _PipelineStage:
    return _PipelineStage(stage_module, stage_index, pipe_info, device, group)




[docs]
class PipelineStage(_PipelineStageBase):
    def __init__(self, submodule: Any, stage_index: int, num_stages: int, device: Any = None, input_args: tuple[Any, ...] | None = None, output_args: Any = None, output_grads: Any = None, input_grads: Any = None, group: Any = None, dw_builder: Callable[[], Callable[..., None]] | None = None, get_mesh: Any = None) -> None:
        super().__init__(submodule, stage_index, num_stages, device, group, dw_builder)
        self._mesh_cache = _MeshCache(get_mesh_cb=get_mesh)
        self._input_example = _normalize_model_output_as_tuple(input_args) if input_args is not None else ()
        self._output_example = _normalize_model_output_as_tuple(output_args) if output_args is not None else None
        input_grad_values = _normalize_model_output_as_tuple(input_grads) if input_grads is not None else None
        output_grad_values = _normalize_model_output_as_tuple(output_grads) if output_grads is not None else None
        self._user_meta = _StageMeta()
        self._user_meta.inputs = extract_tensor_metas(self._input_example) if self._input_example else None
        self._user_meta.outputs = extract_tensor_metas(self._output_example) if self._output_example is not None else None
        self._user_meta.input_grads = extract_tensor_metas(input_grad_values, allow_none=True) if input_grad_values is not None else None
        self._user_meta.output_grads = extract_tensor_metas(output_grad_values, allow_none=True) if output_grad_values is not None else None
        for values in (self._input_example, self._output_example, input_grad_values, output_grad_values):
            if values:
                self._mesh_cache.update_from_tensors(values)
        if self._user_meta.has_dtensors():
            if self._input_example and input_grad_values:
                validate_static_arg_grad_correspondence(
                    self.stage_index,
                    self._input_example,
                    input_grad_values,
                    is_input=True,
                )
            if self._output_example and output_grad_values:
                validate_static_arg_grad_correspondence(
                    self.stage_index,
                    self._output_example,
                    output_grad_values,
                    is_input=False,
                )
        self._inference_mode: InferenceMode | None = None
        self._fwd_outputs_for_bwd_meta: tuple[Any, ...] | None = None
        self._fwd_inputs_for_bwd_meta: tuple[Any, ...] | None = None
        self._fwd_kwargs_tensors_for_bwd_meta: tuple[Any, ...] | None = None
        self._metadata_inference_buffer_backup: list[tuple[Any, Any]] | None = None
        self._inference_mode = None

    def _prepare_forward_infra(self, num_microbatches: int, args: Any, kwargs: Any, has_backward: bool) -> Any:
        self.chunks = int(num_microbatches)
        self.has_backward = bool(has_backward)
        self._inference_mode = (
            InferenceMode.DYNAMIC
            if InferenceMode.needs_dynamic(self._user_meta, has_backward)
            else InferenceMode.STATIC
        )
        source_args = args
        if source_args is None or source_args == ():
            source_args = self._input_example
        fwd_meta_output = None
        if self._inference_mode == InferenceMode.DYNAMIC:
            fwd_meta_output = self._forward_metadata_inference(
                source_args, kwargs, has_backward
            )
        else:
            self._stage_meta.inputs = self._user_meta.inputs
            self._stage_meta.outputs = self._user_meta.outputs
        if self._stage_meta.inputs is None and source_args:
            self._stage_meta.inputs = extract_tensor_metas(tuple(source_args))
        if self._stage_meta.outputs is None and self._output_example is not None:
            self._stage_meta.outputs = extract_tensor_metas(self._output_example)
        self._setup_forward_recv_info(self.chunks, has_backward)
        self._setup_forward_send_info()
        return fwd_meta_output

    def _prepare_backward_infra(
        self,
        num_microbatches: int,
        loss_fn: Any = None,
        target: Any = None,
        received_grad_meta: Any = None,
        loss_kwargs: Any = None,
    ) -> Any:
        self.chunks = int(num_microbatches)
        self.has_backward = True
        if self._inference_mode == InferenceMode.DYNAMIC:
            result = self._backward_metadata_inference(
                loss_fn,
                target,
                received_grad_meta,
                loss_kwargs,
            )
            self._validate_inferred_metadata()
        else:
            result = None
            self._stage_meta.inputs = self._user_meta.inputs
            self._stage_meta.outputs = self._user_meta.outputs
            self._stage_meta.input_grads = self._user_meta.input_grads
            self._stage_meta.output_grads = self._user_meta.output_grads
        if isinstance(received_grad_meta, _StageBackwardMeta):
            self._stage_meta.output_grads = received_grad_meta.input_grad_metas
        if self._stage_meta.output_grads is None:
            if self._stage_meta.outputs is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: output metadata is required for backward inference."
                )
            self._stage_meta.output_grads = _derive_grad_metas(self._stage_meta.outputs)
        if self._stage_meta.input_grads is None:
            if self._stage_meta.inputs is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: input metadata is required for backward inference."
                )
            self._stage_meta.input_grads = _derive_grad_metas(self._stage_meta.inputs)
        self._setup_backward_recv_info(num_microbatches)
        self.grad_send_info = self._create_grad_send_info(self.args_recv_info.get(0, ()))
        return result

    def get_fwd_recv_ops(self, fwd_chunk_id: int) -> list[Any]:
        self._check_chunk_id(fwd_chunk_id)
        if self.is_first:
            return []
        return self._get_recv_ops(
            self.args_recv_info[fwd_chunk_id], self._downstream_group
        )

    def _recv_meta(self, src_stage: int) -> Any:
        objects = [None]
        dist.recv_object_list(
            objects,
            src=self._resolve_peer_global_rank(src_stage),
            group=self.group,
            device=self.device,
        )
        if len(objects) != 1:
            raise PipeliningMetadataError(
                f"expected one metadata object, got {len(objects)}"
            )
        return objects[0]

    def _send_meta(self, meta: Any, dst_stage: int) -> None:
        dist.send_object_list(
            [meta],
            dst=self._resolve_peer_global_rank(dst_stage),
            group=self.group,
            device=self.device,
        )

    def _is_same_rank(self, other_stage: int) -> bool:
        return self.stage_index_to_group_rank[int(other_stage)] == self.group_rank

    def _warmup_forward_vote(self, has_backward: bool, received_acc: Any = None) -> Any:
        my_vote = 0 if InferenceMode.needs_dynamic(self._user_meta, has_backward) else 1
        vote = tp.tensor([my_vote], dtype=tp.int32, device=self.device)
        if self.is_first:
            accumulated = vote
        elif self._is_same_rank(self.stage_index - 1):
            if received_acc is None:
                raise AssertionError("forward vote is missing the accumulated value")
            accumulated = received_acc * vote
        else:
            accumulated = tp.zeros(1, dtype=tp.int32, device=self.device)
            dist.recv(
                accumulated,
                src=self._resolve_peer_global_rank(self.stage_index - 1),
                group=self.group,
            )
            accumulated = accumulated * vote
        if not self.is_last and not self._is_same_rank(self.stage_index + 1):
            dist.send(
                accumulated,
                dst=self._resolve_peer_global_rank(self.stage_index + 1),
                group=self.group,
            )
        return accumulated

    def _warmup_backward_result(self, received_result: Any = None) -> Any:
        if self.is_last or self._is_same_rank(self.stage_index + 1):
            if received_result is None:
                raise AssertionError("backward vote is missing the accumulated value")
            result = received_result
        else:
            result = tp.zeros(1, dtype=tp.int32, device=self.device)
            dist.recv(
                result,
                src=self._resolve_peer_global_rank(self.stage_index + 1),
                group=self.group,
            )
        if not self.is_first and not self._is_same_rank(self.stage_index - 1):
            dist.send(
                result,
                dst=self._resolve_peer_global_rank(self.stage_index - 1),
                group=self.group,
            )
        return result

    def _compute_outputs(self, *args: Any, module: Any = None, **kwargs: Any) -> Any:
        return (self.submod if module is None else module)(*args, **kwargs)

    def _compute_input_grads(
        self,
        outputs: Any,
        all_fwd_inputs: Any,
        grad_outputs: Any = None,
    ) -> tuple[Any, ...]:
        return _autograd_grad_for_inputs(
            tuple(outputs),
            tuple(all_fwd_inputs),
            None if grad_outputs is None else tuple(grad_outputs),
            allow_unused=True,
        )

    def backward_one_chunk(self, bwd_chunk_id: int, loss: Any = None, full_backward: bool = True, last_backward: bool = False) -> Any:
        return super().backward_one_chunk(
            bwd_chunk_id,
            loss=loss,
            full_backward=full_backward,
            last_backward=last_backward,
        )

    def _to_tensor(self, arg: Any) -> Any:
        if isinstance(arg, DTensor):
            local = arg.to_local().detach()
            if getattr(arg, "requires_grad", False) and (
                local.is_floating_point() or local.is_complex()
            ):
                local.requires_grad_(True)
            return DTensor.from_local(
                local,
                device_mesh=arg.device_mesh,
                placements=arg.placements,
                shape=arg.shape,
                stride=arg.stride(),
            )
        if isinstance(arg, tp.Tensor):
            result = arg.detach()
            if arg.requires_grad:
                result.requires_grad_(True)
            return result
        if isinstance(arg, _DTensorMeta):
            mesh = self._mesh_cache.get_mesh(arg.mesh_cache_key)
            local = _make_tensor_from_meta(arg, self.device)
            if arg.requires_grad and (
                local.is_floating_point() or local.is_complex()
            ):
                local.requires_grad_(True)
            return DTensor.from_local(
                local,
                device_mesh=mesh,
                placements=arg.placements,
                shape=arg.global_shape,
                stride=arg.global_stride,
            )
        if isinstance(arg, _TensorMeta):
            result = arg.to_tensor(self.device)
            if arg.requires_grad and (
                result.is_floating_point() or result.is_complex()
            ):
                result.requires_grad_(True)
            return result
        raise PipeliningMetadataError(
            f"unsupported metadata value {type(arg).__name__}"
        )

    def _ones_from_metadata(self, meta: Any) -> Any:
        local = tp.ones(meta.shape, dtype=meta.dtype, device=self.device)
        if isinstance(meta, _DTensorMeta):
            mesh = self._mesh_cache.get_mesh(meta.mesh_cache_key)
            return DTensor.from_local(
                local,
                device_mesh=mesh,
                placements=meta.placements,
                shape=meta.global_shape,
                stride=meta.global_stride,
            )
        return local

    def _pre_metadata_inference_backup(self) -> None:
        if self._inference_mode != InferenceMode.DYNAMIC:
            return
        if self._metadata_inference_buffer_backup is not None:
            raise RuntimeError("metadata inference backup is already active")
        named_buffers = getattr(self.submod, "named_buffers", None)
        if callable(named_buffers):
            self._metadata_inference_buffer_backup = [
                (buffer, buffer.detach().clone())
                for _, buffer in named_buffers(remove_duplicate=False)
            ]

    def _forward_metadata_inference(self, args: Any, kwargs: Any, has_backward: bool) -> Any:
        kwargs = kwargs or {}
        if self.is_first:
            if args is None or isinstance(args, _StageForwardMeta):
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: first stage requires tensor inputs"
                )
            values = tuple(args)
            self._stage_meta.inputs = extract_tensor_metas(values)
            inference_args = tuple(self._to_tensor(value) for value in values)
        elif self._is_same_rank(self.stage_index - 1) or isinstance(args, _StageForwardMeta):
            if not isinstance(args, _StageForwardMeta):
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: forward metadata from the previous stage is required"
                )
            input_metas = args.forward_metas
            self._stage_meta.inputs = tuple(input_metas)
            inference_args = tuple(self._to_tensor(meta) for meta in input_metas)
        else:
            recv_meta = self._recv_meta(self.stage_index - 1)
            if not isinstance(recv_meta, _StageForwardMeta):
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: invalid forward metadata received from the previous stage"
                )
            input_metas = recv_meta.forward_metas
            self._stage_meta.inputs = tuple(input_metas)
            inference_args = tuple(self._to_tensor(meta) for meta in input_metas)
        inference_kwargs = {
            key: self._to_tensor(value) if isinstance(value, tp.Tensor) else value
            for key, value in kwargs.items()
        }
        with (tp.enable_grad() if has_backward else tp.no_grad()):
            output = self._compute_outputs(
                *inference_args,
                module=self.submod,
                **inference_kwargs,
            )
        output_values = _normalize_model_output_as_tuple(output)
        self._stage_meta.outputs = tuple(
            meta
            for meta in (extract_tensor_meta(value) for value in output_values)
            if meta is not None
        )
        self._fwd_outputs_for_bwd_meta = output_values
        self._fwd_inputs_for_bwd_meta = inference_args
        self._fwd_kwargs_tensors_for_bwd_meta = tuple(
            value
            for value in flatten_args(inference_kwargs)
            if isinstance(value, tp.Tensor) or isinstance(value, DTensor)
        )
        fwd_meta = _StageForwardMeta(forward_metas=self._stage_meta.outputs)
        if self.is_last or self._is_same_rank(self.stage_index + 1):
            return fwd_meta
        self._send_meta(fwd_meta, self.stage_index + 1)
        return None

    def _backward_metadata_inference(self, loss_fn: Any, target: Any, received_grad_meta: Any, loss_kwargs: Any) -> Any:
        fwd_outputs = self._fwd_outputs_for_bwd_meta
        fwd_inputs = self._fwd_inputs_for_bwd_meta
        if fwd_outputs is None or fwd_inputs is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: forward metadata inference must run first"
            )
        all_inputs = list(fwd_inputs) + list(self._fwd_kwargs_tensors_for_bwd_meta or ())
        if self.is_last:
            if loss_fn is None or target is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: loss_fn and target are required for backward inference"
                )
            output_value = fwd_outputs[0] if len(fwd_outputs) == 1 else fwd_outputs
            loss = loss_fn(output_value, self._to_tensor(target), **(loss_kwargs or {}))
            input_grads = self._compute_input_grads((loss,), all_inputs)
            self._stage_meta.output_grads = None
        else:
            if self._is_same_rank(self.stage_index + 1) or (
                not dist.is_initialized() and received_grad_meta is not None
            ):
                if not isinstance(received_grad_meta, _StageBackwardMeta):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: backward metadata from the next stage is required"
                    )
                output_grad_metas = received_grad_meta.backward_metas
            else:
                recv_meta = self._recv_meta(self.stage_index + 1)
                if not isinstance(recv_meta, _StageBackwardMeta):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: invalid backward metadata received from the next stage"
                    )
                output_grad_metas = recv_meta.backward_metas
            self._stage_meta.output_grads = output_grad_metas
            if len(fwd_outputs) != len(output_grad_metas):
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: output and gradient metadata counts differ"
                )
            filtered_outputs = []
            filtered_grad_outputs = []
            for index, (output, grad_meta) in enumerate(
                zip(fwd_outputs, output_grad_metas, strict=True)
            ):
                if not isinstance(output, (tp.Tensor, DTensor)):
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: output {index} is not a tensor"
                    )
                if not output.requires_grad and getattr(output, "grad_fn", None) is None:
                    if grad_meta is not None:
                        raise PipeliningMetadataError(
                            f"Stage {self.stage_index}: output {index} has gradient metadata but does not require gradients"
                        )
                    continue
                filtered_outputs.append(output)
                filtered_grad_outputs.append(
                    self._ones_from_metadata(grad_meta) if grad_meta is not None else None
                )
            if filtered_outputs:
                input_grads = self._compute_input_grads(
                    filtered_outputs,
                    all_inputs,
                    filtered_grad_outputs,
                )
            else:
                input_grads = tuple(None for _ in all_inputs)
        input_metas = self._stage_meta.inputs or ()
        if len(input_grads) < len(input_metas):
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: backward returned too few input gradients"
            )
        self._stage_meta.input_grads = tuple(
            extract_tensor_meta(gradient)
            if isinstance(gradient, (tp.Tensor, DTensor))
            else (
                _derive_grad_metas((meta,))[0]
                if meta is not None and meta.requires_grad
                else None
            )
            for meta, gradient in zip(input_metas, input_grads)
        )
        bwd_meta = _StageBackwardMeta(backward_metas=self._stage_meta.input_grads)
        if self.is_first or self._is_same_rank(self.stage_index - 1):
            return bwd_meta
        self._send_meta(bwd_meta, self.stage_index - 1)
        return None

    def _post_metadata_inference_cleanup(self) -> None:
        if self._metadata_inference_buffer_backup is not None:
            with tp.no_grad():
                for buffer, saved in self._metadata_inference_buffer_backup:
                    buffer.copy_(saved)
            self._metadata_inference_buffer_backup = None
        self._fwd_outputs_for_bwd_meta = None
        self._fwd_inputs_for_bwd_meta = None
        self._fwd_kwargs_tensors_for_bwd_meta = None
        self.clear_runtime_states()

    def _validate_inferred_metadata(self) -> None:
        if not self._stage_meta.outputs:
            raise PipeliningMetadataError("stage output metadata is empty")
        for user_meta, inferred_meta, label in (
            (self._user_meta.inputs, self._stage_meta.inputs, "input"),
            (self._user_meta.outputs, self._stage_meta.outputs, "output"),
            (self._user_meta.input_grads, self._stage_meta.input_grads, "input_grad"),
            (self._user_meta.output_grads, self._stage_meta.output_grads, "output_grad"),
        ):
            if user_meta is not None and inferred_meta is not None:
                validate_tensors_metadata(
                    f"Stage {self.stage_index} {label}",
                    user_meta,
                    inferred_meta,
                    raise_on_mismatch=False,
                    warn_on_mismatch=True,
                )

    def _setup_forward_recv_info(self, num_microbatches: int, has_backward: bool) -> None:
        del has_backward
        if self._stage_meta.inputs is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: inputs metadata is required for receive setup."
            )
        self.args_recv_info = {}
        for chunk_id in range(num_microbatches):
            if self.is_first:
                infos = tuple(
                    _RecvInfo(
                        f"root_input_{index}",
                        None,
                        None,
                        meta,
                        True,
                    )
                    for index, meta in enumerate(self._stage_meta.inputs)
                )
            else:
                infos = tuple(
                    _RecvInfo(
                        f"recv_for_{self.stage_index}_from_{self.stage_index - 1}",
                        self.stage_index - 1,
                        self._to_tensor(meta),
                        meta,
                        False,
                    )
                    for meta in self._stage_meta.inputs
                )
            self.args_recv_info[chunk_id] = infos

    def _setup_forward_send_info(self) -> None:
        if self._stage_meta.outputs is None:
            raise PipeliningMetadataError(
                f"Stage {self.stage_index}: outputs metadata is required for send setup."
            )
        self.act_send_info = {
            index: [self.stage_index + 1] if not self.is_last else []
            for index in range(len(self._stage_meta.outputs))
        }

    def _create_grad_recv_info(
        self,
        act_send_info: dict,
    ) -> tuple[_RecvInfo, ...]:
        grad_recv_infos: list[_RecvInfo] = []
        if not self.is_last:
            if self._stage_meta.output_grads is None:
                raise PipeliningMetadataError(
                    f"Stage {self.stage_index}: output_grads metadata is required for creating grad recv info."
                )
            output_grads = self._stage_meta.output_grads
            for index, destinations in act_send_info.items():
                if destinations is None or not destinations:
                    raise PipeliningMetadataError(
                        f"Stage {self.stage_index}: output {index} is not sent to any stage."
                    )
                source = destinations[0]
                grad_meta = output_grads[index]
                grad_recv_infos.append(
                    _RecvInfo(
                        f"recv_grad_for_{self.stage_index}_from_{source}",
                        source,
                        _make_tensor_from_meta(grad_meta, self.device)
                        if grad_meta is not None
                        else None,
                        grad_meta,
                    )
                )
        return tuple(grad_recv_infos)