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

Reworked merge in out of loras with convrot weights for better roundtrip accuracy.

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-02 08:29:17 -06:00
parent a6f6b6b896
commit 73cab2acf5
3 changed files with 54 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
toolkit/lora_special.py
View File

@@ -235,20 +235,33 @@ class FullModule(ToolkitModuleMixin, torch.nn.Module):
def merge_in(self: 'FullModule', merge_weight=1.0):
if not self.can_merge_in:
return
om = self.org_module[0]
if 'weight._data' in om.state_dict():
# quanto quantized weight, can't merge
# a zero diff merges to identity: skip entirely (a quantized base would
# otherwise still get requantized, which is not lossless)
if not self.diff.any() and (self.diff_b is None or not self.diff_b.any()):
return
org_weight = om.weight
orig_dtype = org_weight.dtype
# dequantize torchao weights so we can fold the full precision delta in
merged_weight = _dequantize_if_needed(org_weight).float() + merge_weight * self.diff.float().to(org_weight.device)
om = self.org_module[0]
if getattr(om, "is_ostris_quantized", False):
# fp32 dequant straight from the backend; the bf16 weight property
# would resample the quant scales on every merge cycle
orig_dtype = om.ostris_orig_dtype
base_weight = om.ostris_quantizer.dequantize(om)
weight_device = base_weight.device
else:
if 'weight._data' in om.state_dict():
# quanto quantized weight, can't merge
return
org_weight = om.weight
orig_dtype = org_weight.dtype
base_weight = _dequantize_if_needed(org_weight).float()
weight_device = org_weight.device
# fold the full precision delta in
merged_weight = base_weight + merge_weight * self.diff.float().to(weight_device)
if self.weight_is_quantized:
# re-quantize so the model stays quantized across continuous merge/reset cycles
from toolkit.util.quantize import get_torchao_config, requantize_module_weight
requantize_module_weight(om, merged_weight, orig_dtype, get_torchao_config(self._get_base_qtype()))
else:
om.weight.data = merged_weight.to(org_weight.device, orig_dtype)
om.weight.data = merged_weight.to(weight_device, orig_dtype)
# bias is never quantized
if self.diff_b is not None and getattr(om, 'bias', None) is not None:
om.bias.data = (om.bias.data.float() + merge_weight * self.diff_b.float().to(om.bias.device)).to(om.bias.dtype)