diff --git a/tuning/config/acceleration_configs/fast_moe.py b/tuning/config/acceleration_configs/fast_moe.py index 39eed4f5d..78e65d6b7 100644 --- a/tuning/config/acceleration_configs/fast_moe.py +++ b/tuning/config/acceleration_configs/fast_moe.py @@ -99,7 +99,7 @@ def on_save( def checkpoint(checkpoint_dir, save_dir, is_intermediate: bool = True): if is_intermediate: hf_converted_output_dir = os.path.join( - save_dir, "hf_converted_checkpoint" + save_dir, f"safetensors-{state.global_step}" ) else: hf_converted_output_dir = save_dir diff --git a/tuning/trainercontroller/callback.py b/tuning/trainercontroller/callback.py index 89ab4df93..7e580ddf3 100644 --- a/tuning/trainercontroller/callback.py +++ b/tuning/trainercontroller/callback.py @@ -587,7 +587,7 @@ def on_save( kwargs["is_final"] = False base_path = kwargs["path"] - hf_converted_path = os.path.join(base_path, "hf_converted_checkpoint") + hf_converted_path = os.path.join(base_path, f"safetensors-{state.global_step}") if os.path.isdir(hf_converted_path): kwargs["hf_path"] = hf_converted_path