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

WIP working on convrot offloading

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-11 15:37:28 -06:00
parent 1d1e21177a
commit 4625406093
5 changed files with 176 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
toolkit/lora_special.py
View File

@@ -69,8 +69,16 @@ class LoRAModule(ToolkitModuleMixin, ExtractableModuleMixin, torch.nn.Module):
torch.nn.Module.__init__(self)
self.lora_name = lora_name
self.orig_module_ref = weakref.ref(org_module)
self.scalar = torch.tensor(1.0, device=org_module.weight.device)
# read the device off a param/buffer directly: OstrisLinear.weight is a
# property that dequantizes the whole weight just to answer .device
org_tensor = next(
(t for t in org_module._parameters.values() if t is not None),
next((t for t in org_module._buffers.values() if t is not None), None),
)
self.scalar = torch.tensor(
1.0, device=org_tensor.device if org_tensor is not None else None
)
# if is ara lora module, mark it on the layer so memory manager can handle it
if is_ara:
org_module.ara_lora_ref = weakref.ref(self)
@@ -182,8 +190,9 @@ class FullModule(ToolkitModuleMixin, torch.nn.Module):
# trainable delta, zero initialized so an untrained layer is a no-op (zero diff)
# dequantize first so the delta is full precision and shaped like the real (unpacked) weight
self.weight_is_quantized = _is_quantized_tensor(org_module.weight)
ref_weight = _dequantize_if_needed(org_module.weight)
org_weight = org_module.weight # single access: dequantizes on OstrisLinear
self.weight_is_quantized = _is_quantized_tensor(org_weight)
ref_weight = _dequantize_if_needed(org_weight)
self.diff = torch.nn.Parameter(torch.zeros_like(ref_weight))
# some modules (e.g. Embedding) have no bias attribute at all
org_bias = getattr(org_module, 'bias', None)