Diff Coverage

    def group_context(self, step: Any) -> ContextManager[None]:
        """Enter the run-scoped swap group for one forward leaf."""
        if not self.manages(step):
            return nullcontext()
        key = self._key(step)
        group_name = self._group_names[key]
        self._manager.ensure_group(group_name)
        self._forward_context_keys.add(key)

    def group_name(self, step: Any) -> str:
        """Return the physical group name for a managed step."""
        key = self._key(step)
        if key not in self._group_names:
            raise RuntimeError(f"Pipeline swap does not manage chunk {key}.")
        return self._group_names[key]

    def wait_load(self, step: Any) -> None:
        """Wait for H2D on the scheduler's current compute stream."""
        if self.manages(step):
            self._manager.wait_load(self.group_name(step))

    def protect_aliases(self, step: Any, tensors: Any) -> None:
        """Keep pipeline-owned aliases resident for a managed chunk."""
        if self.manages(step):

        if (
                step is not None
                and step.type == MetaStepType.FSDP_UNSHARD
                and step.stage_index == bwd_leaf.step.stage_index):
            return index
    return bwd_leaf.container_index


def inject_pipeline_swap_steps(order: List[Any]) -> List[Any]:

def swap_wait_offload(step: Any, session: PipelineSwapSession) -> None:
    """Wait for a pipeline swap group's D2H and release device storage."""
    SwapManager().wait_offload(session.group_name(step))


def swap_launch_load(step: Any, session: PipelineSwapSession) -> None:
    """Launch H2D for a pipeline swap group."""
    SwapManager().launch_load(session.group_name(step))


def swap_wait_load(step: Any, session: PipelineSwapSession) -> None:
    """Wait for a pipeline swap group's H2D before its backward container."""

        elif step_type == MetaStepType.FWD_RECV:
            self.recv_fwd(stage, micro_index)

        elif step_type == MetaStepType.FWD:
            self.execute_fwd_leaf(cur_step, arg_mbs, kwarg_mbs, losses)

        elif step_type == MetaStepType.FWD_SEND:
            self.send_fwd(stage, micro_index)

    def _exec_pipeline_swap_step(self, cur_step, arg_mbs, kwarg_mbs):
        """Execute a pipeline activation-swap control step."""
        if self._swap_session is None:
            raise RuntimeError("Pipeline swap step executed without an active run session.")
        if cur_step.type == MetaStepType.SWAP_LAUNCH_OFFLOAD:
            swap_launch_offload(cur_step, self, arg_mbs, kwarg_mbs, self._swap_session)
        elif cur_step.type == MetaStepType.SWAP_WAIT_OFFLOAD:
            swap_wait_offload(cur_step, self._swap_session)
        elif cur_step.type == MetaStepType.SWAP_LAUNCH_LOAD:
            swap_launch_load(cur_step, self._swap_session)
        elif cur_step.type == MetaStepType.SWAP_WAIT_LOAD:
            swap_wait_load(cur_step, self._swap_session)

    def run_microbatches(self, arg_mbs: list, kwarg_mbs: list, losses: list) -> None:

        Returns:
            An iterable of ``(name, buffer)`` pairs, including non-persistent
            buffers and buffers registered by child modules.
        """
        raise NotImplementedError("Platform subclasses must implement buffers_dict")

    @staticmethod
    def get_model_state_dict(model: Any, *, options: Any = None) -> dict[str, Any]:
        """Get the state dictionary of a model.