[BUG]: AssertionError: you are calculating the l2 norm twice

[airlsyn]

🐛 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

[Fridge003]

Hi, could you please provide a detailed error traceback/the code snippet you use.

[airlsyn]

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

[Fridge003]

Thanks, we will try to fix it as soon as possible.

[Fridge003]

Hi @EricXSun, this bug should have been fixed, you can try again.

[airlsyn]

Hi @EricXSun, this bug should have been fixed, you can try again.

Okay, I'm trying, thanks.