wehub-resource-sync
116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
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."""
|
|
|
|
super_offload: bool = False
|
|
""" Enable high performance CPU offloading for Superchips. Only valid with ZeRO Stage 3."""
|
|
|
|
cpuadam_cores_perc: float = Field(0.8, ge=0.0, le=1.0)
|
|
""" Percentage of CPU cores to use for CPU Adam. Only valid with ZeRO Stage 3 and super_offload=True."""
|
|
|
|
@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"
|