diff --git a/comfy/quantization/torchao/__init__.py b/comfy/quantization/torchao/__init__.py new file mode 100644 index 000000000..90da9ad0f --- /dev/null +++ b/comfy/quantization/torchao/__init__.py @@ -0,0 +1,31 @@ +""" +ComfyUI TorchAO INT4 weight-only quantization backend. + +W4A16 scheme: weights quantized to INT4 via torchao, +activations remain FP16. Supports CUDA, Intel XPU, and CPU. + +Includes: +- TINT4Linear: INT4 linear layer with LoRA forward injection + QuaRot +- quantize_model: torchao INT4 quantization entry point +- reconstruct_int4_state_dict: rebuild TINT4Linear from safetensors +- build_hadamard / rotate_weight: QuaRot Hadamard rotation +""" + +from .linear import TINT4Linear +from .quantize import quantize_model, reconstruct_int4_state_dict +from .quarot import build_hadamard, rotate_weight + +# Public API — documented for docstring coverage +__all__ = [ + "TINT4Linear", + "quantize_model", + "reconstruct_int4_state_dict", + "build_hadamard", + "rotate_weight", +] + +# ★ FIX 3: 删除 __doc__ 覆写(5 行) +# 原因:这些类/函数的 .py 文件里已有 docstring,重复赋值会: +# 1. 覆盖原始文档(如果两处不一致会让人困惑) +# 2. 干扰 Sphinx 等文档工具的正确引用 +# 3. 违反 DRY 原则——文档应该只在源码处维护一份 diff --git a/comfy/quantization/torchao/linear.py b/comfy/quantization/torchao/linear.py new file mode 100644 index 000000000..22a2e52e6 --- /dev/null +++ b/comfy/quantization/torchao/linear.py @@ -0,0 +1,234 @@ +""" +TINT4Linear — INT4 weight-only Linear with LoRA + QuaRot. + +Backed by torchao Int4PlainInt32Tensor. Activations stay FP16. +LoRA entries are injected at forward time — no kernel modification needed. +QuaRot activation rotation is applied online; LoRA A/w2 matrices are +rotated into the QuaRot basis inside forward() automatically unless +_lora_prerotated is set (for callers that pre-rotate at injection time). +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchao.quantization.quantize_.workflows.int4.int4_plain_int32_tensor import ( + Int4PlainInt32Tensor, +) + + +class TINT4Linear(nn.Module): + """ + INT4 weight-only linear layer. + + Stores raw int32 qdata + fp16 scales + int8 zero-points. + Int4PlainInt32Tensor is lazily constructed and rebuilt on device change. + + Slots (set externally by loaders): + _use_quarot: bool — enable activation Hadamard rotation + _group_size: int — QuaRot group size + _hadamard_H: Tensor|None — normalized Hadamard matrix + _tint4_lora_entries: dict|None — {lora_name: [(type, ...), ...]} + _lora_prerotated: bool — True if caller already rotated A/w2 + into QuaRot basis (default False). + When False, rotation is done in forward(). + """ + + def __init__(self, in_features: int, out_features: int, + qdata: torch.Tensor, scale: torch.Tensor, + zp: torch.Tensor, block_size: tuple, + bias: torch.Tensor | None = None): + """ + Initialize TINT4Linear from raw torchao INT4 tensors. + + Args: + in_features: Input feature dimension. + out_features: Output feature dimension. + qdata: Packed int32 quantized weight data (out_f, in_f // 8). + scale: Per-block fp16 scales (n_blocks, out_f). + zp: Per-block int8 zero-points (n_blocks, out_f). + block_size: (b0, b1) tuple defining quantization block shape. + bias: Optional bias tensor (out_f,). + """ + super().__init__() + self.in_features = in_features + self.out_features = out_features + + if bias is not None: + self.bias = nn.Parameter(bias) + else: + self.register_parameter('bias', None) + + self.register_buffer('_qdata', qdata) + self.register_buffer('_scale', scale) + self.register_buffer('_zp', zp) + self._block_size = block_size + self._qt: Int4PlainInt32Tensor | None = None + + # QuaRot (optional) + self._use_quarot: bool = False + self._group_size: int = 128 + self._hadamard_H: torch.Tensor | None = None + + # LoRA entries slot + self._tint4_lora_entries: dict | None = None + + # Pre-rotation flag + self._lora_prerotated: bool = False + + @property + def weight(self) -> Int4PlainInt32Tensor: + """Lazily construct and return the Int4PlainInt32Tensor view.""" + if self._qt is None: + self._qt = Int4PlainInt32Tensor( + self._qdata, self._scale, self._zp, + self._block_size, [self.out_features, self.in_features], + ) + return self._qt + + @weight.setter + def weight(self, value: Int4PlainInt32Tensor) -> None: + """Replace the underlying Int4PlainInt32Tensor.""" + if isinstance(value, Int4PlainInt32Tensor): + self._qt = value + + def _rotate_into_quarot(self, tensor: torch.Tensor) -> torch.Tensor: + """ + Rotate a LoRA A or w2 matrix into the QuaRot Hadamard basis. + + Skipped when _lora_prerotated is True (caller already rotated). + Returns unchanged tensor if QuaRot is disabled or dimensions + are not divisible by _group_size. + """ + if self._lora_prerotated: + return tensor + if not self._use_quarot or self._hadamard_H is None: + return tensor + if tensor.shape[1] % self._group_size != 0: + return tensor + Ht = self._hadamard_H.T.to(device=tensor.device, dtype=tensor.dtype) + ng = tensor.shape[1] // self._group_size + return (tensor.reshape(tensor.shape[0], ng, self._group_size) + @ Ht).reshape(tensor.shape[0], -1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass with optional QuaRot activation rotation and LoRA injection. + + Args: + x: Input tensor, shape (..., in_features). + + Returns: + Output tensor, shape (..., out_features). + + Raises: + ValueError: If activation dim is not divisible by QuaRot group_size, + or if LoRA output shapes are incompatible. + """ + x_flat = x.reshape(-1, x.shape[-1]) + + # QuaRot activation rotation (online) + if self._use_quarot and self._hadamard_H is not None: + 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}" + ) + 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 + + # Rebuild Int4PlainInt32Tensor on device change + if self._qt is None or self._qt.device != dev: + self._qt = Int4PlainInt32Tensor( + self._qdata.to(dev), self._scale.to(dev), self._zp.to(dev), + self._block_size, [self.out_features, self.in_features], + ) + + out = F.linear(x_flat, self._qt, None) + + # LoRA forward injection + entries = self._tint4_lora_entries + if entries is not None and entries: + cd = x.dtype + 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 + se = e[6] if len(e) > 6 else None + + w2 = self._rotate_into_quarot(w2) + w1d = w1.to(device=dev, dtype=cd) + w2d = w2.to(device=dev, dtype=cd) + + # kron 展开为完整 delta(代替 repeat_interleave) + dw = torch.kron(w1d, w2d) + + # Pad/trim to layer (or slice) dimensions + tgt_out = (se - sl) if sl is not None else self.out_features + if dw.shape[0] < tgt_out: + dw = dw.repeat( + (tgt_out + dw.shape[0] - 1) // dw.shape[0], 1) + if dw.shape[0] > tgt_out: + dw = dw[:tgt_out, :] + if dw.shape[1] < self.in_features: + dw = dw.repeat( + 1, (self.in_features + dw.shape[1] - 1) // dw.shape[1]) + if dw.shape[1] > self.in_features: + dw = dw[:, :self.in_features] + + dw = dw.contiguous().mul_(mult) + lo = x_flat.to(cd) @ dw.T + if sl is not None: + 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 path + A, B, mult = e[:3] + sl = e[3] if len(e) > 3 else None + se = e[4] if len(e) > 4 else None + A = self._rotate_into_quarot(A) + Ad = A.to(device=dev, dtype=cd) + 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): + 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) + + return out.reshape(*x.shape[:-1], out.shape[-1]) + + def release(self) -> None: + """Drop cached Int4PlainInt32Tensor to free VRAM.""" + self._qt = None diff --git a/comfy/quantization/torchao/quantize.py b/comfy/quantization/torchao/quantize.py new file mode 100644 index 000000000..868a945c1 --- /dev/null +++ b/comfy/quantization/torchao/quantize.py @@ -0,0 +1,115 @@ +""" +TorchAO INT4 quantization helpers for ComfyUI. + +- quantize_model: in-place INT4 quantization via torchao +- reconstruct_int4_state_dict: rebuild TINT4Linear layers from safetensors +""" + +import torch +import torch.nn as nn +from torchao.quantization import Int4WeightOnlyConfig, quantize_ + +from .linear import TINT4Linear + + +def quantize_model( + model: nn.Module, + group_size: int = 128, + filter_fn=None, +) -> nn.Module: + """ + Quantize all nn.Linear layers in-place via torchao INT4 weight-only. + + Args: + model: PyTorch model to quantize. + group_size: Quantization group size (32/64/128/256). + filter_fn: Optional callable(module, name) → bool. + Return True to quantize that module. + + Returns: + Same model (quantized in-place). + """ + config = Int4WeightOnlyConfig( + group_size=group_size, + int4_packing_format="plain_int32", + ) + quantize_(model, config, filter_fn=filter_fn) + return model + + +_QUANT_SUFFIXES = ( + ".weight_scale", ".weight_zp", + ".weight_b0", ".weight_b1", + ".weight_sh0", ".weight_sh1", + ".comfy_quant", +) + + +def _is_quantized_weight(key: str, sd: dict) -> bool: + """Check whether a safetensors key is a torchao-quantized weight tensor.""" + if not key.endswith(".weight"): + return False + base = key[:-len(".weight")] + return f"{base}.weight_scale" in sd and f"{base}.weight_zp" in sd + + +def reconstruct_int4_state_dict( + sd: dict, + device: torch.device = torch.device("cpu"), +) -> dict[str, TINT4Linear]: + """ + Extract TINT4Linear layers from a torchao-quantized safetensors dict. + + Pops quantized weight tensors from sd and replaces them with + fully-constructed TINT4Linear modules keyed by base layer name. + + Caller is responsible for placing the modules into the target model. + + Args: + sd: Safetensors state dict with torchao int4 packed weights. + device: Target device for the reconstructed layers. + + Returns: + Dict mapping base key (e.g. "blocks.0.attn.qkv") → TINT4Linear. + """ + replaced: dict[str, TINT4Linear] = {} + + for key in list(sd.keys()): + if not _is_quantized_weight(key, sd): + continue + + base = key[:-len(".weight")] + qdata = sd.pop(key) + scale = sd.pop(f"{base}.weight_scale") + zp = sd.pop(f"{base}.weight_zp") + b0 = sd.pop(f"{base}.weight_b0") + b1 = sd.pop(f"{base}.weight_b1") + + for suffix in _QUANT_SUFFIXES: + sd.pop(f"{base}{suffix}", None) + + bias = sd.pop(f"{base}.bias", None) + + in_f = qdata.shape[1] * 8 + out_f = qdata.shape[0] + + layer = TINT4Linear( + in_features=in_f, + out_features=out_f, + qdata=qdata.to(torch.int32), + scale=scale.to(torch.float16), + zp=zp.to(torch.int8), + block_size=(b0.item(), b1.item()), + bias=bias, + ) + + if device.type != "cpu": + layer._qdata = layer._qdata.to(device) + layer._scale = layer._scale.to(device) + layer._zp = layer._zp.to(device) + if bias is not None: + layer.bias.data = layer.bias.data.to(device) + + replaced[base] = layer + + return replaced diff --git a/comfy/quantization/torchao/quarot.py b/comfy/quantization/torchao/quarot.py new file mode 100644 index 000000000..3fa03d9fb --- /dev/null +++ b/comfy/quantization/torchao/quarot.py @@ -0,0 +1,87 @@ +""" +QuaRot: Group-wise Hadamard rotation for INT quantization quality. + +Spreads activation outliers across channels via orthogonal Hadamard matrices. +Based on QuaRot (2024). Pure PyTorch — scipy optional for large matrices. +""" + +import torch + +try: + from scipy.linalg import hadamard as _scipy_hadamard + _SCIPY_AVAILABLE = True +except ImportError: + _SCIPY_AVAILABLE = False + +_HADAMARD_CACHE: dict = {} +_HADAMARD_MAX_SIZE = 64 # Maximum number of cached (size, device, dtype) entries + + +def build_hadamard( + size: int, + device: str = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """ + Build normalized orthogonal Hadamard matrix (size must be power of 2). + Returns H such that H @ H^T = I. + Uses scipy if available, otherwise pure-PyTorch Sylvester construction. + """ + if size & (size - 1) != 0: + raise ValueError(f"Hadamard size must be a power of 2, got {size}") + + cache_key = (size, str(device), dtype) + if cache_key in _HADAMARD_CACHE: + return _HADAMARD_CACHE[cache_key] + + if _SCIPY_AVAILABLE: + import numpy as np + 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: + 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 + + +def rotate_weight( + weight: torch.Tensor, + H: torch.Tensor, + group_size: int, +) -> torch.Tensor: + """ + Rotate weight matrix offline: W_rot = W @ H_block^T. + + For Linear(in, out) with weight (out_f, in_f): + each row is split into groups of group_size and rotated by H^T. + + Args: + weight: (out_features, in_features) + H: Normalized Hadamard matrix (group_size, group_size) + group_size: Group size for block-diagonal rotation + + Returns: + Rotated weight, same shape. + """ + out_f, in_f = weight.shape + if in_f % group_size != 0: + raise ValueError( + f"in_features {in_f} not divisible by group_size {group_size}") + n_groups = in_f // group_size + W_grouped = weight.view(out_f, n_groups, group_size) + H_t = H.T.to(dtype=weight.dtype, device=weight.device) + W_rot = torch.matmul(W_grouped, H_t) + return W_rot.reshape(out_f, in_f)