fix FP16 optimizer and adapted torch amp with tensor parallel (#18)

Author: ver217

colossalai/engine/schedule/_no_pipeline.py

@@ -16,11 +16,21 @@
 import torch.nn as nn
 from torch.optim import Optimizer
 
+<<<<<<< HEAD
 from colossalai.nn import (ZeroRedundancyOptimizer_Level_2,
                            ZeroRedundancyOptimizer_Level_3)
 from colossalai.nn.optimizer._utils import clip_grad_norm_fp32
 from ._base_schedule import BaseSchedule
 from ._utils import convert_to_fp16, convert_to_fp32
+=======
+from colossalai.context import ParallelMode
+from colossalai.core import global_context as gpc
+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
+>>>>>>> c8cb9f9... fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
 from ..amp import AMP_TYPE, GradScaler
 
 
@@ -191,10 +201,14 @@ def forward_backward_step(self,
     def optimizer_step(self, model: nn.Module, optimizer: Optimizer, grad_clipping: float = 0.0):
         # step optimizer
         if self.fp16 and self.amp_type == AMP_TYPE.TORCH:
-            if grad_clipping > 0.0:
-                self._torch_amp_scaler.unscale_(optimizer)
-                clip_grad_norm_fp32(model.parameters(), grad_clipping)
-            self._torch_amp_scaler.step(optimizer)
+            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()

colossalai/engine/schedule/_pipeline.py

@@ -163,8 +163,7 @@ def forward_step(self, model, criterion, input_tensor, return_tensors,
         if gpc.is_last_rank(ParallelMode.PIPELINE):
             if return_loss:
                 input_tensor, label = self.load_micro_batch()
-                loss_reduced = self.criterion(output_tensor, *
-                label) / (self.num_microbatches * self.grad_accum)
+                loss_reduced = self.criterion(output_tensor, *label) / (self.num_microbatches * self.grad_accum)
                 return_tensors.append(
                     tuple((output_tensor, label[0], loss_reduced)))
                 return loss_reduced