Merge 879833ee37 into 1d1099bea0

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 原则——文档应该只在源码处维护一份

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

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

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)