Add adapter_to() method for moving & converting adapter weights

notebooks/QLoRA_Llama_Finetuning.ipynb

@@ -261,9 +261,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# for _, v in model.get_adapter(\"assistant_adapter\").items():\n",
-    "#     for _, module in v.items():\n",
-    "#         module.to(\"cuda\")"
+    "# model.adapter_to(\"assistant_adapter\", device=\"cuda\")"
    ]
   },
   {

src/adapters/methods/bottleneck.py

@@ -176,6 +176,13 @@ def get_adapter(self, adapter_name: str):
         else:
             return None
 
+    def get_adapter_fusion(self, adapter_names: Union[List, str]):
+        adapter_names = adapter_names if isinstance(adapter_names, str) else ",".join(adapter_names)
+        if adapter_names in self.adapter_fusion_layer:
+            return self.adapter_fusion_layer[adapter_names]
+        else:
+            return None
+
     def pre_block(self, adapter_setup: Union[AdapterCompositionBlock, str], state: BottleneckState) -> BottleneckState:
         if isinstance(adapter_setup, AdapterCompositionBlock):
             adapter_name = adapter_setup.first()

src/adapters/model_mixin.py

@@ -1013,6 +1013,42 @@ def get_adapter(self, name) -> dict:
 
         return dict(destination)
 
+    def adapter_to(
+        self, name: str, device: Optional[Union[torch.device, str]] = None, dtype: Optional[torch.dtype] = None
+    ):
+        """
+        Moves the adapter with the given name to the specified device and data type.
+
+        Args:
+            name (str): The name of the adapter to be moved.
+            device (torch.device or str, optional): The device on which the adapter should be moved.
+            dtype (torch.dtype, optional): The data type to which the adapter should be cast.
+        """
+        for _, v in self.get_adapter(name).items():
+            for _, module in v.items():
+                module.to(device=device, dtype=dtype)
+
+    def adapter_fusion_to(
+        self,
+        adapter_names: Union[Fuse, list, str],
+        device: Optional[Union[torch.device, str]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        """
+        Moves the adapter fusion layer with the given name to the specified device and data type.
+
+        Args:
+            adapter_names (Union[Fuse, list, str]): The name of the adapter fusion layer to be moved.
+            device (torch.device or str, optional): The device on which the adapter fusion layer should be moved.
+            dtype (torch.dtype, optional): The data type to which the adapter fusion layer should be cast.
+        """
+        for _, layer in self.iter_layers():
+            for module in layer.modules():
+                if isinstance(module, BottleneckLayer):
+                    fusion = module.get_adapter_fusion(adapter_names)
+                    if fusion is not None:
+                        fusion.to(device=device, dtype=dtype)
+
     def adapter_summary(self, as_dict=False) -> Union[str, dict]:
         """
         Returns a string summary of all adapters currently added to the model. Each entry in the summary table has the

src/adapters/trainer.py

@@ -212,6 +212,7 @@ def _load_best_model(self):
             adapter_dir = os.path.join(self.state.best_model_checkpoint, adapter)
             if os.path.exists(adapter_dir):
                 model.load_adapter(adapter_dir)
+                model.adapter_to(adapter, device=self.args.device)
         if self.train_adapter_fusion:
             logger.info(
                 f"Loading best adapter fusion(s) from {self.state.best_model_checkpoint} (score:"
@@ -222,7 +223,7 @@ def _load_best_model(self):
                 fusion_dir = os.path.join(self.state.best_model_checkpoint, fusion)
                 if os.path.exists(fusion_dir):
                     model.load_adapter_fusion(fusion_dir)
-        model.to(self.args.device)
+                    model.adapter_fusion_to(fusion, device=self.args.device)
 
 
 class AdapterTrainerCallback(TrainerCallback):