Repository: microsoft/DeepSpeed
File: deepspeed/runtime/engine.py (Python)

Code before:
                       device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
                       pin_memory: bool = True,
                       non_blocking: bool = False) -> None:
        """Move the ZeRO optimizer buffers to the specified device.

        Arguments:
            include: Optional. The set of states to offload. If not provided, all states are offloaded.
            device: Optional. The device to move the ZeRO optimizer buffers to.
            pin_memory: Optional. Whether to pin the memory of the offloaded states.
            non_blocking: Optional. Whether to offload the states asynchronously.
        """
        assert self.zero_optimization_stage(
        ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3."

        assert not self.zero_offload_param(), "Moving states across devices is not supported for offloaded parameters."

        if device == OffloadDeviceEnum.none:
            logger.warning("No device specified for offloading states.")
            return

        if device == OffloadDeviceEnum.nvme:
            raise ValueError("NVMe offload is not supported for offloading states.")

        self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking)

    def offload_states_back(self, non_blocking: bool = False) -> None:
        """Move the ZeRO optimizer buffers back to the original device.

        Arguments:
            non_blocking: Optional. Whether to offload the states asynchronously.
        """
        assert self.zero_optimization_stage(
        ) == ZeroStageEnum.weights, "Moving buffers back is supported only for ZeRO stage 3."
        self.optimizer.offload_states_back(non_blocking=non_blocking)


Review comment (suggestion):
Let's improve name?

```suggestion
    def restore_offloaded_states(self, non_blocking: bool = False) -> None:
```

Code after:
                       device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
                       pin_memory: bool = True,
                       non_blocking: bool = False) -> None:
        """Offload the engine's states to the specified device.

        Arguments:
            include: Optional. The set of states to offload. If not provided, all states are offloaded.
            device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported.
            pin_memory: Optional. Whether to pin the memory of the offloaded states.
            non_blocking: Optional. Whether to offload the states asynchronously.
        """
        assert self.zero_optimization_stage(
        ) == ZeroStageEnum.weights, "Moving buffers across devices is supported only for ZeRO stage 3."

        assert not self.zero_offload_param(), "Moving states across devices is not supported for offloaded parameters."

        if device == OffloadDeviceEnum.none:
            logger.warning("No device specified for offloading states.")
            return

        if device == OffloadDeviceEnum.nvme:
            raise ValueError("NVMe offload is not supported for offloading states.")

        self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking)

    def reload_states(self, non_blocking: bool = False) -> None:
        """Reload the engine states to the original device.

        Arguments:
            non_blocking: Optional. Whether to offload the states asynchronously.
        """
        assert self.zero_optimization_stage(
        ) == ZeroStageEnum.weights, "Moving buffers back is supported only for ZeRO stage 3."
        self.optimizer.reload_states(non_blocking=non_blocking)