[legacy] remove cli of benchmark and update optim

colossalai/amp/apex_amp/apex_amp.py

@@ -10,11 +10,11 @@
 
 from torch import Tensor
 
-from colossalai.nn.optimizer import ColossalaiOptimizer
+from colossalai.interface import OptimizerWrapper
 from colossalai.utils import clip_grad_norm_fp32
 
 
-class ApexAMPOptimizer(ColossalaiOptimizer):
+class ApexAMPOptimizer(OptimizerWrapper):
     """ A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm
     methods
     """

colossalai/amp/naive_amp/naive_amp.py

@@ -13,12 +13,12 @@
 
 from colossalai.context import ParallelMode
 from colossalai.core import global_context as gpc
-from colossalai.nn.optimizer import ColossalaiOptimizer
+from colossalai.interface import OptimizerWrapper
 
 from ._fp16_optimizer import FP16Optimizer
 
 
-class NaiveAMPOptimizer(ColossalaiOptimizer):
+class NaiveAMPOptimizer(OptimizerWrapper):
     """A wrapper class for optimizer to cast all parameters to fp16
 
     Args:

colossalai/amp/torch_amp/torch_amp.py

@@ -7,13 +7,13 @@
 from torch.nn.modules.loss import _Loss
 from torch.optim import Optimizer
 
-from colossalai.nn.optimizer import ColossalaiOptimizer
+from colossalai.interface import OptimizerWrapper
 from colossalai.utils import clip_grad_norm_fp32
 
 from ._grad_scaler import GradScaler
 
 
-class TorchAMPOptimizer(ColossalaiOptimizer):
+class TorchAMPOptimizer(OptimizerWrapper):
     """A wrapper class which integrate Pytorch AMP with an optimizer
 
     Args:

colossalai/auto_parallel/offload/amp_optimizer.py

@@ -5,8 +5,8 @@
 from torch.optim import Optimizer
 
 from colossalai.amp.naive_amp.grad_scaler import DynamicGradScaler
+from colossalai.interface import OptimizerWrapper
 from colossalai.logging import get_dist_logger
-from colossalai.nn.optimizer import ColossalaiOptimizer
 from colossalai.utils import get_current_device
 
 from .base_offload_module import BaseOffloadModule
@@ -19,7 +19,7 @@ class OptimState(Enum):
     UNSCALED = 1
 
 
-class AMPOptimizer(ColossalaiOptimizer):
+class AMPOptimizer(OptimizerWrapper):
     """
     A wrapper for Optimizer.
     Code reference: https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/optimizer/zero_optimizer.py

colossalai/checkpoint_io/utils.py

@@ -13,7 +13,6 @@
 from torch.optim import Optimizer
 
 from colossalai.interface import ModelWrapper, OptimizerWrapper
-from colossalai.nn.optimizer import ColossalaiOptimizer
 from colossalai.tensor.d_tensor import (
     is_customized_distributed_tensor,
     is_distributed_tensor,
@@ -130,10 +129,7 @@ def unwrap_optimizer(optimizer: OptimizerWrapper):
     This method should be used before saving/loading it to/from sharded checkpoints.
     '''
 
-    # TODO(Baizhou): ColossalaiOptimizer will be replaced with OptimizerWrapper in the future
     unwrapped_optim = optimizer.optim
-    if isinstance(unwrapped_optim, ColossalaiOptimizer):
-        unwrapped_optim = unwrapped_optim.optim
     return unwrapped_optim