fix FP16 optimizer and adapted torch amp with tensor parallel

Author: ver217

colossalai/engine/__init__.py

@@ -1,7 +1,7 @@
-from .amp_type import AMP_TYPE
 from ._base_engine import Engine
 from .gradient_handler import *
 from .schedule import *
+from .amp import *
 
 
 __all__ = ['Engine']

colossalai/engine/amp/__init__.py

@@ -0,0 +1,2 @@
+from .grad_scaler import GradScaler
+from .amp_type import AMP_TYPE

colossalai/engine/amp_type.py renamed to colossalai/engine/amp/amp_type.py

@@ -12,11 +12,12 @@
 
 from colossalai.context import ParallelMode
 from colossalai.core import global_context as gpc
-from colossalai.engine.amp_type import AMP_TYPE
 from colossalai.nn import (ZeroRedundancyOptimizer_Level_2,
                            ZeroRedundancyOptimizer_Level_3)
+from colossalai.nn.optimizer._utils import clip_grad_norm_fp32
 from ._utils import convert_to_fp16
 from ._base_schedule import BaseSchedule
+from ..amp import AMP_TYPE, GradScaler
 
 
 class NoPipelineSchedule(BaseSchedule):
@@ -30,6 +31,7 @@ class NoPipelineSchedule(BaseSchedule):
     :type amp_type: AMP_TYPE
     :type amp_config: dict
     """
+
     def __init__(
             self,
             amp_type: AMP_TYPE = None,
@@ -101,7 +103,7 @@ def initialize(self,
 
         if self.fp16:
             if self.amp_type == AMP_TYPE.TORCH:
-                self._torch_amp_scaler = torch_amp.GradScaler(**self.amp_cfg)
+                self._torch_amp_scaler = GradScaler(**self.amp_cfg)
             elif self.amp_type == AMP_TYPE.APEX:
                 self.model, self.optimizer = apex_amp.initialize(
                     self.model, self.optimizer, **self.amp_cfg)
@@ -175,9 +177,16 @@ def forward_backward_step(self, forward_only=False, return_loss=True):
     def step(self):
         # step optimizer
         if self.fp16 and self.amp_type == AMP_TYPE.TORCH:
+            if getattr(gpc.config, 'clip_grad', 0.0) > 0.0:
+                self._torch_amp_scaler.unscale_(self.optimizer)
+                clip_grad_norm_fp32(self.model.parameters(),
+                                    gpc.config.clip_grad)
             self._torch_amp_scaler.step(self.optimizer)
             self._torch_amp_scaler.update()
         else:
+            if not self.fp16 and not self.use_zero_level_2_3 and getattr(gpc.config, 'clip_grad', 0.0) > 0.0:
+                clip_grad_norm_fp32(self.model.parameters(),
+                                    gpc.config.clip_grad)
             self.optimizer.step()
 
         # update lr scheduler

colossalai/engine/amp/grad_scaler.py

@@ -15,7 +15,7 @@
 from colossalai.utils import get_current_device
 from ._base_schedule import BaseSchedule
 from ._utils import convert_to_fp16
-from ..amp_type import AMP_TYPE
+from ..amp import AMP_TYPE
 
 
 def squeeze(x: Union[Tensor, tuple, list]):
@@ -163,7 +163,7 @@ def forward_step(self, input_tensor, return_tensors, return_loss=True):
             if return_loss:
                 input_tensor, label = self.load_micro_batch()
                 loss_reduced = self.criterion(output_tensor, *
-                label) / self.num_microbatches
+                                              label) / self.num_microbatches
                 return_tensors.append(
                     tuple((output_tensor, label[0], loss_reduced)))
                 return loss_reduced
@@ -200,7 +200,7 @@ def backward_step(self, input_tensor, output_tensor, output_tensor_grad):
     def forward_backward_step(self, forward_only=True, return_loss=True):
         """Runs non-interleaved 1F1B schedule, with communication between pipeline stages.
         Returns a tuple with losses if the last stage, an empty tuple otherwise.
-        
+
         :return: (output, label, loss)
         """
 

colossalai/engine/schedule/_no_pipeline.py

@@ -106,7 +106,7 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
             tensor_parallel_norm = _calc_lp(tensor_parallel_grads, norm_type)
             no_tensor_parallel_grads = _calc_lp(
                 no_tensor_parallel_grads, norm_type)
-        if gpc.is_initialized(ParallelMode.TENSOR):
+        if gpc.is_initialized(ParallelMode.TENSOR) and len(tensor_parallel_grads) > 0:
             # Sum across all model-parallel GPUs.
             torch.distributed.all_reduce(tensor_parallel_norm,
                                          op=torch.distributed.ReduceOp.SUM,

colossalai/engine/schedule/_pipeline.py

@@ -0,0 +1,5 @@
+colossalai.engine.amp.amp\_type
+===============================
+
+.. automodule:: colossalai.engine.amp.amp_type
+   :members:

colossalai/nn/optimizer/_utils.py

@@ -0,0 +1,5 @@
+colossalai.engine.amp.grad\_scaler
+==================================
+
+.. automodule:: colossalai.engine.amp.grad_scaler
+   :members:

docs/colossalai/colossalai.engine.amp.amp_type.rst

@@ -0,0 +1,12 @@
+colossalai.engine.amp
+=====================
+
+.. automodule:: colossalai.engine.amp
+   :members:
+
+
+.. toctree::
+   :maxdepth: 2
+
+   colossalai.engine.amp.amp_type
+   colossalai.engine.amp.grad_scaler