diff --git a/testing/test_lora_compile_scalars.py b/testing/test_lora_compile_scalars.py new file mode 100644 index 00000000..86797330 --- /dev/null +++ b/testing/test_lora_compile_scalars.py @@ -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() diff --git a/toolkit/kohya_lora.py b/toolkit/kohya_lora.py index fa7b0675..a48f7cd4 100644 --- a/toolkit/kohya_lora.py +++ b/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( diff --git a/toolkit/lora_special.py b/toolkit/lora_special.py index bf8309db..86994dbc 100644 --- a/toolkit/lora_special.py +++ b/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 diff --git a/toolkit/lycoris_special.py b/toolkit/lycoris_special.py index 8bafb6d9..1bce9dfe 100644 --- a/toolkit/lycoris_special.py +++ b/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 diff --git a/toolkit/models/DoRA.py b/toolkit/models/DoRA.py index 1d402778..b08377f8 100644 --- a/toolkit/models/DoRA.py +++ b/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) diff --git a/toolkit/models/lokr.py b/toolkit/models/lokr.py index 468127a9..34d507f6 100644 --- a/toolkit/models/lokr.py +++ b/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: diff --git a/toolkit/network_mixins.py b/toolkit/network_mixins.py index 2e3bfb3f..cacd69d2 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 = 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)