enable selective unfreeze (#7326)

* wip Signed-off-by: arendu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * wip Signed-off-by: arendu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * avoid PTL method conflicts Signed-off-by: arendu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update Signed-off-by: arendu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update Signed-off-by: arendu <[email protected]> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: arendu <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
NVIDIA · Sep 2, 2023 · 1c10b36 · 1c10b36
1 parent bc81900
commit 1c10b36
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/nemo/collections/common/parts/adapter_modules.py b/nemo/collections/common/parts/adapter_modules.py
@@ -61,6 +61,15 @@ def get_default_strategy_config(self) -> 'dataclass':
         """
         return adapter_mixin_strategies.ResidualAddAdapterStrategyConfig()
 
+    def adapter_unfreeze(self,):
+        """
+        Sets the requires grad for all parameters in the adapter to True.
+        This method should be overridden for any custom unfreeze behavior that is required.
+        For example, if not all params of the adapter should be unfrozen.
+        """
+        for param in self.parameters():
+            param.requires_grad_(True)
+
 
 class LinearAdapter(nn.Module, AdapterModuleUtil):
 

diff --git a/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py b/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py
@@ -28,6 +28,7 @@
 from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, init_method_const, init_method_normal
 from nemo.core.classes.mixins import adapter_mixin_strategies
 
+
 try:
     from apex.normalization.fused_layer_norm import MixedFusedLayerNorm
 
@@ -206,6 +207,12 @@ def _get_init_fn(self, init_method: str):
             raise NotImplementedError("out_init_method should be zero, normal or xavier")
         return init_fn
 
+    def adapter_unfreeze(self,):
+        """
+        Can be customized to allow for selective training of only some params in the PEFT.
+        """
+        super().adapter_unfreeze()
+
     def forward(self, x):
 
         if self.norm_position == 'pre':

diff --git a/nemo/core/classes/mixins/adapter_mixins.py b/nemo/core/classes/mixins/adapter_mixins.py
@@ -409,12 +409,12 @@ def unfreeze_enabled_adapters(self, freeze_batchnorm: bool = True) -> None:
 
                     # Check if adapter is enabled or not
                     if self.adapter_cfg[name]['enabled'] and name in module.adapter_layer:
+
                         # Recursively set training mode of submodules
                         module.adapter_layer[name].train()
 
                         # Recursively set grad required for submodules
-                        for pname, param in module.adapter_layer[name].named_parameters():
-                            param.requires_grad_(True)
+                        module.adapter_layer[name].adapter_unfreeze()
 
                         # unfreeze batch norm if any in the adapter submodules
                         for mname, module_ in module.adapter_layer[name].named_modules():