Fix adapter scalar handling under torch.compile (#946)

* Fix adapter scalar handling under torch.compile

* Fix instance where error could happen when merging in a lora to the base model

---------

Co-authored-by: Rydén Johan <johan.ryden@bostad.uppsala.se>
Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>

toolkit/kohya_lora.py

@@ -63,9 +63,9 @@ class LoRAModule(torch.nn.Module):
             self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
 
         if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+            alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
+        self.scale = float(alpha) / self.lora_dim
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's

toolkit/lora_special.py

@@ -69,15 +69,6 @@ class LoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
         torch.nn.Module.__init__(self)
         self.lora_name = lora_name
         self.orig_module_ref = weakref.ref(org_module)
-        # read the device off a param/buffer directly: OstrisLinear.weight is a
-        # property that dequantizes the whole weight just to answer .device
-        org_tensor = next(
-            (t for t in org_module._parameters.values() if t is not None),
-            next((t for t in org_module._buffers.values() if t is not None), None),
-        )
-        self.scalar = torch.tensor(
-            1.0, device=org_tensor.device if org_tensor is not None else None
-        )
 
         # if is ara lora module, mark it on the layer so memory manager can handle it
         if is_ara:
@@ -120,9 +111,9 @@ class LoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
                 self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=use_bias)
 
         if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+            alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
+        self.scale = float(alpha) / self.lora_dim
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's

toolkit/lycoris_special.py

@@ -94,9 +94,9 @@ class LoConSpecialModule(ToolkitModuleMixin, LoConModule, ExtractableModuleMixin
         self.module_dropout = module_dropout
 
         if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+            alpha = float(alpha.detach().float().item())
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
+        self.scale = float(alpha) / self.lora_dim
         self.register_buffer('alpha', torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's

toolkit/models/DoRA.py

@@ -54,7 +54,6 @@ class DoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
         ToolkitModuleMixin.__init__(self, network=network)
         torch.nn.Module.__init__(self)
         self.lora_name = lora_name
-        self.scalar = torch.tensor(1.0)
 
         self.lora_dim = lora_dim
 
@@ -62,9 +61,9 @@ class DoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
             raise NotImplementedError("Convolutional layers are not supported yet")
 
         if type(alpha) == torch.Tensor:
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+            alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = alpha / self.lora_dim
+        self.scale = float(alpha) / self.lora_dim
         # self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える eng: treat as constant
 
         self.multiplier: Union[float, List[float]] = multiplier

toolkit/models/lokr.py

@@ -207,12 +207,12 @@ class LokrModule(ToolkitModuleMixin, nn.Module):
         self.module_dropout = module_dropout
 
         if isinstance(alpha, torch.Tensor):
-            alpha = alpha.detach().float().numpy()  # without casting, bf16 causes error
+            alpha = float(alpha.detach().float().item())
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
         if self.use_w2 and self.use_w1:
             # use scale = 1
             alpha = lora_dim
-        self.scale = alpha / self.lora_dim
+        self.scale = float(alpha) / self.lora_dim
         self.register_buffer('alpha', torch.tensor(alpha))  # treat as constant
 
         if self.use_w2:
@@ -236,7 +236,7 @@ class LokrModule(ToolkitModuleMixin, nn.Module):
             (self.lokr_w2 if self.use_w2
              else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp
              else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.multiplier * self.scale)
+            self.multiplier * self.scale
         )
         assert torch.sum(torch.isnan(weight)) == 0, "weight is nan"
 
@@ -251,7 +251,7 @@ class LokrModule(ToolkitModuleMixin, nn.Module):
             (self.lokr_w2 if self.use_w2
              else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp
              else self.lokr_w2_a@self.lokr_w2_b),
-            torch.tensor(self.scale)
+            self.scale
         )
         if orig_weight is not None:
             weight = weight.reshape(orig_weight.shape)

toolkit/network_mixins.py

@@ -158,7 +158,7 @@ class ExtractableModuleMixin:
 
         # set up alphas
         self.alpha = (self.alpha * 0) + down_weight.shape[0]
-        self.scale = self.alpha / self.lora_dim
+        self.scale = float(self.alpha.detach().float().item()) / self.lora_dim
 
         # assign them
 
@@ -387,7 +387,7 @@ class ToolkitModuleMixin:
         weight_device = weight.device
         if weight.device != down_weight.device:
             weight = weight.to(down_weight.device)
-        if scale.device != down_weight.device:
+        if isinstance(scale, torch.Tensor) and scale.device != down_weight.device:
             scale = scale.to(down_weight.device)
         # merge weight
         if self.full_rank: