Skip to content

GitLab

Unverified Commit 80d50076 authored by Sourab Mangrulkar's avatar Sourab Mangrulkar Committed by GitHub
Browse files

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

support FSDP+PEFT
parent f4977959
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
src/transformers/trainer.py
View file @ 80d50076
...@@ -1697,6 +1697,8 @@ class Trainer: ...@@ -1697,6 +1697,8 @@ class Trainer:
use_accelerator_prepare = True if model is self.model else False use_accelerator_prepare = True if model is self.model else False
if delay_optimizer_creation: 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) self.create_optimizer_and_scheduler(num_training_steps=max_steps)
# prepare using `accelerator` prepare # prepare using `accelerator` prepare
......
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