fix: prevent errors during model reset from interrupting export (#565)

Fixes #563
This commit is contained in:
Philipp Emanuel Weidmann authored and GitHub committed 2026-11-05 16:20:16 +05:30
1 parent 6fc6074c61
commit b4864fdd59
1 file changed
+10 -10
+10 -10
View File
@@ -1005,7 +1005,6 @@ 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()
@@ -1018,10 +1017,8 @@ 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...")
@@ -1048,6 +1045,9 @@ def run():
f"[bold]{filename}:[/] [red]File not found[/]"
)
if strategy != ExportStrategy.ADAPTER:
reset_trial_model()
case "upload":
# We don't use huggingface_hub.login() because that stores the token on disk,
# and since this program will often be run on rented or shared GPU servers,
@@ -1175,7 +1175,6 @@ 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()
@@ -1198,10 +1197,6 @@ 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):
card = ModelCard.load(settings.model)
@@ -1269,6 +1264,8 @@ 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...")
@@ -1320,6 +1317,9 @@ def run():
f"[bold]{filename}:[/] [red]File not found[/]"
)
if strategy != ExportStrategy.ADAPTER:
reset_trial_model()
case "chat":
print()
print(