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

Fix adapter scalar handling under torch.compile (#946)

Browse Source 
* Fix adapter scalar handling under torch.compile

* Fix instance where error could happen when merging in a lora to the base model

---------

Co-authored-by: Rydén Johan <johan.ryden@bostad.uppsala.se>
Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>
This commit is contained in:
Zironic
2026-07-15 19:07:55 +02:00
committed by GitHub
parent e44c34a955
commit 4ece17b71f
6 changed files with 14 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
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 = self.alpha / self.lora_dim
self.scale = float(self.alpha.detach().float().item()) / self.lora_dim
# assign them
@@ -387,7 +387,7 @@ class ToolkitModuleMixin:
weight_device = weight.device
if weight.device != down_weight.device:
weight = weight.to(down_weight.device)
if scale.device != down_weight.device:
if isinstance(scale, torch.Tensor) and scale.device != down_weight.device:
scale = scale.to(down_weight.device)
# merge weight
if self.full_rank: