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

[haofanwang]

🐛 Describe the bug

Describe the bug

I meet AssertionError: you are calculating the l2 norm twice, which looks similar to another issue.

I guess it is related to set_l2_norm. But strangely, this error doesn't raise at the very beginning, but appear after several steps. Once this happens, the loss becomes nan.

To Reproduce

python -m torch.distributed.run --nproc_per_node=$GPU_NUM --nnodes=$WORLD_SIZE \
  --node_rank=$RANK --master_addr=$MASTER_ADDR --master_port=$MASTER_PORT ./train_dreambooth_colossalai.py

I'm working with this dreambooth example. Both single machine and multiple machine cannot work well in the training.

Environment

I install the lastest ColossalAI from the source as instructed by the new readme.

[haofanwang]

Is there any help with this issue?

[1SAA]

Hi @haofanwang

I can run properly on 2 A100. could you show more information about your configuration file and running environment?

[haofanwang]

Thanks, but I don't think it is related to configuration or environment. @1SAA

Please note that this error is raised suddenly after several steps (16K steps in my cases), and intermediate checkpoints have been successfully save every 5K steps. Could you try to run with more steps? Also, I train on my own dataset, is it possible that there exists some problems with the data itself? But I haven't met any error without using colossalai.

I do try to add skip once I meet this AssertError, but all following steps will be skipped. So it looks more like a problem of overflow. If I set L2 norm to None manually, it leads to "Found overflow. Skip step".

FYI, my Python is 3.8.12, the Colossalai is installed from the source.

Colossal-AI version: 0.2.0
----------------------------
PyTorch Version: 1.13.1
PyTorch Version required by Colossal-AI: 1.13
PyTorch version match: ✓
----------------------------
System CUDA Version: 11.6
CUDA Version required by PyTorch: 11.6
CUDA Version required by Colossal-AI: 11.6
CUDA Version Match: ✓
----------------------------
CUDA Extension: ✓

Much appreciated for help.

[haofanwang]

Full error message is attached for reference.

