Repository: hpcaitech/ColossalAI

Issue:
[BUG]: unscaled_grads don't work in sharded_optim_v2.py
### 🐛 Describe the bug

Hi, I have great respect for this work! But when i debug the tests/test_zero/test_zero_engine.py, i found that on line 230 of sharded_optim_v2.py, the p.grad always equal None, scaling loss never works。And i also debug Colossalai-Examples/language/gpt/train_gpt.py,  unscale_grads function also don't work!

### Environment

cossalai==v0.1.7
torch==1.11.0+cu113


Discussion:
And i found that on line 285 in sharded_model_v2.py, all parameter's grad set to None，so unscaled_grad function will works abnormlly.
@[feifeibear](https://github.com/hpcaitech/ColossalAI/commits?author=feifeibear)
@[BoxiangW](https://github.com/BoxiangW)
Maybe self._prepare_grads() should be called before unscaled_grads?Likes bellow:
`if self._prepare_grads_flag == False`
`    self._prepare_grads()`
`do unscaled grads`

Resolving patch:
diff --git a/colossalai/utils/common.py b/colossalai/utils/common.py
--- a/colossalai/utils/common.py
+++ b/colossalai/utils/common.py
@@ -195,7 +195,7 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
             # Make sure the grads are in fp32
             assert param.grad.dtype == torch.float, \
                 f'expected gradient to be dtype torch.float, but got {param.grad.type()}'
-            if hasattr(param, 'zero_is_sharded'):
+            if hasattr(param, 'colo_attr') and param.colo_attr.sharded_data_tensor.is_sharded:
                 has_zero_shared_param = True
             params.append(param)
 
@@ -234,7 +234,7 @@ def clip_grad_norm_fp32(parameters, max_norm, norm_type=2):
             if is_model_parallel_parameter(p):
                 reductor = (gpc.get_world_size(ParallelMode.TENSOR) / getattr(p, NUM_PARTITIONS))**(1 / norm_type)
                 tensor_parallel_grads.append(p.grad.data / reductor)
-            elif hasattr(p, 'zero_is_sharded'):
+            elif hasattr(p, 'colo_attr') and p.colo_attr.sharded_data_tensor.is_sharded:
                 zero_sharded_grads.append(p.grad.data)
             else:
                 no_tensor_parallel_grads.append(p.grad.data)
diff --git a/colossalai/zero/sharded_optim/sharded_optim_v2.py b/colossalai/zero/sharded_optim/sharded_optim_v2.py
--- a/colossalai/zero/sharded_optim/sharded_optim_v2.py
+++ b/colossalai/zero/sharded_optim/sharded_optim_v2.py
@@ -169,21 +169,27 @@ def backward(self, loss: Tensor) -> None:
         self.model.backward(loss)
 
     def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None:
+        # This function is called except the last stage of pipeline parallel
+        # It receives the scaled grad from the previous rank
+        # No need to scale the grad again
+        # Need to unscale when optimizing
+        self.optim_state = OptimState.SCALED
         self.model.backward_by_grad(tensor, grad)
 
     def clip_grad_norm(self, model: nn.Module, max_norm: float):
         if self.optim_state == OptimState.SCALED:
+            self._prepare_grads()
             self._unscale_grads()
         return super().clip_grad_norm(model, max_norm)
 
     def step(self, *args, **kwargs):
-        self._prepare_grads()
-        self._maybe_move_fp32_shards()
 
         # unscale grads if scaled
         if self.optim_state == OptimState.SCALED:
+            self._prepare_grads()
             self._unscale_grads()
 
+        self._maybe_move_fp32_shards()
         found_inf = self._check_overflow()
         self.grad_scaler.update(found_inf)
 


Test patch: