Restore device runtime scales for compiled adapters (#963)

Co-authored-by: Rydén Johan <johan.ryden@bostad.uppsala.se>

toolkit/kohya_lora.py

@@ -65,7 +65,7 @@ class LoRAModule(torch.nn.Module):
         if type(alpha) == torch.Tensor:
             alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = float(alpha) / self.lora_dim
+        self._set_runtime_scale(float(alpha) / self.lora_dim)
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's
@@ -78,6 +78,21 @@ class LoRAModule(torch.nn.Module):
         self.rank_dropout = rank_dropout
         self.module_dropout = module_dropout
 
+    def _set_runtime_scale(self, value) -> None:
+        """Keep float metadata while using a device tensor in compiled math."""
+        self.scale = float(value)
+        runtime_scale = getattr(self, "_runtime_scale", None)
+        if runtime_scale is None:
+            reference = next(self.parameters(), None)
+            if reference is None:
+                runtime_scale = torch.tensor(self.scale, dtype=torch.float32)
+            else:
+                runtime_scale = reference.new_tensor(self.scale, dtype=torch.float32)
+            self.register_buffer("_runtime_scale", runtime_scale, persistent=False)
+        else:
+            with torch.no_grad():
+                runtime_scale.fill_(self.scale)
+
     def apply_to(self):
         self.org_forward = self.org_module.forward
         self.org_module.forward = self.forward
@@ -108,9 +123,9 @@ class LoRAModule(torch.nn.Module):
 
             # scaling for rank dropout: treat as if the rank is changed
             # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる
-            scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
+            scale = self._runtime_scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
         else:
-            scale = self.scale
+            scale = self._runtime_scale
 
         lx = self.lora_up(lx)
 
@@ -219,7 +234,7 @@ class LoRAInfModule(LoRAModule):
 
     def default_forward(self, x):
         # print("default_forward", self.lora_name, x.size())
-        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+        return self.org_forward(x) + self.lora_up(self.lora_down(x)) * self.multiplier * self._runtime_scale
 
     def forward(self, x):
         if not self.enabled:
@@ -258,7 +273,7 @@ class LoRAInfModule(LoRAModule):
             return self.default_forward(x)
 
         # apply mask for LoRA result
-        lx = self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+        lx = self.lora_up(self.lora_down(x)) * self.multiplier * self._runtime_scale
         mask = self.get_mask_for_x(lx)
         # print("regional", self.lora_name, self.network.sub_prompt_index, lx.size(), mask.size())
         lx = lx * mask
@@ -302,7 +317,7 @@ class LoRAInfModule(LoRAModule):
 
         # apply sub prompt of X
         lx = x[emb_idx :: self.network.num_sub_prompts]
-        lx = self.lora_up(self.lora_down(lx)) * self.multiplier * self.scale
+        lx = self.lora_up(self.lora_down(lx)) * self.multiplier * self._runtime_scale
 
         # print("sub_prompt_forward", self.lora_name, x.size(), lx.size(), emb_idx)
 
@@ -322,7 +337,7 @@ class LoRAInfModule(LoRAModule):
 
         # call own LoRA
         x1 = x[self.network.batch_size + self.network.sub_prompt_index :: self.network.num_sub_prompts]
-        lx1 = self.lora_up(self.lora_down(x1)) * self.multiplier * self.scale
+        lx1 = self.lora_up(self.lora_down(x1)) * self.multiplier * self._runtime_scale
 
         if self.network.is_last_network:
             lx = torch.zeros(

toolkit/lora_special.py

@@ -113,7 +113,7 @@ class LoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
         if type(alpha) == torch.Tensor:
             alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = float(alpha) / self.lora_dim
+        self._set_runtime_scale(float(alpha) / self.lora_dim)
         self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's

toolkit/lycoris_special.py

@@ -96,7 +96,7 @@ class LoConSpecialModule(ToolkitModuleMixin, LoConModule, ExtractableModuleMixin
         if type(alpha) == torch.Tensor:
             alpha = float(alpha.detach().float().item())
         alpha = lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = float(alpha) / self.lora_dim
+        self._set_runtime_scale(float(alpha) / self.lora_dim)
         self.register_buffer('alpha', torch.tensor(alpha))  # 定数として扱える
 
         # same as microsoft's

toolkit/models/DoRA.py

@@ -63,7 +63,7 @@ class DoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
         if type(alpha) == torch.Tensor:
             alpha = float(alpha.detach().float().item())
         alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
-        self.scale = float(alpha) / self.lora_dim
+        scale = float(alpha) / self.lora_dim
         # self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える eng: treat as constant
 
         self.multiplier: Union[float, List[float]] = multiplier
@@ -89,6 +89,8 @@ class DoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
         # self.lora_down.weight.data = torch.zeros_like(self.lora_down.weight.data)
         self.lora_down.weight.data = torch.randn_like(self.lora_down.weight.data) * std_dev
 
+        self._set_runtime_scale(scale)
+
         # m = Magnitude column-wise across output dimension
         weight = self.get_orig_weight()
         weight = weight.to(self.lora_up.weight.device, dtype=self.lora_up.weight.dtype)

toolkit/models/lokr.py

@@ -212,7 +212,7 @@ class LokrModule(ToolkitModuleMixin, nn.Module):
         if self.use_w2 and self.use_w1:
             # use scale = 1
             alpha = lora_dim
-        self.scale = float(alpha) / self.lora_dim
+        self._set_runtime_scale(float(alpha) / self.lora_dim)
         self.register_buffer('alpha', torch.tensor(alpha))  # treat as constant
 
         if self.use_w2:
@@ -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),
-            self.scale
+            self._runtime_scale
         )
         if orig_weight is not None:
             weight = weight.reshape(orig_weight.shape)
@@ -382,7 +382,7 @@ class LokrModule(ToolkitModuleMixin, nn.Module):
 
         # scale folded into A (not applied to the reduction output) to avoid
         # an inductor lowering bug under torch.compile
-        delta = torch.einsum('...qo,pq->...po', tmp, A * self.scale)  # (..., out_l, out_k)
+        delta = torch.einsum('...qo,pq->...po', tmp, A * self._runtime_scale)  # (..., out_l, out_k)
         delta = delta.flatten(-2, -1)
 
         if self.training and self.rank_dropout:

toolkit/network_mixins.py

@@ -158,7 +158,7 @@ class ExtractableModuleMixin:
 
         # set up alphas
         self.alpha = (self.alpha * 0) + down_weight.shape[0]
-        self.scale = float(self.alpha.detach().float().item()) / self.lora_dim
+        self._set_runtime_scale(float(self.alpha.detach().float().item()) / self.lora_dim)
 
         # assign them
 
@@ -179,6 +179,21 @@ class ToolkitModuleMixin:
         self.is_checkpointing = False
         self._multiplier: Union[float, list, torch.Tensor] = None
 
+    def _set_runtime_scale(self: Module, value) -> None:
+        """Keep float metadata while using a device tensor in compiled math."""
+        self.scale = float(value)
+        runtime_scale = getattr(self, "_runtime_scale", None)
+        if runtime_scale is None:
+            reference = next(self.parameters(), None)
+            if reference is None:
+                runtime_scale = torch.tensor(self.scale, dtype=torch.float32)
+            else:
+                runtime_scale = reference.new_tensor(self.scale, dtype=torch.float32)
+            self.register_buffer("_runtime_scale", runtime_scale, persistent=False)
+        else:
+            with torch.no_grad():
+                runtime_scale.fill_(self.scale)
+
     def _call_forward(self: Module, x):
         # module dropout
         if self.module_dropout is not None and self.training:
@@ -211,9 +226,9 @@ class ToolkitModuleMixin:
 
             # scaling for rank dropout: treat as if the rank is changed
             # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる
-            scale = self.scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
+            scale = self._runtime_scale * (1.0 / (1.0 - self.rank_dropout))  # redundant for readability
         else:
-            scale = self.scale
+            scale = self._runtime_scale
 
         lx = self.lora_up(lx)