fix save issue of deepspeed zero3 (#828)

Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
intel · Nov 30, 2023 · cf5ff82 · cf5ff82
1 parent 464962e
commit cf5ff82
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 9 deletions.
diff --git a/intel_extension_for_transformers/llm/finetuning/finetuning.py b/intel_extension_for_transformers/llm/finetuning/finetuning.py
@@ -553,15 +553,11 @@ def concatenate_data(dataset, max_seq_length):
                 )
 
             trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
-            with training_args.main_process_first(desc="save model"):
-                if is_main_process(training_args.local_rank):
-                    unwrapped_model = unwrap_model(model)
-                    unwrapped_model.save_pretrained(
-                        training_args.output_dir, state_dict=unwrapped_model.state_dict()
-                    )
+            trainer.save_model()
         if finetune_args.do_lm_eval and finetune_args.task == "code-generation":
             tokenizer.padding_side = "right" # padding on the right is needed to cut off padding in `complete_code`
             tokenizer.truncation_side = "left"
+            unwrapped_model = unwrap_model(model)
             unwrapped_model.eval()
             class Eval_Args:
                 n_samples = 20

diff --git a/..._for_transformers/neural_chat/examples/finetuning/instruction/llama2_ds_zero3_config.json b/..._for_transformers/neural_chat/examples/finetuning/instruction/llama2_ds_zero3_config.json
@@ -10,6 +10,7 @@
     "zero_optimization": {
         "stage": 3,
         "overlap_comm": false,
-        "contiguous_gradients": false
+        "contiguous_gradients": false,
+        "stage3_gather_16bit_weights_on_model_save": true
     }
 }
diff --git a/...xtension_for_transformers/neural_chat/examples/finetuning/ppo_pipeline/reward_modeling.py b/...xtension_for_transformers/neural_chat/examples/finetuning/ppo_pipeline/reward_modeling.py
@@ -475,5 +475,4 @@ def compute_metrics(eval_pred):
     trainer.train()
 
     trainer.model = trainer.model.merge_and_unload()
-    if trainer.is_world_process_zero():
-        trainer.save_model()
+    trainer.save_model()