nodes/bbc-s-float-increment-per-latent
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fuse direct LoKr adapters in aggregate hooks

Browse Source
This commit is contained in:
bbc-s
2026-05-14 18:38:24 +02:00
parent 794100e05f
commit 1fc9b386e1
1 changed files with 110 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
float_increment_per_latent.py
+110 -3
View File
@@ -28,6 +28,7 @@ class AggregateBypassForwardHook:
for adapter in adapters
]
self.fused_adapters = []
self.fused_lokr_adapters = []
self.fallback_adapters = adapters[:]
self.fused_down = None
self.fused_up = None
@@ -36,11 +37,14 @@ class AggregateBypassForwardHook:
self.fused_is_conv = False
self.fused_conv_dim = 0
self.fused_kw_dict = {}
self.fused_lokr_groups = []
def _forward(self, x, *args, **kwargs):
out = self.original_forward(x, *args, **kwargs)
if self.fused_adapters:
out = out + self._fused_lora_delta(x, out)
for group in self.fused_lokr_groups:
out = out + self._fused_lokr_delta(x, out, group)
for adapter in self.fallback_adapters:
out = adapter.g(out + adapter.h(x, out))
return out
@@ -83,6 +87,35 @@ class AggregateBypassForwardHook:
hidden = hidden * torch.cat(rank_scales, dim=1).view(*scale_shape)
return F.linear(hidden, self.fused_up.to(dtype=x.dtype))
def _fused_lokr_delta(self, x, out, group):
adapters = group["adapters"]
scales = torch.stack([
self._adapter_multiplier(adapter, x, out).to(dtype=x.dtype)
for adapter in adapters
], dim=1)
if group["is_conv"]:
conv_fn = (F.conv1d, F.conv2d, F.conv3d)[group["conv_dim"] - 1]
b, _c, *spatial = x.shape
uq = group["uq"]
h_in_group = x.reshape(b * uq, -1, *spatial)
hb = conv_fn(h_in_group, group["w2"].to(dtype=x.dtype), **group["kw_dict"])
out_k = group["out_k"]
spatial_out = hb.shape[2:]
hb = hb.view(b, uq, len(adapters), out_k, *spatial_out)
hc = torch.einsum("bqnk...,nlq->bnlk...", hb, group["w1"].to(dtype=x.dtype))
hc = hc * scales.view(b, len(adapters), 1, 1, *([1] * len(spatial_out)))
return hc.sum(dim=1).reshape(b, -1, *spatial_out)
uq = group["uq"]
h_in_group = x.reshape(*x.shape[:-1], uq, -1)
hb = F.linear(h_in_group, group["w2"].to(dtype=x.dtype))
out_k = group["out_k"]
hb = hb.view(*h_in_group.shape[:-1], len(adapters), out_k)
hc = torch.einsum("...qnk,nlq->...nlk", hb, group["w1"].to(dtype=x.dtype))
hc = hc * scales.view(*([x.shape[0]] + [1] * (hc.ndim - 4) + [len(adapters), 1, 1]))
return hc.sum(dim=-3).reshape(*x.shape[:-1], -1)
def _prepare_fused_lora(self, dtype):
self.fused_is_conv = getattr(self.fused_adapters[0], "is_conv", False)
self.fused_conv_dim = getattr(self.fused_adapters[0], "conv_dim", 0)
@@ -122,6 +155,49 @@ class AggregateBypassForwardHook:
self.fused_ranks = ranks
self.fused_alpha_scales = alpha_scales
def _prepare_fused_lokr(self, dtype):
groups = {}
for adapter in self.fused_lokr_adapters:
if not self._is_lokr_direct_like(adapter):
continue
w1, w2 = adapter.weights[0], adapter.weights[1]
key = (
tuple(w1.shape),
tuple(w2.shape),
getattr(adapter, "is_conv", False),
getattr(adapter, "conv_dim", 0),
tuple(sorted(getattr(adapter, "kw_dict", {}).items())),
)
groups.setdefault(key, []).append(adapter)
self.fused_lokr_groups = []
for adapters in groups.values():
first = adapters[0]
is_conv = getattr(first, "is_conv", False)
conv_dim = getattr(first, "conv_dim", 0)
w1_list = []
w2_list = []
for adapter in adapters:
w1, w2 = adapter.weights[0], adapter.weights[1]
if dtype is not None:
w1 = w1.to(dtype=dtype)
w2 = w2.to(dtype=dtype)
if is_conv and w2.dim() == 2:
w2 = w2.view(*w2.shape, *([1] * conv_dim))
w1_list.append(w1)
w2_list.append(w2)
self.fused_lokr_groups.append({
"adapters": adapters,
"w1": torch.stack(w1_list, dim=0).contiguous(),
"w2": torch.cat(w2_list, dim=0).contiguous(),
"uq": int(w1_list[0].shape[1]),
"out_k": int(w2_list[0].shape[0]),
"is_conv": is_conv,
"conv_dim": conv_dim,
"kw_dict": getattr(first, "kw_dict", {}) if is_conv else {},
})
def _can_fuse_lora(self):
fused, _fallback, _reasons = self._partition_adapters_for_fusion()
return len(fused) > 0
@@ -139,6 +215,20 @@ class AggregateBypassForwardHook:
and hasattr(down, "to")
)
@staticmethod
def _is_lokr_direct_like(adapter):
weights = getattr(adapter, "weights", None)
if type(adapter).__name__ != "LoKrAdapter":
return False
if not isinstance(weights, (list, tuple)) or len(weights) != 9:
return False
w1, w2 = weights[0], weights[1]
if w1 is None or w2 is None:
return False
if any(weights[index] is not None for index in (3, 4, 5, 6, 7)):
return False
return hasattr(w1, "shape") and hasattr(w1, "to") and hasattr(w2, "shape") and hasattr(w2, "to")
@staticmethod
def _adapter_debug_name(adapter):
weights = getattr(adapter, "weights", None)
@@ -171,12 +261,20 @@ class AggregateBypassForwardHook:
def _partition_adapters_for_fusion(self):
ref_adapter = next((adapter for adapter in self.adapters if self._is_lora_like(adapter)), None)
if ref_adapter is None:
return [], self.adapters[:], {"non_lora_adapter": len(self.adapters)}
fused = [adapter for adapter in self.adapters if self._is_lokr_direct_like(adapter)]
fallback = [adapter for adapter in self.adapters if not self._is_lokr_direct_like(adapter)]
reasons = {}
if fallback:
reasons["non_lora_adapter"] = len(fallback)
return fused, fallback, reasons
fused = []
fallback = []
reasons = {}
for adapter in self.adapters:
if self._is_lokr_direct_like(adapter):
fused.append(adapter)
continue
reason = self._adapter_fuse_blocker_reason(adapter, ref_adapter)
if reason is None:
fused.append(adapter)
@@ -199,9 +297,12 @@ class AggregateBypassForwardHook:
for setup_hook in self.setup_hooks:
setup_hook._move_adapter_weights_to_device(device, dtype)
self.fused_adapters, self.fallback_adapters, _reasons = self._partition_adapters_for_fusion()
fused_adapters, self.fallback_adapters, _reasons = self._partition_adapters_for_fusion()
self.fused_adapters = [adapter for adapter in fused_adapters if self._is_lora_like(adapter)]
self.fused_lokr_adapters = [adapter for adapter in fused_adapters if self._is_lokr_direct_like(adapter)]
if self.fused_adapters:
self._prepare_fused_lora(dtype)
self._prepare_fused_lokr(dtype)
self.original_forward = self.module.forward
self.module.forward = self._forward
@@ -215,7 +316,9 @@ class AggregateBypassForwardHook:
self.fused_ranks = None
self.fused_alpha_scales = None
self.fused_adapters = []
self.fused_lokr_adapters = []
self.fallback_adapters = self.adapters[:]
self.fused_lokr_groups = []
class FloatIncrementPerLatent:
@@ -548,7 +651,11 @@ class PerSampleLoraLoader:
model_sd_keys = set(new_model.model.state_dict().keys())
for key, adapter in bypass_patches.items():
if key in model_sd_keys:
if fallback_adapter_policy == "static_base_only" and not AggregateBypassForwardHook._is_lora_like(adapter):
if (
fallback_adapter_policy == "static_base_only"
and not AggregateBypassForwardHook._is_lora_like(adapter)
and not AggregateBypassForwardHook._is_lokr_direct_like(adapter)
):
skipped_static_adapters += 1
continue
self._attach_per_sample_multiplier(