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

Restore device runtime scales for compiled adapters (#963)

Browse Source 
Co-authored-by: Rydén Johan <johan.ryden@bostad.uppsala.se>
This commit is contained in:
Zironic
2026-07-18 16:19:57 +02:00
committed by GitHub
parent cfdc9033a6
commit a0224793ce
7 changed files with 202 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/models/DoRA.py
View File

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

6
toolkit/models/lokr.py
View File

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