fix model state_dict retrieving in zero3 (#576)

Co-authored-by: 聂靖入 <niejingru@bytedance.com>
CarperAI · Oct 30, 2023 · c26d450 · c26d450
1 parent 91a0f43
commit c26d450
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trlx/trainer/accelerate_ppo_trainer.py b/trlx/trainer/accelerate_ppo_trainer.py
@@ -539,7 +539,7 @@ def save_pretrained(self, directory: Optional[str] = None, **kwargs):
 
         # Save only the base model, so that is could be loaded directly
         # with Hugging Face's `from_pretrained` method
-        state_dict = self.accelerator.get_state_dict(self.model.base_model)
+        state_dict = self.accelerator.get_state_dict(self.model, unwrap=True)
 
         self.accelerator.unwrap_model(self.model).save_pretrained(
             directory,