Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[BUG]: AssertionError: you are calculating the l2 norm twice
### 🐛 Describe the bug

following the section `Gradient Accumulation on GeminiPlugin` in https://colossalai.org/docs/features/gradient_accumulation_with_booster/ ,  I got the exception `AssertionError: you are calculating the l2 norm twice` at the second step.

### Environment

Pytorch: 2.0.1
Colossalai: 0.3.4
[tensor] fix some unittests



Discussion:

Hi, could you please provide a detailed error traceback/the code snippet you use.
Here is the traceback:


```
[2023-11-01 17:05:43] [I][unknown_node][pretrain.py:237:main] - step: 1, loss: 1.17188, cost[ms][total: 2575, dataloader: 991, forward: 1463, backward: 118, optim: 2], batch token passed: 611, global token passed: 611
[2023-11-01 17:05:43] [I][unknown_node][pretrain.py:272:main] - -> [start]empty cuda cache at first forward+backward
[2023-11-01 17:05:43] [I][unknown_node][pretrain.py:274:main] - -> [end]empty cuda cache at first forward+backward
[2023-11-01 17:05:43] [E][unknown_node][pretrain.py:205:main] -
╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ pretrain.py:283 in <module>                                                                      │
│                                                                                                  │
│   280                                                                                            │
│   281                                                                                            │
│   282 if __name__ == "__main__":                                                                 │
│ ❱ 283 │   main()                                                                                 │
│   284                                                                                            │
│                                                                                                  │
│                                                                                                  │
│ pretrain.py:183 in main                                                                          │
│                                                                                                  │
│   180 │   │   │   │                                                                              │
│   181 │   │   │   │   loss = outputs.loss / args.gradient_accum_step                             │
│   182 │   │   │   │                                                                              │
│ ❱ 183 │   │   │   │   booster.backward(loss, optimizer)                                          │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/booster/booster.py:167 in backward                │
│                                                                                                  │
│   164 │   │   │   optimizer (Optimizer): The optimizer to be updated.                            │
│   165 │   │   """                                                                                │
│   166 │   │   # TODO(frank lee): implement this method with plugin                               │
│ ❱ 167 │   │   optimizer.backward(loss)                                                           │
│   168 │                                                                                          │
│   169 │   def execute_pipeline(                                                                  │
│   170 │   │   self,                                                                              │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/zero/gemini/gemini_optimizer.py:274               │
│ in backward                                                                                      │
│                                                                                                  │
│   271 │                                                                                          │
│   272 │   def backward(self, loss: torch.Tensor):                                                │
│   273 │   │   loss = self.mix_precision_mixin.pre_backward(loss)                                 │
│ ❱ 274 │   │   self.module.backward(loss)                                                         │
│   275 │                                                                                          │
│   276 │   def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tensor):                  │
│   277 │   │   # This function is called except the last stage of pipeline parallel               │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/zero/gemini/gemini_ddp.py:317 in                  │
│ backward                                                                                         │
│                                                                                                  │
│   314 │   def backward(self, loss: torch.Tensor):                                                │
│   315 │   │   self._pre_backward()                                                               │
│   316 │   │   with self.param_op_hook.switch_to_backward(), ColoParamOpHookManager.use_hooks(s   │
│ ❱ 317 │   │   │   loss.backward()                                                                │
│   318 │   │   self._post_backward()                                                              │
│   319 │                                                                                          │
│   320 │   def backward_by_grad(self, tensor, grad):                                              │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/torch/_tensor.py:478 in backward                             │
│                                                                                                  │
│    475 │   │   │   │   used to compute the attr::tensors.                                        │
│    476 │   │   """                                                                               │
│    477 │   │   if has_torch_function_unary(self):                                                │
│ ❱  478 │   │   │   return handle_torch_function(                                                 │
│    479 │   │   │   │   Tensor.backward,                                                          │
│    480 │   │   │   │   (self,),                                                                  │
│    481 │   │   │   │   self,                                                                     │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/torch/overrides.py:1551 in                                   │
│ handle_torch_function                                                                            │
│                                                                                                  │
│   1548 │   │                                                                                     │
│   1549 │   │   # Use `public_api` instead of `implementation` so __torch_function__              │
│   1550 │   │   # implementations can do equality/identity comparisons.                           │
│ ❱ 1551 │   │   result = torch_func_method(public_api, types, args, kwargs)                       │
│   1552 │   │                                                                                     │
│   1553 │   │   if result is not NotImplemented:                                                  │
│   1554 │   │   │   return result                                                                 │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/tensor/colo_tensor.py:81 in                       │
│ __torch_function__                                                                               │
│                                                                                                  │
│    78 │   │   │   │   assert len(args) == 1  # only has 1 parameter                              │
│    79 │   │   │   │   backward_tensor = torch.Tensor(args[0])                                    │
│    80 │   │   │   │   tensor_kwargs = {k: torch.Tensor(v) if torch.is_tensor(v) else v for k,    │
│ ❱  81 │   │   │   │   return backward_tensor.backward(**tensor_kwargs)                           │
│    82 │   │                                                                                      │
│    83 │   │   # replace the in-place function                                                    │
│    84 │   │   if func in INPALCE_MAPPING:                                                        │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/torch/_tensor.py:487 in backward                             │
│                                                                                                  │
│    484 │   │   │   │   create_graph=create_graph,                                                │
│    485 │   │   │   │   inputs=inputs,                                                            │
│    486 │   │   │   )                                                                             │
│ ❱  487 │   │   torch.autograd.backward(                                                          │
│    488 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs                     │
│    489 │   │   )                                                                                 │
│    490                                                                                           │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/torch/autograd/__init__.py:200 in backward                   │
│                                                                                                  │
│   197 │   # The reason we repeat same the comment below is that                                  │
│   198 │   # some Python versions print out the first line of a multi-line function               │
│   199 │   # calls in the traceback and some print out the last line                              │
│ ❱ 200 │   Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the bac   │
│   201 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,                        │
│   202 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to ru   │
│   203                                                                                            │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/zero/gemini/gemini_ddp.py:374 in                  │
│ grad_handle                                                                                      │
│                                                                                                  │
│   371 │   │   │   │   self.overflow_counter += grad_chunk.has_inf_or_nan                         │
│   372 │   │   │   │   # record l2 norm for gradient clipping. flag is bound to fp16 chunk        │
│   373 │   │   │   │   if chunk.l2_norm_flag:                                                     │
│ ❱ 374 │   │   │   │   │   grad_chunk.set_l2_norm()                                               │
│   375 │   │   │   │   self.chunk_manager.move_chunk(grad_chunk, self.grads_device[p], force_co   │
│   376 │   │   │   │   if not (self.master_weights) or (self.enable_gradient_accumulation):       │
│   377 │   │   │   │   │   self.chunk_manager.move_chunk(chunk, self.grads_device[p], force_cop   │
│                                                                                                  │
│ .local/lib/python3.10/site-packages/colossalai/zero/gemini/chunk/chunk.py:251 in                 │
│ set_l2_norm                                                                                      │
│                                                                                                  │
│   248 │                                                                                          │
│   249 │   def set_l2_norm(self) -> None:                                                         │
│   250 │   │   """Record l2 norm of this chunks on CUDA."""                                       │
│ ❱ 251 │   │   assert self.l2_norm is None, "you are calculating the l2 norm twice"               │
│   252 │   │   if self.is_gathered:                                                               │
│   253 │   │   │   valid_tensor = self.cuda_global_chunk[: self.utilized_size]                    │
│   254 │   │   else:                                                                              │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
AssertionError: you are calculating the l2 norm twice

```
Thanks, we will try to fix it as soon as possible.


