Repository: hpcaitech/ColossalAI

Issue:
[BUG]: colossalai  0.3.3 + torch 2.0.1 + baichuan-2 7b 训练保存 lr_scheduler 时会报  NotImplementedError 错
### 🐛 Describe the bug

用 colossalai  0.3.3 + torch 2.0.1 + baichuan-2 7b 训练保存 lr_scheduler 时 colossalai/nn/lr_scheduler/delayed.py 会报  NotImplementedError 错。

In [25]: lr_scheduler
Out[25]: <[colossalai.nn.lr](http://colossalai.nn.lr/)_scheduler.cosine.CosineAnnealingWarmupLR at 0x7f01cd616e00>
In [26]: booster.save_lr_scheduler(lr_scheduler, "/data/checkpoint/lr_scheduler")

```
 in <module>:1                                                                                    
                                                                                                  
 python3.10/site-packages/colossalai/booster/booster.py:308 in            
 save_lr_scheduler                                                                                
                                                                                                  
   305          lr_scheduler (LRScheduler): A lr scheduler boosted by Booster.                 
   306          checkpoint (str): Path to the checkpoint. It must be a local file path.        
   307       """                                                                                
 ❱ 308       self.checkpoint_io.save_lr_scheduler(lr_scheduler, checkpoint)                     
   309                                                                                           
   310    def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str) -> None:       
   311       """Load lr scheduler from checkpoint.                                              
                                                                                                  
 python3.10/site-packages/colossalai/booster/plugin/gemini_plugin.py:225  
 in save_lr_scheduler                                                                             
                                                                                                  
   222       Save model to checkpoint but only on master process.                               
   223       """                                                                                
   224       if self.coordinator.is_master():                                                   
 ❱ 225          super().save_lr_scheduler(lr_scheduler, checkpoint)                            
   226                                                                                            
   227                                                                                            
   228 class GeminiPlugin(DPPluginBase):                                                          
                                                                                                  
 python3.10/site-packages/colossalai/checkpoint_io/checkpoint_io_base.py: 
 318 in save_lr_scheduler                                                                         
                                                                                                  
   315          lr_scheduler (LRScheduler): lr scheduler to be saved.                          
   316          checkpoint: checkpoint path. The checkpoint path can only be a file path.      
   317       """                                                                                
 ❱ 318       torch.save(lr_scheduler.state_dict(), checkpoint)                                  
   319                                                                                           
   320    def load_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):               
   321       """                                                                                
                                                                                                  
 python3.10/site-packages/colossalai/nn/lr_scheduler/delayed.py:93 in     
 state_dict                                                                                       
                                                                                                  
    90          state_dict["after_scheduler_dict"] = state_dict["after_scheduler"].state_dic   
    91          del state_dict["after_scheduler"]                                              
    92       else:                                                                              
 ❱  93          raise NotImplementedError()                                                    
    94       return state_dict                                                                  
    95                                                                                           
    96    def get_lr(self):
```

进一步分析 lr_scheduler 里的信息
```
state_dict = {key: value for key, value in lr_scheduler.__dict__.items() if key not in "optimizer"}

# =>
{
  'warmup_epochs': 2000,
 'after_scheduler': <torch.optim.lr_scheduler.CosineAnnealingLR at 0x7f01cd6173a0>,
 'finished': False,
 'base_lrs': [0.0003],
 'last_epoch': 1,
 'verbose': False,
 '_step_count': 2,
 '_get_lr_called_within_step': False,
 '_last_lr': [3e-07]
}
```

- 其中 after_scheduler 是 torch.optim.lr_scheduler.CosineAnnealingLR 的实例，而 torch.optim.lr_scheduler.CosineAnnealingLR 是继承的 LRScheduler，那么 after_scheduler 的父类是 LRScheduler

- _LRScheduler 是继承了 LRScheduler

- 而在 [save lr scheduler 时（delayed.py) 中](https://github.com/hpcaitech/ColossalAI/blob/822051d8884a46d4d8626330e21adfd6427c99a0/colossalai/nn/lr_scheduler/delayed.py#L88)，是 `isinstance(state_dict['after_scheduler'], _LRScheduler)`

```
from torch.optim.lr_scheduler import _LRScheduler, LRScheduler

isinstance(state_dict['after_scheduler'], LRScheduler)

# => True

isinstance(state_dict['after_scheduler'], _LRScheduler)

# => False

```

**那这样，是否说明 应该用 `LRScheduler` 而不是 `_LRScheduler` 呢？**


注：baichuan-2 依赖 torch 2.0+，不能降到 2.0 以下（用 1.13 会报 TypeError: sdp_kernel() got an unexpected keyword argument 'enable_mem_efficient'）

### Environment

- colossalai  0.3.3
- torch 2.0.1
- baichuan-2 7b 
[tensor] fix some unittests

[tensor] fix some unittests



Discussion:
 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿      
----  
**Title:** [BUG]: colossalai 0.3.3 + torch 2.0.1 + baichuan-2 7b will report a NotImplementedError error when saving lr_scheduler during training    
      
hi，您可以将pytorch版本降到1.13.1试试
 
> Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿         
----    
Hi, you can try downgrading the pytorch version to 1.13.1  
      
> hi，您可以将pytorch版本降到1.13.1试试

没办法降，baichuan-2 依赖 torch 2.0+，降到 1.13 会报 TypeError: sdp_kernel() got an unexpected keyword argument 'enable_mem_efficient'




Resolving patch:
diff --git a/colossalai/nn/lr_scheduler/delayed.py b/colossalai/nn/lr_scheduler/delayed.py
--- a/colossalai/nn/lr_scheduler/delayed.py
+++ b/colossalai/nn/lr_scheduler/delayed.py
@@ -1,4 +1,10 @@
-from torch.optim.lr_scheduler import _LRScheduler
+import torch
+from packaging.version import Version
+
+if Version(torch.__version__) >= Version("2.0.0"):
+    from torch.optim.lr_scheduler import LRScheduler as _LRScheduler
+else:
+    from torch.optim.lr_scheduler import _LRScheduler
 
 
 class _enable_get_lr_call:


Test patch:
diff --git a/tests/test_checkpoint_io/test_general_checkpoint_io.py b/tests/test_checkpoint_io/test_general_checkpoint_io.py
--- a/tests/test_checkpoint_io/test_general_checkpoint_io.py
+++ b/tests/test_checkpoint_io/test_general_checkpoint_io.py
@@ -6,6 +6,7 @@
 from torchvision.models import resnet18
 
 from colossalai.checkpoint_io import GeneralCheckpointIO
+from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR
 from colossalai.testing import check_state_dict_equal, clear_cache_before_run, parameterize
 
 # ========
@@ -22,6 +23,7 @@ def test_unsharded_checkpoint(use_safetensors: bool):
     # create a model and optimizer
     model = resnet18()
     optimizer = Adam(model.parameters(), lr=0.001)
+    lr_scheduler = CosineAnnealingWarmupLR(optimizer, total_steps=10)
 
     # create test data sample
     x = torch.randn(1, 3, 224, 224)
@@ -31,6 +33,7 @@ def test_unsharded_checkpoint(use_safetensors: bool):
     loss = y.sum()
     loss.backward()
     optimizer.step()
+    lr_scheduler.step()
 
     # create a temp file for checkpoint
     if use_safetensors:
@@ -39,19 +42,23 @@ def test_unsharded_checkpoint(use_safetensors: bool):
         suffix = ".bin"
     model_ckpt_tempfile = tempfile.NamedTemporaryFile(suffix=suffix)
     optimizer_ckpt_tempfile = tempfile.NamedTemporaryFile()
+    lr_scheduler_ckpt_tempfile = tempfile.NamedTemporaryFile()
 
-    # save the model and optimizer
+    # save the model, optimizer, lr_scheduler
     ckpt_io = GeneralCheckpointIO()
     ckpt_io.save_model(model, model_ckpt_tempfile.name, use_safetensors=use_safetensors)
     ckpt_io.save_optimizer(optimizer, optimizer_ckpt_tempfile.name)
+    ckpt_io.save_lr_scheduler(lr_scheduler, lr_scheduler_ckpt_tempfile.name)
 
     # create new model
     new_model = resnet18()
     new_optimizer = Adam(new_model.parameters(), lr=0.001)
+    new_lr_scheduler = CosineAnnealingWarmupLR(optimizer, total_steps=10)
 
-    # load the model and optimizer
+    # load the model, optimizer, lr_scheduler
     ckpt_io.load_model(new_model, model_ckpt_tempfile.name)
     ckpt_io.load_optimizer(new_optimizer, optimizer_ckpt_tempfile.name)
+    ckpt_io.load_lr_scheduler(new_lr_scheduler, lr_scheduler_ckpt_tempfile.name)
 
     # check for model and optimizer state dict recursively
     check_state_dict_equal(model.state_dict(), new_model.state_dict())
diff --git a/tests/test_zero/test_gemini/test_zerooptim_state_dict.py b/tests/test_zero/test_gemini/test_zerooptim_state_dict.py
--- a/tests/test_zero/test_gemini/test_zerooptim_state_dict.py
+++ b/tests/test_zero/test_gemini/test_zerooptim_state_dict.py
@@ -72,6 +72,7 @@ def run_dist(rank, world_size, port):
     exam_zero_optim_state_dict()
 
 
+@pytest.mark.skip
 @pytest.mark.dist
 @pytest.mark.parametrize("world_size", [1, 4])
 @rerun_if_address_is_in_use()