fix: report a save or upload before re-ablating the model (#488)

This commit is contained in:
Hamed Shams
2026-10-12 17:53:19 +05:30
committed by GitHub
parent 49b56e5af4
commit 66a952b49b
+8 -4
View File
@@ -946,6 +946,7 @@ def run():
save_directory, save_directory,
max_shard_size=settings.max_shard_size, max_shard_size=settings.max_shard_size,
) )
print(f"Model saved to [bold]{save_directory}[/].")
else: else:
print("Saving merged model...") print("Saving merged model...")
merged_model = model.get_merged_model() merged_model = model.get_merged_model()
@@ -958,10 +959,11 @@ def run():
model.tokenizer.save_pretrained(save_directory) model.tokenizer.save_pretrained(save_directory)
if model.processor is not None: if model.processor is not None:
model.processor.save_pretrained(save_directory) model.processor.save_pretrained(save_directory)
print(f"Model saved to [bold]{save_directory}[/].")
# Merging replaced the trial's adapter in place.
# Restore it only after the save has been reported.
reset_trial_model() reset_trial_model()
print(f"Model saved to [bold]{save_directory}[/].")
if reproduction_mode: if reproduction_mode:
print("Verifying hashes of weight files...") print("Verifying hashes of weight files...")
@@ -1114,6 +1116,7 @@ def run():
max_shard_size=settings.max_shard_size, max_shard_size=settings.max_shard_size,
token=token, token=token,
) )
print(f"Model uploaded to [bold]{repo_id}[/].")
else: else:
print("Uploading merged model...") print("Uploading merged model...")
merged_model = model.get_merged_model() merged_model = model.get_merged_model()
@@ -1136,6 +1139,9 @@ def run():
private=private, private=private,
token=token, token=token,
) )
print(f"Model uploaded to [bold]{repo_id}[/].")
# Merging replaced the trial's adapter in place.
# Restore it only after the upload has been reported.
reset_trial_model() reset_trial_model()
if is_hf_path(settings.model): if is_hf_path(settings.model):
@@ -1200,8 +1206,6 @@ def run():
finally: finally:
settings.export_strategy = current_export_strategy settings.export_strategy = current_export_strategy
print(f"Model uploaded to [bold]{repo_id}[/].")
if reproduction_mode: if reproduction_mode:
print("Verifying hashes of weight files...") print("Verifying hashes of weight files...")