support asyncio for 3d

flybird11111 · Nov 26, 2024 · 29eb97e · 29eb97e
1 parent ab856fd
commit 29eb97e
Show file tree

Hide file tree

Showing 4 changed files with 253 additions and 49 deletions.
diff --git a/colossalai/checkpoint_io/general_checkpoint_io.py b/colossalai/checkpoint_io/general_checkpoint_io.py
@@ -13,7 +13,7 @@
 from .checkpoint_io_base import CheckpointIO
 from .index_file import CheckpointIndexFile
 from .utils import (
-    async_save_state_dict_shards,
+    async_move_save_state_dict_shards,
     create_pinned_state_dict,
     get_model_base_filenames,
     get_optimizer_base_filenames,
@@ -189,7 +189,7 @@ def save_sharded_model(
 
         if use_async:
             pinned_state_dict = self.pinned_state_dicts.get(id(model), None)
-            total_size, new_pinned_state_dict, writers = async_save_state_dict_shards(
+            total_size, new_pinned_state_dict, writers = async_move_save_state_dict_shards(
                 sharded_state_dict=state_dict_shard,
                 checkpoint=checkpoint_path,
                 index_file=index_file,