github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix

ckpt_api
wangbluo 2024-11-20 14:12:01 +08:00
parent 2aa6e44355
commit cf2e9ed345
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py
View File

@ -232,7 +232,7 @@ class HybridParallelCheckpointIO(GeneralCheckpointIO):
index_file=index_file,
base_filename=weights_name,
is_master=control_saving,
use_safetensors=use_safetensors,
pinned_state_dict=pinned_state_dict,
n_write_entries=self.N_WRITE_ENTRIES,
)
self.pinned_state_dicts[id(model)] = pinned_state_dict
@ -278,7 +278,7 @@ class HybridParallelCheckpointIO(GeneralCheckpointIO):
index_file=index_file,
base_filename=weights_name,
is_master=control_saving,
use_safetensors=use_safetensors,
pinned_state_dict=pinned_state_dict,
n_write_entries=self.N_WRITE_ENTRIES,
)
self.pinned_state_dicts[id(model)] = pinned_state_dict