github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-21 01:24:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

[ckpt] Add async ckpt api (#6136)

Browse Source 
* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix

* fix
This commit is contained in:
Wang Binluo
2024-11-15 18:19:16 +08:00
committed by Hongxin Liu
parent d4a436051d
commit 8e08c27e19
12 changed files with 174 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/checkpoint_io/general_checkpoint_io.py
View File

@@ -61,8 +61,8 @@ class GeneralCheckpointIO(CheckpointIO):
self.pinned_state_dicts[id(model)] = create_pinned_state_dict(state_dict)
self.async_writers.append(writer)
move_and_save(writer, state_dict, self.pinned_state_dicts[id(model)])
else:
else:
# save the checkpoint
save_state_dict(state_dict, checkpoint, use_safetensors)