import torch from torch import Tensor, nn import torch.nn.functional as F import logging import comfy.model_patcher import comfy.memory_management import comfy.model_management import comfy.lora import comfy.utils try: import comfy_aimdo.host_buffer import comfy_aimdo.torch _AIMDO_FILE_SLICE_LOAD = True except Exception: _AIMDO_FILE_SLICE_LOAD = False # Add this at the top of your file try: from .int8_fused_kernel import triton_int8_linear from .int8_fused_kernel import triton_int8_linear_per_row from .int8_fused_kernel import triton_quantize_rowwise _TRITON_AVAILABLE = True except ImportError: _TRITON_AVAILABLE = False print("Triton not found, falling back to torch._int_mm") # Runtime toggle — set by Int8TensorwiseOps.use_triton via the loader node _use_triton = True # ConvRot Configuration CONVROT_GROUP_SIZE = 256 # Must be a power of 4 for Regular Hadamard (e.g. 16, 64, 256) # --- Quantization Utils --- def quantize_int8(x: Tensor, scale: float | Tensor) -> Tensor: return x.float().mul(1.0 / scale).round_().clamp_(-128.0, 127.0).to(torch.int8) def quantize_int8_tensorwise(x: Tensor) -> tuple[Tensor, Tensor]: abs_max = x.abs().max() scale = (abs_max.float() / 127.0).clamp(min=1e-30) return quantize_int8(x, scale), scale def quantize_int8_axiswise(x: Tensor, dim: int) -> tuple[Tensor, Tensor]: abs_max = x.abs().amax(dim=dim, keepdim=True) scale = (abs_max.float() / 127.0).clamp(min=1e-30) return quantize_int8(x, scale), scale def dequantize(q: Tensor, scale: float | Tensor) -> Tensor: return q.float() * scale def tensor_to_device_file_slice(tensor: Tensor, device: torch.device) -> Tensor: if ( not _AIMDO_FILE_SLICE_LOAD or tensor.device.type != "cpu" or device is None or device.type != "cuda" ): return tensor.to(device, non_blocking=True) size = tensor.numel() * tensor.element_size() if size == 0: return tensor.to(device, non_blocking=True) hostbuf = comfy_aimdo.host_buffer.HostBuffer(size) host_tensor = comfy_aimdo.torch.hostbuf_to_tensor(hostbuf) host_view = host_tensor[:size].view(dtype=tensor.dtype).view(tensor.shape) if comfy.memory_management.read_tensor_file_slice_into(tensor, host_view): out = torch.empty_like(tensor, device=device) out.copy_(host_view, non_blocking=False) return out return tensor.to(device, non_blocking=True) def stochastic_round_int8_delta(x: Tensor, scale: float | Tensor, seed: int = 0) -> Tensor: """ Quantize a delta tensor to INT8 using stochastic rounding. Used for LoRA deltas to minimize quantization error. """ generator = torch.Generator(device=x.device) generator.manual_seed(seed) # Scale to INT8 range — move scale to x's device to handle CPU-stored scales if isinstance(scale, torch.Tensor): scale = scale.to(x.device) x_scaled = x / scale # Stochastic rounding x_floor = torch.floor(x_scaled) fraction = x_scaled - x_floor del x_scaled # High-precision input no longer needed # Speed optimization: Create random values directly on the target device random_vals = torch.rand(x_floor.shape, generator=generator, device=x.device, dtype=x_floor.dtype) x_rounded = torch.where(random_vals < fraction, x_floor + 1, x_floor) del random_vals del fraction del x_floor return torch.clamp(x_rounded, -128, 127).to(torch.int8) # --- LinearW8A8 Core --- @torch.no_grad() def int8_forward_dynamic(x: Tensor, weight: Tensor, weight_scale: float | Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: """Forward with dynamic per-token activation quantization.""" # --- FAST PATH: Triton Fused Kernel --- if _TRITON_AVAILABLE and _use_triton and x.is_cuda: return triton_int8_linear(x, weight, weight_scale, bias, compute_dtype) # --- SLOW PATH: Standard PyTorch --- # Quantize activations per row (dynamic) x_8, x_scale = quantize_int8_axiswise(x, dim=-1) # INT8 Matmul (Outputs Int32) res = torch._int_mm(x_8, weight.T) # Dequantize: (res * weight_scale * x_scale) # Note: Creating intermediate Float tensors here is VRAM heavy res_scaled = res.float().mul_(weight_scale * x_scale).to(compute_dtype) if bias is not None: res_scaled = res_scaled + bias.to(compute_dtype) return res_scaled @torch.no_grad() def int8_forward_dynamic_per_row(x: Tensor, weight: Tensor, weight_scale: Tensor, bias: Tensor | None, compute_dtype: torch.dtype) -> Tensor: """Forward with dynamic per-token activation quantization and per-row weight quantization. Args: x: Input activations [batch, in_features] weight: INT8 weight matrix [out_features, in_features] weight_scale: Per-row weight scales [out_features, 1] bias: Optional bias compute_dtype: Output dtype """ # --- FAST PATH: Triton Fused Kernel (per-row) --- if _TRITON_AVAILABLE and _use_triton and x.is_cuda: return triton_int8_linear_per_row(x, weight, weight_scale, bias, compute_dtype) # --- SLOW PATH: Standard PyTorch --- x_8, x_scale = quantize_int8_axiswise(x, dim=-1) # INT8 Matmul (Outputs Int32) res = torch._int_mm(x_8, weight.T) # [batch, out_features] # Dequantize with per-row weight scales # res[i,j] = sum_k(x_8[i,k] * weight[j,k]) * x_scale[i] * weight_scale[j] # Broadcasting: res * x_scale * weight_scale.T res_scaled = res.float().mul_(x_scale).mul_(weight_scale.T).to(compute_dtype) if bias is not None: res_scaled = res_scaled + bias.to(compute_dtype) return res_scaled # ============================================================================= # Int8TensorwiseOps - ComfyUI Custom Operations # ============================================================================= try: from comfy.ops import manual_cast, cast_bias_weight, uncast_bias_weight _COMFY_OPS_AVAILABLE = True except ImportError: _COMFY_OPS_AVAILABLE = False if _COMFY_OPS_AVAILABLE: class Int8TensorwiseOps(manual_cast): """ Custom ComfyUI operations for INT8 tensorwise quantization. """ excluded_names = [] dynamic_quantize = False # Manual toggle for on-the-fly quantization enable_convrot = False # Toggle for ConvRot Hadamard rotation use_triton = True # Toggle for Triton fused kernel (mirrors _use_triton) compute_dtype = None # Optional override for INT8 activation/output compute dtype _is_prequantized = False # Keep this as a status flag, but don't use for detection lora_mode = "None" # None/Stochastic bake into INT8 weights; Dynamic applies LoRA at inference dynamic_lora = False # If True, apply LoRA dynamically at inference; if False, bake into INT8 weights at load time lora_patches = {} # Map of model_key -> patch list (from load_lora) lora_strength = 1.0 dynamic_load_device = None # Set by the loader when Aimdo should avoid a full CPU staging copy skeleton_meta_init = False # Temporary mode for LoRA key-map discovery _auto_compute_dtype_by_device = {} @staticmethod def _default_compute_dtype(x: Tensor) -> torch.dtype: if x.dtype in (torch.float16, torch.bfloat16): return x.dtype if x.dtype == torch.float32 and x.is_cuda: device_index = x.device.index if device_index is None: device_index = torch.cuda.current_device() cached = Int8TensorwiseOps._auto_compute_dtype_by_device.get(device_index) if cached is not None: return cached compute_dtype = torch.float32 try: capability = torch.cuda.get_device_capability(device_index) name = torch.cuda.get_device_name(device_index).lower() if capability == (7, 5) and ("rtx" in name or "t4" in name): compute_dtype = torch.float16 except Exception: pass Int8TensorwiseOps._auto_compute_dtype_by_device[device_index] = compute_dtype return compute_dtype if x.dtype == torch.float32: return torch.float32 return torch.float16 class Linear(manual_cast.Linear): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): if getattr(Int8TensorwiseOps, "skeleton_meta_init", False): nn.Module.__init__(self) self.in_features = in_features self.out_features = out_features tensor_kwargs = {"device": "meta"} if dtype is not None: tensor_kwargs["dtype"] = dtype self.weight = nn.Parameter(torch.empty((out_features, in_features), **tensor_kwargs), requires_grad=False) self.bias = nn.Parameter(torch.empty((out_features,), **tensor_kwargs), requires_grad=False) if bias else None self.weight_comfy_model_dtype = dtype self.bias_comfy_model_dtype = dtype # Preserve ComfyUI's Windows/Aimdo lazy-init path. The base # disable_weight_init.Linear only takes this path for classes # that do not override _load_from_state_dict; this INT8 class # does override it, so calling super() would allocate full # skeleton weights during Pre-LoRA key-map discovery. elif comfy.model_management.WINDOWS and comfy.memory_management.aimdo_enabled: nn.Module.__init__(self) self.in_features = in_features self.out_features = out_features self.weight = None self.bias = None self.comfy_need_lazy_init_bias = bias self.weight_comfy_model_dtype = dtype self.bias_comfy_model_dtype = dtype else: super().__init__(in_features, out_features, bias, device, dtype) self.register_buffer('weight_scale', None) self._is_quantized = False self._is_per_row = False # Track quantization granularity self._use_convrot = False # Track if ConvRot was applied self._weight_scale_scalar = None # For scalar (non-tensor) scales self.compute_dtype = None self.comfy_cast_weights = False self.lora_patches = [] # List of (down_scaled, up, start, size) set by INT8ModelPatcher def reset_parameters(self): return None @staticmethod def _normalize_lora_key(key): if not isinstance(key, str): return key for p in ["diffusion_model.", "model.diffusion_model.", "model.", "transformer."]: if key.startswith(p): return key[len(p):] return key @staticmethod def _is_bias_key(key): return isinstance(key, str) and key.endswith(".bias") @staticmethod def _format_lora_patches(patches): formatted = [] for patch in patches or []: if len(patch) == 4: v, offset, function, strength = patch else: v, offset, function = patch strength = getattr(Int8TensorwiseOps, "lora_strength", 1.0) formatted.append((strength, v, 1.0, offset, function)) return formatted def _apply_int8_lora_patches(self, tensor, key, patches, device): if not patches or tensor.dtype == torch.int8: return tensor temp_dtype = comfy.model_management.lora_compute_dtype(device) tensor_temp = tensor_to_device_file_slice(tensor, device).to(dtype=temp_dtype) return comfy.lora.calculate_weight(self._format_lora_patches(patches), tensor_temp, key) def finalize_pending_int8(self): pending = getattr(self, "_pending_int8_finalize", None) if pending is None: return False weight_key = pending["weight_key"] device = pending.get("device") if device is None: device = torch.device("cuda") if torch.cuda.is_available() else self.weight.device weight_tensor = self.weight.detach() weight_tensor = self._apply_int8_lora_patches(weight_tensor, weight_key, pending.get("lora_patches"), device) if pending["quantize"]: if not hasattr(Int8TensorwiseOps, '_logged_otf'): print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {pending.get('enable_convrot', False)})") Int8TensorwiseOps._logged_otf = True w_gpu = tensor_to_device_file_slice(weight_tensor, device).float() self._use_convrot = False if pending.get("enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0: try: from .convrot import build_hadamard, rotate_weight H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype) w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE) self._use_convrot = True # Stamp the active groupsize so INT8ModelSave can # round-trip it deterministically. self._convrot_groupsize = CONVROT_GROUP_SIZE except ImportError as e: logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}") q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1) self.weight = nn.Parameter(q_weight.cpu(), requires_grad=False) self.register_buffer('weight_scale', q_scale.cpu()) self._weight_scale_scalar = None self._is_quantized = True self._is_per_row = True del w_gpu, q_weight, q_scale else: self.weight = nn.Parameter(weight_tensor.cpu(), requires_grad=False) self.weight_comfy_model_dtype = self.weight.dtype if self.weight_scale is not None: self.weight_scale_comfy_model_dtype = self.weight_scale.dtype if self.bias is not None: self.bias_comfy_model_dtype = self.bias.dtype delattr(self, "_pending_int8_finalize") return True def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): weight_key = prefix + "weight" # Utility to normalize keys by stripping common prefixes def normalize_key(key): return self._normalize_lora_key(key) def apply_lora_patches(tensor, key): if self._is_bias_key(key) or not Int8TensorwiseOps.lora_patches or tensor.dtype == torch.int8: return tensor nk = normalize_key(key) patches = Int8TensorwiseOps.lora_patches.get(nk) if patches: # Track applied patches if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'): Int8TensorwiseOps.applied_lora_patches = set() Int8TensorwiseOps.applied_lora_patches.add(nk) # Print only if multiple sub-patches map to the same layer # if "weight" in key and len(patches) > 1: # print(f"INT8 Fast: Baking multiple LoRA parts into {nk} ({len(patches)} sub-patches)") # ComfyUI dynamically patches during inference using lora_compute_dtype() # On most modern GPUs, this evaluates to torch.float16. # We simulate that exact intermediate cast here to achieve a 1:1 binary match. device = getattr(Int8TensorwiseOps, "dynamic_load_device", None) if device is None: device = tensor.device result_temp = self._apply_int8_lora_patches(tensor, key, patches, device) return result_temp.to(tensor.dtype) return tensor def source_tensor(tensor): if tensor is not None and getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None: return tensor.cpu() return tensor scale_key = prefix + "weight_scale" input_scale_key = prefix + "input_scale" bias_key = prefix + "bias" def pop_metadata(sd, p, k): v = sd.pop(p + k, None) if v is not None: return v v = sd.pop("model." + p + k, None) if v is not None: return v if p.startswith("model."): v = sd.pop(p[6:] + k, None) if v is not None: return v if p.startswith("diffusion_model."): v = sd.pop("diffusion_model." + p + k, None) if v is not None: return v return None weight_scale = pop_metadata(state_dict, prefix, "weight_scale") comfy_quant_tensor = pop_metadata(state_dict, prefix, "comfy_quant") weight_tensor = state_dict.pop(weight_key, None) bias_tensor = state_dict.pop(bias_key, None) # Pop input_scale to clean state_dict, but ignore it _ = state_dict.pop(input_scale_key, None) quant_conf_parsed = None if comfy_quant_tensor is not None: try: import json quant_conf_parsed = json.loads(bytes(comfy_quant_tensor.tolist()).decode('utf-8')) if quant_conf_parsed.get("convrot", False): self._use_convrot = True Int8TensorwiseOps.enable_convrot = True # Propagate globally for LoRA if "convrot_groupsize" in quant_conf_parsed: self._convrot_groupsize = int(quant_conf_parsed["convrot_groupsize"]) Int8TensorwiseOps._global_convrot_groupsize = self._convrot_groupsize except Exception: pass pending_weight_lora_patches = None if weight_tensor is not None and weight_tensor.dtype != torch.int8: pending_weight_lora_patches = Int8TensorwiseOps.lora_patches.get(normalize_key(weight_key)) defer_weight_lora = ( getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None and pending_weight_lora_patches ) # Apply LoRA patches to weight and bias once. With Aimdo, large # weight patches are deferred until KSampler/model load time so # the loader node stays cheap and VBAR geometry is finalized once. if weight_tensor is not None and not defer_weight_lora: weight_tensor = apply_lora_patches(weight_tensor, weight_key) if bias_tensor is not None: bias_tensor = apply_lora_patches(bias_tensor, bias_key) if weight_tensor is not None: if weight_tensor.dtype == torch.int8 and weight_scale is not None: # Load Quantized self._is_quantized = True self.weight = nn.Parameter(weight_tensor, requires_grad=False) Int8TensorwiseOps._is_prequantized = True # Found a quantized layer # Optional explicit hint from saved comfy_quant per_row_hint = None if isinstance(quant_conf_parsed, dict) and "per_row" in quant_conf_parsed: per_row_hint = bool(quant_conf_parsed["per_row"]) if isinstance(weight_scale, torch.Tensor): if weight_scale.numel() == 1: # Scalar scale — store as float for speed self._weight_scale_scalar = weight_scale.float().item() self.register_buffer('weight_scale', weight_scale.float().reshape(1)) self._weight_scale_scalar = None elif weight_scale.dim() == 2 and weight_scale.shape[1] == 1: self.register_buffer('weight_scale', weight_scale.float()) self._weight_scale_scalar = None self._is_per_row = True if per_row_hint is None else per_row_hint else: self.register_buffer('weight_scale', weight_scale.float()) self._weight_scale_scalar = None self._is_per_row = False if per_row_hint is None else per_row_hint else: self.weight_scale = nn.Parameter( torch.tensor(float(weight_scale), dtype=torch.float32), requires_grad=False ) self.weight_scale = None self._is_per_row = False if per_row_hint is None else per_row_hint elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32): # Load High-Precision is_excluded = any(ex in prefix for ex in Int8TensorwiseOps.excluded_names) is_dim1 = self.in_features == 1 or self.out_features == 1 or weight_tensor.ndim == 1 should_quantize = not (is_excluded or is_dim1 or not Int8TensorwiseOps.dynamic_quantize) defer_finalize = ( getattr(Int8TensorwiseOps, "dynamic_load_device", None) is not None and (should_quantize or pending_weight_lora_patches) ) if defer_finalize: self._is_quantized = False self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False) self._pending_int8_finalize = { "weight_key": weight_key, "quantize": should_quantize, "lora_patches": pending_weight_lora_patches, "device": getattr(Int8TensorwiseOps, "dynamic_load_device", None), "enable_convrot": getattr(Int8TensorwiseOps, "enable_convrot", False), } if pending_weight_lora_patches: if not hasattr(Int8TensorwiseOps, 'applied_lora_patches'): Int8TensorwiseOps.applied_lora_patches = set() Int8TensorwiseOps.applied_lora_patches.add(normalize_key(weight_key)) elif not should_quantize: self._is_quantized = False self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False) else: # Quantize on the fly device = getattr(Int8TensorwiseOps, "dynamic_load_device", None) if device is None: device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device # Log the first time we quantize in this loader pass if not hasattr(Int8TensorwiseOps, '_logged_otf'): print(f"INT8 Fast: Quantizing on-the-fly (ConvRot: {getattr(Int8TensorwiseOps, 'enable_convrot', False)})") Int8TensorwiseOps._logged_otf = True # Cast to float32 before rotation and scale computation w_gpu = weight_tensor.to(device, non_blocking=True).float() self._use_convrot = False if getattr(Int8TensorwiseOps, "enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0: try: import logging from .convrot import build_hadamard, rotate_weight H = build_hadamard(CONVROT_GROUP_SIZE, device=w_gpu.device, dtype=w_gpu.dtype) w_gpu = rotate_weight(w_gpu, H, group_size=CONVROT_GROUP_SIZE) self._use_convrot = True # Stamp the active groupsize on the module so # INT8ModelSave can serialize it deterministically # (instead of relying on the loader's default). self._convrot_groupsize = CONVROT_GROUP_SIZE except ImportError as e: import logging logging.warning(f"INT8 Fast: ConvRot enabled but convrot module error: {e}") q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1) q_weight = q_weight.cpu() q_scale = q_scale.cpu() self.weight = nn.Parameter(q_weight, requires_grad=False) self.register_buffer('weight_scale', q_scale) self._weight_scale_scalar = None self._is_quantized = True self._is_per_row = True del w_gpu, q_weight, q_scale else: self._is_quantized = False self.weight = nn.Parameter(source_tensor(weight_tensor), requires_grad=False) else: missing_keys.append(weight_key) # Assign bias if it exists (already patched if needed) if bias_tensor is not None: self.bias = nn.Parameter(source_tensor(bias_tensor), requires_grad=False) else: self.bias = None # Update archived model dtypes so VBAR geometry uses the correct # sizes. archive_model_dtypes runs before state_dict loading, so # weight_comfy_model_dtype is stale (e.g. bfloat16 instead of int8). # Without this, VBAR allocates 2x the needed memory and the cast # buffer path misinterprets int8 data as bfloat16. if self.weight is not None: self.weight_comfy_model_dtype = self.weight.dtype if self.weight_scale is not None: self.weight_scale_comfy_model_dtype = self.weight_scale.dtype if self.bias is not None: self.bias_comfy_model_dtype = self.bias.dtype def _get_weight_scale(self): return self.weight_scale def convert_weight(self, _weight, inplace=False): if not self._is_quantized: return _weight return self.weight def set_weight(self, out_weight, inplace_update=False, seed=0, return_weight=False, **kwargs): if not self._is_quantized: new_weight = out_weight.to(self.weight.dtype) if return_weight: return new_weight if inplace_update: self.weight.data.copy_(new_weight) else: self.weight = nn.Parameter(new_weight, requires_grad=False) return if out_weight.dtype == torch.int8: if return_weight: return out_weight if inplace_update: self.weight.data.copy_(out_weight) else: self.weight = nn.Parameter(out_weight, requires_grad=False) return # Re-quantize if fallback occurred new_weight = quantize_int8(out_weight, self._get_weight_scale()) if return_weight: return new_weight if inplace_update: self.weight.data.copy_(new_weight) else: self.weight = nn.Parameter(new_weight, requires_grad=False) def set_bias(self, out_bias, inplace_update=False, seed=0, return_weight=False, **kwargs): if out_bias is None: return None new_bias = out_bias if return_weight: return new_bias if inplace_update: if self.bias is not None: self.bias.data.copy_(new_bias) else: self.bias = nn.Parameter(new_bias, requires_grad=False) def forward(self, x: Tensor) -> Tensor: """Fast forward using torch._int_mm for quantized weights.""" # Check if ComfyUI needs to manage weight transfer (VBAR, offloading, LoRA patches, etc.) # This mirrors the base class check in disable_weight_init.Linear.forward() need_cast = self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0 if not self._is_quantized: if need_cast: weight, bias, offload_stream = cast_bias_weight(self, x, offloadable=True) out = F.linear(x, weight, bias) uncast_bias_weight(self, weight, bias, offload_stream) return out else: if x.device != self.weight.device or x.dtype != self.weight.dtype: weight = self.weight.to(device=x.device, dtype=x.dtype) bias = self.bias.to(device=x.device, dtype=x.dtype) if self.bias is not None else None return F.linear(x, weight, bias) return F.linear(x, self.weight, self.bias) # INT8 quantized path if need_cast: # VBAR / offload / lowvram path weight, bias, offload_stream = cast_bias_weight( self, input=None, dtype=torch.int8, device=x.device, bias_dtype=x.dtype, offloadable=True ) else: # Fast path: weights already on GPU, no functions to apply weight = self.weight bias = self.bias offload_stream = None w_scale = self._get_weight_scale() if isinstance(w_scale, torch.Tensor) and w_scale.device != x.device: w_scale = w_scale.to(x.device, non_blocking=True) compute_dtype = Int8TensorwiseOps.compute_dtype if compute_dtype is None: compute_dtype = Int8TensorwiseOps._default_compute_dtype(x) x_shape = x.shape x_2d = x.reshape(-1, x_shape[-1]) if x_2d.dtype != compute_dtype: x_2d = x_2d.to(compute_dtype) if getattr(self, "_use_convrot", False): from .convrot import build_hadamard, rotate_activation group_size = getattr(self, "_convrot_groupsize", CONVROT_GROUP_SIZE) H = build_hadamard(group_size, device=x.device, dtype=x_2d.dtype) x_2d = rotate_activation(x_2d, H, group_size=group_size) # Sync the loader toggle to the module-level flag read by the forward fns import sys as _sys _mod = _sys.modules[__name__] _mod._use_triton = Int8TensorwiseOps.use_triton if x_2d.shape[0] > 16: if self._is_per_row: y = int8_forward_dynamic_per_row(x_2d, weight, w_scale, bias, compute_dtype) else: y = int8_forward_dynamic(x_2d, weight, w_scale, bias, compute_dtype) else: # Small batch fallback w_float = dequantize(weight, w_scale).to(x_2d.dtype) bias_typed = bias.to(x_2d.dtype) if bias is not None else None y = F.linear(x_2d, w_float, bias_typed) # Dynamic LoRA Path — handles split QKV via per-patch offsets for lora_down, lora_up, lora_start, lora_size in self.lora_patches: lD = lora_down.to(x.device, non_blocking=True) lU = lora_up.to(x.device, non_blocking=True) lora_x = F.linear(x_2d.to(lD.dtype), lD) lora_y = F.linear(lora_x, lU) # [batch, slice_size or full_out] if lora_start is not None: y[:, lora_start:lora_start + lora_size] = ( y[:, lora_start:lora_start + lora_size] + lora_y.to(y.dtype) ) else: y = y + lora_y.to(y.dtype) if need_cast: uncast_bias_weight(self, weight, bias, offload_stream) return y.reshape(*x_shape[:-1], y.shape[-1]) # Pass-through for other layers class GroupNorm(manual_cast.GroupNorm): pass class LayerNorm(manual_cast.LayerNorm): pass class Conv2d(manual_cast.Conv2d): pass class Conv3d(manual_cast.Conv3d): pass class ConvTranspose2d(manual_cast.ConvTranspose2d): pass class Embedding(manual_cast.Embedding): pass @classmethod def conv_nd(cls, dims, *args, **kwargs): if dims == 2: return cls.Conv2d(*args, **kwargs) elif dims == 3: return cls.Conv3d(*args, **kwargs) else: raise ValueError(f"unsupported dimensions: {dims}") # ============================================================================= # INT8 Model Patcher - Unified LoRA Handling # ============================================================================= import inspect try: _prefetch_sig = inspect.signature(comfy.lora.prefetch_prepared_value) _use_new_prefetch = len(_prefetch_sig.parameters) == 5 except Exception: _use_new_prefetch = False class INT8LowVramPatch: is_lowvram_patch = True def __init__(self, key, patches, module, lora_mode): self.key = key self.patches = patches self.module = module self.lora_mode = lora_mode self.prepared_patches = None def memory_required(self): if not _use_new_prefetch: return 0 counter = [0] for patch in self.patches[self.key]: comfy.lora.prefetch_prepared_value(patch[1], counter, None, None, False) return counter[0] def prepare(self, *args, **kwargs): if _use_new_prefetch: # 0.22.0+ signature: prepare(self, destination, stream, copy=True, commit=True) destination = args[0] if len(args) > 0 else kwargs.get("destination") stream = args[1] if len(args) > 1 else kwargs.get("stream") copy = args[2] if len(args) > 2 else kwargs.get("copy", True) commit = args[3] if len(args) > 3 else kwargs.get("commit", True) counter = [0] prepared_patches = [ (patch[0], comfy.lora.prefetch_prepared_value(patch[1], counter, destination, stream, copy), patch[2], patch[3], patch[4]) for patch in self.patches[self.key] ] if commit: self.prepared_patches = prepared_patches return prepared_patches else: # 0.21.1- signature: prepare(self, allocate_buffer, stream) allocate_buffer = args[0] if len(args) > 0 else kwargs.get("allocate_buffer") stream = args[1] if len(args) > 1 else kwargs.get("stream") self.prepared_patches = [ (patch[0], comfy.lora.prefetch_prepared_value(patch[1], allocate_buffer, stream), patch[2], patch[3], patch[4]) for patch in self.patches[self.key] ] return self.prepared_patches def clear_prepared(self): self.prepared_patches = None def __call__(self, weight): patches = self.prepared_patches if self.prepared_patches is not None else self.patches[self.key] scale = self.module._get_weight_scale() if isinstance(scale, torch.Tensor): scale = scale.to(weight.device) weight_float = dequantize(weight, scale) use_convrot = getattr(self.module, "_use_convrot", False) if use_convrot: group_size = getattr(self.module, "_convrot_groupsize", CONVROT_GROUP_SIZE) try: from .convrot import build_hadamard, rotate_weight H = build_hadamard(group_size, device=weight.device, dtype=weight_float.dtype) weight_float = rotate_weight(weight_float, H, group_size=group_size) except ImportError: use_convrot = False patched_weight_float = comfy.lora.calculate_weight( patches, weight_float, self.key, intermediate_dtype=weight_float.dtype, ) if use_convrot: patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size) if self.lora_mode == "Stochastic": return stochastic_round_int8_delta( patched_weight_float, scale, seed=comfy.utils.string_to_seed(self.key), ) return quantize_int8(patched_weight_float, scale) class INT8ModelPatcher(comfy.model_patcher.ModelPatcher): """ Custom ModelPatcher that intercepts patching for INT8 layers. Routes patching through either a bake-in path (dequant-patch-requant) or a dynamic path (runtime injection), depending on the dynamic_lora toggle. """ def finalize_pending_int8(self): finalized = 0 for module in self.model.modules(): finalize = getattr(module, "finalize_pending_int8", None) if finalize is not None and finalize(): finalized += 1 if finalized > 0: self.size = 0 #logging.info(f"INT8 Fast: Finalized {finalized} deferred INT8 layer(s) at model load time.") if torch.cuda.is_available(): torch.cuda.empty_cache() def patch_weight_to_device(self, key, device_to=None, inplace_update=False, return_weight=False, force_cast=False): if key not in self.patches and not force_cast: return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast) # Check if this is one of our INT8 modules module_path = key.rsplit('.', 1)[0] try: module = comfy.utils.get_attr(self.model, module_path) except AttributeError: module = None is_int8_module = hasattr(module, "_is_quantized") and module._is_quantized patches = self.patches.get(key, []) if is_int8_module and Int8TensorwiseOps.Linear._is_bias_key(key): return comfy.utils.get_attr(self.model, key) if return_weight else None if is_int8_module: if not Int8TensorwiseOps.dynamic_lora: # --- BAKE-IN LORA PATH (Dequant → Patch → Quant) --- # Works with the native ComfyUI LoRA Loader (and also INT8LoraLoader). # All patches are applied in float space via ComfyUI's standard mechanism, # then the result is re-quantized back to INT8. # Identify current weight in the model current_weight = comfy.utils.get_attr(self.model, key) scale = module._get_weight_scale() if device_to is None: device_to = current_weight.device # ALWAYS use the weight from backup as the source if it exists to prevent additive stacking. # If it doesn't exist, this is the first patch, so create it from the current model weight. if key not in self.backup: import collections BackupEntry = collections.namedtuple('Dimension', ['weight', 'inplace_update']) self.backup[key] = BackupEntry( weight=current_weight.to(device=self.offload_device, copy=inplace_update), inplace_update=inplace_update, ) source_weight = current_weight else: # Use existing backup as source source_weight = self.backup[key].weight # 1. Dequantize to float (move scale to device_to since it lives on CPU) if isinstance(scale, torch.Tensor): scale = scale.to(device_to) weight_float = dequantize(source_weight.to(device_to), scale) # 2. Handle ConvRot: de-rotate into weight space before patching use_convrot = getattr(module, "_use_convrot", False) if use_convrot: group_size = getattr(module, "_convrot_groupsize", CONVROT_GROUP_SIZE) try: from .convrot import build_hadamard, rotate_weight H = build_hadamard(group_size, device=device_to, dtype=weight_float.dtype) weight_float = rotate_weight(weight_float, H, group_size=group_size) except ImportError: pass # 3. Patch in float space using ComfyUI's standard mechanism. # calculate_weight handles LoRA, LoHA, LoKR, DoRA, etc. patches_list = self.patches.get(key, []) patched_weight_float = comfy.lora.calculate_weight(patches_list, weight_float, key) # 4. Handle ConvRot: re-rotate if use_convrot: patched_weight_float = rotate_weight(patched_weight_float, H, group_size=group_size) # 5. Re-quantize back to INT8 using the original scale if getattr(Int8TensorwiseOps, "lora_mode", "None") == "Stochastic": patched_weight_int8 = stochastic_round_int8_delta(patched_weight_float, scale) else: patched_weight_int8 = quantize_int8(patched_weight_float, scale) # 6. Move back to original device and store patched_weight_int8 = patched_weight_int8.to(current_weight.device) if return_weight: return patched_weight_int8 if inplace_update: current_weight.data.copy_(patched_weight_int8) else: comfy.utils.set_attr(self.model, key, nn.Parameter(patched_weight_int8, requires_grad=False)) return else: # --- DYNAMIC LORA PATH --- # Build a list of (down_scaled, up, start, size) per patch. # Keeping patches separate preserves the offset info needed for # fused QKV layers where each of Q/K/V targets a different output slice. weight = comfy.utils.get_attr(self.model, key) device = weight.device if weight is not None else self.offload_device lora_patches = [] for p in patches: strength_patch = p[0] # float adapter = p[1] # the LoRA adapter object strength_model = p[2] # float offset = p[3] if len(p) > 3 else None # (dim, start, size) or None if not hasattr(adapter, "weights"): continue strength = strength_patch * strength_model weights = adapter.weights # Standard LoRA: (up, down, alpha, mid, dora_scale, reshape) if len(weights) == 6: up, down, alpha, mid, dora, reshape = weights rank = down.shape[0] if down.ndim >= 2 else 1 scale = (alpha / rank) * strength if alpha is not None else strength down_scaled = down.flatten(1) * scale if mid is not None: down_scaled = torch.mm(mid.flatten(1), down.flatten(1)) * scale # If this layer has ConvRot applied, rotate the 'down' matrix # so the LoRA delta is coherent with the rotated weight basis: # W_rot = W @ H^T => ΔW_rot = ΔW @ H^T => rotate down only if getattr(module, "_use_convrot", False) and down_scaled.shape[1] % CONVROT_GROUP_SIZE == 0: try: from .convrot import build_hadamard, rotate_weight group_size = getattr(module, "_convrot_groupsize", CONVROT_GROUP_SIZE) H = build_hadamard(group_size, device=down_scaled.device, dtype=down_scaled.dtype) down_scaled = rotate_weight(down_scaled, H, group_size=group_size) except ImportError: pass # Extract offset: which output rows this patch targets start, size = None, None if offset is not None: _dim, start, size = offset # dim is always 0 for linear weights lora_patches.append((down_scaled.to(device), up.flatten(1).to(device), start, size)) module.lora_patches = lora_patches if return_weight: return weight return # Skip standard weight-merging path # --- NON-INT8 MODULE PATH --- return super().patch_weight_to_device(key, device_to, inplace_update, return_weight, force_cast) def load(self, *args, **kwargs): self.finalize_pending_int8() save_materialized = bool(getattr(self, "_int8_save_materialized_lora", False)) if not Int8TensorwiseOps.dynamic_lora and not save_materialized: for k in list(self.backup): if k in self.patches: try: module = comfy.utils.get_attr(self.model, k.rsplit('.', 1)[0]) except AttributeError: module = None if hasattr(module, "_is_quantized") and module._is_quantized: bk = self.backup.pop(k) if bk.inplace_update: dest = comfy.utils.get_attr(self.model, k) dest.data.copy_(bk.weight) else: comfy.utils.set_attr(self.model, k, bk.weight) # Cleanup: Revert any keys that are in backup but no longer in patches (stale patches) # This ensures that when a LoRA is disabled, the model returns to its base state. stale_keys = [k for k in self.backup if k not in self.patches] for k in stale_keys: bk = self.backup.pop(k) if bk.inplace_update: dest = comfy.utils.get_attr(self.model, k) dest.data.copy_(bk.weight) else: comfy.utils.set_attr(self.model, k, bk.weight) # Cleanup: Clear stale dynamic LoRA patches. # This prevents LoRA from "sticking" when dynamic_lora is toggled or LoRAs are disabled. for name, module in self.model.named_modules(): if hasattr(module, "lora_patches") and module.lora_patches: # If dynamic LoRA is disabled globally, or if this module has no active patches, clear them. if not Int8TensorwiseOps.dynamic_lora or (name + ".weight") not in self.patches: module.lora_patches = [] res = super().load(*args, **kwargs) if hasattr(super(), "load") else None device_to = kwargs.get("device_to", args[0] if len(args) > 0 else self.model.device) for name, module in self.model.named_modules(): if hasattr(module, "_is_quantized") and module._is_quantized: weight_key = name + ".weight" if weight_key in self.patches: if save_materialized: if hasattr(module, "weight_lowvram_function"): module.weight_lowvram_function = None if hasattr(module, "weight_function"): module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"] elif Int8TensorwiseOps.dynamic_lora: if hasattr(module, "weight_lowvram_function"): module.weight_lowvram_function = None if hasattr(module, "weight_function"): module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"] self.patch_weight_to_device(weight_key, device_to=device_to) else: if hasattr(module, "weight_function"): module.weight_function = [f for f in getattr(module, "weight_function", []) if type(f).__name__ != "LowVramPatch"] lowvram_patch = INT8LowVramPatch( weight_key, self.patches, module, getattr(Int8TensorwiseOps, "lora_mode", "None"), ) pin_state = getattr(self.model, "dynamic_pins", {}).get(self.load_device, None) if pin_state is not None: lowvram_patch._pin_state = pin_state module.weight_lowvram_function = lowvram_patch return res def unpatch_model(self, device_to=None, unpatch_weights=True): if unpatch_weights: for name, module in self.model.named_modules(): if hasattr(module, "lora_patches"): module.lora_patches = [] return super().unpatch_model(device_to, unpatch_weights) def clone(self, *args, **kwargs): src_cls = self.__class__ if src_cls is INT8ModelPatcher: return super().clone(*args, **kwargs) if not issubclass(src_cls, INT8ModelPatcher): name = f"INT8_{src_cls.__name__}" dynamic_cls = type(name, (INT8ModelPatcher, src_cls), {}) else: dynamic_cls = src_cls self.__class__ = dynamic_cls # Static clones do not need a disk reload factory. Dynamic-to-static # delegates do: sharing the dynamic model object makes ComfyUI treat # the static copy as a replacement instead of an independent model. if not self.is_dynamic() and getattr(self, "cached_patcher_init", None) is None: self.cached_patcher_init = (lambda *a, **kw: self, ()) n = super().clone(*args, **kwargs) # If disable_dynamic is True, the core strips dynamic wrappers. We must re-apply INT8! disable_dyn = kwargs.get("disable_dynamic", False) if len(args) > 0: disable_dyn = args[0] if disable_dyn and not issubclass(n.__class__, INT8ModelPatcher): new_cls = type(f"INT8_{n.__class__.__name__}", (INT8ModelPatcher, n.__class__), {}) n.__class__ = new_cls self.__class__ = src_cls return n