""" INT8 Quantizer — Per-Channel Symmetric Quantization for Legacy GPUs Provides symmetric INT8 quantization optimized for Turing/Pascal GPUs where FP8 is emulated and slow. Uses per-channel scales to minimize precision loss. All functions are stateless pure functions operating on tensors and keys. This module mirrors the API shape of engine/fp8_quantizer.py. """ from typing import Dict, Optional, Tuple import torch try: from ..config import INT8_CLIP_PERCENTILE except ImportError: from config import INT8_CLIP_PERCENTILE # Optional ConvRot (Hadamard rotation) support for INT8 quantization quality. # Requires ComfyUI-Flux2-INT8 custom node to be installed. # ConvRot applies an orthogonal Hadamard transform before quantization, # spreading per-channel outliers uniformly to reduce quantization error. _HAS_CONVROT = False _HADAMARD_CACHE = {} # (size, device_str, dtype_str) -> H matrix try: import importlib # Use importlib.import_module() because the module name contains # hyphens ('ComfyUI-Flux2-INT8') which are not valid in Python's # from X import Y syntax (SyntaxError on '-' identifier). _convrot_mod = importlib.import_module("custom_nodes.ComfyUI-Flux2-INT8.convrot") build_hadamard = _convrot_mod.build_hadamard rotate_weight = _convrot_mod.rotate_weight _HAS_CONVROT = True except (ImportError, AttributeError): pass # ── INT8 skip patterns (tensors too small to benefit from INT8) ────── # LoRA weights are tiny (typically 1-50K elements each); the per-channel # scale overhead (~2 bytes per output channel) is comparable to the data # savings, making INT8 quantization counterproductive. INT8_SKIP_PATTERNS: tuple = ( 'lora.', # All LoRA weight patterns start with a module containing 'lora' ) # 🔥 INT8 IMPROVEMENT: Patterns that should be preserved in FP16 instead of quantized to INT8 # # INT8 has zero mantissa bits — step size is 1.0 in integer space. After scaling, # effective precision is `scale` per channel. Input/output projections, modulation # layers, and norm weights are the most sensitive to quantization error because they # handle signal entry/exit from the transformer blocks. # # This mirrors FP8_PRESERVE_BF16_PATTERNS from engine/fp8_quantizer.py. The key # difference: FP8 uses BF16 (3 mantissa bits preserved), while INT8 uses FP16 # (10 mantissa bits). For INT8 dequant math (int8 * scale), FP16's higher mantissa # precision produces more accurate results than BF16 — so FP16 is universally correct # regardless of GPU generation (Pascal/Turing/Ampere+). # # Patterns support two key formats: # - Diffusers-style: uses '/' path separators (model/diffusion_model/...) # - FLUX-style: uses '.' dot notation (double_blocks.0.img_attn...) INT8_PRESERVE_FP16_PATTERNS: tuple = ( # ── Diffusers-style path separators ── '/final_layer/', '/time_in/', '/txt_in/', '/vector_in/', '/guidance_in/', '/norm.', # ── FLUX dot-notation keys ── 'final_layer.', 'time_in.', 'txt_in.', 'vector_in.', 'guidance_in.', 'img_in.', 'single_stream_modulation.', # FLUX modulation tensors use underscore suffixes (_img, _txt), not dots 'double_stream_modulation_img.', 'double_stream_modulation_txt.', # ── SDXL/SD1.5 dot-notation keys ── # Time embedding projection — SDXL/SD1.5 equivalent of Flux's time_in. # Controls how the denoising timestep signal enters the model. # Errors here compound through all denoising steps. 'time_embed.', # Label embedding (SDXL only) — conditioning signal entry point. # SDXL uses pooled text embeddings from both CLIP models. 'label_emb.', # First input block — the very first convolution processing the # noisy latent. Signal entry point — errors propagate everywhere. 'input_blocks.0.', # UNet middle block (bottleneck) — all features pass through. # Small number of tensors but critical for global coherence. 'middle_block.', # Output projection — final convolution before VAE decode. # Signal exit point — errors directly affect pixel output. 'out.', # ── SDXL diffusers-style path equivalents ── '/time_embed/', '/label_emb/', '/input_blocks/0/', '/middle_block/', '/out/', # ── Anima architecture patterns ── # Input patch embedding — the very first layer processing the noisy latent. # Equivalent of Flux's img_in. and SDXL's input_blocks.0. # Signal entry — errors here propagate through all 35 DiT blocks. # Matches: net.x_embedder.proj.weight 'net.x_embedder.', # Time embedding projection — timestep conditioning signal entry. # Equivalent of Flux's time_in. and SDXL's time_embed. # Errors here compound through all 35 denoising steps. # Matches: net.t_embedder.0.weight, net.t_embedder.2.weight 'net.t_embedder.', # First DiT block — signal entry into the transformer core. # All embedded signals (latent, timestep, text conditioning) flow # through blocks.0 first. Errors here corrupt every subsequent block. # Matches: net.blocks.0.adaln_modulation_cross_attn.1.weight 'net.blocks.0.', # Runtime format variant — ComfyUI wraps UNet under diffusion_model. # When a baked/converted checkpoint is re-loaded, keys use the # diffusion_model. prefix instead of net. — but the same first-block # signal entry rationale applies. # Matches: diffusion_model.blocks.0.adaln_modulation_cross_attn.1.weight 'diffusion_model.blocks.0.', # Output projection — final linear layer before VAE decode. # Equivalent of Flux's final_layer. and SDXL's out. # Signal exit — errors directly affect pixel output. # Matches: net.final_layer.linear.weight 'net.final_layer.', # ── LLM Adapter — Anima's T5 cross-attention conditioning interface ── # First llm_adapter block — signal entry for T5 cross-attention conditioning. # Similar rationale to net.blocks.0.: the first block processes raw T5 # hidden states before they propagate to downstream blocks. Errors here # corrupt all cross-attention conditioning across every DiT block. # Matches: net.llm_adapter.blocks.0.cross_attn.k_proj.weight 'net.llm_adapter.blocks.0.', # LLM Adapter input embedding — projects T5 hidden states into the # adapter's feature space. Signal entry point. # Matches: net.llm_adapter.embed.weight 'net.llm_adapter.embed.', # LLM Adapter output norm — final normalization before the output # projection. Should be preserved alongside other norm parameters. # Matches: net.llm_adapter.norm.weight 'net.llm_adapter.norm.', # LLM Adapter output projection — projects adapter features back into # the main DiT block's feature space. Signal exit point. # Matches: net.llm_adapter.out_proj.weight, net.llm_adapter.out_proj.bias 'net.llm_adapter.out_proj.', # ── SDXL CLIP Text Encoder — embeddings & first encoder layer ── # Token/position embedding tables encode semantic vocabulary. # Matches: conditioner.embedders.0.transformer.text_model.embeddings.* # Note: We do NOT add 'conditioner.embedders.' here — that broad pattern # would match BOTH CLIP models (embedders.0 ~123M, embedders.1 ~694M), # forcing the entire 694M-param OpenCLIP ViT-bigG to FP16 (+670 MB). # The two specific patterns below already protect the sensitive layers # (token embeddings + first encoder layer) in BOTH CLIP models. 'text_model.embeddings.', # First text encoder layer — signal entry for text processing. 'text_model.encoder.layers.0.', # ── Output block embedding layers (time embedding projections) ── # Similar criticality to time_embed. but on the output side. # Matches: output_blocks.*.0.emb_layers.* 'emb_layers.', ) # ── Public API ──────────────────────────────────────────────────────── def should_skip_int8(key: str) -> bool: """Return True if this key should skip INT8 quantization. Args: key: The checkpoint key to check. Returns: True if the key matches a skip pattern (e.g. LoRA weights). """ key_lower = key.lower() for pattern in INT8_SKIP_PATTERNS: if pattern in key_lower: return True return False def should_preserve_fp16_for_int8(key: str) -> bool: """Return True if this key should stay in FP16 when target precision is INT8. INT8 has zero mantissa bits — input projections (img_in, time_in, txt_in), output layers (final_layer), and modulation tensors are the most sensitive to quantization error. This mirrors FP8's ``should_preserve_bf16()`` but uses FP16 preservation instead of BF16. For INT8 dequant math (``int8 * scale``), FP16's 10 mantissa bits produce more accurate results than BF16's 7 mantissa bits — so FP16 is universally better regardless of GPU generation. Args: key: The checkpoint key to check. Returns: True if this key should be preserved in FP16 instead of INT8. """ key_lower = key.lower() for pattern in INT8_PRESERVE_FP16_PATTERNS: if pattern in key_lower: return True return False def compute_int8_scale( tensor: torch.Tensor, percentile: Optional[float] = None, ) -> torch.Tensor: """Per-channel INT8 scale factor per output channel. Computes ``threshold / 127.0`` where ``threshold`` is the per-channel absolute value at the given ``percentile`` (default: module-level :data:`INT8_CLIP_PERCENTILE`). Clipping outliers via percentile (e.g. 99.9th) prevents a single extreme outlier from compressing the effective INT8 range for the bulk of weights. Symmetric quantization maps the range [-threshold, +threshold] to [-127, +127] per output channel. Scale is clamped to a minimum of 1e-12 to prevent division by zero for all-zero tensors. Handles 1D tensors (``(C,)``) by using a per-tensor (scalar) scale — there is no "output channel" dimension to quantize per-channel. This is a defensive guard; callers should route 1D weights through float16, but if they don't, the function still produces correct output instead of crashing. Args: tensor: Float tensor, shape ``(out_channels, ...)`` or ``(C,)``. percentile: If set to < 1.0, clips outliers at this percentile before computing scale. E.g. 0.999 = 99.9th percentile. Default (None) uses :data:`INT8_CLIP_PERCENTILE`. Set to 1.0 to use max (disable clipping). Returns: Scale tensor, dtype ``float32``. Shape is ``(out_channels,)`` for 2D+ input, or scalar (0-dim) for 1D input. """ # Use module-level default if not specified if percentile is None: percentile = INT8_CLIP_PERCENTILE # Handle 1D tensors: use per-tensor (scalar) scale — flatten(1) would crash if tensor.dim() <= 1: flat = tensor.float() # shape (C,) is_1d = True else: # Flatten all dimensions except the first (output channel) flat = tensor.float().flatten(1) # shape (C, K*...) is_1d = False # Guard: sanitize NaN/Inf -> 0 to prevent corrupted scale factors if not torch.isfinite(flat).all(): flat = torch.nan_to_num(flat, nan=0.0, posinf=0.0, neginf=0.0) if is_1d: # 1D: per-tensor (scalar) scale — no channel dimension to reduce abs_vals = flat.abs() if percentile < 1.0: k = max(1, int(abs_vals.numel() * percentile)) threshold = abs_vals.kthvalue(k).values flat = flat.clamp(-threshold, threshold) # Recompute abs after clipping (outliers now saturated at threshold) abs_vals = flat.abs() abs_max = abs_vals.max() else: # 2D+: per-channel scale — max across all flattened non-channel dims abs_vals = flat.abs() # shape (C, K*...) if percentile < 1.0: # Per-channel percentile: independent threshold per output channel k = max(1, int(abs_vals.size(-1) * percentile)) threshold = abs_vals.kthvalue(k, dim=-1).values # shape (C,) # Clamp each channel independently flat = flat.clamp( -threshold.unsqueeze(-1), threshold.unsqueeze(-1), ) # Recompute abs after clipping (outliers now saturated at threshold) abs_vals = flat.abs() abs_max = abs_vals.max(dim=1).values scale = abs_max / 127.0 scale = torch.clamp(scale, min=1e-12) return scale.to(dtype=torch.float32) def quantize_to_int8( tensor: torch.Tensor, percentile: Optional[float] = None, use_convrot: bool = False, convrot_group_size: int = 256, ) -> Tuple[torch.Tensor, torch.Tensor]: """Symmetric per-channel INT8 quantization. Steps: 1. (Optional) Apply Hadamard rotation if ``use_convrot=True``, spreading outlier channels to reduce quantization error. 2. Compute per-channel scale = threshold / 127.0 (with optional percentile clipping for outlier reduction) 3. Divide tensor by scale (broadcasted across non-channel dims) 4. Round and clamp to [-127, +127] 5. Cast to ``torch.int8`` Args: tensor: Float tensor, shape ``(out_channels, ...)``. percentile: Passed through to :func:`compute_int8_scale`. If < 1.0, clips outliers at this percentile before computing scale. Default (None) uses :data:`INT8_CLIP_PERCENTILE`. use_convrot: If True, apply Hadamard rotation before quantization. Requires ComfyUI-Flux2-INT8 custom node to be installed. convrot_group_size: Group size for Hadamard matrix (must be power of 4). Only used when ``use_convrot=True``. Returns: Tuple of ``(quantized_int8_tensor, scale_f32_1d)``. Raises: ImportError: If ``use_convrot=True`` but ComfyUI-Flux2-INT8 is not installed. """ if use_convrot: if not _HAS_CONVROT: raise ImportError( "ConvRot requires ComfyUI-Flux2-INT8 custom node. " "Install from: https://github.com/BobJohnson24/ComfyUI-INT8-Fast" ) # Skip ConvRot for layers where in_features is not divisible by # group_size (e.g. small embedding projections like [64, 64]). # rotate_weight raises ValueError otherwise. in_features = tensor.shape[-1] if in_features % convrot_group_size != 0: print(f" ⚠️ ConvRot: skipping layer with in_features={in_features} " f"(not divisible by group_size={convrot_group_size})") else: # Cache Hadamard matrix per (size, device, dtype) to avoid recomputation cache_key = (convrot_group_size, str(tensor.device), str(tensor.dtype)) if cache_key not in _HADAMARD_CACHE: H = build_hadamard(convrot_group_size, device=tensor.device, dtype=tensor.dtype) _HADAMARD_CACHE[cache_key] = H else: H = _HADAMARD_CACHE[cache_key] # Rotate weight: W_rot = W @ H^T (orthogonal, preserves magnitude) tensor = rotate_weight(tensor, H, group_size=convrot_group_size) scale = compute_int8_scale(tensor, percentile=percentile) # Reshape scale for broadcasting: [C] -> [C, 1, 1, ...] scale_view = scale.view(-1, *([1] * (tensor.dim() - 1))) q = (tensor / scale_view).round().clamp(-127, 127).to(torch.int8) return q, scale def quantize_weight_to_int8_with_scales( tensor: torch.Tensor, percentile: Optional[float] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Convenience wrapper: quantize a weight tensor to INT8. Identical to :func:`quantize_to_int8` but named to mirror the FP8 API convention from :mod:`engine.fp8_quantizer.quantize_weight_to_fp8_with_scales`. Args: tensor: Original float tensor (before INT8 quantization). percentile: Passed through to :func:`quantize_to_int8`. If < 1.0, clips outliers at this percentile before computing scale. Default (None) uses :data:`INT8_CLIP_PERCENTILE`. Returns: Tuple of ``(quantized_int8_tensor, weight_scale_f32_1d)``. """ return quantize_to_int8(tensor, percentile=percentile) def dequant_int8_tensor( tensor: torch.Tensor, key: str, target_dtype: torch.dtype, scale_store: Dict[str, torch.Tensor], ) -> torch.Tensor: """Dequantize an INT8 tensor to target dtype using companion scale factors. INT8 tensors store per-channel quantization scales as companion tensors (``.weight_scale`` suffix, matching the FP8 convention). This function: 1. Casts INT8 -> target dtype (e.g. fp16) — gives scaled values 2. Looks up the per-channel scale factor from ``scale_store`` 3. Broadcasts and multiplies: out = int8_as_target * scale Args: tensor: INT8 tensor (``torch.int8``). key: Checkpoint key for the tensor (used to derive scale key names). target_dtype: Target dtype for dequantization (e.g. ``torch.float16``). scale_store: Dict-like containing companion scale tensors. Supports plain dict and ``_LazyCheckpointMapping`` (both have ``__contains__``). Returns: Dequantized tensor at ``target_dtype``. Returns tensor unchanged (cast to target_dtype) if no companion scale is found. """ if tensor.dtype != torch.int8: return tensor.to(dtype=target_dtype) if tensor.dtype != target_dtype else tensor result = tensor.to(dtype=target_dtype) # Derive the base prefix for scale key lookup base = key[:-len('.weight')] if key.endswith('.weight') else key scale_key = base + '.weight_scale' if scale_key in scale_store: scale = scale_store[scale_key].to(dtype=target_dtype) # Broadcast: weight [C, K1, K2, ...] * scale [C, 1, 1, ...] for _ in range(result.dim() - scale.dim()): scale = scale.unsqueeze(-1) result = result * scale return result