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 s
(self,
include: Container[OffloadStateTypeEnum] = None,
device: OffloadDeviceEnum = OffloadDeviceEnum.cpu,
pin_memory: bool = True,
non_blocking: bool = False)
| 5571 | return include |
| 5572 | |
| 5573 | def offload_states(self, |
| 5574 | include: Container[OffloadStateTypeEnum] = None, |
| 5575 | device: OffloadDeviceEnum = OffloadDeviceEnum.cpu, |
| 5576 | pin_memory: bool = True, |
| 5577 | non_blocking: bool = False) -> None: |
| 5578 | """Offload the engine's states to the specified device. |
| 5579 | |
| 5580 | Arguments: |
| 5581 | include: Optional. The set of states to offload. If not provided, all states are offloaded. |
| 5582 | device: Optional. The device to move the ZeRO optimizer buffers to. Currently only `OffloadDeviceEnum.cpu` is supported. |
| 5583 | pin_memory: Optional. Whether to pin the memory of the offloaded states. |
| 5584 | non_blocking: Optional. Whether to offload the states asynchronously. |
| 5585 | """ |
| 5586 | include = self._refine_include_states(include) |
| 5587 | param_offload_config = self.zero_offload_param() |
| 5588 | assert param_offload_config is None or param_offload_config.device == OffloadDeviceEnum.none, "Moving states across devices is not supported for offloaded parameters." |
| 5589 | |
| 5590 | assert not isinstance( |
| 5591 | self.optimizer, |
| 5592 | DeepSpeedZeRoOffload), "Moving states across devices is not supported without an optimizer." |
| 5593 | |
| 5594 | if device == OffloadDeviceEnum.none: |
| 5595 | logger.warning("No device specified for offloading states.") |
| 5596 | return |
| 5597 | |
| 5598 | if device == OffloadDeviceEnum.nvme: |
| 5599 | raise ValueError("NVMe offload is not supported for offloading states.") |
| 5600 | |
| 5601 | self.optimizer.offload_states(include=include, device=device, pin_memory=pin_memory, non_blocking=non_blocking) |
| 5602 | |
| 5603 | def reload_states(self, non_blocking: bool = False) -> None: |
| 5604 | """Reload the engine states to the original device. |