Fix multiple deletions of the same files in save_pretrained (#16947)

* Fix multiple deletions of the same files in save_pretrained

* Add is_main_process argument
This commit is contained in:
Sylvain Gugger
2022-04-27 12:28:42 -04:00
committed by GitHub
parent bfbec17765
commit c79bbc3ba5
14 changed files with 105 additions and 35 deletions

View File

@@ -669,7 +669,9 @@ def main():
if (args.push_to_hub and epoch < args.num_train_epochs - 1) or args.output_dir is not None:
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.output_dir, save_function=accelerator.save)
unwrapped_model.save_pretrained(
args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
)
if (args.push_to_hub and epoch < args.num_train_epochs - 1) and accelerator.is_main_process:
repo.push_to_hub(
@@ -720,7 +722,9 @@ def main():
if args.output_dir is not None:
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(args.output_dir, save_function=accelerator.save)
unwrapped_model.save_pretrained(
args.output_dir, is_main_process=accelerator.is_main_process, save_function=accelerator.save
)
if accelerator.is_main_process:
if args.push_to_hub:
repo.push_to_hub(commit_message="End of training", auto_lfs_prune=True)