github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-19 20:23:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
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, index_file=index_file,
base_filename=weights_name, base_filename=weights_name,
is_master=control_saving, is_master=control_saving,
use_safetensors=use_safetensors, pinned_state_dict=pinned_state_dict,
n_write_entries=self.N_WRITE_ENTRIES, n_write_entries=self.N_WRITE_ENTRIES,
) )
self.pinned_state_dicts[id(model)] = pinned_state_dict self.pinned_state_dicts[id(model)] = pinned_state_dict
@ -278,7 +278,7 @@ class HybridParallelCheckpointIO(GeneralCheckpointIO):
index_file=index_file, index_file=index_file,
base_filename=weights_name, base_filename=weights_name,
is_master=control_saving, is_master=control_saving,
use_safetensors=use_safetensors, pinned_state_dict=pinned_state_dict,
n_write_entries=self.N_WRITE_ENTRIES, n_write_entries=self.N_WRITE_ENTRIES,
) )
self.pinned_state_dicts[id(model)] = pinned_state_dict self.pinned_state_dicts[id(model)] = pinned_state_dict