RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Restore device runtime scales for compiled adapters (#963) 

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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 202 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

154
testing/test_lora_compile_scalars.py Normal file
View File

@ -0,0 +1,154 @@
import unittest
from unittest import mock
import torch
from toolkit.kohya_lora import LoRAModule as KohyaLoRAModule
from toolkit.lora_special import LoRAModule
from toolkit.lycoris_special import LoConSpecialModule
from toolkit.models.DoRA import DoRAModule
from toolkit.models.lokr import LokrModule
class _Network:
network_type = "lora"
is_lorm = False
is_active = True
is_merged_in = False
_multiplier = 1.0
def _linear(device=None, dtype=None):
return torch.nn.Linear(8, 8, bias=False, device=device, dtype=dtype)
class AdapterScaleTest(unittest.TestCase):
def test_adapters_keep_float_metadata_and_nonpersistent_runtime_buffer(self):
network = _Network()
modules = [
LoRAModule(
"lora_scale",
_linear(),
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
network=network,
),
KohyaLoRAModule(
"kohya_scale",
_linear(),
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
),
LoConSpecialModule(
"locon_scale",
_linear(),
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
network=network,
),
DoRAModule(
"dora_scale",
_linear(),
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
network=network,
),
LokrModule(
"lokr_scale",
_linear(),
lora_dim=2,
alpha=torch.tensor(4, dtype=torch.bfloat16),
network=network,
),
]
for module in modules:
with self.subTest(module=type(module).__name__):
self.assertIs(type(module.scale), float)
self.assertEqual(module._runtime_scale.item(), module.scale)
self.assertNotIn("_runtime_scale", module.state_dict())
self.assertFalse(module._runtime_scale.requires_grad)
def test_extract_weight_synchronizes_runtime_scale(self):
module = LoRAModule(
"extract_scale",
_linear(),
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
network=_Network(),
)
runtime_scale = module._runtime_scale
down = torch.randn(2, 8)
up = torch.randn(8, 2)
with mock.patch(
"toolkit.network_mixins.extract_linear",
return_value=(down, up, 2, None),
):
module.extract_weight(extract_mode="fixed", extract_mode_param=2)
self.assertIs(module._runtime_scale, runtime_scale)
self.assertEqual(module.scale, 1.0)
self.assertEqual(module._runtime_scale.item(), 1.0)
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required")
def test_dynamic_compile_stays_cuda_and_scale_updates_do_not_recompile(self):
torch.manual_seed(0)
torch._dynamo.reset()
from torch._inductor import metrics
metrics.reset()
network = _Network()
network.torch_multiplier = torch.ones(1, device="cuda")
original = _linear(device="cuda", dtype=torch.bfloat16)
original.requires_grad_(False)
module = LoRAModule(
"compiled_scale",
original,
lora_dim=4,
alpha=torch.tensor(8, dtype=torch.bfloat16),
network=network,
).to("cuda")
module.org_forward = original.forward
with torch.no_grad():
module.lora_up.weight.normal_()
value = torch.randn(2, 8, device="cuda", dtype=torch.bfloat16)
eager = module(value)
eager.square().mean().backward()
eager_down_grad = module.lora_down.weight.grad.detach().clone()
eager_up_grad = module.lora_up.weight.grad.detach().clone()
module.zero_grad(set_to_none=True)
compiled = torch.compile(module, fullgraph=False, dynamic=True)
actual = compiled(value)
actual.square().mean().backward()
torch.cuda.synchronize()
torch.testing.assert_close(actual, eager, rtol=2e-2, atol=5e-2)
torch.testing.assert_close(
module.lora_down.weight.grad, eager_down_grad, rtol=2e-2, atol=5e-2
)
torch.testing.assert_close(
module.lora_up.weight.grad, eager_up_grad, rtol=2e-2, atol=5e-2
)
self.assertEqual(
getattr(metrics, "generated_cpp_vec_kernel_count", 0), 0
)
self.assertEqual(module._runtime_scale.device.type, "cuda")
base = original(value)
original_delta = actual - base
generated_kernels = metrics.generated_kernel_count
module._set_runtime_scale(0.5)
updated = compiled(value)
torch.cuda.synchronize()
torch.testing.assert_close(
updated - base, original_delta * 0.25, rtol=2e-2, atol=5e-2
)
self.assertEqual(metrics.generated_kernel_count, generated_kernels)
if __name__ == "__main__":
unittest.main()

29
toolkit/kohya_lora.py
View File

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

2
toolkit/lora_special.py
View File

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

2
toolkit/lycoris_special.py
View File

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

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:

21
toolkit/network_mixins.py
View File

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