only main process should call _save on deepspeed zero3 (#25959)

only main process should call _save when deepspeed zero3
huggingface · Sep 11, 2023 · 7fd2d68 · 7fd2d68
1 parent 95b3749
commit 7fd2d68
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -2850,7 +2850,8 @@ def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = Fa
                     " stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
                     " zero_to_fp32.py to recover weights"
                 )
-                self._save(output_dir, state_dict={})
+                if self.args.should_save:
+                    self._save(output_dir, state_dict={})
                 # remove the dummy state_dict
                 remove_dummy_checkpoint(self.args.should_save, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME])
                 self.model_wrapped.save_checkpoint(output_dir)