rename DSOptimWrapper to DeepSpeedOptimWrapper

open-mmlab · Jun 20, 2023 · 8436d00 · 8436d00
1 parent 9088abc
commit 8436d00
Show file tree

Hide file tree

Showing 5 changed files with 8 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -7,7 +7,7 @@
 
 
 @OPTIM_WRAPPERS.register_module()
-class DSOptimWrapper(BaseOptimWrapper):
+class DeepSpeedOptimWrapper(BaseOptimWrapper):
 
     def __init__(self, optimizer):
         self.optimizer = optimizer