|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from enum import Enum |
|
|
from pathlib import Path |
|
|
from pydantic import Field, model_validator |
|
|
from typing import Optional |
|
|
|
|
|
from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int |
|
|
|
|
|
|
|
|
class OffloadDeviceEnum(str, Enum): |
|
|
""" Enum for valid offload devices """ |
|
|
none = "none" |
|
|
cpu = "cpu" |
|
|
nvme = "nvme" |
|
|
|
|
|
|
|
|
class DeepSpeedZeroOffloadParamConfig(DeepSpeedConfigModel): |
|
|
""" Set options for parameter offload. Valid only with stage 3. """ |
|
|
|
|
|
device: OffloadDeviceEnum = "none" |
|
|
""" |
|
|
Device memory to offload model parameters. Supported options are `cpu` and |
|
|
`nvme`. |
|
|
""" |
|
|
|
|
|
nvme_path: Optional[Path] = None |
|
|
""" Filesystem path for NVMe device for parameter offloading. """ |
|
|
|
|
|
buffer_count: int = Field(5, ge=0) |
|
|
""" Number of buffers in buffer pool for parameter offloading to NVMe. """ |
|
|
|
|
|
buffer_size: int = Field(pp_int(1e8), ge=0) |
|
|
""" Size of buffers in buffer pool for parameter offloading to NVMe. """ |
|
|
|
|
|
max_in_cpu: int = Field(pp_int(1e9), ge=0) |
|
|
""" |
|
|
Number of parameter elements to maintain in CPU memory when offloading to |
|
|
NVMe is enabled. |
|
|
""" |
|
|
|
|
|
pin_memory: bool = False |
|
|
""" |
|
|
Offload to page-locked CPU memory. This could boost throughput at the cost |
|
|
of extra memory overhead. |
|
|
""" |
|
|
|
|
|
|
|
|
class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel): |
|
|
""" Set options for optimizer offload. Valid with stage 1, 2, and 3. """ |
|
|
|
|
|
device: OffloadDeviceEnum = "none" |
|
|
""" |
|
|
Device memory to offload optimizer state. Supported options are `cpu` and |
|
|
`nvme`. Optimizer computation is offload to CPU regardless of device option. |
|
|
""" |
|
|
|
|
|
nvme_path: Optional[Path] = None |
|
|
""" Filesystem path for NVMe device for optimizer state offloading. """ |
|
|
|
|
|
buffer_count: int = Field(4, ge=0) |
|
|
""" |
|
|
Number of buffers in buffer pool for optimizer state offloading to NVMe. |
|
|
This should be at least the number of states maintained per parameter by |
|
|
the optimizer. For example, Adam optimizer has 4 states (parameter, |
|
|
gradient, momentum, and variance). |
|
|
""" |
|
|
|
|
|
pin_memory: bool = False |
|
|
""" |
|
|
Offload to page-locked CPU memory. This could boost throughput at the cost |
|
|
of extra memory overhead. |
|
|
""" |
|
|
|
|
|
pipeline_read: bool = False |
|
|
""" |
|
|
For tile-based optimizer step processing, overlap read of next tile with |
|
|
computation of current tile. Used in ZeRO-Infinity. |
|
|
""" |
|
|
|
|
|
pipeline_write: bool = False |
|
|
""" |
|
|
For tile-based optimizer step processing, overlap write of previous tile |
|
|
with computation of current tile. |
|
|
""" |
|
|
|
|
|
fast_init: bool = False |
|
|
""" Enable fast optimizer initialization when offloading to NVMe. """ |
|
|
|
|
|
ratio: float = Field(1.0, ge=0.0, le=1.0) |
|
|
""" Percentage of offloaded optimizer states to CPU Adam. Only valid with ZeRO Stage 3.""" |
|
|
|
|
|
@model_validator(mode="after") |
|
|
def set_pipeline(self): |
|
|
pipeline = self.pipeline_read or self.pipeline_write |
|
|
self.__dict__["pipeline"] = pipeline |
|
|
return self |
|
|
|
|
|
|
|
|
class OffloadStateTypeEnum(str, Enum): |
|
|
""" Enum for internal buffer types """ |
|
|
optim_states = "optim_states" |
|
|
hp_params = "hp_params" |
|
|
lp_params = "lp_params" |
|
|
lp_grads = "lp_grads" |
|
|
contiguous_grad_buffer = "contiguous_grad_buffer" |
|
|
|
