Merge pull request LianjiaTech#335 from Flow3rDown/main

修正final model保存问题
dellinbcg · May 6, 2023 · 58a71c2 · 58a71c2
2 parents 3212ee8 + b015f9c
commit 58a71c2
Showing 1 changed file with 13 additions and 13 deletions.
diff --git a/train/main.py b/train/main.py
@@ -386,19 +386,19 @@ def evaluation(model, eval_dataloader):
 
         model.tput_timer.update_epoch_count()
 
-        if args.output_dir is not None:
-            print_rank_0('saving the final model ...', args.global_rank)#It will overwrite the last epoch model
-            model = convert_lora_to_linear_layer(model)
-
-            if args.global_rank == 0:
-                save_hf_format(model, tokenizer, args)
-
-            if args.zero_stage == 3:
-                # For zero stage 3, each gpu only has a part of the model, so we need a special save function
-                save_zero_three_model(model,
-                                    args.global_rank,
-                                    args.output_dir,
-                                    zero_stage=args.zero_stage)
+    if args.output_dir is not None:
+        print_rank_0('saving the final model ...', args.global_rank)#It will overwrite the last epoch model
+        model = convert_lora_to_linear_layer(model)
+
+        if args.global_rank == 0:
+            save_hf_format(model, tokenizer, args)
+
+        if args.zero_stage == 3:
+            # For zero stage 3, each gpu only has a part of the model, so we need a special save function
+            save_zero_three_model(model,
+                                args.global_rank,
+                                args.output_dir,
+                                zero_stage=args.zero_stage)
 
 
 if __name__ == "__main__":