From 4a2d8d085417bee3dcc9c2a560f7603443b8c1c3 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Wed, 26 Feb 2025 14:47:45 -0800 Subject: [PATCH] enable fsdp flag for FSDP2Strategy Signed-off-by: Alexandros Koumparoulis --- nemo/lightning/pytorch/callbacks/peft.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nemo/lightning/pytorch/callbacks/peft.py b/nemo/lightning/pytorch/callbacks/peft.py index cc7407e5a785..17aaf38ae8f0 100644 --- a/nemo/lightning/pytorch/callbacks/peft.py +++ b/nemo/lightning/pytorch/callbacks/peft.py @@ -140,7 +140,7 @@ def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) super().setup(trainer, pl_module, stage=stage) - self._is_fsdp_v1 = type(trainer.strategy).__name__ == 'FSDPStrategy' + self._is_fsdp_v1 = type(trainer.strategy).__name__ in ['FSDPStrategy', 'FSDP2Strategy'] trainer.strategy.trainer = trainer wrapped_io = partial(WrappedAdapterIO, peft=self)