fix: report a save or upload before re-ablating the model (#488)
Some checks failed
CI / Check and build (Python 3.10) (push) Has been cancelled
CI / Check and build (Python 3.11) (push) Has been cancelled
CI / Check and build (Python 3.12) (push) Has been cancelled
CI / Check and build (Python 3.13) (push) Has been cancelled

This commit is contained in:
Hamed Shams 2026-10-12 15:23:19 +03:00 • committed by GitHub
parent 49b56e5af4
commit 66a952b49b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -946,6 +946,7 @@ def run():
save_directory,
max_shard_size=settings.max_shard_size,
)
print(f"Model saved to [bold]{save_directory}[/].")
else:
print("Saving merged model...")
merged_model = model.get_merged_model()
@ -958,10 +959,11 @@ def run():
model.tokenizer.save_pretrained(save_directory)
if model.processor is not None:
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()
print(f"Model saved to [bold]{save_directory}[/].")
if reproduction_mode:
print("Verifying hashes of weight files...")
@ -1114,6 +1116,7 @@ def run():
max_shard_size=settings.max_shard_size,
token=token,
)
print(f"Model uploaded to [bold]{repo_id}[/].")
else:
print("Uploading merged model...")
merged_model = model.get_merged_model()
@ -1136,6 +1139,9 @@ def run():
private=private,
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()
if is_hf_path(settings.model):
@ -1200,8 +1206,6 @@ def run():
finally:
settings.export_strategy = current_export_strategy
print(f"Model uploaded to [bold]{repo_id}[/].")
if reproduction_mode:
print("Verifying hashes of weight files...")