github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
* 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>
This commit is contained in:
Zironic
2026-07-15 19:07:55 +02:00
committed by GitHub
parent e44c34a955
commit 4ece17b71f
6 changed files with 14 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
toolkit/models/DoRA.py
View File

@@ -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

8
toolkit/models/lokr.py
View File

@@ -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)