From df0e1fb59266c9903ddd6dbfe7339bcd2068d150 Mon Sep 17 00:00:00 2001
From: Younes Belkada <49240599+younesbelkada@users.noreply.github.com>
Date: Tue, 14 Mar 2023 13:06:33 +0100
Subject: [PATCH] [`core`] Fix peft multi-gpu issue (#145)

* add multi-gpu support

* rm deepcopy

* tryo to comment

* style

* fix nits
---
 src/peft/tuners/lora.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/peft/tuners/lora.py b/src/peft/tuners/lora.py
index 22d1a1e..0f65cbf 100644
--- a/src/peft/tuners/lora.py
+++ b/src/peft/tuners/lora.py
@@ -200,6 +200,11 @@ class LoraModel(torch.nn.Module):
             new_module.state = old_module.state
             new_module.to(old_module.weight.device)
 
+        # dispatch to correct device
+        for name, module in new_module.named_modules():
+            if "lora_" in name:
+                module.to(old_module.weight.device)
+
     def __getattr__(self, name: str):
         """Forward missing attributes to the wrapped module."""
         try:
@@ -345,6 +350,7 @@ class Linear(nn.Linear, LoraLayer):
                     transpose(self.lora_B.weight @ self.lora_A.weight, self.fan_in_fan_out) * self.scaling
                 )
                 self.merged = False
+
             return F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)
         elif self.r > 0 and not self.merged:
             result = F.linear(x, transpose(self.weight, self.fan_in_fan_out), bias=self.bias)