auto_wrap_policy for PEFT with FSDP

src/accelerate/utils/dataclasses.py

@@ -967,6 +967,7 @@ def get_module_class_from_name(module, name):
                     return module_class
 
     def set_auto_wrap_policy(self, model):
+        from peft import PeftModel
         from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy
 
         default_transformer_cls_names_to_wrap = (
@@ -986,11 +987,36 @@ def set_auto_wrap_policy(self, model):
                     else:
                         transformer_cls_to_wrap.add(transformer_cls)
 
-                self.auto_wrap_policy = functools.partial(
+                auto_wrap_policy = functools.partial(
                     transformer_auto_wrap_policy,
                     # Transformer layer class to wrap
                     transformer_layer_cls=transformer_cls_to_wrap,
                 )
+
+                # In an FSDP setting PEFT models require individually wrapping unfrozen parameters
+
+                if isinstance(model, PeftModel):
+                    print("PEFT wrapping")
+                    from torch.distributed.fsdp.wrap import (
+                        _or_policy,
+                        lambda_auto_wrap_policy,
+                        transformer_auto_wrap_policy,
+                    )
+
+                    def lambda_policy_fn(module):
+                        if (
+                            len(list(module.named_children())) == 0
+                            and getattr(module, "weight", None) is not None
+                            and module.weight.requires_grad
+                        ):
+                            return True
+                        return False
+
+                    lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn)
+                    auto_wrap_policy = functools.partial(_or_policy, policies=[lambda_policy, auto_wrap_policy])
+
+                self.auto_wrap_policy = auto_wrap_policy
+
             elif auto_wrap_policy == FSDP_AUTO_WRAP_POLICY[1]:
                 min_num_params = int(os.environ.get("FSDP_MIN_NUM_PARAMS", 0))
                 if min_num_params > 0:

tests/fsdp/test_fsdp.py

@@ -103,6 +103,32 @@ def test_state_dict_type(self):
                     self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu)
                     self.assertTrue(fsdp_plugin.state_dict_config.rank0_only)
 
+    def test_auto_wrap_policy_peft(self):
+        from peft import LoraConfig, TaskType, get_peft_model
+        from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
+
+        peft_config = LoraConfig(
+            task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
+        )
+        model = AutoModel.from_pretrained(BERT_BASE_CASED)
+        model = get_peft_model(model, peft_config)
+
+        env = self.dist_env.copy()
+        env["FSDP_AUTO_WRAP_POLICY"] = "TRANSFORMER_BASED_WRAP"
+        env["FSDP_TRANSFORMER_CLS_TO_WRAP"] = "BertLayer"
+        env["FSDP_USE_ORIG_PARAMS"] = "false"
+        env["RANK"] = "0"
+        with mockenv_context(**env):  #
+            fsdp_plugin = FullyShardedDataParallelPlugin()
+            fsdp_plugin.set_auto_wrap_policy(model)
+            kwargs = {
+                "sharding_strategy": fsdp_plugin.sharding_strategy,
+                "auto_wrap_policy": fsdp_plugin.auto_wrap_policy,
+                "use_orig_params": fsdp_plugin.use_orig_params,
+            }
+            torch.distributed.init_process_group(backend="nccl")
+            model = FSDP(model, **kwargs)
+
     def test_auto_wrap_policy(self):
         model = AutoModel.from_pretrained(BERT_BASE_CASED)
         for policy in FSDP_AUTO_WRAP_POLICY: