diff --git a/backend/loader.py b/backend/loader.py index efa23a2d..4ce34886 100644 --- a/backend/loader.py +++ b/backend/loader.py @@ -379,12 +379,8 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p override_dtype = backend.args.dynamic_args.forge_unet_storage_dtype if override_dtype is torch.int8: - if state_dict_dtype is torch.bfloat16: - override_dtype = torch.bfloat16 - try_int8 = True - else: - override_dtype = None - logger.warning("int8 only supports bfloat16 models...") + override_dtype = torch.bfloat16 + try_int8 = True if guess.nunchaku: storage_dtype = torch.bfloat16 diff --git a/backend/operations.py b/backend/operations.py index 77ae44e6..b0655b45 100644 --- a/backend/operations.py +++ b/backend/operations.py @@ -62,9 +62,11 @@ def get_weight_and_bias(layer: torch.nn.Module) -> tuple[torch.Tensor, torch.Ten def weights_manual_cast( layer: Union[torch.nn.Module, "ForgeWeights"], - x: torch.Tensor, + x: torch.Tensor = None, *, dtype: torch.dtype = None, + device: torch.device = None, + bias_dtype: torch.dtype = None, weight_fn: Callable = None, bias_fn: Callable = None, skip_weight_dtype: bool = False, @@ -75,7 +77,10 @@ def weights_manual_cast( * Reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/ops.py#L210 """ - target_dtype, target_device = x.dtype, x.device + if x is not None: + target_dtype, target_device = x.dtype, x.device + else: + target_dtype, target_device = dtype, device non_blocking = memory_management.device_supports_non_blocking(target_device) weight, bias = None, None @@ -87,7 +92,7 @@ def weights_manual_cast( if skip_weight_dtype or weight_has_function: weight_args.pop("dtype") - bias_args = dict(device=target_device, dtype=target_dtype, non_blocking=non_blocking) + bias_args = dict(device=target_device, dtype=bias_dtype or target_dtype, non_blocking=non_blocking) if skip_bias_dtype or bias_has_function: bias_args.pop("dtype") @@ -362,46 +367,133 @@ class ForgeOperations: # region Int8 + from backend.operations_int8 import ( + CONVROT_GROUP_SIZE, dequantize, int8_forward_dynamic, - quantize_int8_tensorwise, - stochastic_round_int8_delta, + int8_forward_dynamic_per_row, + quantize_int8, + quantize_int8_axiswise, ) +from backend.patcher.lora import merge_lora_to_weight +from backend.quant_rotation import build_hadamard, rotate_activation, rotate_weight class ForgeOperationsInt8(ForgeOperations): """Custom operations for INT8 tensorwise quantization""" excluded_names = [] - _is_prequantized = None + dynamic_quantize = True # Toggle for on-the-fly quantization + enable_convrot = True # Toggle for ConvRot Hadamard rotation + + _is_prequantized = False # status flag (not used for detection) + + applied_lora_patches = set() + lora_patches = {} # Map of model_key -> patch list (from load_lora) + lora_strength = 1.0 class Linear(torch.nn.Linear, ForgeWeights): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.weight_scale = None + 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 = torch.bfloat16 - self.lora_A = None - self.lora_B = None - self.lora_alpha = None + self.lora_patches = [] # List of (down_scaled, up, start, size) set by INT8ModelPatcher def reset_parameters(self): return None def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): weight_key = prefix + "weight" - scale_key = prefix + "weight_scale" + + # Utility to normalize keys by stripping common prefixes + def normalize_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 + + def apply_lora_patches(tensor, key): + if not ForgeOperationsInt8.lora_patches or tensor.dtype == torch.int8: + return tensor + nk = normalize_key(key) + patches = ForgeOperationsInt8.lora_patches.get(nk) + if patches: + # calculate_weight expects: [(strength, v, strength_model, offset, function)] + formatted = [] + for patch in patches: + if len(patch) == 4: + v, offset, function, strength = patch + else: + v, offset, function = patch + strength = getattr(ForgeOperationsInt8, "lora_strength", 1.0) + formatted.append((strength, v, 1.0, offset, function)) + + # Track applied patches + ForgeOperationsInt8.applied_lora_patches.add(nk) + + device = torch.device("cuda") if torch.cuda.is_available() else tensor.device + temp_dtype = memory_management.lora_compute_dtype(device) + + tensor_temp = tensor.to(temp_dtype) + result_temp = merge_lora_to_weight(formatted, tensor_temp, key) + return result_temp.to(tensor.dtype) + return tensor + input_scale_key = prefix + "input_scale" bias_key = prefix + "bias" - weight_scale = state_dict.pop(scale_key, None) - state_dict.pop(prefix + "comfy_quant", None) + 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) + if comfy_quant_tensor is not None: + try: + import json + + quant_conf = json.loads(bytes(comfy_quant_tensor.tolist()).decode("utf-8")) + if quant_conf.get("convrot", False): + self._use_convrot = True + ForgeOperationsInt8.enable_convrot = True # Propagate globally for LoRA + if "convrot_groupsize" in quant_conf: + self._convrot_groupsize = quant_conf["convrot_groupsize"] + except Exception: + pass + + # Apply LoRA patches to weight and bias once + if weight_tensor is not None: + 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 @@ -410,72 +502,94 @@ class ForgeOperationsInt8(ForgeOperations): ForgeOperationsInt8._is_prequantized = True # Found a quantized layer if isinstance(weight_scale, torch.Tensor): - self.weight_scale = weight_scale.float().item() if weight_scale.numel() == 1 else weight_scale.float() + if weight_scale.numel() == 1: + # Scalar scale — store as float for speed + self._weight_scale_scalar = weight_scale.float().item() + self.weight_scale = None + self._is_per_row = False + 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 + else: + self.register_buffer("weight_scale", weight_scale.float()) + self._weight_scale_scalar = None + self._is_per_row = False else: - self.weight_scale = float(weight_scale) + self._weight_scale_scalar = float(weight_scale) + self.weight_scale = None + self._is_per_row = False - elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32): + elif weight_tensor.dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float8_e4m3fn): # Load High-Precision - # Detect if the model is pre-quantized if we don't know yet - if ForgeOperationsInt8._is_prequantized is None: - # Robust detection: scan keys and a sample of values - is_prequant = False - for k in state_dict.keys(): - if "weight_scale" in k or "comfy_quant" in k: - is_prequant = True - break - - if not is_prequant: - # Fallback: scan a sample of values for int8 tensors - for i, v in enumerate(state_dict.values()): - if i > 200: - break # Safety limit - if getattr(v, "dtype", None) == torch.int8: - is_prequant = True - break - ForgeOperationsInt8._is_prequantized = is_prequant - is_excluded = any(ex in prefix for ex in ForgeOperationsInt8.excluded_names) is_dim1 = self.in_features == 1 or self.out_features == 1 or weight_tensor.ndim == 1 - if is_excluded or is_dim1 or ForgeOperationsInt8._is_prequantized: + if is_excluded or is_dim1 or not ForgeOperationsInt8.dynamic_quantize: self._is_quantized = False self.weight = torch.nn.Parameter(weight_tensor, requires_grad=False) else: # Quantize on the fly device = torch.device("cuda") if torch.cuda.is_available() else weight_tensor.device - w_gpu = weight_tensor.to(device, non_blocking=True) - q_weight, q_scale = quantize_int8_tensorwise(w_gpu) + + # Cast to float32 before rotation and scale computation + w_gpu = weight_tensor.to(device, non_blocking=True).float() + + self._use_convrot = False + if getattr(ForgeOperationsInt8, "enable_convrot", False) and self.in_features % CONVROT_GROUP_SIZE == 0: + try: + 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 + except ImportError as e: + memory_management.logger.warning(f"[INT8 Fast] ConvRot Error: {e}") + + q_weight, q_scale = quantize_int8_axiswise(w_gpu, dim=1) self.weight = torch.nn.Parameter(q_weight.cpu(), requires_grad=False) - self.weight_scale = q_scale.cpu() if isinstance(q_scale, torch.Tensor) else q_scale + self.register_buffer("weight_scale", q_scale.cpu()) + self._weight_scale_scalar = None self._is_quantized = True + self._is_per_row = True else: self._is_quantized = False self.weight = torch.nn.Parameter(weight_tensor, requires_grad=False) else: missing_keys.append(weight_key) - bias_tensor = state_dict.pop(bias_key, None) + # Assign bias if it exists (already patched if needed) if bias_tensor is not None: self.bias = torch.nn.Parameter(bias_tensor, requires_grad=False) else: self.bias = None + def _get_weight_scale(self): + """Get weight scale, preferring scalar if available.""" + if self._weight_scale_scalar is not None: + return self._weight_scale_scalar + 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): + 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_(out_weight) + self.weight.data.copy_(new_weight) else: - self.weight = torch.nn.Parameter(out_weight.to(self.weight.dtype), requires_grad=False) + self.weight = torch.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: @@ -483,35 +597,56 @@ class ForgeOperationsInt8(ForgeOperations): return # Re-quantize if fallback occurred - new_weight = stochastic_round_int8_delta(out_weight, self.weight_scale, seed) + 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 = torch.nn.Parameter(new_weight, requires_grad=False) - def set_bias(self, out_bias, inplace_update=False, seed=0): + def set_bias(self, out_bias, inplace_update=False, seed=0, return_weight=False, **kwargs): if out_bias is None: - return + 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_(out_bias) + self.bias.data.copy_(new_bias) else: - self.bias = torch.nn.Parameter(out_bias, requires_grad=False) + self.bias = torch.nn.Parameter(new_bias, requires_grad=False) def forward(self, x: torch.Tensor) -> torch.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.parameters_manual_cast or len(self.weight_function) > 0 or len(self.bias_function) > 0 + if not self._is_quantized: - weight, bias, signal = weights_manual_cast(self, x) - with main_stream_worker(weight, bias, signal): - return torch.nn.functional.linear(x, weight, bias) + if need_cast: + weight, bias, signal = weights_manual_cast(self, x) + with main_stream_worker(weight, bias, signal): + return torch.nn.functional.linear(x, weight, bias) + else: + return torch.nn.functional.linear(x, self.weight, self.bias) - # 1. Move weight/bias/scale to device (non_blocking) - weight = self.weight.to(x.device, non_blocking=True) - bias = self.bias.to(x.device, non_blocking=True) if self.bias is not None else None + # INT8 quantized path + if need_cast: + # VBAR / offload / lowvram path + weight, bias, signal = weights_manual_cast(self, x=None, dtype=torch.int8, device=x.device, bias_dtype=x.dtype) + else: + # Fast path: weights already on GPU, no functions to apply + weight = self.weight + bias = self.bias - w_scale = self.weight_scale - if isinstance(w_scale, torch.Tensor): + 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 = x.dtype if x.dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 @@ -519,26 +654,36 @@ class ForgeOperationsInt8(ForgeOperations): x_shape = x.shape x_2d = x.reshape(-1, x_shape[-1]) + if getattr(self, "_use_convrot", False): + group_size = getattr(self, "_convrot_groupsize", CONVROT_GROUP_SIZE) + H = build_hadamard(group_size, device=x.device, dtype=x.dtype) + x_2d = rotate_activation(x_2d, H, group_size=group_size) + if x_2d.shape[0] > 16: - y = int8_forward_dynamic(x_2d, weight, w_scale, bias, compute_dtype) + 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.dtype) - y = torch.nn.functional.linear(x_2d, w_float, bias) + bias_typed = bias.to(x.dtype) if bias is not None else None + y = torch.nn.functional.linear(x_2d, w_float, bias_typed) - # Dynamic LoRA Path - if self.lora_A is not None and self.lora_B is not None: - # Ensure LoRA tensors are on the same device as x - lA = self.lora_A.to(x.device, non_blocking=True) - lB = self.lora_B.to(x.device, non_blocking=True) + # 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 = torch.nn.functional.linear(x_2d.to(lD.dtype), lD) + lora_y = torch.nn.functional.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) - lora_x = torch.nn.functional.linear(x_2d.to(lA.dtype), lA) - lora_y = torch.nn.functional.linear(lora_x, lB) - - if self.lora_alpha is not None: - lora_y = lora_y * self.lora_alpha - - y = y + lora_y.to(y.dtype) + if need_cast: + with main_stream_worker(weight, bias, signal): + pass return y.reshape(*x_shape[:-1], y.shape[-1]) @@ -834,14 +979,14 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast if isinstance(bnb_dtype, str): # https://github.com/BobJohnson24/ComfyUI-Flux2-INT8/blob/main/int8_unet_loader.py - ForgeOperationsInt8._is_prequantized = None + ForgeOperationsInt8._is_prequantized = False match bnb_dtype: case "Flux2K4B" | "Flux2K9B": - ForgeOperationsInt8.excluded_names = ["img_in", "time_in", "guidance_in", "txt_in", "final_layer", "double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation"] + ForgeOperationsInt8.excluded_names = ["img_in", "time_in", "guidance_in", "txt_in", "double_stream_modulation_img", "double_stream_modulation_txt", "single_stream_modulation"] operations = ForgeOperationsInt8 case "ZImage": - ForgeOperationsInt8.excluded_names = ["cap_embedder", "t_embedder", "x_embedder", "cap_pad_token", "context_refiner", "final_layer", "noise_refiner", "adaLN", "x_pad_token"] + ForgeOperationsInt8.excluded_names = ["cap_embedder", "t_embedder", "x_embedder", "cap_pad_token", "context_refiner", "final_layer", "noise_refiner", "adaLN", "x_pad_token", "layers.0."] operations = ForgeOperationsInt8 case "Chroma": ForgeOperationsInt8.excluded_names = ["distilled_guidance_layer", "final_layer", "img_in", "txt_in", "nerf_image_embedder", "nerf_blocks", "nerf_final_layer_conv", "__x0__", "nerf_final_layer_conv"] @@ -849,8 +994,14 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast case "QwenImage": ForgeOperationsInt8.excluded_names = ["time_text_embed", "img_in", "norm_out", "proj_out", "txt_in"] operations = ForgeOperationsInt8 + case "ErnieImage": + ForgeOperationsInt8.excluded_names = ["time", "x_embedder", "text_proj", "adaLN"] + operations = ForgeOperationsInt8 + case "Anima": + ForgeOperationsInt8.excluded_names = ["embed", "adaln"] + operations = ForgeOperationsInt8 case "WAN21_T2V" | "WAN21_I2V": - ForgeOperationsInt8.excluded_names = ["patch_embedding", "text_embedding", "time_embedding", "time_projection" "head", "img_emb"] + ForgeOperationsInt8.excluded_names = ["patch_embedding", "text_embedding", "time_embedding", "time_projection", "head", "img_emb"] operations = ForgeOperationsInt8 elif isinstance(bnb_dtype, dict): # https://github.com/Comfy-Org/ComfyUI/blob/v0.16.4/comfy/ops.py#L950 diff --git a/backend/operations_int8.py b/backend/operations_int8.py index 07100ae5..1f426836 100644 --- a/backend/operations_int8.py +++ b/backend/operations_int8.py @@ -1,9 +1,9 @@ -# https://github.com/BobJohnson24/ComfyUI-Flux2-INT8/blob/main/int8_quant.py +# https://github.com/BobJohnson24/ComfyUI-INT8-Fast/blob/main/int8_quant.py import torch try: - from backend.operations_triton import triton_int8_linear + from backend.operations_triton import triton_int8_linear, triton_int8_linear_per_row except ImportError: # Triton not found, fall back to torch._int_mm _TRITON_AVAILABLE = False @@ -11,6 +11,9 @@ else: _TRITON_AVAILABLE = True +CONVROT_GROUP_SIZE = 256 + + # region Quantization Utils @@ -34,28 +37,6 @@ def dequantize(q: torch.Tensor, scale: float | torch.Tensor) -> torch.Tensor: return q.float() * scale -def stochastic_round_int8_delta(x: torch.Tensor, scale: float | torch.Tensor, seed: int = 0) -> torch.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 - x_scaled = x / scale - - # Stochastic rounding - x_floor = torch.floor(x_scaled) - fraction = x_scaled - x_floor - - # Speed optimization: Create random values directly on the target device - random_vals = torch.rand(x_scaled.shape, generator=generator, device=x.device, dtype=x_scaled.dtype) - x_rounded = torch.where(random_vals < fraction, x_floor + 1, x_floor) - - return torch.clamp(x_rounded, -128, 127).to(torch.int8) - - # region LinearW8A8 Core @@ -83,104 +64,35 @@ def int8_forward_dynamic(x: torch.Tensor, weight: torch.Tensor, weight_scale: fl return res_scaled -# region Dynamic LoRA Hook +@torch.no_grad() +def int8_forward_dynamic_per_row(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, bias: torch.Tensor | None, compute_dtype: torch.dtype) -> torch.Tensor: + """Forward with dynamic per-token activation quantization and per-row weight quantization. - -class DynamicLoRAHook: - """ - Hook registered on the diffusion_model to synchronize dynamic LoRA attributes - with the current ModelPatcher context at the start of each forward pass. + 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 x.is_cuda: + return triton_int8_linear_per_row(x, weight, weight_scale, bias, compute_dtype) - def __init__(self): - self.current_lora_id = None + # --- SLOW PATH: Standard PyTorch --- + x_8, x_scale = quantize_int8_axiswise(x, dim=-1) - def pre_forward(self, module, input_args, input_kwargs): - # 1. Try to find transformer_options - transformer_options = input_kwargs.get("transformer_options", {}) - if not transformer_options: - # Fallback for models that pass it in context - context = input_args[2] if len(input_args) > 2 else None - if isinstance(context, dict) and "transformer_options" in context: - transformer_options = context["transformer_options"] + # INT8 Matmul (Outputs Int32) + res = torch._int_mm(x_8, weight.T) # [batch, out_features] - dynamic_loras = transformer_options.get("dynamic_loras", []) + # 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) - # 2. Generate a unique ID for this set of LoRAs - # We use handles/strengths to detect changes - lora_id = hash(tuple((id(d["patches"]), d["strength"]) for d in dynamic_loras)) if dynamic_loras else None - - if lora_id == self.current_lora_id: - return None # Already synchronized - - # 3. Synchronize all linear layers - self.apply_composition(module, dynamic_loras) - self.current_lora_id = lora_id - return None - - def apply_composition(self, diffusion_model, dynamic_loras): - # Pre-group patches by layer - layer_patches = {} - if dynamic_loras: - for entry in dynamic_loras: - strength = entry["strength"] - for key, adapter in entry["patches"].items(): - if key not in layer_patches: - layer_patches[key] = [] - layer_patches[key].append((adapter, strength)) - - # Update all modules - for name, module in diffusion_model.named_modules(): - if not hasattr(module, "lora_A"): - continue - - # Find patches for this module - # ComfyUI keys are often 'diffusion_model.path.to.weight' or 'path.to.weight' - possible_keys = [f"diffusion_model.{name}.weight", f"{name}.weight"] - patches = None - for pk in possible_keys: - if pk in layer_patches: - patches = layer_patches[pk] - break - - if not patches: - module.lora_A = None - module.lora_B = None - module.lora_alpha = None - continue - - # Compose - all_A = [] - all_B = [] - for adapter, strength in patches: - v = adapter.weights - up, down, alpha, mid = v[0], v[1], v[2], v[3] - rank = down.shape[0] if down.ndim >= 2 else 1 - scale = (alpha / rank) * strength if alpha is not None else strength - - curr_A = down - if mid is not None: - curr_A = torch.mm(mid.flatten(1), down.flatten(1)).reshape(down.shape) - - all_A.append(curr_A * scale) - all_B.append(up) - - if all_A: - device = getattr(module, "weight", torch.tensor(0)).device - module.lora_A = torch.cat(all_A, dim=0).to(device) - module.lora_B = torch.cat(all_B, dim=1).to(device) - module.lora_alpha = None - else: - module.lora_A = None - module.lora_B = None - - @classmethod - def register(cls, diffusion_model): - if not hasattr(diffusion_model, "_dynamic_lora_hook"): - hook = cls() - diffusion_model._dynamic_lora_hook = hook - diffusion_model.register_forward_pre_hook(hook.pre_forward, with_kwargs=True) - return diffusion_model._dynamic_lora_hook + if bias is not None: + res_scaled = res_scaled + bias.to(compute_dtype) + return res_scaled # region load_lora_int8 @@ -188,38 +100,152 @@ class DynamicLoRAHook: from typing import TYPE_CHECKING if TYPE_CHECKING: - from backend.patcher.unet import UnetPatcher + from backend.operations import ForgeOperationsInt8 + +from backend.patcher.lora import merge_lora_to_weight +from backend.patcher.unet import UnetPatcher +from backend.quant_rotation import build_hadamard, rotate_weight +from backend.utils import get_attr, set_attr -from backend.patcher.lora import load_lora, model_lora_keys_unet +class INT8ModelPatcher(UnetPatcher): + def _process_online_loras(self): + for key in self.online_patches: + self.patch_weight_to_device(key, online=True) -def load_lora_int8(model: "UnetPatcher", lora: dict[str, torch.Tensor], strength: float, lora_name: str) -> "UnetPatcher": - if strength == 0: - return model + def patch_weight_to_device(self, key, device_to=None, inplace_update=False, *, online=False): + if (not online) and (key not in self.patches): + return + if online and (key not in self.online_patches): + return - model_patcher = model.clone() + # Check if this is one of our INT8 modules + module_path = key.rsplit(".", 1)[0] - # 1. Get Patch Map - key_map = model_lora_keys_unet(model_patcher.model) + try: + module: "ForgeOperationsInt8.Linear" = get_attr(self.model, module_path) + except Exception: + return super().patch_weight_to_device(key, device_to, inplace_update) - patch_dict, _unmatch = load_lora(lora, key_map) + if not getattr(module, "_is_quantized", False): + return super().patch_weight_to_device(key, device_to, inplace_update) - # 2. Register Global Hook (if not exists) - DynamicLoRAHook.register(model_patcher.model.diffusion_model) + if online: + # --- 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. + patches = self.online_patches[key] + weight = 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 - # 3. Add to Dynamic LoRA list in transformer_options - # This ensures ComfyUI's cloning handles everything and it's non-sticky - if "transformer_options" not in model_patcher.model_options: - model_patcher.model_options["transformer_options"] = {} + if not hasattr(adapter, "weights"): + continue - opts = model_patcher.model_options["transformer_options"] - if "dynamic_loras" not in opts: - opts["dynamic_loras"] = [] - else: - # Shallow copy the list to avoid modifying the parent patcher's list - opts["dynamic_loras"] = opts["dynamic_loras"].copy() + 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 - opts["dynamic_loras"].append({"name": lora_name, "strength": strength, "patches": patch_dict}) + down_scaled = down.flatten(1) * scale + if mid is not None: + down_scaled = torch.mm(mid.flatten(1), down.flatten(1)) * scale - return model_patcher + # 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: + 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 Exception: + 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 + + else: + # --- 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. + patches = self.patches[key] + weight_int8 = get_attr(self.model, key) + scale = module._get_weight_scale() + + if device_to is None: + device_to = weight_int8.device + + # Save original weight so unpatch_model can restore it. + # Must use the same namedtuple format as ComfyUI's base patcher + # (collections.namedtuple('Dimension', ['weight', 'inplace_update'])) + # otherwise unpatch_model crashes with AttributeError on bk.inplace_update. + if key not in self.backup: + import collections + + BackupEntry = collections.namedtuple("Dimension", ["weight", "inplace_update"]) + self.backup[key] = BackupEntry( + weight=weight_int8.to(device=self.offload_device, copy=inplace_update), + inplace_update=inplace_update, + ) + + # 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(weight_int8.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: + 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 = merge_lora_to_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 + patched_weight_int8 = quantize_int8(patched_weight_float, scale) # stochastic_round_int8_delta(patched_weight_float, scale) + # I'm not really sure whether to stochastic round or not, results seem to depend on a per-lora basis. + # If quality is of the utmost importance, I recommend Pre-Lora instead of worrying about this. + + # 6. Move back to original device and store + patched_weight_int8 = patched_weight_int8.to(weight_int8.device) + + if inplace_update: + weight_int8.data.copy_(patched_weight_int8) + else: + set_attr(self.model, key, patched_weight_int8) + + 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) diff --git a/backend/operations_triton.py b/backend/operations_triton.py index 2c0f0e24..96ebc3af 100644 --- a/backend/operations_triton.py +++ b/backend/operations_triton.py @@ -1,4 +1,4 @@ -# https://github.com/BobJohnson24/ComfyUI-Flux2-INT8/blob/main/int8_fused_kernel.py +# https://github.com/BobJohnson24/ComfyUI-INT8-Fast/blob/main/int8_fused_kernel.py import torch import triton @@ -257,3 +257,142 @@ def triton_int8_linear(x: torch.Tensor, weight: torch.Tensor, weight_scale, bias # 6. Reshape output return output.reshape(x_shape_orig[:-1] + (N,)) + + +# ============================================================================= +# Kernel 3: INT8 GEMM + Fused Dequant with Per-Row Weight Scales +# ============================================================================= + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256, "BLOCK_K": 64, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 256, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 128, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 32, "BLOCK_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + ], + key=["M", "N", "K"], +) +@triton.jit +def _int8_matmul_dequant_per_row_kernel( + # Pointers + a_ptr, + b_ptr, + c_ptr, + a_scale_ptr, + b_scale_ptr, + bias_ptr, + # Matrix Dimensions + M, + N, + K, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + # Meta-parameters + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + HAS_BIAS: tl.constexpr, +): + """ + Computes: C = ((A * B) * (scale_a[:, None] * scale_b[None, :])) + bias + A: [M, K] int8, scale_a: [M, 1] per-row activation scales + B: [N, K] int8, scale_b: [N, 1] per-row weight scales + """ + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_M) + num_pid_n = tl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # 1. Prepare Pointers for A and B + offs_am = (pid_m * BLOCK_M + tl.arange(0, BLOCK_M)) % M + offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N)) % N + offs_k = tl.arange(0, BLOCK_K) + + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # 2. Main Loop (Accumulate in Int32) + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32) + + for k in range(0, tl.cdiv(K, BLOCK_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0) + accumulator += tl.dot(a, b) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + + # 3. Fused Epilogue (Dequantize & Bias) + + # A Scale is per-row [M, 1] + scale_a = tl.load(a_scale_ptr + offs_am) # Vector [BLOCK_M] + + # B Scale is per-row [N, 1] (the key difference from the scalar kernel) + scale_b = tl.load(b_scale_ptr + offs_bn) # Vector [BLOCK_N] + + c = accumulator.to(tl.float32) + + # Outer product of scales: [BLOCK_M, 1] * [1, BLOCK_N] + total_scale = scale_a[:, None] * scale_b[None, :] + + c = c * total_scale + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_bn) + c = c + bias[None, :] + + # 4. Store Result + c_ptrs = c_ptr + stride_cm * offs_am[:, None] + stride_cn * offs_bn[None, :] + c_mask = (offs_am[:, None] < M) & (offs_bn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +# ============================================================================= +# Python Wrapper (Per-Row Weight Scales) +# ============================================================================= + + +def triton_int8_linear_per_row(x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor, bias=None, compute_dtype=torch.float16): + """ + Fused pipeline for W8A8 Linear Layer with per-row weight quantization. + weight_scale: [N, 1] per-row scales + """ + # 1. Flatten inputs if 3D + x_shape_orig = x.shape + x_2d = x.reshape(-1, x_shape_orig[-1]) + + M, K = x_2d.shape + N = weight.shape[0] + + # 2. Dynamic Activation Quantization + x_int8, x_scale = triton_quantize_rowwise(x_2d) + + # 3. Allocate Output + output = torch.empty((M, N), device=x.device, dtype=compute_dtype) + + # 4. Prepare weight scales - flatten [N, 1] -> [N] for kernel + ws = weight_scale.reshape(N).contiguous() + + # 5. Fused GEMM + Per-Row Dequant + grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]) * triton.cdiv(N, META["BLOCK_N"]),) + + has_bias = bias is not None + bias_ptr = bias if has_bias else x # Dummy pointer if None + + _int8_matmul_dequant_per_row_kernel[grid](a_ptr=x_int8, b_ptr=weight, c_ptr=output, a_scale_ptr=x_scale, b_scale_ptr=ws, bias_ptr=bias_ptr, M=M, N=N, K=K, stride_am=x_int8.stride(0), stride_ak=x_int8.stride(1), stride_bk=weight.stride(1), stride_bn=weight.stride(0), stride_cm=output.stride(0), stride_cn=output.stride(1), HAS_BIAS=has_bias) + + # 6. Reshape output + return output.reshape(x_shape_orig[:-1] + (N,)) diff --git a/backend/patcher/unet.py b/backend/patcher/unet.py index 0857ac1b..567fc93d 100644 --- a/backend/patcher/unet.py +++ b/backend/patcher/unet.py @@ -1,5 +1,6 @@ import torch +from backend.args import dynamic_args from backend.modules.k_model import KModel from backend.patcher.base import ModelPatcher @@ -16,6 +17,11 @@ class UnetPatcher(ModelPatcher): if isinstance(model.diffusion_model, NunchakuModelMixin): return NunchakuPatcher(model, load_device=model.diffusion_model.load_device, offload_device=model.diffusion_model.offload_device, current_device=model.diffusion_model.initial_device) else: + if dynamic_args.ops.endswith("Int8"): + from backend.operations_int8 import INT8ModelPatcher + + return INT8ModelPatcher(model, load_device=model.diffusion_model.load_device, offload_device=model.diffusion_model.offload_device, current_device=model.diffusion_model.initial_device) + return UnetPatcher(model, load_device=model.diffusion_model.load_device, offload_device=model.diffusion_model.offload_device, current_device=model.diffusion_model.initial_device) def __init__(self, *args, **kwargs): diff --git a/backend/quant_rotation.py b/backend/quant_rotation.py new file mode 100644 index 00000000..8a261f3c --- /dev/null +++ b/backend/quant_rotation.py @@ -0,0 +1,112 @@ +# https://github.com/BobJohnson24/ComfyUI-INT8-Fast/blob/main/convrot.py + +# Group-wise Hadamard rotation for INT8 quantization quality improvement +# reference: https://github.com/newgrit1004/ComfyUI-ZImage-Triton +# License: MIT + + +import torch + +# Cache Hadamard matrices by (size, device, dtype) to avoid recomputation +_HADAMARD_CACHE: dict[tuple[int, str, torch.dtype], torch.Tensor] = {} + + +def build_hadamard( + size: int, + device: str | torch.device = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Build a normalized REGULAR orthogonal Hadamard matrix (ConvRot). + + Size must be a power of 4 (e.g., 4, 16, 64, 256, 1024...). + Uses the Kronecker construction from Theorem 3.3 to avoid the all-1s + column of standard Sylvester Hadamard matrices, which amplifies + row-wise outliers in diffusion models. + """ + import math + + cache_key = (size, str(device), dtype) + if cache_key in _HADAMARD_CACHE: + return _HADAMARD_CACHE[cache_key] + + if size < 4 or (size & (size - 1)) != 0 or math.log(size, 4) % 1 != 0: + raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") + + # Base H4 from Theorem 3.3 (Eq 9 in the paper) + # Notice how every row and column sums to exactly 2 + H4 = torch.tensor([[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], dtype=dtype, device=device) + + H = H4 + current_size = 4 + + # Kronecker construction for larger sizes: H_{4^{k+1}} = H_{4^k} \otimes H_4 + while current_size < size: + H = torch.kron(H, H4) + current_size *= 4 + + # Normalize to make it orthogonal + H_normalized = H / (size**0.5) + _HADAMARD_CACHE[cache_key] = H_normalized + + return H_normalized + + +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 shape (out, in): + Each row of W is split into groups of group_size and rotated by H^T. + + Args: + weight: Shape (out_features, in_features). + H: Normalized Hadamard matrix, shape (group_size, group_size). + group_size: Group size for block-diagonal rotation. + + Returns: + Rotated weight, same shape as input. + """ + 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 + + # (out, in) → (out, n_groups, group_size) + W_grouped = weight.view(out_f, n_groups, group_size) + # Apply H^T to each group: (..., group_size) @ (group_size, 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) + + +def rotate_activation( + x: torch.Tensor, + H: torch.Tensor, + group_size: int, +) -> torch.Tensor: + """Rotate activation online: x_rot = x @ H_block. + + Group-wise Hadamard spreads outliers across channels within each group. + + Args: + x: Shape (..., features). Last dim must be divisible by group_size. + H: Normalized Hadamard matrix, shape (group_size, group_size). + group_size: Group size for block-diagonal rotation. + + Returns: + Rotated activation, same shape as input. + """ + orig_shape = x.shape + features = orig_shape[-1] + if features % group_size != 0: + raise ValueError(f"features {features} not divisible by group_size {group_size}") + n_groups = features // group_size + + # (..., features) → (..., n_groups, group_size) + x_grouped = x.view(*orig_shape[:-1], n_groups, group_size) + H_dev = H.to(dtype=x.dtype, device=x.device) + x_rot = torch.matmul(x_grouped, H_dev) + return x_rot.view(orig_shape) diff --git a/extensions-builtin/sd_forge_lora/networks.py b/extensions-builtin/sd_forge_lora/networks.py index 19d05776..0bf04ad8 100644 --- a/extensions-builtin/sd_forge_lora/networks.py +++ b/extensions-builtin/sd_forge_lora/networks.py @@ -37,11 +37,6 @@ def load_lora_for_models(model: "UnetPatcher", clip: "CLIP", lora: dict[str, tor if dynamic_args.nunchaku: model.model.diffusion_model.loras.append((filename, strength_model)) return model, clip - if dynamic_args.ops.endswith("Int8"): - from backend.operations_int8 import load_lora_int8 - - model = load_lora_int8(model, lora, strength_model, filename) - return model, clip model_flag: str = type(model.model).__name__ if model is not None else "default" @@ -127,8 +122,6 @@ def load_networks(names: list[str], te_multipliers: list[float] = None, unet_mul online_mode = dynamic_args.online_lora or False - if current_sd.forge_objects.unet.model.storage_dtype in [torch.float32, torch.float16, torch.bfloat16]: - online_mode = False if dynamic_args.ops.startswith("Mixed") or dynamic_args.ops.endswith("FP8"): online_mode = False