Repository: hpcaitech/ColossalAI

Issue:
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests

[FEATURE]: Adapt the booster plugins with the refactored chekcpointIO API
### Describe the feature

With changes introduced in #3427 , it is important to align the booster plugins with the new APIs.


Discussion:

Resolving patch:
diff --git a/colossalai/booster/plugin/gemini_plugin.py b/colossalai/booster/plugin/gemini_plugin.py
--- a/colossalai/booster/plugin/gemini_plugin.py
+++ b/colossalai/booster/plugin/gemini_plugin.py
@@ -13,6 +13,7 @@
 from torch.utils.data.distributed import DistributedSampler
 
 from colossalai.checkpoint_io import CheckpointIO, GeneralCheckpointIO
+from colossalai.checkpoint_io.utils import save_state_dict
 from colossalai.cluster import DistCoordinator
 from colossalai.interface import ModelWrapper, OptimizerWrapper
 from colossalai.tensor.colo_parameter import ColoParameter
@@ -83,7 +84,7 @@ def load_unsharded_model(self, model: GeminiDDP, checkpoint: str, strict: bool =
         # the model should be unwrapped in self.load_model via ModelWrapper.unwrap
         return super().load_unsharded_model(model, checkpoint, strict=strict)
 
-    def save_unsharded_model(self, model: GeminiDDP, checkpoint: str):
+    def save_unsharded_model(self, model: GeminiDDP, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):
         """
         Save model to checkpoint but only on master process.
         """
@@ -91,14 +92,14 @@ def save_unsharded_model(self, model: GeminiDDP, checkpoint: str):
         # as there is communication when get state dict, this must be called on all processes
         state_dict = model.state_dict(only_rank_0=True)
         if self.coordinator.is_master():
-            self.save_checkpoint(state_dict, checkpoint)
+            save_state_dict(state_dict, checkpoint, use_safetensors)
 
-    def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str):
+    def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):
         """
         Save optimizer to checkpoint but only on master process.
         """
         # TODO(ver217): optimizer state dict is sharded
-        super().save_unsharded_optimizer(optimizer, checkpoint)
+        super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)
 
     def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
         """
diff --git a/colossalai/booster/plugin/torch_ddp_plugin.py b/colossalai/booster/plugin/torch_ddp_plugin.py
--- a/colossalai/booster/plugin/torch_ddp_plugin.py
+++ b/colossalai/booster/plugin/torch_ddp_plugin.py
@@ -33,20 +33,20 @@ def load_unsharded_model(self, model: nn.Module, checkpoint: str, strict: bool =
         # the model should be unwrapped in self.load_model via ModelWrapper.unwrap
         return super().load_unsharded_model(model, checkpoint, strict=strict)
 
-    def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool):
+    def save_unsharded_model(self, model: nn.Module, checkpoint: str, gather_dtensor: bool, use_safetensors: bool):
         """
         Save model to checkpoint but only on master process.
         """
         # the model should be unwrapped in self.load_model via ModelWrapper.unwrap
         if self.coordinator.is_master():
-            super().save_unsharded_model(model, checkpoint)
+            super().save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors)
 
     def save_unsharded_optimizer(self, optimizer: Optimizer, checkpoint: str, gather_dtensor: bool):
         """
         Save optimizer to checkpoint but only on master process.
         """
         if self.coordinator.is_master():
-            super().save_unsharded_optimizer(optimizer, checkpoint)
+            super().save_unsharded_optimizer(optimizer, checkpoint, gather_dtensor)
 
     def save_lr_scheduler(self, lr_scheduler: LRScheduler, checkpoint: str):
         """


Test patch: