From 4ece17b71fccfcac453a8e786907192017d210aa Mon Sep 17 00:00:00 2001 From: Zironic Date: Wed, 15 Jul 2026 19:07:55 +0200 Subject: [PATCH] Fix adapter scalar handling under torch.compile (#946) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 Co-authored-by: Jaret Burkett --- toolkit/kohya_lora.py | 4 ++-- toolkit/lora_special.py | 13 ++----------- toolkit/lycoris_special.py | 4 ++-- toolkit/models/DoRA.py | 5 ++--- toolkit/models/lokr.py | 8 ++++---- toolkit/network_mixins.py | 4 ++-- 6 files changed, 14 insertions(+), 24 deletions(-) diff --git a/toolkit/kohya_lora.py b/toolkit/kohya_lora.py index b085748a..fa7b0675 100644 --- a/toolkit/kohya_lora.py +++ b/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 diff --git a/toolkit/lora_special.py b/toolkit/lora_special.py index 2b733f9b..bf8309db 100644 --- a/toolkit/lora_special.py +++ b/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 diff --git a/toolkit/lycoris_special.py b/toolkit/lycoris_special.py index e3da1e1f..8bafb6d9 100644 --- a/toolkit/lycoris_special.py +++ b/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 diff --git a/toolkit/models/DoRA.py b/toolkit/models/DoRA.py index 3cfeb4a2..1d402778 100644 --- a/toolkit/models/DoRA.py +++ b/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 diff --git a/toolkit/models/lokr.py b/toolkit/models/lokr.py index 7865725e..8f11d6c3 100644 --- a/toolkit/models/lokr.py +++ b/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) diff --git a/toolkit/network_mixins.py b/toolkit/network_mixins.py index 941e09d6..2e3bfb3f 100644 --- a/toolkit/network_mixins.py +++ b/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: