Add PyDantic v2 support

.github/workflows/python.yml

@@ -20,7 +20,7 @@ jobs:
   version-check:
     strategy:
       matrix:
-        pyVersion: ["3.6", "3.7", "3.8", "3.9", "3.10"]
+        pyVersion: ["3.7", "3.8", "3.9", "3.10"]
       fail-fast: false
 
     runs-on: ubuntu-20.04

deepspeed/comm/config.py

@@ -3,17 +3,12 @@
 
 # DeepSpeed Team
 
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 from .constants import *
 
 
 class CommsConfig(BaseModel):
-
-    class Config:
-        validate_all = True
-        validate_assignment = True
-        use_enum_values = True
-        extra = 'forbid'
+    model_config = ConfigDict(validate_default=True, validate_assignment=True, use_enum_values=True, extra='forbid')
 
 
 class CommsLoggerConfig(CommsConfig):

deepspeed/inference/config.py

@@ -7,9 +7,8 @@
 import deepspeed
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
 from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
-from pydantic import Field
-from pydantic import validator
-from typing import Dict, Union
+from pydantic import ConfigDict, Field, FieldValidationInfo, field_validator
+from typing import Dict, Union, Optional
 from enum import Enum
 
 
@@ -92,24 +91,24 @@ class QuantTypeEnum(str, Enum):
 
 
 class BaseQuantConfig(DeepSpeedConfigModel):
-    enabled = True
-    num_bits = 8
+    enabled: bool = True
+    num_bits: int = 8
     q_type: QuantTypeEnum = QuantTypeEnum.sym
     q_groups: int = 1
 
 
 class WeightQuantConfig(BaseQuantConfig):
-    enabled = True
+    enabled: bool = True
     quantized_initialization: Dict = {}
     post_init_quant: Dict = {}
 
 
 class ActivationQuantConfig(BaseQuantConfig):
-    enabled = True
+    enabled: bool = True
 
 
 class QKVQuantConfig(DeepSpeedConfigModel):
-    enabled = True
+    enabled: bool = True
 
 
 class QuantizationConfig(DeepSpeedConfigModel):
@@ -121,9 +120,9 @@ class QuantizationConfig(DeepSpeedConfigModel):
 
 # todo: brainstorm on how to do ckpt loading for DS inference
 class InferenceCheckpointConfig(DeepSpeedConfigModel):
-    checkpoint_dir: str = None
-    save_mp_checkpoint_path: str = None
-    base_dir: str = None
+    checkpoint_dir: Optional[str] = None
+    save_mp_checkpoint_path: Optional[str] = None
+    base_dir: Optional[str] = None
 
 
 class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
@@ -199,12 +198,12 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
     """
 
     #todo: refactor the following 3 into the new checkpoint_config
-    checkpoint: Union[str, Dict] = None
+    checkpoint: Optional[Union[str, Dict]] = None
     """
     Path to deepspeed compatible checkpoint or path to JSON with load policy.
     """
 
-    base_dir: str = ""
+    base_dir: Optional[str] = ""
     """
     This shows the root directory under which all the checkpoint files exists.
     This can be passed through the json config too.
@@ -215,7 +214,7 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
     specifying whether the inference-module is created with empty or real Tensor
     """
 
-    save_mp_checkpoint_path: str = None
+    save_mp_checkpoint_path: Optional[str] = None
     """
     The path for which we want to save the loaded model with a checkpoint. This
     feature is used for adjusting the parallelism degree to help alleviate the
@@ -247,16 +246,16 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
         deprecated=True,
         deprecated_msg="This parameter is no longer needed, please remove from your call to DeepSpeed-inference")
 
-    injection_policy: Dict = Field(None, alias="injection_dict")
+    injection_policy: Optional[Dict] = Field(None, alias="injection_dict")
     """
     Dictionary mapping a client nn.Module to its corresponding injection
     policy. e.g., `{BertLayer : deepspeed.inference.HFBertLayerPolicy}`
     """
 
-    injection_policy_tuple: tuple = None
+    injection_policy_tuple: Optional[tuple] = None
     """ TODO: Add docs """
 
-    config: Dict = Field(None, alias="args")  # todo: really no need for this field if we can refactor
+    config: Optional[Dict] = Field(None, alias="args")  # todo: really no need for this field if we can refactor
 
     max_out_tokens: int = Field(1024, alias="max_tokens")
     """
@@ -281,25 +280,27 @@ class DeepSpeedInferenceConfig(DeepSpeedConfigModel):
     Deprecated, please use the ``tensor_parallel` config to control model
     parallelism.
     """
-    mpu: object = Field(None, deprecated=True, new_param="tensor_parallel.mpu")
+    mpu: Optional[object] = Field(None, deprecated=True, new_param="tensor_parallel.mpu")
     ep_size: int = Field(1, deprecated=True, new_param="moe.ep_size")
-    ep_group: object = Field(None, alias="expert_group", deprecated=True, new_param="moe.ep_group")
-    ep_mp_group: object = Field(None, alias="expert_mp_group", deprecated=True, new_param="moe.ep_mp_group")
+    ep_group: Optional[object] = Field(None, alias="expert_group", deprecated=True, new_param="moe.ep_group")
+    ep_mp_group: Optional[object] = Field(None, alias="expert_mp_group", deprecated=True, new_param="moe.ep_mp_group")
     moe_experts: list = Field([1], deprecated=True, new_param="moe.moe_experts")
     moe_type: MoETypeEnum = Field(MoETypeEnum.standard, deprecated=True, new_param="moe.type")
 
-    @validator("moe")
-    def moe_backward_compat(cls, field_value, values):
-        if isinstance(field_value, bool):
-            return DeepSpeedMoEConfig(moe=field_value)
-        return field_value
-
-    @validator("use_triton")
-    def has_triton(cls, field_value, values):
-        if field_value and not deepspeed.HAS_TRITON:
+    @field_validator("moe")
+    @classmethod
+    def moe_backward_compat(cls, v: Union[bool, DeepSpeedMoEConfig], info: FieldValidationInfo) -> DeepSpeedMoEConfig:
+        if isinstance(v, bool):
+            return DeepSpeedMoEConfig(moe=v)
+        return v
+
+    @field_validator("use_triton")
+    @classmethod
+    def has_triton(cls, v: bool, info: FieldValidationInfo) -> bool:
+        if v and not deepspeed.HAS_TRITON:
             raise ValueError('Triton needs to be installed to use deepspeed with triton kernels')
-        return field_value
+        return v
 
-    class Config:
-        # Get the str representation of the datatype for serialization
-        json_encoders = {torch.dtype: lambda x: str(x)}
+    # TODO[pydantic]: The following keys were removed: `json_encoders`.
+    # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
+    model_config = ConfigDict(json_encoders={torch.dtype: lambda x: str(x)})

deepspeed/monitor/config.py

@@ -3,7 +3,7 @@
 
 # DeepSpeed Team
 
-from pydantic import root_validator
+from typing import Optional
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel
 
 
@@ -34,10 +34,10 @@ class WandbConfig(DeepSpeedConfigModel):
     enabled: bool = False
     """ Whether logging to WandB is enabled. Requires `wandb` package is installed. """
 
-    group: str = None
+    group: Optional[str] = None
     """ Name for the WandB group. This can be used to group together runs. """
 
-    team: str = None
+    team: Optional[str] = None
     """ Name for the WandB team. """
 
     project: str = "deepspeed"
@@ -63,6 +63,10 @@ class CSVConfig(DeepSpeedConfigModel):
 class DeepSpeedMonitorConfig(DeepSpeedConfigModel):
     """Sets parameters for various monitoring methods."""
 
+    @property
+    def enabled(self) -> bool:
+        return any((self.tensorboard.enabled, self.wandb.enabled, self.csv_monitor.enabled))
+
     tensorboard: TensorBoardConfig = {}
     """ TensorBoard monitor, requires `tensorboard` package is installed. """
 
@@ -71,9 +75,3 @@ class DeepSpeedMonitorConfig(DeepSpeedConfigModel):
 
     csv_monitor: CSVConfig = {}
     """ Local CSV output of monitoring data. """
-
-    @root_validator
-    def check_enabled(cls, values):
-        values["enabled"] = values.get("tensorboard").enabled or values.get("wandb").enabled or values.get(
-            "csv_monitor").enabled
-        return values

deepspeed/runtime/config_utils.py

@@ -9,7 +9,7 @@
 import collections
 import collections.abc
 from functools import reduce
-from pydantic import BaseModel
+from pydantic import ConfigDict, BaseModel
 from deepspeed.utils import logger
 
 
@@ -56,11 +56,11 @@ def __init__(self, strict=False, **data):
         super().__init__(**data)
         self._deprecated_fields_check(self)
 
-    def _process_deprecated_field(self, pydantic_config, field):
+    def _process_deprecated_field(self, pydantic_config, field_name, field):
         # Get information about the deprecated field
-        fields_set = pydantic_config.__fields_set__
-        dep_param = field.name
-        kwargs = field.field_info.extra
+        fields_set = pydantic_config.model_fields_set
+        dep_param = field_name
+        kwargs = field.json_schema_extra
         new_param_fn = kwargs.get("new_param_fn", lambda x: x)
         param_value = new_param_fn(getattr(pydantic_config, dep_param))
         new_param = kwargs.get("new_param", "")
@@ -96,25 +96,26 @@ def _process_deprecated_field(self, pydantic_config, field):
                     raise e
 
     def _deprecated_fields_check(self, pydantic_config):
-        fields = pydantic_config.__fields__
-        for field in fields.values():
-            if field.field_info.extra.get("deprecated", False):
-                self._process_deprecated_field(pydantic_config, field)
+        fields = pydantic_config.model_fields
+        for field_name, field in fields.items():
+            if getattr(field, "json_schema_extra") is None:
+                continue
+            if field.json_schema_extra.get("deprecated", False):
+                self._process_deprecated_field(pydantic_config, field_name, field)
 
