155 lines
5.1 KiB
Python
155 lines
5.1 KiB
Python
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()
|