support distribute checkpoint io

Author: flybird11111

colossalai/booster/plugin/hybrid_parallel_plugin.py

@@ -78,6 +78,9 @@ def __init__(
         self.require_grad_sync = True
         self.overlap_allgather = overlap_allgather
         self.use_fp8 = use_fp8
+        self.param_origin_shape = {}
+        for name, param in module.named_parameters():
+            self.param_origin_shape[name] = param.shape
 
         shardformer = ShardFormer(shard_config)
         if custom_policy is not None:

colossalai/checkpoint_io/__init__.py

@@ -1,6 +1,8 @@
 from .checkpoint_io_base import CheckpointIO
 from .general_checkpoint_io import GeneralCheckpointIO
 from .hybrid_parallel_checkpoint_io import HybridParallelCheckpointIO
+
+from .distributed_checkpoint_io import DistributedCheckpointIO
 from .index_file import CheckpointIndexFile
 from .moe_checkpoint import MoECheckpointIO
 
@@ -10,4 +12,5 @@
     "GeneralCheckpointIO",
     "HybridParallelCheckpointIO",
     "MoECheckpointIO",
+    "DistributedCheckpointIO",
 ]