From 8436d006120f3d7ead9eba48b4b3cc21a68f49e0 Mon Sep 17 00:00:00 2001
From: zhouzaida <zhouzaida@163.com>
Date: Tue, 20 Jun 2023 16:22:04 +0800
Subject: [PATCH] rename DSOptimWrapper to DeepSpeedOptimWrapper

---
 examples/distributed_training_with_flexible_runner.py | 2 +-
 mmengine/_strategy/deepspeed.py                       | 2 ++
 mmengine/model/wrappers/_deepspeed.py                 | 4 ++--
 mmengine/optim/optimizer/__init__.py                  | 4 ++--
 mmengine/optim/optimizer/_deepspeed.py                | 2 +-
 5 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/examples/distributed_training_with_flexible_runner.py b/examples/distributed_training_with_flexible_runner.py
index 730158b6f1..d14f1eb76b 100644
--- a/examples/distributed_training_with_flexible_runner.py
+++ b/examples/distributed_training_with_flexible_runner.py
@@ -103,7 +103,7 @@ def main():
                 contiguous_gradients=True,
                 cpu_offload=False))
         optim_wrapper = dict(
-            type='DSOptimWrapper',
+            type='DeepSpeedOptimWrapper',
             optimizer=dict(type=SGD, lr=0.001, momentum=0.9))
     else:
         strategy = None
diff --git a/mmengine/_strategy/deepspeed.py b/mmengine/_strategy/deepspeed.py
index a0f7778997..4ddc5811b4 100644
--- a/mmengine/_strategy/deepspeed.py
+++ b/mmengine/_strategy/deepspeed.py
@@ -26,6 +26,8 @@ class DeepSpeedStrategy(BaseStrategy):
         https://www.deepspeed.ai/docs/config-json/.
 
     Args:
+        config (str or dict, optional): If it is a string, it is a path to load
+            config for deepspeed. Defaults to None.
         zero_optimization (dict, optional): Enabling and configuring ZeRO
             memory optimizations. Defaults to None.
         gradient_clipping (float): Enable gradient clipping with value.
diff --git a/mmengine/model/wrappers/_deepspeed.py b/mmengine/model/wrappers/_deepspeed.py
index fe0e449928..587bec879c 100644
--- a/mmengine/model/wrappers/_deepspeed.py
+++ b/mmengine/model/wrappers/_deepspeed.py
@@ -4,7 +4,7 @@
 import torch
 from deepspeed.runtime.engine import DeepSpeedEngine
 
-from mmengine.optim.optimizer._deepspeed import DSOptimWrapper
+from mmengine.optim.optimizer._deepspeed import DeepSpeedOptimWrapper
 from mmengine.registry import MODEL_WRAPPERS
 
 
@@ -29,7 +29,7 @@ def __getattr__(self, name):
     def train_step(
         self,
         data: Union[dict, tuple, list],
-        optim_wrapper: DSOptimWrapper,
+        optim_wrapper: DeepSpeedOptimWrapper,
     ) -> Dict[str, torch.Tensor]:
         data = self.model.module.data_preprocessor(data, training=True)
         data = self._cast_inputs_half(data)
diff --git a/mmengine/optim/optimizer/__init__.py b/mmengine/optim/optimizer/__init__.py
index a551f1de27..b9674ee05d 100644
--- a/mmengine/optim/optimizer/__init__.py
+++ b/mmengine/optim/optimizer/__init__.py
@@ -18,5 +18,5 @@
 ]
 
 if is_installed('deepspeed'):
-    from ._deepspeed import DSOptimWrapper  # noqa:F401
-    __all__.append('DSOptimWrapper')
+    from ._deepspeed import DeepSpeedOptimWrapper  # noqa:F401
+    __all__.append('DeepSpeedOptimWrapper')
diff --git a/mmengine/optim/optimizer/_deepspeed.py b/mmengine/optim/optimizer/_deepspeed.py
index 589776492f..d81531cc55 100644
--- a/mmengine/optim/optimizer/_deepspeed.py
+++ b/mmengine/optim/optimizer/_deepspeed.py
@@ -7,7 +7,7 @@
 
 
 @OPTIM_WRAPPERS.register_module()
-class DSOptimWrapper(BaseOptimWrapper):
+class DeepSpeedOptimWrapper(BaseOptimWrapper):
 
     def __init__(self, optimizer):
         self.optimizer = optimizer