diff --git a/comfy/model_base.py b/comfy/model_base.py index 52804e477..cfd67cae2 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -130,7 +130,7 @@ class BaseModel(torch.nn.Module): if not unet_config.get("disable_unet_model_creation", False): if model_config.custom_operations is None: fp8 = model_config.optimizations.get("fp8", False) - operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype) + operations = comfy.ops.pick_operations(model_config.target_dtype, self.manual_cast_dtype, fast_fp8=fp8) else: operations = model_config.custom_operations self.diffusion_model = unet_model(**unet_config, device=device, operations=operations) diff --git a/comfy/ops.py b/comfy/ops.py index 142fd215f..342513ba5 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -19,11 +19,10 @@ import torch import logging import comfy.model_management -from comfy.cli_args import args, PerformanceFeature import comfy.float import comfy.rmsnorm import contextlib -from comfy.quant_tensor import Q_TYPES, tensor_quantizer, tensor_dequantizer, dynamic_tensor_quantizer, woq_fwd, quantized_fwd +from comfy.quant_tensor import Q_TYPES, tensor_quantizer, tensor_dequantizer, dynamic_tensor_quantizer, woq_fwd, quantized_fwd, get_quantizer_with_constraints import types import inspect @@ -128,7 +127,7 @@ class disable_weight_init: def _set_quantizer_fn(self, scale_weight, scale_input): if scale_weight.ndim != 0 and scale_weight.shape[0] != 1: raise ValueError("Blockwise quantization is not supported") - if self.use_dynamic_quantizer: + if scale_input is None and self.use_dynamic_quantizer: setattr(self, "quantizer", dynamic_tensor_quantizer) else: setattr(self, "quantizer", tensor_quantizer) @@ -147,14 +146,22 @@ class disable_weight_init: def _init_parameters_from_sd(self, state_dict, prefix): if not state_dict: - logging.warning("No state dict provided.") + logging.warning(f"No state dict provided for {prefix}.") weight = torch.nn.Parameter( torch.empty((self.out_features, self.in_features), dtype=self.compute_dtype, device=self.device) ) self.register_buffer('weight', weight) return - weight_dtype = state_dict[f"{prefix}weight"].dtype + scale_weight = None + _w = state_dict.pop(f"{prefix}weight") + if len(self.weight_function): + _w, scale_weight = self.weight_function[0](_w) + state_dict[f"{prefix}weight"] = _w + if scale_weight is not None: + state_dict[f"{prefix}scale_weight"] = scale_weight + + weight_dtype = _w.dtype weight = torch.nn.Parameter( torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype)) @@ -164,20 +171,21 @@ class disable_weight_init: scale_weight = state_dict.get(f"{prefix}scale_weight", None) if scale_weight is None: - logging.warning("Using quantized Weights requires a scale to be present! Falling back to 1.0") + logging.warning("Using quantized weights without a scale can result in low accuracy.") scale_weight = torch.ones(1) state_dict[f"{prefix}scale_weight"] = scale_weight self.register_buffer('scale_weight', scale_weight.to(device=self.device)) scale_input = state_dict.get(f"{prefix}scale_input", None) - if scale_input is None and self.use_dynamic_quantizer: - scale_input = torch.ones(1) # Placeholder for API if scale_input is not None: self.register_buffer('scale_input', scale_input.to(device=self.device)) + elif not self.use_dynamic_quantizer: + # Fallback to WoQ + self.fp8_compute = False if self.bias is not None: # WAR not really nice, but Qwen VL has an input scale but uses f32 intermediates and quantized bias - self.fp8_compute = not self.bias.dtype in Q_TYPES + self.fp8_compute = self.fp8_compute and (self.bias.dtype in [torch.float16, torch.bfloat16]) self._set_quantizer_fn(scale_weight, scale_input) self._set_dequantizer_fn(scale_weight) @@ -416,9 +424,12 @@ def operator_factory(**factory_kwargs): # TODO might be nicer to have a unified interface to the factory # TODO logic might not be 1-1 match to original implementation -def pick_operations(weight_dtype=None, compute_dtype=None, load_device=None, disable_fast_fp8=False): +def pick_operations(weight_dtype=None, compute_dtype=None, load_device=None, fast_fp8=False, disable_fast_fp8=False): fp8_compute = (comfy.model_management.supports_fp8_compute(load_device) and not disable_fast_fp8) - use_dynamic_quantizer = PerformanceFeature.DynamicQuantizer in args.fast + use_dynamic_quantizer = fast_fp8 manual_cast = not((weight_dtype == compute_dtype) or use_dynamic_quantizer or fp8_compute) - return operator_factory(comfy_cast_weights=manual_cast, use_dynamic_quantizer=use_dynamic_quantizer, fp8_compute=fp8_compute) + weight_function = [] + if weight_dtype is not None and compute_dtype is not None and not manual_cast: + weight_function = [get_quantizer_with_constraints(weight_dtype)] + return operator_factory(comfy_cast_weights=manual_cast, use_dynamic_quantizer=use_dynamic_quantizer, fp8_compute=fp8_compute, weight_function=weight_function) diff --git a/comfy/quant_tensor.py b/comfy/quant_tensor.py index cb01349cc..8f4cac1de 100644 --- a/comfy/quant_tensor.py +++ b/comfy/quant_tensor.py @@ -1,73 +1,54 @@ import torch +import logging +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Tuple +from torch.utils._triton import has_triton +from typing import Dict -""" -x: 512x1024 w:1024x1024 -- For TensorWise scaling, a and b should be float8, scales should be float and singletons. -- For RowWise scaling, a and b should be float8, scales should be float, scale_a should be (512, 1) and scale_b should be (1, 1024), and both should be contiguous. -- For BlockWise 1x128 scaling, a and b should be float8, scales should be float, scale_a should be (512, 8) and scale_b should be (8, 1024), and both should be outer-dim-major. -- For BlockWise 128x128 scaling, a and b should be float8, scales should be float, scale_a should be (4, 8) and scale_b should be (8, 8), and both should be near-inner-dim-major (with 16-byte aligned strides). -- For Blockwise 1x32 scaling, a and b should be float8, scales should be float8_e8m0fnu, scale_a should have 16384 elements and scale_b should have 32768 elements, and both should be contiguous. -- For Blockwise 1x16 scaling, a and b should be float4 (packed 2x), scales should be float8_e4m3fn, scale_a should have 65536 elements and scale_b should have 131072 elements, and both should be contiguous. -""" -Q_TYPES = [torch.float8_e4m3fn, torch.float4_e2m1fn_x2] +Q_TYPES = [torch.float8_e4m3fn] -def dynamic_tensor_quantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): +if has_triton(): + q_compile_decorator = torch.compile() +else: + q_compile_decorator = lambda func: func + +def get_quantizer_with_constraints(target_dtype: torch.dtype): + if target_dtype == torch.float8_e4m3fn: + q_fn = dynamic_tensor_quantizer + else: + raise ValueError(f"Unsupported dtype {target_dtype}") + + alignment_check_fn = lambda x: x.shape[0] % 16 or x.shape[1] % 16 + + def fn(x, **kwargs): + if alignment_check_fn(x): + return x, None + return q_fn(x, **kwargs) + + return fn + +@q_compile_decorator +def dynamic_tensor_quantizer(x: torch.Tensor, dtype=torch.dtype, *args, **kwargs): input_scale = torch.abs(x).max() / torch.finfo(dtype).max x = (x / input_scale).clamp(torch.finfo(dtype).min, torch.finfo(dtype).max).to(dtype=dtype) return x, input_scale.float() -def mxfp8_quantizer(x: torch.Tensor, dtype: torch.dtype): - block_size = 32 - orig_shape = x.shape - x = x.reshape(-1, block_size) - scale = (torch.amax(torch.abs(x), dim=-1) / torch.finfo(dtype).max) - x = (x / scale.unsqueeze(-1)).clamp(torch.finfo(dtype).min, torch.finfo(dtype).max).to(dtype=dtype).contiguous() - x = x.view(orig_shape) - - return x, scale.to(dtype=torch.float8_e8m0fnu).contiguous() - +@q_compile_decorator def tensor_quantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): x = (x / scale).clamp(torch.finfo(dtype).min, torch.finfo(dtype).max).to(dtype=dtype).contiguous() return x, scale.float() -def nvfp4_quantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): - """ - orig_shape = x.shape - x = x.reshape(orig_shape[0], -1, block_size) - max_abs = torch.amax(torch.abs(x), dim=-1) - block_scale = (max_abs / torch.finfo(torch.float4_e2m1fn_x2.max))-float() - scaled_block_scales = block_scale / scale - scaled_block_scales_fp8 = torch.clamp( - scaled_block_scales, min=E4M3_EPS, max=F8E4M3_MAX - ).to(torch.float8_e4m3fn) - scaled_block_scales_fp32 = scaled_block_scales_fp8.to(torch.float32) - # We "temporarily" dequant the scaled_block_scales_fp32 to get the per_tensor_scale - # To apply to data - total_scale = scale * scaled_block_scales_fp32 - data_scaled = x / total_scale.unsqueeze(-1) - out_scales = scaled_block_scales_fp8 - - data_scaled = torch.clamp(data_scaled, -F4_E2M1_MAX, F4_E2M1_MAX) - data_scaled = data_scaled.view(orig_shape) - data_lp = f32_to_f4_unpacked(data_scaled) - # TODO: NotImplementedError: "copy_kernel" not implemented for 'Float4_e2m1fn_x2' - # data_lp = pack_uint4(data_lp).view(torch.float4_e2m1fn_x2) - data_lp = pack_uint4(data_lp) - return out_scales, data_lp - """ - block_size: int = 16 - raise NotImplementedError - - +@q_compile_decorator def tensor_dequantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): - x = (x.to(dtype=scale.dtype) * scale).to(dtype=dtype) + x = x.to(dtype=dtype) * scale.to(dtype=dtype) return x def woq_fwd(self, x): - dq_weight = self.dequantizer(self.weight, self.scale_weight, x.dtype) + dq_weight = self.dequantizer(self.weight, scale=self.scale_weight, dtype=x.dtype) bias = self.bias if bias is not None and bias.dtype == self.weight.dtype: - bias = self.dequantizer(bias, self.scale_weight, x.dtype) + bias = self.dequantizer(bias, torch.ones_like(self.scale_weight), x.dtype) return torch.nn.functional.linear(x, dq_weight, bias) def quantized_fwd(self, input): @@ -80,9 +61,10 @@ def quantized_fwd(self, input): input_dtype = input.dtype assert len(input_shape) == 3, "input must be 3D" - q_input, input_scale = self.quantizer(input, self.scale_input, self.weight.dtype) + scale_input = getattr(self, "scale_input", None) + q_input, scale_input = self.quantizer(input, scale=scale_input, dtype=self.weight.dtype) q_input = q_input.reshape(-1, input_shape[2]) - o = torch._scaled_mm(q_input, self.weight.T, scale_a=input_scale, scale_b=self.scale_weight.float(), + o = torch._scaled_mm(q_input, self.weight.T, scale_a=scale_input, scale_b=self.scale_weight.float(), bias=self.bias, out_dtype=input_dtype) if isinstance(o, tuple): o = o[0] diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py index 54573abb1..bb3ae5940 100644 --- a/comfy/supported_models_base.py +++ b/comfy/supported_models_base.py @@ -48,6 +48,7 @@ class BASE: memory_usage_factor = 2.0 manual_cast_dtype = None + target_dtype = None custom_operations = None scaled_fp8 = None optimizations = {"fp8": False} @@ -115,5 +116,6 @@ class BASE: return utils.state_dict_prefix_replace(state_dict, replace_prefix) def set_inference_dtype(self, dtype, manual_cast_dtype): - self.unet_config['dtype'] = dtype + self.unet_config['dtype'] = dtype if manual_cast_dtype is None else manual_cast_dtype + self.target_dtype = dtype self.manual_cast_dtype = manual_cast_dtype