Restore device runtime scales for compiled adapters (#963)

Co-authored-by: Rydén Johan <johan.ryden@bostad.uppsala.se>

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