wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-20 05:18:15 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

fix: register_buffer for int4 tensors, kron-based LoKr path with ceiling-div padding, remove __doc__ overrides

Browse Source
This commit is contained in:
JWLHS 2026-07-16 01:31:11 +08:00
parent 9c576caa7f
commit 2990416508
1 changed files with 202 additions and 202 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

404
comfy/quantization/torchao/linear.py
View File

@ -12,240 +12,240 @@ import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.functional as F import torch.nn.functional as F
from torchao.quantization.quantize_.workflows.int4.int4_plain_int32_tensor import ( from torchao.quantization.quantize_.workflows.int4.int4_plain_int32_tensor import (
Int4PlainInt32Tensor, Int4PlainInt32Tensor,
) )
class TINT4Linear(nn.Module): class TINT4Linear(nn.Module):
""" """
INT4 weight-only linear layer. INT4 weight-only linear layer.
Stores raw int32 qdata + fp16 scales + int8 zero-points. Stores raw int32 qdata + fp16 scales + int8 zero-points.
Int4PlainInt32Tensor is lazily constructed and rebuilt on device change. Int4PlainInt32Tensor is lazily constructed and rebuilt on device change.
Slots (set externally by loaders): Slots (set externally by loaders):
_use_quarot: bool — enable activation Hadamard rotation _use_quarot: bool — enable activation Hadamard rotation
_group_size: int — QuaRot group size _group_size: int — QuaRot group size
_hadamard_H: Tensor|None — normalized Hadamard matrix _hadamard_H: Tensor|None — normalized Hadamard matrix
_tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]} _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]}
_lora_prerotated: bool — True if caller already rotated A/w2 _lora_prerotated: bool — True if caller already rotated A/w2
into QuaRot basis (default False). into QuaRot basis (default False).
When False, rotation is done in forward(). When False, rotation is done in forward().
""" """
def __init__(self, in_features: int, out_features: int, def __init__(self, in_features: int, out_features: int,
qdata: torch.Tensor, scale: torch.Tensor, qdata: torch.Tensor, scale: torch.Tensor,
zp: torch.Tensor, block_size: tuple, zp: torch.Tensor, block_size: tuple,
bias: torch.Tensor | None = None): bias: torch.Tensor | None = None):
""" """
Initialize TINT4Linear from raw torchao INT4 tensors. Initialize TINT4Linear from raw torchao INT4 tensors.
Args: Args:
in_features: Input feature dimension. in_features: Input feature dimension.
out_features: Output feature dimension. out_features: Output feature dimension.
qdata: Packed int32 quantized weight data (out_f, in_f // 8). qdata: Packed int32 quantized weight data (out_f, in_f // 8).
scale: Per-block fp16 scales (n_blocks, out_f). scale: Per-block fp16 scales (n_blocks, out_f).
zp: Per-block int8 zero-points (n_blocks, out_f). zp: Per-block int8 zero-points (n_blocks, out_f).
block_size: (b0, b1) tuple defining quantization block shape. block_size: (b0, b1) tuple defining quantization block shape.
bias: Optional bias tensor (out_f,). bias: Optional bias tensor (out_f,).
""" """
super().__init__() super().__init__()
self.in_features = in_features self.in_features = in_features
self.out_features = out_features self.out_features = out_features
if bias is not None: if bias is not None:
self.bias = nn.Parameter(bias) self.bias = nn.Parameter(bias)
else: else:
self.register_parameter('bias', None) self.register_parameter('bias', None)
# ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值 # ★ FIX 1: 使用 register_buffer 而不是直接用 = 赋值
# 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理 # 好处：model.to(device) / state_dict() / load_state_dict() 能自动管理
# 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据 # 旧代码：self._qdata = qdata → PyTorch 不知道这是持久化数据
self.register_buffer('_qdata', qdata) self.register_buffer('_qdata', qdata)
self.register_buffer('_scale', scale) self.register_buffer('_scale', scale)
self.register_buffer('_zp', zp) self.register_buffer('_zp', zp)
self._block_size = block_size self._block_size = block_size
self._qt: Int4PlainInt32Tensor | None = None self._qt: Int4PlainInt32Tensor | None = None
# QuaRot (optional) # QuaRot (optional)
self._use_quarot: bool = False self._use_quarot: bool = False
self._group_size: int = 128 self._group_size: int = 128
self._hadamard_H: torch.Tensor | None = None self._hadamard_H: torch.Tensor | None = None
# LoRA entries slot # LoRA entries slot
self._tint4_lora_entries: dict | None = None self._tint4_lora_entries: dict | None = None
# Pre-rotation flag # Pre-rotation flag
self._lora_prerotated: bool = False self._lora_prerotated: bool = False
@property @property
def weight(self) -> Int4PlainInt32Tensor: def weight(self) -> Int4PlainInt32Tensor:
"""Lazily construct and return the Int4PlainInt32Tensor view.""" """Lazily construct and return the Int4PlainInt32Tensor view."""
if self._qt is None: if self._qt is None:
self._qt = Int4PlainInt32Tensor( self._qt = Int4PlainInt32Tensor(
self._qdata, self._scale, self._zp, self._qdata, self._scale, self._zp,
self._block_size, [self.out_features, self.in_features], self._block_size, [self.out_features, self.in_features],
) )
return self._qt return self._qt
@weight.setter @weight.setter
def weight(self, value: Int4PlainInt32Tensor) -> None: def weight(self, value: Int4PlainInt32Tensor) -> None:
"""Replace the underlying Int4PlainInt32Tensor.""" """Replace the underlying Int4PlainInt32Tensor."""
if isinstance(value, Int4PlainInt32Tensor): if isinstance(value, Int4PlainInt32Tensor):
self._qt = value self._qt = value
def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor: def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor:
""" """
Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis. Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis.
Skipped when _lora_prerotated is True (caller already rotated). Skipped when _lora_prerotated is True (caller already rotated).
Returns unchanged tensor if QuaRot is disabled or dimensions Returns unchanged tensor if QuaRot is disabled or dimensions
are not divisible by _group_size. are not divisible by _group_size.
""" """
if self._lora_prerotated: if self._lora_prerotated:
return tensor return tensor
if not self._use_quarot or self._hadamard_H is None: if not self._use_quarot or self._hadamard_H is None:
return tensor return tensor
if tensor.shape[1] % self._group_size != 0: if tensor.shape[1] % self._group_size != 0:
return tensor return tensor
Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype) Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype)
ng = tensor.shape[1] // self._group_size ng = tensor.shape[1] // self._group_size
return (tensor.reshape(tensor.shape[0], ng, self._group_size) return (tensor.reshape(tensor.shape[0], ng, self._group_size)
@ Ht).reshape(tensor.shape[0], -1) @ Ht).reshape(tensor.shape[0], -1)
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
""" """
Forward pass with optional QuaRot activation rotation and LoRA injection. Forward pass with optional QuaRot activation rotation and LoRA injection.
Args: Args:
x: Input tensor, shape (..., in_features). x: Input tensor, shape (..., in_features).
Returns: Returns:
Output tensor, shape (..., out_features). Output tensor, shape (..., out_features).
Raises: Raises:
ValueError: If activation dim is not divisible by QuaRot group_size, ValueError: If activation dim is not divisible by QuaRot group_size,
or if LoRA output shapes are incompatible. or if LoRA output shapes are incompatible.
""" """
x_flat = x.reshape(-1, x.shape[-1]) x_flat = x.reshape(-1, x.shape[-1])
# QuaRot activation rotation (online) # QuaRot activation rotation (online)
if self._use_quarot and self._hadamard_H is not None: if self._use_quarot and self._hadamard_H is not None:
if x_flat.shape[-1] % self._group_size != 0: if x_flat.shape[-1] % self._group_size != 0:
raise ValueError( raise ValueError(
f"activation dim {x_flat.shape[-1]} not divisible " f"activation dim {x_flat.shape[-1]} not divisible "
f"by QuaRot group_size {self._group_size}" f"by QuaRot group_size {self._group_size}"
) )
H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype) H_dev = self._hadamard_H.to(device=x.device, dtype=x.dtype)
n_groups = x_flat.shape[-1] // self._group_size n_groups = x_flat.shape[-1] // self._group_size
x_g = x_flat.view(-1, n_groups, 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) x_flat = torch.matmul(x_g, H_dev).view(-1, n_groups * self._group_size)
dev = x.device dev = x.device
# Rebuild Int4PlainInt32Tensor on device change # Rebuild Int4PlainInt32Tensor on device change
if self._qt is None or self._qt.device != dev: if self._qt is None or self._qt.device != dev:
self._qt = Int4PlainInt32Tensor( self._qt = Int4PlainInt32Tensor(
self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev), self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev),
self._block_size, [self.out_features, self.in_features], self._block_size, [self.out_features, self.in_features],
) )
out = F.linear(x_flat, self._qt, None) out = F.linear(x_flat, self._qt, None)
# LoRA forward injection # LoRA forward injection
entries = self._tint4_lora_entries entries = self._tint4_lora_entries
if entries is not None and entries: if entries is not None and entries:
cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16) cd = (x.dtype if x.dtype in (torch.float16, torch.bfloat16)
else torch.float16) else torch.float16)
for lora_entries in entries.values(): for lora_entries in entries.values():
for e in lora_entries: for e in lora_entries:
# ── LoKr path ── # ── LoKr path ──
# ★ FIX 2: 用 torch.kron 替代 repeat_interleave # ★ FIX 2: 用 torch.kron 替代 repeat_interleave
# 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错） # 旧方案要求 w2 维度能被 factor 整除（例如 64÷60=不整除→报错）
# 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解 # 新方案 kron 展开 + pad/trim 到目标维度，兼容任意分解
if isinstance(e[0], str) and e[0] == "lokr": if isinstance(e[0], str) and e[0] == "lokr":
_, w1, w2, mult, factor = e[:5] _, w1, w2, mult, factor = e[:5]
sl = e[5] if len(e) > 5 else None sl = e[5] if len(e) > 5 else None
se = e[6] if len(e) > 6 else None se = e[6] if len(e) > 6 else None
# QuaRot: 先旋转 w2（小矩阵，转发负担小） # QuaRot: 先旋转 w2（小矩阵，转发负担小）
# 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的） # 数学上等价于旋转完整 kron 展开（Hadamard 是分块对角的）
w2 = self._rotate_into_quarot(w2) w2 = self._rotate_into_quarot(w2)
w1d = w1.to(device=dev, dtype=cd) w1d = w1.to(device=dev, dtype=cd)
w2d = w2.to(device=dev, dtype=cd) w2d = w2.to(device=dev, dtype=cd)
# kron 展开为完整 delta（代替 repeat_interleave） # kron 展开为完整 delta（代替 repeat_interleave）
dw = torch.kron(w1d, w2d) dw = torch.kron(w1d, w2d)
# pad/trim 到层实际维度 # Pad/trim to layer (or slice) dimensions
# QKV 融合层：kron 覆盖 1 个 head，模块持有 3 个 tgt_out = (se - sl) if sl is not None else self.out_features
if dw.shape[0] < self.out_features: if dw.shape[0] < tgt_out:
dw = dw.repeat( dw = dw.repeat(
self.out_features // dw.shape[0], 1) (tgt_out + dw.shape[0] - 1) // dw.shape[0], 1)
elif dw.shape[0] > self.out_features: if dw.shape[0] > tgt_out:
dw = dw[:self.out_features, :] dw = dw[:tgt_out, :]
if dw.shape[1] < self.in_features: if dw.shape[1] < self.in_features:
dw = dw.repeat( dw = dw.repeat(
1, self.in_features // dw.shape[1]) 1, (self.in_features + dw.shape[1] - 1) // dw.shape[1])
elif dw.shape[1] > self.in_features: if dw.shape[1] > self.in_features:
dw = dw[:, :self.in_features] dw = dw[:, :self.in_features]
dw = dw.contiguous().mul_(mult) dw = dw.contiguous().mul_(mult)
lo = x_flat.to(cd) @ dw.T lo = x_flat.to(cd) @ dw.T
if sl is not None: if sl is not None:
if lo.shape[1] != (se - sl): if lo.shape[1] != (se - sl):
raise ValueError( raise ValueError(
f"LoKr output width {lo.shape[1]} " f"LoKr output width {lo.shape[1]} "
f"doesn't match target slice width " f"doesn't match target slice width "
f"{se - sl}" f"{se - sl}"
) )
out[:, sl:se] += lo out[:, sl:se] += lo
elif lo.shape == out.shape: elif lo.shape == out.shape:
out += lo out += lo
else: else:
raise ValueError( raise ValueError(
f"LoKr output shape {lo.shape} " f"LoKr output shape {lo.shape} "
f"incompatible with {out.shape}" f"incompatible with {out.shape}"
) )
continue continue
# Standard LoRA path # Standard LoRA path
A, B, mult = e[:3] A, B, mult = e[:3]
sl = e[3] if len(e) > 3 else None sl = e[3] if len(e) > 3 else None
se = e[4] if len(e) > 4 else None se = e[4] if len(e) > 4 else None
A = self._rotate_into_quarot(A) A = self._rotate_into_quarot(A)
Ad = A.to(device=dev, dtype=cd) Ad = A.to(device=dev, dtype=cd)
Bd = B.to(device=dev, dtype=cd) Bd = B.to(device=dev, dtype=cd)
lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult lo = (x_flat.to(cd) @ Ad.T) @ Bd.T * mult
if sl is not None: if sl is not None:
if lo.shape[1] != (se - sl): if lo.shape[1] != (se - sl):
raise ValueError( raise ValueError(
f"LoRA output width {lo.shape[1]} " f"LoRA output width {lo.shape[1]} "
f"doesn't match target slice width " f"doesn't match target slice width "
f"{se - sl}" f"{se - sl}"
) )
out[:, sl:se] += lo out[:, sl:se] += lo
elif lo.shape[1] == out.shape[1]: elif lo.shape[1] == out.shape[1]:
out += lo out += lo
else: else:
raise ValueError( raise ValueError(
f"LoRA output shape {lo.shape} " f"LoRA output shape {lo.shape} "
f"incompatible with {out.shape}" f"incompatible with {out.shape}"
) )
if self.bias is not None: if self.bias is not None:
out += self.bias.to(device=dev, dtype=out.dtype) out += self.bias.to(device=dev, dtype=out.dtype)
return out.reshape(*x.shape[:-1], out.shape[-1]) return out.reshape(*x.shape[:-1], out.shape[-1])
def release(self) -> None: def release(self) -> None:
"""Drop cached Int4PlainInt32Tensor to free VRAM.""" """Drop cached Int4PlainInt32Tensor to free VRAM."""
self._qt = None self._qt = None
def __del__(self): def __del__(self):
"""Cleanup raw tensor references on deletion.""" """Cleanup raw tensor references on deletion."""
self._qdata = None self._qdata = None
self._scale = None self._scale = None
self._zp = None self._zp = None
self._qt = None self._qt = None
self._tint4_lora_entries = None self._tint4_lora_entries = None