"git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "99b189df6de71b2f01d6f72e6b1f4aa74455275b"
Unverified Commit a6f37f88 authored by Sourab Mangrulkar's avatar Sourab Mangrulkar Committed by GitHub
Browse files

fixes issue when saving fsdp via accelerate's FSDP plugin (#24446)

parent 2898fd39
...@@ -2322,7 +2322,7 @@ class Trainer: ...@@ -2322,7 +2322,7 @@ class Trainer:
torch.save(self.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME)) torch.save(self.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME))
elif self.args.should_save and not self.is_deepspeed_enabled: elif self.args.should_save and not self.is_deepspeed_enabled:
# deepspeed.save_checkpoint above saves model/optim/sched # deepspeed.save_checkpoint above saves model/optim/sched
if self.fsdp: if self.fsdp and not self.is_fsdp_enabled:
torch.save(full_osd, os.path.join(output_dir, OPTIMIZER_NAME)) torch.save(full_osd, os.path.join(output_dir, OPTIMIZER_NAME))
else: else:
torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME)) torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment