[legacy] clean up legacy code (#4743)

* [legacy] remove outdated codes of pipeline (#4692)

* [legacy] remove cli of benchmark and update optim (#4690)

* [legacy] remove cli of benchmark and update optim

* [doc] fix cli doc test

* [legacy] fix engine clip grad norm

* [legacy] remove outdated colo tensor (#4694)

* [legacy] remove outdated colo tensor

* [test] fix test import

* [legacy] move outdated zero to legacy (#4696)

* [legacy] clean up utils (#4700)

* [legacy] clean up utils

* [example] update examples

* [legacy] clean up amp

* [legacy] fix amp module

* [legacy] clean up gpc (#4742)

* [legacy] clean up context

* [legacy] clean core, constants and global vars

* [legacy] refactor initialize

* [example] fix examples ci

* [example] fix examples ci

* [legacy] fix tests

* [example] fix gpt example

* [example] fix examples ci

* [devops] fix ci installation

* [example] fix examples ci

colossalai/legacy/amp/apex_amp/apex_amp.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import torch.nn as nn
+
+try:
+    import apex.amp as apex_amp
+except ImportError:
+    pass
+
+from torch import Tensor
+
+from colossalai.interface import OptimizerWrapper
+from colossalai.legacy.utils import clip_grad_norm_fp32
+
+
+class ApexAMPOptimizer(OptimizerWrapper):
+    """ A wrapper class for APEX optimizer and it implements apex-specific backward and clip_grad_norm
+    methods
+    """
+
+    def backward(self, loss: Tensor):
+        """Backward pass to get all gradients
+
+        Args:
+            loss (torch.Tensor): Loss computed by a loss function
+        """
+        with apex_amp.scale_loss(loss, self.optim) as scaled_loss:
+            scaled_loss.backward()
+
+    def clip_grad_norm(self, model: nn.Module, max_norm: float):
+        """Clip gradients by norm
+
+        Args:
+            model (torch.nn.Module): Your model object
+            max_norm (float): The max norm value for gradient clipping
+        """
+        if max_norm > 0:
+            clip_grad_norm_fp32(apex_amp.master_params(self.optim), max_norm)