Reduce GPU memory usage when using FSDP+PEFT (#28830)

support FSDP+PEFT
huggingface · Feb 2, 2024 · 80d5007 · 80d5007
1 parent f497795
commit 80d5007
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -1697,6 +1697,8 @@ def _inner_training_loop(
         use_accelerator_prepare = True if model is self.model else False
 
         if delay_optimizer_creation:
+            if use_accelerator_prepare:
+                self.model = self.accelerator.prepare(self.model)
             self.create_optimizer_and_scheduler(num_training_steps=max_steps)
 
         # prepare using `accelerator` prepare