Resolving patch:
diff --git a/colossalai/zero/gemini/chunk/chunk.py b/colossalai/zero/gemini/chunk/chunk.py
--- a/colossalai/zero/gemini/chunk/chunk.py
+++ b/colossalai/zero/gemini/chunk/chunk.py
@@ -637,6 +637,7 @@ def init_grad_chunk(self) -> "Chunk":
             # grad chunk is initialized, just reallocate cuda global chunk
             self.grad_chunk.cuda_shard = None
             self.grad_chunk.is_gathered = True
+            self.grad_chunk.l2_norm = None
             alloc_storage(self.grad_chunk.cuda_global_chunk)
 
         return self.grad_chunk
diff --git a/colossalai/zero/gemini/gemini_ddp.py b/colossalai/zero/gemini/gemini_ddp.py
--- a/colossalai/zero/gemini/gemini_ddp.py
+++ b/colossalai/zero/gemini/gemini_ddp.py
@@ -343,6 +343,7 @@ def grad_handle(self, p, grad):
                         grad_chunk = self.chunk_manager.rearrange_accumulated_grad_chunk(chunk)
                     else:
                         grad_chunk = chunk.grad_chunk
+                        chunk.grad_chunk.l2_norm = None
 
                 # hold -> compute -> hold after bwd
                 grad_chunk.tensor_trans_state(p, TensorState.COMPUTE)


Test patch:
diff --git a/tests/test_zero/test_gemini/test_grad_accum.py b/tests/test_zero/test_gemini/test_grad_accum.py
--- a/tests/test_zero/test_gemini/test_grad_accum.py
+++ b/tests/test_zero/test_gemini/test_grad_accum.py
@@ -49,7 +49,10 @@ def check_grad(model: GeminiDDP, torch_model: torch.nn.Module):
 @parameterize("keep_gathered", [False, True])
 @parameterize("model_name", ["transformers_gpt_lm"])
 @parameterize("master_weights", [False, True])
-def exam_gemini_grad_acc(placement_config, keep_gathered: bool, model_name: str, master_weights: bool):
+@parameterize("use_grad_checkpoint", [False, True])
+def exam_gemini_grad_acc(
+    placement_config, keep_gathered: bool, model_name: str, master_weights: bool, use_grad_checkpoint: bool
+):
     init_device = get_current_device()
     model_builder, data_gen_fn, output_transform_fn, loss_fn, *_ = next(
         iter(model_zoo.get_sub_registry(model_name).values())
@@ -63,6 +66,10 @@ def exam_gemini_grad_acc(placement_config, keep_gathered: bool, model_name: str,
     for torch_p, p in zip(torch_model.parameters(), gemini_model.parameters()):
         torch_p.data.copy_(p.data)
 
+    if use_grad_checkpoint:
+        gemini_model.gradient_checkpointing_enable()
+        torch_model.gradient_checkpointing_enable()
+
     world_size = torch.distributed.get_world_size()
     config_dict, *_ = search_chunk_configuration(gemini_model, search_range_m=1, search_interval=100)
     config_dict[world_size]["chunk_size"] = 5000
@@ -77,7 +84,7 @@ def exam_gemini_grad_acc(placement_config, keep_gathered: bool, model_name: str,
         **placement_config,
     )
     optimizer = HybridAdam(gemini_model.parameters(), lr=1e-3)
-    gemini_optim = GeminiOptimizer(optimizer, gemini_model, initial_scale=1)
+    gemini_optim = GeminiOptimizer(optimizer, gemini_model, initial_scale=1, max_norm=1.0)
 
     rank = dist.get_rank()
 
@@ -112,6 +119,7 @@ def exam_gemini_grad_acc(placement_config, keep_gathered: bool, model_name: str,
         check_grad(gemini_model, torch_model)
 
         if (i + 1) % accum_iter == 0:
+            torch.nn.utils.clip_grad_norm_(amp.master_params(torch_optim), 1.0)
             torch_optim.step()
             gemini_optim.step()
             torch_optim.zero_grad()
diff --git a/tests/test_zero/test_gemini/test_grad_clip.py b/tests/test_zero/test_gemini/test_grad_clip.py
--- a/tests/test_zero/test_gemini/test_grad_clip.py
+++ b/tests/test_zero/test_gemini/test_grad_clip.py
@@ -88,7 +88,7 @@ def exam_grad_clipping(placement_config, model_name: str, master_weights: bool):
     )
 
     optimizer = HybridAdam(model.parameters(), lr=1e-3)
-    zero_optim = GeminiOptimizer(optimizer, model, initial_scale=32, clipping_norm=1.0)
+    zero_optim = GeminiOptimizer(optimizer, model, initial_scale=32, max_norm=1.0)
 
     model.train()
     torch_model.train()