@@ -6,3 +6,3 @@
 
 from ...utils import set_mode_by_env_config, SKIP_ENV_GRAPH_MODE, SKIP_ENV_PYNATIVE_MODE
-set_mode_by_env_config()
+set_mode_by_env_config(), enable_backward

@@ -41,2 +47,3 @@
                 self.modules_buffers.append(buffer)
-        self.broadcast = ops.Broadcast(0, pg_name)
+
+        self.broadcast_bucket_size = int(250 * 1024 * 1024)

@@ -258,0 +260,4 @@
+                raise ValueError("Under GraphMode, has to pass `grads` to optimizer.step.")
+            grads = [param.grad if param.grad is not None
+                     else ms.ops.zeros_like(param) for param in self.parameters]
+            grads = tuple(grads)

@@ -88,6 +88,209 @@ def test_sgd():
     torch_result2 = opt.param_groups[0]['params'][0].detach().numpy()
     assert np.allclose(ms_result2, torch_result2)
 
+@SKIP_ENV_GRAPH_MODE("backward not support graph mode.")

@@ -56,0 +64,4 @@
+        for p, g in zip(l.parameters(), ms_grads):
+            p.grad = g
+        ms_total_norm = ms_torch.nn.utils.clip_grad_norm_(l.parameters(), max_norm,
+                                                        norm_type=norm_type)

@@ -64,4 +71,3 @@
     def scale(self, outputs):
         if not self._enabled:
             return outputs
-        return DynamicLossScaler.scale(self, outputs)

@@ -79,3 +103,2 @@
+
         optimizer_state['found_inf_per_device'] = self._check_inf(grads)
-        if graph_mode_condition():
-            raise RuntimeError("Under graph mode, GradScalar not support unscale_(), please use unscale(). "

@@ -43,0 +50,4 @@
+        # TODO: not support 'parameters_to_ignore' now, because it is used by 'delay_all_reduce_named_params',
+        # but 'delay_all_reduce_named_params' relies on Parameter's hook, which is not support yet.
+        self.parameters_to_ignore = set()
+        _sync_module_states(

@@ -63,3 +80,3 @@
         if self.will_sync_module_buffers():
             self._sync_buffers()
-        return self.network(*inputs, **kwargs)
+        return self.module(*inputs, **kwargs)

@@ -93,15 +114,15 @@ class GradScaler(DynamicLossScaler):
         optimizer_state["stage"] = OptState.UNSCALED
         return DynamicLossScaler.unscale(self, grads)
 
-    def _maybe_opt_step(self, optimizer, grads, optimizer_state, *args, **kwargs):

@@ -138,2 +159,3 @@
         if optimizer_state["stage"] is OptState.READY:
-            self.unscale_(optimizer, grads)
+            # To see if grads is pass in.
+            if len(args) > 0 and isinstance(args[0], tuple) and \

@@ -1042,0 +1042,4 @@
+            return
+
+        for p in self.parameters():
+            if p.grad is not None:

@@ -51,3 +52,3 @@
     return new_grads, cast_to_adapter_tensor(total_norm)
 
-def clip_grad_norm_(parameters, max_norm, grads, norm_type=2.0, error_if_nonfinite=False, foreach=None):
+def clip_grad_norm_(parameters, max_norm, grads=None, norm_type=2.0,