Merge branch 'main' into 3d-checkpoint-async-2

hpcaitech · Dec 16, 2024 · 611dcb5 · 611dcb5
2 parents dbc4ada + e994c64
commit 611dcb5
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/colossalai/checkpoint_io/general_checkpoint_io.py b/colossalai/checkpoint_io/general_checkpoint_io.py
@@ -8,8 +8,6 @@
 import torch.nn as nn
 from torch.optim import Optimizer
 
-from colossalai.utils.safetensors import load_flat
-
 from .checkpoint_io_base import CheckpointIO
 from .index_file import CheckpointIndexFile
 from .utils import (

diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py
@@ -19,7 +19,6 @@
     to_global,
     to_global_for_customized_distributed_tensor,
 )
-from colossalai.utils.safetensors import _flatten_optim_state_dict, move_and_save
 
 SAFE_WEIGHTS_NAME = "model.safetensors"
 WEIGHTS_NAME = "pytorch_model.bin"
@@ -347,6 +346,7 @@ def async_move_save_state_dict_shards(
     Returns:
         int: the total size of shards
     """
+    from colossalai.utils.safetensors import move_and_save
 
     total_size = 0
     shard_filenames = []