│   585 if __name__ == "__main__":                                             │
│ ❱ 586 │   main()                                                             │
│   587                                                                        │
│                                                                              │
│ /mnt/mmtech01/usr/yanzhen1/multimodal-intelligence/BLIP/./train_text_to_imag │
│ e_colossalai.py:554 in main                                                  │
│                                                                              │
│   551 │   │   │   model_pred = unet(noisy_latents, timesteps, encoder_hidden │
│   552 │   │   │   loss = F.mse_loss(model_pred.float(), target.float(), redu │
│   553 │   │   │                                                              │
│ ❱ 554 │   │   │   optimizer.backward(loss)                                   │
│   555 │   │   │                                                              │
│   556 │   │   │   optimizer.step()                                           │
│   557 │   │   │   lr_scheduler.step()                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/optimizer/zero_optimize │
│ r.py:225 in backward                                                         │
│                                                                              │
│   222 │   def backward(self, loss: torch.Tensor):                            │
│   223 │   │   loss = self.loss_scale * loss                                  │
│   224 │   │   self.optim_state = OptimState.SCALED                           │
│ ❱ 225 │   │   self.module.backward(loss)                                     │
│   226 │                                                                      │
│   227 │   def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tenso │
│   228 │   │   # This function is called except the last stage of pipeline pa │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel. │
│ py:297 in backward                                                           │
│                                                                              │
│   294 │                                                                      │
│   295 │   def backward(self, loss: torch.Tensor):                            │
│   296 │   │   with self.param_op_hook.switch_to_backward(), ColoParamOpHookM │
│ ❱ 297 │   │   │   loss.backward()                                            │
│   298 │   │   self._post_backward()                                          │
│   299 │                                                                      │
│   300 │   def backward_by_grad(self, tensor, grad):                          │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:479 in backward      │
│                                                                              │
│    476 │   │   │   │   used to compute the attr::tensors.                    │
│    477 │   │   """                                                           │
│    478 │   │   if has_torch_function_unary(self):                            │
│ ❱  479 │   │   │   return handle_torch_function(                             │
│    480 │   │   │   │   Tensor.backward,                                      │
│    481 │   │   │   │   (self,),                                              │
│    482 │   │   │   │   self,                                                 │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/overrides.py:1534 in            │
│ handle_torch_function                                                        │
│                                                                              │
│   1531 │   │                                                                 │
│   1532 │   │   # Use `public_api` instead of `implementation` so __torch_fun │
│   1533 │   │   # implementations can do equality/identity comparisons.       │
│ ❱ 1534 │   │   result = torch_func_method(public_api, types, args, kwargs)   │
│   1535 │   │                                                                 │
│   1536 │   │   if result is not NotImplemented:                              │
│   1537 │   │   │   return result                                             │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/tensor/colo_tensor.py:180  │
│ in __torch_function__                                                        │
│                                                                              │
│   177 │   │   │   │   assert len(args) == 1    # only has 1 paramter         │
│   178 │   │   │   │   backward_tensor = torch.Tensor(args[0])                │
│   179 │   │   │   │   tensor_kwargs = {k: torch.Tensor(v) if torch.is_tensor │
│ ❱ 180 │   │   │   │   return backward_tensor.backward(**tensor_kwargs)       │
│   181 │   │                                                                  │
│   182 │   │   with torch._C.DisableTorchFunction():                          │
│   183 │   │   │   ret = func(*args, **kwargs)                                │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:488 in backward      │
│                                                                              │
│    485 │   │   │   │   create_graph=create_graph,                            │
│    486 │   │   │   │   inputs=inputs,                                        │
│    487 │   │   │   )                                                         │
│ ❱  488 │   │   torch.autograd.backward(                                      │
│    489 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs │
│    490 │   │   )                                                             │
│    491                                                                       │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py:197 in     │
│ backward                                                                     │
│                                                                              │
│   194 │   # The reason we repeat same the comment below is that              │
│   195 │   # some Python versions print out the first line of a multi-line fu │
│   196 │   # calls in the traceback and some print out the last line          │
│ ❱ 197 │   Variable._execution_engine.run_backward(  # Calls into the C++ eng │
│   198 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,    │
│   199 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th │
│   200                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/function.py:267 in     │
│ apply                                                                        │
│                                                                              │
│   264 │   │   │   │   │   │   │      "Function is not allowed. You should on │
│   265 │   │   │   │   │   │   │      "of them.")                             │
│   266 │   │   user_fn = vjp_fn if vjp_fn is not Function.vjp else backward_f │
│ ❱ 267 │   │   return user_fn(self, *args)                                    │
│   268 │                                                                      │
│   269 │   def apply_jvp(self, *args):                                        │
│   270 │   │   # _forward_cls is defined by derived class                     │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/utils/checkpoint.py:157 in      │
│ backward                                                                     │
│                                                                              │
│   154 │   │   │   raise RuntimeError(                                        │
│   155 │   │   │   │   "none of output has requires_grad=True,"               │
│   156 │   │   │   │   " this checkpoint() is not necessary")                 │
│ ❱ 157 │   │   torch.autograd.backward(outputs_with_grad, args_with_grad)     │
│   158 │   │   grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else N │
│   159 │   │   │   │   │     for inp in detached_inputs)                      │
│   160                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py:197 in     │
│ backward                                                                     │
│                                                                              │
│   194 │   # The reason we repeat same the comment below is that              │
│   195 │   # some Python versions print out the first line of a multi-line fu │
│   196 │   # calls in the traceback and some print out the last line          │
│ ❱ 197 │   Variable._execution_engine.run_backward(  # Calls into the C++ eng │
│   198 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,    │
│   199 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th │
│   200                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel. │
│ py:323 in grad_handle                                                        │
│                                                                              │
│   320 │   │   │   │   self.overflow_counter += chunk.has_inf_or_nan          │
│   321 │   │   │   │   # record l2 norm for gradient clipping                 │
│   322 │   │   │   │   if chunk.l2_norm_flag:                                 │
│ ❱ 323 │   │   │   │   │   chunk.set_l2_norm()                                │
│   324 │   │   │   │   self.chunk_manager.move_chunk(chunk, self.grads_device │
│   325 │   │   return empty_grad                                              │
│   326                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py:242  │
│ in set_l2_norm                                                               │
│                                                                              │
│   239 │   def set_l2_norm(self) -> None:                                     │
│   240 │   │   """Record l2 norm of this chunks on CUDA.                      │
│   241 │   │   """                                                            │
│ ❱ 242 │   │   assert self.l2_norm is None, "you are calculating the l2 norm  │
│   243 │   │   if self.is_gathered:                                           │
│   244 │   │   │   valid_tensor = self.cuda_global_chunk[:self.utilized_size] │
│   245 │   │   else:                                                          │
╰──────────────────────────────────────────────────────────────────────────────╯
AssertionError: you are calculating the l2 norm twice
╭───────────────────── Traceback (most recent call last) ──────────────────────╮
│ /mnt/mmtech01/usr/yanzhen1/multimodal-intelligence/BLIP/./train_text_to_imag │
│ e_colossalai.py:586 in <module>                                              │
│                                                                              │
│   583 │   │   │   │   break                                                  │
│   584                                                                        │
│   585 if __name__ == "__main__":                                             │
│ ❱ 586 │   main()                                                             │
│   587                                                                        │
│                                                                              │
│ /mnt/mmtech01/usr/yanzhen1/multimodal-intelligence/BLIP/./train_text_to_imag │
│ e_colossalai.py:554 in main                                                  │
│                                                                              │
│   551 │   │   │   model_pred = unet(noisy_latents, timesteps, encoder_hidden │
│   552 │   │   │   loss = F.mse_loss(model_pred.float(), target.float(), redu │
│   553 │   │   │                                                              │
│ ❱ 554 │   │   │   optimizer.backward(loss)                                   │
│   555 │   │   │                                                              │
│   556 │   │   │   optimizer.step()                                           │
│   557 │   │   │   lr_scheduler.step()                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/optimizer/zero_optimize │
│ r.py:225 in backward                                                         │
│                                                                              │
│   222 │   def backward(self, loss: torch.Tensor):                            │
│   223 │   │   loss = self.loss_scale * loss                                  │
│   224 │   │   self.optim_state = OptimState.SCALED                           │
│ ❱ 225 │   │   self.module.backward(loss)                                     │
│   226 │                                                                      │
│   227 │   def backward_by_grad(self, tensor: torch.Tensor, grad: torch.Tenso │
│   228 │   │   # This function is called except the last stage of pipeline pa │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel. │
│ py:297 in backward                                                           │
│                                                                              │
│   294 │                                                                      │
│   295 │   def backward(self, loss: torch.Tensor):                            │
│   296 │   │   with self.param_op_hook.switch_to_backward(), ColoParamOpHookM │
│ ❱ 297 │   │   │   loss.backward()                                            │
│   298 │   │   self._post_backward()                                          │
│   299 │                                                                      │
│   300 │   def backward_by_grad(self, tensor, grad):                          │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:479 in backward      │
│                                                                              │
│    476 │   │   │   │   used to compute the attr::tensors.                    │
│    477 │   │   """                                                           │
│    478 │   │   if has_torch_function_unary(self):                            │
│ ❱  479 │   │   │   return handle_torch_function(                             │
│    480 │   │   │   │   Tensor.backward,                                      │
│    481 │   │   │   │   (self,),                                              │
│    482 │   │   │   │   self,                                                 │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/overrides.py:1534 in            │
│ handle_torch_function                                                        │
│                                                                              │
│   1531 │   │                                                                 │
│   1532 │   │   # Use `public_api` instead of `implementation` so __torch_fun │
│   1533 │   │   # implementations can do equality/identity comparisons.       │
│ ❱ 1534 │   │   result = torch_func_method(public_api, types, args, kwargs)   │
│   1535 │   │                                                                 │
│   1536 │   │   if result is not NotImplemented:                              │
│   1537 │   │   │   return result                                             │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/tensor/colo_tensor.py:180  │
│ in __torch_function__                                                        │
│                                                                              │
│   177 │   │   │   │   assert len(args) == 1    # only has 1 paramter         │
│   178 │   │   │   │   backward_tensor = torch.Tensor(args[0])                │
│   179 │   │   │   │   tensor_kwargs = {k: torch.Tensor(v) if torch.is_tensor │
│ ❱ 180 │   │   │   │   return backward_tensor.backward(**tensor_kwargs)       │
│   181 │   │                                                                  │
│   182 │   │   with torch._C.DisableTorchFunction():                          │
│   183 │   │   │   ret = func(*args, **kwargs)                                │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/_tensor.py:488 in backward      │
│                                                                              │
│    485 │   │   │   │   create_graph=create_graph,                            │
│    486 │   │   │   │   inputs=inputs,                                        │
│    487 │   │   │   )                                                         │
│ ❱  488 │   │   torch.autograd.backward(                                      │
│    489 │   │   │   self, gradient, retain_graph, create_graph, inputs=inputs │
│    490 │   │   )                                                             │
│    491                                                                       │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py:197 in     │
│ backward                                                                     │
│                                                                              │
│   194 │   # The reason we repeat same the comment below is that              │
│   195 │   # some Python versions print out the first line of a multi-line fu │
│   196 │   # calls in the traceback and some print out the last line          │
│ ❱ 197 │   Variable._execution_engine.run_backward(  # Calls into the C++ eng │
│   198 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,    │
│   199 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th │
│   200                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/function.py:267 in     │
│ apply                                                                        │
│                                                                              │
│   264 │   │   │   │   │   │   │      "Function is not allowed. You should on │
│   265 │   │   │   │   │   │   │      "of them.")                             │
│   266 │   │   user_fn = vjp_fn if vjp_fn is not Function.vjp else backward_f │
│ ❱ 267 │   │   return user_fn(self, *args)                                    │
│   268 │                                                                      │
│   269 │   def apply_jvp(self, *args):                                        │
│   270 │   │   # _forward_cls is defined by derived class                     │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/utils/checkpoint.py:157 in      │
│ backward                                                                     │
│                                                                              │
│   154 │   │   │   raise RuntimeError(                                        │
│   155 │   │   │   │   "none of output has requires_grad=True,"               │
│   156 │   │   │   │   " this checkpoint() is not necessary")                 │
│ ❱ 157 │   │   torch.autograd.backward(outputs_with_grad, args_with_grad)     │
│   158 │   │   grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else N │
│   159 │   │   │   │   │     for inp in detached_inputs)                      │
│   160                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/torch/autograd/__init__.py:197 in     │
│ backward                                                                     │
│                                                                              │
│   194 │   # The reason we repeat same the comment below is that              │
│   195 │   # some Python versions print out the first line of a multi-line fu │
│   196 │   # calls in the traceback and some print out the last line          │
│ ❱ 197 │   Variable._execution_engine.run_backward(  # Calls into the C++ eng │
│   198 │   │   tensors, grad_tensors_, retain_graph, create_graph, inputs,    │
│   199 │   │   allow_unreachable=True, accumulate_grad=True)  # Calls into th │
│   200                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/nn/parallel/data_parallel. │
│ py:323 in grad_handle                                                        │
│                                                                              │
│   320 │   │   │   │   self.overflow_counter += chunk.has_inf_or_nan          │
│   321 │   │   │   │   # record l2 norm for gradient clipping                 │
│   322 │   │   │   │   if chunk.l2_norm_flag:                                 │
│ ❱ 323 │   │   │   │   │   chunk.set_l2_norm()                                │
│   324 │   │   │   │   self.chunk_manager.move_chunk(chunk, self.grads_device │
│   325 │   │   return empty_grad                                              │
│   326                                                                        │
│                                                                              │
│ /opt/conda/lib/python3.8/site-packages/colossalai/gemini/chunk/chunk.py:242  │
│ in set_l2_norm                                                               │
│                                                                              │
│   239 │   def set_l2_norm(self) -> None:                                     │
│   240 │   │   """Record l2 norm of this chunks on CUDA.                      │
│   241 │   │   """                                                            │
│ ❱ 242 │   │   assert self.l2_norm is None, "you are calculating the l2 norm  │
│   243 │   │   if self.is_gathered:                                           │
│   244 │   │   │   valid_tensor = self.cuda_global_chunk[:self.utilized_size] │
│   245 │   │   else:                                                          │
╰──────────────────────────────────────────────────────────────────────────────╯
AssertionError: you are calculating the l2 norm twice

Steps:   3%|▎         | 16024/500000 [3:08:37<94:57:08,  1.42it/s, loss=0.409, lr=0.0001]

[1SAA]

Hi @haofanwang

I believe this problem is caused by the overflow skip. I'll fix this bug as soon as possible.

[haofanwang]

Thanks, it seems that the overflow leads to this issue. Once the overflow happens, all following steps will goes wrong. Please let me know once you fix it. @1SAA

[haofanwang]

There are still some problems. The recent PR solves the AssertionError and works similar to what I have done before by setting L2 norm to None manually, but this raises another new issue where almost all following steps are skipped. The loss also becomes nan. @1SAA

Does the gradient clipping works?

[01/10/23 08:27:14] INFO     colossalai - colossalai - INFO:                    
                             /opt/conda/lib/python3.8/site-packages/colossalai/n
                             n/optimizer/zero_optimizer.py:205 step             
                    INFO     colossalai - colossalai - INFO: Found overflow.    
                             Skip step                                          
                    INFO     colossalai - colossalai - INFO: Found overflow.    
                             Skip step                            

[1SAA]

Hi @haofanwang

I reckon that this problem is cacused by fp16 training. Your lose is too small to be represented as a fp16. You should tune your gradient scaler in the ZeroOptimizer. I recommend you log all your lose and scaled loss to check whether you have encoutered overflow or underflow. You can find more information here.

[haofanwang]

It makes sense. I decrease the learning rate, and the problem is alleviated. Will try to adjust gradient scale. Thanks.