-    class Config:
-        validate_all = True
-        validate_assignment = True
-        use_enum_values = True
-        allow_population_by_field_name = True
-        extra = "forbid"
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(validate_default=True,
+                              validate_assignment=True,
+                              use_enum_values=True,
+                              populate_by_name=True,
+                              extra="forbid",
+                              arbitrary_types_allowed=True)
 
 
 def get_config_default(config, field_name):
-    assert field_name in config.__fields__, f"'{field_name}' is not a field in {config}"
-    assert not config.__fields__.get(
-        field_name).required, f"'{field_name}' is a required field and does not have a default value"
-    return config.__fields__.get(field_name).default
+    assert field_name in config.model_fields, f"'{field_name}' is not a field in {config}"
+    assert not config.model_fields.get(
+        field_name).is_required(), f"'{field_name}' is a required field and does not have a default value"
+    return config.model_fields.get(field_name).default
 
 
 class pp_int(int):

deepspeed/runtime/zero/config.py

@@ -3,9 +3,9 @@
 
 # DeepSpeed Team
 
-from pydantic import Field, validator
+from pydantic import Field, model_validator
 import sys
-from typing import Optional
+from typing import Any, Optional
 from enum import Enum
 from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel
 from deepspeed.utils import logger
@@ -119,7 +119,7 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     the allgather for large model sizes
     """
 
-    overlap_comm: bool = None  # None for dynamic default value (see validator `overlap_comm_valid` below)
+    overlap_comm: Optional[bool] = None  # None for dynamic default value (see validator `overlap_comm_valid` below)
     """
     Attempts to overlap the reduction of the gradients with backward computation
     """
@@ -159,23 +159,23 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     parameters). Used by ZeRO3-Offload and ZeRO-Infinity
     """
 
-    cpu_offload_param: bool = Field(
+    cpu_offload_param: Optional[bool] = Field(
         None,
         deprecated=True,
         new_param="offload_param",
         new_param_fn=(lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu) if val else None),
     )
     """ Deprecated, please use ``offload_param`` """
 
-    cpu_offload_use_pin_memory: bool = Field(
+    cpu_offload_use_pin_memory: Optional[bool] = Field(
         None,
         deprecated=True,
         new_param="offload_param or offload_optimizer",
         set_new_param=False,
     )
     """ Deprecated, please use ``offload_param`` or ``offload_optimizer`` """
 
-    cpu_offload: bool = Field(
+    cpu_offload: Optional[bool] = Field(
         None,
         deprecated=True,
         new_param="offload_optimizer",
@@ -196,7 +196,7 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     latency-bound messages).
     """
 
-    model_persistence_threshold: int = Field(pp_int(sys.maxsize, "sys.maxsize"),
+    model_persistence_threshold: int = Field(int(pp_int(sys.maxsize, "sys.maxsize")),
                                              ge=0,
                                              alias="stage3_model_persistence_threshold")
     """
@@ -294,9 +294,9 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
     """
 
     # Validators
-    @validator("overlap_comm")
-    def overlap_comm_valid(cls, field_value, values):
-        if field_value is None:
-            assert ("stage" in values), "DeepSpeedZeroConfig: 'stage' must be defined before 'overlap_comm'"
-            field_value = values["stage"] == ZeroStageEnum.weights
-        return field_value
+    @model_validator(mode="after")
+    @classmethod
+    def overlap_comm_valid(cls, data: Any) -> Any:
+        if data.overlap_comm is None:
+            data.overlap_comm = (data.stage == ZeroStageEnum.weights)
+        return data

deepspeed/runtime/zero/offload_config.py

@@ -3,7 +3,7 @@
 
 # DeepSpeed Team
 
-from pydantic import Field, validator
+from pydantic import Field, field_validator
 from enum import Enum
 from pathlib import Path
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int
@@ -88,7 +88,7 @@ class DeepSpeedZeroOffloadOptimizerConfig(DeepSpeedConfigModel):
     fast_init: bool = False
     """ Enable fast optimizer initialization when offloading to NVMe. """
 
-    @validator("pipeline_read", "pipeline_write", always=True)
+    @field_validator("pipeline_read", "pipeline_write")
     def set_pipeline(cls, field_value, values):
         values["pipeline"] = field_value or values.get("pipeline", False)
         return field_value

requirements/requirements-readthedocs.txt

@@ -1,9 +1,9 @@
-autodoc_pydantic<2.0.0
+autodoc_pydantic>=2.0.0
 docutils<0.18
 hjson
 packaging
 psutil
 py-cpuinfo
-pydantic<2.0.0
+pydantic>=2.0.0
 torch
 tqdm

requirements/requirements.txt

@@ -4,6 +4,6 @@ numpy
 packaging>=20.0
 psutil
 py-cpuinfo
-pydantic<2.0.0
+pydantic>=2.0.0
 torch
 tqdm