mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-21 05:48:16 +08:00
fix: register_buffer for int4 tensors, kron-based LoKr path with ceiling-div padding, remove __doc__ overrides
This commit is contained in:
JWLHS
parent
9c576caa7f
commit
2990416508
@ -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
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user