From 894cba09f31d4385adb2921ada55ed88cf49bd12 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Sun, 30 Jul 2023 21:46:44 -0400 Subject: [PATCH] fix FSDP save of final model (#329) --- scripts/finetune.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/scripts/finetune.py b/scripts/finetune.py index 9eb65037c..ddf1992d6 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -344,7 +344,9 @@ def train( # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file - if cfg.local_rank == 0: + if cfg.fsdp: + model.save_pretrained(cfg.output_dir) + elif cfg.local_rank == 0: if cfg.flash_optimum: model = BetterTransformer.reverse(model) model.save_pretrained(cfg.output_dir)