fix lora ckpt save format (ColoTensor to Tensor)

Author: BurkeHulk

colossalai/booster/plugin/low_level_zero_plugin.py

@@ -290,7 +290,8 @@ def save_lora_as_pretrained(self, model, checkpoint, use_safetensors):
         assert isinstance(
             peft_model, PeftModel
         ), "The model doesn't have lora adapters, please enable lora before saving."
-        return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors)
+        return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors,
+                                          state_dict = tree_map(lambda x: x.data if torch.is_tensor(x) else x, peft_model.state_dict()))
 
 
 class LowLevelZeroPlugin(DPPluginBase):

colossalai/booster/plugin/torch_ddp_plugin.py

@@ -1,10 +1,12 @@
 from typing import Callable, Dict, Iterator, List, Optional, Tuple, Union
 
+import torch
 import torch.nn as nn
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
 from torch.utils.data import DataLoader
+from torch.utils._pytree import tree_map
 
 from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO
 from colossalai.cluster import DistCoordinator
@@ -134,7 +136,9 @@ def save_lora_as_pretrained(
             assert isinstance(
                 peft_model, PeftModel
             ), "The model doesn't have lora adapters, please enable lora before saving."
-            peft_model.save_pretrained(save_directory=checkpoint, safe_serialization=use_safetensors)
+            return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors,
+                                              state_dict=tree_map(lambda x: x.data if torch.is_tensor(x) else x,
+                                                                  peft_model.state_dict()))
 
 
 class TorchDDPModel(ModelWrapper):

colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py

@@ -11,6 +11,7 @@
 import torch.nn as nn
 from torch.distributed import ProcessGroup
 from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+from torch.utils._pytree import tree_map
 
 from colossalai.cluster import DistCoordinator
 from colossalai.interface import ModelWrapper, OptimizerWrapper
@@ -956,4 +957,6 @@ def save_lora_as_pretrained(self, model, checkpoint, use_safetensors):
         assert isinstance(
             peft_model, PeftModel
         ), "The model doesn't have lora adapters, please enable lora before saving."
-        return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors)
+        return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors,
+                                          state_dict=tree_map(lambda x: x.data if torch.is_tensor(x) else x,
+                                                              peft_model.state_dict()))