Update disable/enable logic

Signed-off-by: Jingyu Xin <[email protected]>

Author: jingyu-ml

modelopt/torch/peft/conversion.py

@@ -185,13 +185,11 @@ def add_adapter(model, config: PEFTConfig):
                         continue
                 else:
                     raise NotImplementedError(f"Unsupported type {type(wildcard_or_filter_func)}")
-                if adapter_setting.enable:  # type: ignore[union-attr]
-                    module.update_layer_lora(
-                        adapter_name,
-                        adapter_setting,
-                    )
+                module.update_layer_lora(
+                    adapter_name,
+                    adapter_setting,
+                )
 
-    _update_peft_metadata_in_state(model)
     return model
 
 

modelopt/torch/peft/lora/layer.py

@@ -31,7 +31,6 @@ class LoRAModule(DynamicModule):
     def _setup(self) -> None:
         """Initialize LoRA-specific attributes."""
         self._lora_adapters: dict[str, dict[str, Any]] = {}
-        self._active_adapters: set = set()
 
     @property
     def adapter_names(self) -> set:
@@ -43,39 +42,14 @@ def active_adapters(self) -> set:
         """Return the set of currently active adapter names."""
         return self._active_adapters.copy()
 
-    def activate_adapter(self, adapter_name: str) -> None:
-        """Activate a specific adapter.
-
-        Args:
-            adapter_name: Name of the adapter to activate
-
-        Raises:
-            ValueError: If adapter_name is not registered
-        """
-        if adapter_name not in self._lora_adapters:
-            raise ValueError(
-                f"Adapter '{adapter_name}' not found. Available: {list(self._lora_adapters.keys())}"
-            )
-        self._active_adapters.add(adapter_name)
-
-    def deactivate_adapter(self, adapter_name: str) -> None:
-        """Deactivate a specific adapter.
-
-        Args:
-            adapter_name: Name of the adapter to deactivate
-        """
-        self._active_adapters.discard(adapter_name)
-
-    def activate_all_adapters(self) -> None:
-        """Activate all registered adapters."""
-        self._active_adapters = self.adapter_names.copy()
-
-    def deactivate_all_adapters(self) -> None:
-        """Deactivate all adapters."""
-        self._active_adapters.clear()
-
     def _register_adapter(
-        self, adapter_name: str, lora_a: nn.Module, lora_b: nn.Module, rank: int, scale: float = 1.0
+        self,
+        adapter_name: str,
+        lora_a: nn.Module,
+        lora_b: nn.Module,
+        rank: int,
+        scale: float = 1.0,
+        enable: bool = True,
     ) -> None:
         """Register a new LoRA adapter with explicit rank tracking.
 
@@ -86,21 +60,18 @@ def _register_adapter(
             rank: Rank of the LoRA decomposition
             scale: Scale factor for the LoRA output
         """
-        # Add as submodules for proper parameter registration
         self.add_module(f"lora_a_{adapter_name}", lora_a)
         self.add_module(f"lora_b_{adapter_name}", lora_b)
 
         # Store in adapter dictionary with explicit rank
         self._lora_adapters[adapter_name] = {
             "lora_a": lora_a,
             "lora_b": lora_b,
-            "rank": rank,  # Store rank explicitly for reliability
+            "rank": rank,
             "scale": scale,
+            "enable": enable,
         }
 
-        # Automatically activate new adapters
-        self.activate_adapter(adapter_name)
-
     @abstractmethod
     def update_layer_lora(
         self,
@@ -156,14 +127,11 @@ def get_peft_state(self) -> dict[str, Any]:
 
             adapters_config[adapter_name] = {
                 "rank": rank,
-                "is_active": adapter_name in self._active_adapters,
-                "lora_a_type": type(lora_a).__name__,
-                "lora_b_type": type(lora_b).__name__,
+                "enable": adapter_modules.get("enable", True),
                 "scale": adapter_modules.get("scale", 1.0),
             }
 
         modelopt_state["adapters"] = adapters_config
-        modelopt_state["active_adapters"] = list(self._active_adapters)
 
         return modelopt_state
 
@@ -246,41 +214,29 @@ def forward(self, x: torch.Tensor, *args, **kwargs) -> Any:
         Returns:
             Output from the base layer plus active LoRA adaptations
         """
-        # Call the base layer's forward method
         output = super().forward(x, *args, **kwargs)
 
-        # Handle different output types from base layer
         if isinstance(output, tuple):
-            # If output is a tuple, assume first element is the main result
             result = output[0]
             other_outputs = output[1:]
         else:
-            # If output is a single tensor
             result = output
             other_outputs = ()
 
-        # Apply active LoRA adapters
-        if self._active_adapters and self._lora_adapters:
-            for adapter_name in self._active_adapters:
-                if adapter_name in self._lora_adapters:
-                    adapter = self._lora_adapters[adapter_name]
-                    # LoRA computation: result = result + B(A(x))
-                    lora_a = adapter["lora_a"]
-                    lora_b = adapter["lora_b"]
-
-                    # Handle different forward signatures
-                    lora_a_output = lora_a(x)
-                    if isinstance(lora_a_output, tuple):
-                        lora_a_output = lora_a_output[0]
-
-                    lora_b_output = lora_b(lora_a_output)
-                    if isinstance(lora_b_output, tuple):
-                        lora_b_output = lora_b_output[0]
-
-                    scale = adapter.get("scale", 1.0)
-                    result = result + scale * lora_b_output
+        for adapter_name in self._lora_adapters:
+            adapter = self._lora_adapters[adapter_name]
+            if adapter["enable"]:
+                lora_a = adapter["lora_a"]
+                lora_b = adapter["lora_b"]
+                lora_a_output = lora_a(x)
+                if isinstance(lora_a_output, tuple):
+                    lora_a_output = lora_a_output[0]
+                lora_b_output = lora_b(lora_a_output)
+                if isinstance(lora_b_output, tuple):
+                    lora_b_output = lora_b_output[0]
+                scale = adapter["scale"]
+                result = result + scale * lora_b_output
 
-        # Return output in the same format as the base layer
         if other_outputs:
             return (result, *other_outputs)
         else:

modelopt/torch/peft/lora/tp_layer.py

@@ -46,7 +46,13 @@ def _get_init_methods(self, lora_a_init, lora_b_init) -> tuple[Callable, Callabl
         return lora_a_init, lora_b_init
 
     def _register_adapter_with_device(
-        self, adapter_name: str, lora_a: nn.Module, lora_b: nn.Module, rank: int, scale: float
+        self,
+        adapter_name: str,
+        lora_a: nn.Module,
+        lora_b: nn.Module,
+        rank: int,
+        scale: float,
+        enable: bool,
     ) -> None:
         """Register LoRA adapter modules and ensure correct device placement.
 
@@ -78,7 +84,7 @@ def _register_adapter_with_device(
             lora_a = lora_a.to(dtype)
             lora_b = lora_b.to(dtype)
 
-        super()._register_adapter(adapter_name, lora_a, lora_b, rank, scale)
+        super()._register_adapter(adapter_name, lora_a, lora_b, rank, scale, enable)
 
 
 @LoRAModuleRegistry.register({ColumnParallelLinear: "megatron_ColumnParallelLinear"})
@@ -120,7 +126,7 @@ def update_layer_lora(
         )
 
         self._register_adapter_with_device(
-            adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale
+            adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale, attr_config.enable
         )
 
 
@@ -163,7 +169,7 @@ def update_layer_lora(
         )
 
         self._register_adapter_with_device(
-            adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale
+            adapter_name, lora_a, lora_b, attr_config.rank, attr_config.scale, attr_config.enable
         )