github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-01 15:54:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix for async io

Browse Source
This commit is contained in:
flybird11111 2025-02-13 14:06:57 +08:00
parent ec73f1b5e2
commit 6fc6a059a0
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
colossalai/checkpoint_io/utils.py
View File

@ -309,12 +309,13 @@ def async_save_state_dict_shards(
checkpoint_file_path = os.path.join(checkpoint, shard_file) checkpoint_file_path = os.path.join(checkpoint, shard_file)
if state_preprocess: if state_preprocess:
state_dict, _ = _flatten_optim_state_dict(state_dict=shard, seperator=".") state_dict, metadata = _flatten_optim_state_dict(state_dict=shard, seperator=".")
else: else:
state_dict = shard state_dict = shard
metadata = None
# Only save on master rank. # Only save on master rank.
writer = save(checkpoint_file_path, state_dict=state_dict) writer = save(checkpoint_file_path, state_dict=state_dict, metadata=metadata)
writers.append(writer) writers.append(writer)
shard_filenames.append(shard_file) shard_filenames.append(shard_file)
del shard del shard
@ -371,9 +372,10 @@ def async_move_save_state_dict_shards(
checkpoint_file_path = os.path.join(checkpoint, shard_file) checkpoint_file_path = os.path.join(checkpoint, shard_file)
if state_preprocess: if state_preprocess:
state_dict, _ = _flatten_optim_state_dict(state_dict=shard) state_dict, metadata = _flatten_optim_state_dict(state_dict=shard)
else: else:
state_dict = shard state_dict = shard
metadata = None
if pinned_state_dict is not None: if pinned_state_dict is not None:
sub_pinned_state_dict = {k: pinned_state_dict[k] for k in state_dict.keys()} sub_pinned_state_dict = {k: pinned_state_dict[k] for k in state_dict.keys()}
@ -382,7 +384,7 @@ def async_move_save_state_dict_shards(
returned_state_dict.update(sub_pinned_state_dict) returned_state_dict.update(sub_pinned_state_dict)
# Only save on master rank. # Only save on master rank.
writer = move_and_save(checkpoint_file_path, state_dict, sub_pinned_state_dict) writer = move_and_save(checkpoint_file_path, state_dict, sub_pinned_state_dict, metadata)
writers.append(writer) writers.append(writer)
shard_filenames.append(shard_file) shard_filenames.append(shard_file)
del shard del shard