Skip to content

Commit c307497

Browse files
fix: add global step number in hf moe chpt
Signed-off-by: yashasvi <yashasvi@ibm.com>
1 parent 3a57f8a commit c307497

2 files changed

Lines changed: 4 additions & 2 deletions

File tree

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"hf_converted_checkpoint-{state.global_step}"
103103
)
104104
else:
105105
hf_converted_output_dir = save_dir

tuning/trainercontroller/callback.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,9 @@ 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(
591+
base_path, f"hf_converted_checkpoint-{state.global_step}"
592+
)
591593

592594
if os.path.isdir(hf_converted_path):
593595
kwargs["hf_path"] = hf_converted_path

0 commit comments

Comments
 (0)