Skip to content

Commit facbce8

Browse files
fix: add global step number in hf moe chpt (#643)
Signed-off-by: yashasvi <yashasvi@ibm.com>
1 parent 87d90b4 commit facbce8

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

tuning/config/acceleration_configs/fast_moe.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ def on_save(
9999
def checkpoint(checkpoint_dir, save_dir, is_intermediate: bool = True):
100100
if is_intermediate:
101101
hf_converted_output_dir = os.path.join(
102-
save_dir, "hf_converted_checkpoint"
102+
save_dir, f"safetensors-{state.global_step}"
103103
)
104104
else:
105105
hf_converted_output_dir = save_dir

tuning/trainercontroller/callback.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,7 @@ def on_save(
587587
kwargs["is_final"] = False
588588

589589
base_path = kwargs["path"]
590-
hf_converted_path = os.path.join(base_path, "hf_converted_checkpoint")
590+
hf_converted_path = os.path.join(base_path, f"safetensors-{state.global_step}")
591591

592592
if os.path.isdir(hf_converted_path):
593593
kwargs["hf_path"] = hf_converted_path

0 commit comments

Comments
 (0)