From 3e5738da2463e23e83026b14ca20211668151a90 Mon Sep 17 00:00:00 2001 From: JWLHS <847135749@qq.com> Date: Fri, 10 Jul 2026 11:07:08 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20address=20CodeRabbit=20feedback=20?= =?UTF-8?q?=E2=80=94=20QuaRot=20error=20surface,=20LoRA=20shape=20validati?= =?UTF-8?q?on,=20Sylvester=20bug=20fix,=20cache=20cap?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- comfy/quantization/torchao/linear.py | 50 +++++++++++++++++++--------- comfy/quantization/torchao/quarot.py | 11 ++++-- 2 files changed, 43 insertions(+), 18 deletions(-) diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py index de6b3b3d2..32a1c7c61 100644 --- a/comfy/quantization/torchao/linear.py +++ b/comfy/quantization/torchao/linear.py @@ -74,18 +74,15 @@ class TINT4Linear(nn.Module): # ── QuaRot activation rotation (online) ────────────────── if self._use_quarot and self._hadamard_H is not None: - try: - H_dev = self._hadamard_H.to( - device=x.device, - dtype=(x.dtype if x.dtype in ( - torch.float16, torch.bfloat16) - else torch.float16), + if x_flat.shape[-1] % self._group_size != 0: + raise ValueError( + f"activation dim {x_flat.shape[-1]} not divisible " + f"by QuaRot group_size {self._group_size}" ) - n_groups = x_flat.shape[-1] // self._group_size - x_g = x_flat.view(-1, n_groups, self._group_size) - x_flat = torch.matmul(x_g, H_dev).view(-1, x_g.shape[2] * n_groups) - except Exception: - pass + H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype) + n_groups = x_flat.shape[-1] // self._group_size + x_g = x_flat.view(-1, n_groups, self._group_size) + x_flat = torch.matmul(x_g, H_dev).view(-1, n_groups * self._group_size) dev = x.device @@ -105,6 +102,7 @@ class TINT4Linear(nn.Module): else torch.float16) for lora_entries in entries.values(): for e in lora_entries: + # ── LoKr path ──────────────────────────────── if isinstance(e[0], str) and e[0] == "lokr": _, w1, w2, mult, factor = e[:5] sl = e[5] if len(e) > 5 else None @@ -118,13 +116,23 @@ class TINT4Linear(nn.Module): dw = (w1x * w2d).mul_(mult) lo = x_flat.to(cd) @ dw.T if sl is not None: - if lo.shape[1] == (se - sl): - out[:, sl:se] += lo + if lo.shape[1] != (se - sl): + raise ValueError( + f"LoKr output width {lo.shape[1]} " + f"doesn't match target slice width " + f"{se - sl}" + ) + out[:, sl:se] += lo elif lo.shape == out.shape: out += lo + else: + raise ValueError( + f"LoKr output shape {lo.shape} " + f"incompatible with {out.shape}" + ) continue - # standard LoRA + # ── Standard LoRA path ──────────────────────── A, B, mult = e[:3] sl = e[3] if len(e) > 3 else None se = e[4] if len(e) > 4 else None @@ -132,10 +140,20 @@ class TINT4Linear(nn.Module): Bd = B.to(device=dev, dtype=cd) lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult if sl is not None: - if lo.shape[1] == (se - sl): - out[:, sl:se] += lo + if lo.shape[1] != (se - sl): + raise ValueError( + f"LoRA output width {lo.shape[1]} " + f"doesn't match target slice width " + f"{se - sl}" + ) + out[:, sl:se] += lo elif lo.shape[1] == out.shape[1]: out += lo + else: + raise ValueError( + f"LoRA output shape {lo.shape} " + f"incompatible with {out.shape}" + ) if self.bias is not None: out += self.bias.to(device=dev, dtype=out.dtype) diff --git a/comfy/quantization/torchao/quarot.py b/comfy/quantization/torchao/quarot.py index 798dd89b3..3fa03d9fb 100644 --- a/comfy/quantization/torchao/quarot.py +++ b/comfy/quantization/torchao/quarot.py @@ -14,6 +14,7 @@ except ImportError: _SCIPY_AVAILABLE = False _HADAMARD_CACHE: dict = {} +_HADAMARD_MAX_SIZE = 64 # Maximum number of cached (size, device, dtype) entries def build_hadamard( @@ -38,14 +39,20 @@ def build_hadamard( H_np = _scipy_hadamard(size) H = torch.tensor(H_np, dtype=dtype, device=device) / (size ** 0.5) else: + # Pure-PyTorch Sylvester construction H = torch.tensor([[1.0]], dtype=dtype, device=device) n = 1 while n < size: - H = torch.cat([torch.cat([H, H], dim=1), - torch.cat([H, -H], dim=0)], dim=0) + top = torch.cat([H, H], dim=1) + bottom = torch.cat([H, -H], dim=1) + H = torch.cat([top, bottom], dim=0) n *= 2 H = H / (size ** 0.5) + # Enforce cache size cap (evict oldest entry) + if len(_HADAMARD_CACHE) >= _HADAMARD_MAX_SIZE: + _HADAMARD_CACHE.pop(next(iter(_HADAMARD_CACHE))) + _HADAMARD_CACHE[cache_key] = H return H