diff --git a/comfy/cli_args.py b/comfy/cli_args.py index de3e85c08..759a1c36b 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -143,6 +143,7 @@ class PerformanceFeature(enum.Enum): Fp16Accumulation = "fp16_accumulation" Fp8MatrixMultiplication = "fp8_matrix_mult" CublasOps = "cublas_ops" + DynamicQuantizer = "dynamic_quantizer" parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: fp16_accumulation fp8_matrix_mult cublas_ops") diff --git a/comfy/gligen.py b/comfy/gligen.py index 1d7b6c2f4..9a3db0f5c 100644 --- a/comfy/gligen.py +++ b/comfy/gligen.py @@ -3,7 +3,7 @@ import torch from torch import nn from .ldm.modules.attention import CrossAttention, FeedForward import comfy.ops -ops = comfy.ops.manual_cast +ops = comfy.ops.operator_factory(comfy_cast_weights=True) class GatedCrossAttentionDense(nn.Module): diff --git a/comfy/ldm/flux/redux.py b/comfy/ldm/flux/redux.py index 527e83164..46b966831 100644 --- a/comfy/ldm/flux/redux.py +++ b/comfy/ldm/flux/redux.py @@ -1,7 +1,7 @@ import torch import comfy.ops -ops = comfy.ops.manual_cast +ops = comfy.ops.operator_factory(comfy_cast_weights=True) class ReduxImageEncoder(torch.nn.Module): def __init__( diff --git a/comfy/ops.py b/comfy/ops.py index 6aa2ff71d..142fd215f 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -23,62 +23,9 @@ 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 import types - - -Q_TYPES = [torch.float8_e4m3fn, torch.float4_e2m1fn_x2] - -def dynamic_quantizer(x: torch.Tensor, dtype: torch.dtype): - input_scale = 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 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 - -def dequantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): - x = (x.to(dtype=scale.dtype) * scale).to(dtype=dtype) - return x - -def woq_fwd(self, x): - dq_weight = self.dequantizer(self.weight, self.scale_weight, x.dtype) - bias = self.bias - if bias is not None and bias.dtype != x.dtype: - bias = self.dequantizer(bias, self.scale_weight, x.dtype) - return torch.nn.functional.linear(x, dq_weight, bias) - -def quantized_fwd(self, input): - tensor_2d = False - if len(input.shape) == 2: - tensor_2d = True - input = input.unsqueeze(1) - - input_shape = input.shape - input_dtype = input.dtype - assert len(input_shape) == 3, "input must be 3D" - - q_input = self.quantizer(input, self.scale_input, self.weight.dtype) - q_input = q_input.reshape(-1, input_shape[2]) - o = torch._scaled_mm(q_input, self.weight.T, scale_a=self.scale_input, scale_b=self.scale_weight, bias=self.bias, out_dtype=input_dtype) - if isinstance(o, tuple): - o = o[0] - if tensor_2d: - return o.reshape(input_shape[0], -1) - return o.reshape((-1, input_shape[1], self.weight.shape[0])) - -def get_quantized_forward(scale_weight, scale_input): - if scale_input is None: - return woq_fwd - else: - return quantized_fwd - -def get_quantizer_fn(scale_weight, scale_input): - # TODO Block Scaling, MX Scaling, Double-Q-NVFP4 - return quantizer - -def get_dequantizer_fn(scale_weight, scale_input): - return dequantizer +import inspect def scaled_dot_product_attention(q, k, v, *args, **kwargs): return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs) @@ -149,6 +96,8 @@ class CastWeightBiasOp: comfy_cast_weights = False weight_function = [] bias_function = [] + fp8_compute: bool = False + use_dynamic_quantizer: bool = False class disable_weight_init: class Linear(torch.nn.Module, CastWeightBiasOp): @@ -165,8 +114,8 @@ class disable_weight_init: self.in_features = in_features self.out_features = out_features - self.device = device - self.compute_dtype = dtype + self.device = torch.device("cpu") if device is None else device + self.compute_dtype = torch.float32 if dtype is None else dtype if bias: self.bias = torch.nn.Parameter(torch.empty(out_features, **factory_kwargs)) @@ -176,6 +125,26 @@ class disable_weight_init: def reset_parameters(self): return None + 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: + setattr(self, "quantizer", dynamic_tensor_quantizer) + else: + setattr(self, "quantizer", tensor_quantizer) + + def _set_dequantizer_fn(self, scale_weight): + if scale_weight.ndim != 0 and scale_weight.shape[0] != 1: + raise ValueError("Blockwise quantization is not supported") + setattr(self, "dequantizer", tensor_dequantizer) + + def _set_quantized_forward(self): + if not self.fp8_compute: + q_fwd = woq_fwd + else: + q_fwd = quantized_fwd + self.forward = types.MethodType(q_fwd, self) + def _init_parameters_from_sd(self, state_dict, prefix): if not state_dict: logging.warning("No state dict provided.") @@ -187,8 +156,7 @@ class disable_weight_init: weight_dtype = state_dict[f"{prefix}weight"].dtype weight = torch.nn.Parameter( - torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype) - ) + torch.empty((self.out_features, self.in_features), device=self.device, dtype=weight_dtype)) self.register_buffer('weight', weight) if weight_dtype not in Q_TYPES: @@ -199,17 +167,21 @@ class disable_weight_init: logging.warning("Using quantized Weights requires a scale to be present! Falling back to 1.0") scale_weight = torch.ones(1) state_dict[f"{prefix}scale_weight"] = scale_weight - self.register_buffer('scale_weight', scale_weight.to(device=self.device, dtype=torch.float32)) + 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: - if scale_input != 1: # TODO not really nice but e.g. Qwen VL has an input scale but does not use it? - self.register_buffer('scale_input', scale_input.to(device=self.device, dtype=torch.float32)) - else: - scale_input = None - self.forward = types.MethodType(get_quantized_forward(scale_weight, scale_input), self) - setattr(self, "quantizer", get_quantizer_fn(scale_weight, scale_input)) - setattr(self, "dequantizer", get_dequantizer_fn(scale_weight, scale_input)) + self.register_buffer('scale_input', scale_input.to(device=self.device)) + + 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._set_quantizer_fn(scale_weight, scale_input) + self._set_dequantizer_fn(scale_weight) + self._set_quantized_forward() def _load_from_state_dict( self, @@ -240,7 +212,6 @@ class disable_weight_init: weight, bias = cast_bias_weight(self, input) return torch.nn.functional.linear(input, weight, bias) - class Conv1d(torch.nn.Conv1d, CastWeightBiasOp): def reset_parameters(self): return None @@ -405,152 +376,6 @@ class disable_weight_init: else: raise ValueError(f"unsupported dimensions: {dims}") - -class manual_cast(disable_weight_init): - class Linear(disable_weight_init.Linear): - comfy_cast_weights = True - - class Conv1d(disable_weight_init.Conv1d): - comfy_cast_weights = True - - class Conv2d(disable_weight_init.Conv2d): - comfy_cast_weights = True - - class Conv3d(disable_weight_init.Conv3d): - comfy_cast_weights = True - - class GroupNorm(disable_weight_init.GroupNorm): - comfy_cast_weights = True - - class LayerNorm(disable_weight_init.LayerNorm): - comfy_cast_weights = True - - class ConvTranspose2d(disable_weight_init.ConvTranspose2d): - comfy_cast_weights = True - - class ConvTranspose1d(disable_weight_init.ConvTranspose1d): - comfy_cast_weights = True - - class RMSNorm(disable_weight_init.RMSNorm): - comfy_cast_weights = True - - class Embedding(disable_weight_init.Embedding): - comfy_cast_weights = True - - -def fp8_linear(self, input): - dtype = self.weight.dtype - if dtype not in [torch.float8_e4m3fn]: - return None - - tensor_2d = False - if len(input.shape) == 2: - tensor_2d = True - input = input.unsqueeze(1) - - input_shape = input.shape - input_dtype = input.dtype - if len(input.shape) == 3: - w, bias = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype) - w = w.t() - - scale_weight = self.scale_weight - scale_input = self.scale_input - if scale_weight is None: - scale_weight = torch.ones((), device=input.device, dtype=torch.float32) - else: - scale_weight = scale_weight.to(input.device) - - if scale_input is None: - scale_input = torch.ones((), device=input.device, dtype=torch.float32) - input = torch.clamp(input, min=-448, max=448, out=input) - input = input.reshape(-1, input_shape[2]).to(dtype).contiguous() - else: - scale_input = scale_input.to(input.device) - input = (input * (1.0 / scale_input).to(input_dtype)).reshape(-1, input_shape[2]).to(dtype).contiguous() - - if bias is not None: - o = torch._scaled_mm(input, w, out_dtype=input_dtype, bias=bias, scale_a=scale_input, scale_b=scale_weight) - else: - o = torch._scaled_mm(input, w, out_dtype=input_dtype, scale_a=scale_input, scale_b=scale_weight) - - if isinstance(o, tuple): - o = o[0] - - if tensor_2d: - return o.reshape(input_shape[0], -1) - - return o.reshape((-1, input_shape[1], self.weight.shape[0])) - - return None - -class fp8_ops(manual_cast): - class Linear(manual_cast.Linear): - def reset_parameters(self): - self.scale_weight = None - self.scale_input = None - return None - - def forward_comfy_cast_weights(self, input): - try: - out = fp8_linear(self, input) - if out is not None: - return out - except Exception as e: - logging.info("Exception during fp8 op: {}".format(e)) - - weight, bias = cast_bias_weight(self, input) - return torch.nn.functional.linear(input, weight, bias) - -def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None): - logging.info("Using scaled fp8: fp8 matrix mult: {}, scale input: {}".format(fp8_matrix_mult, scale_input)) - class scaled_fp8_op(manual_cast): - class Linear(manual_cast.Linear): - def __init__(self, *args, **kwargs): - if override_dtype is not None: - kwargs['dtype'] = override_dtype - super().__init__(*args, **kwargs) - - def reset_parameters(self): - if not hasattr(self, 'scale_weight'): - self.scale_weight = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False) - - if not scale_input: - self.scale_input = None - - if not hasattr(self, 'scale_input'): - self.scale_input = torch.nn.parameter.Parameter(data=torch.ones((), device=self.weight.device, dtype=torch.float32), requires_grad=False) - return None - - def forward_comfy_cast_weights(self, input): - if fp8_matrix_mult: - out = fp8_linear(self, input) - if out is not None: - return out - - weight, bias = cast_bias_weight(self, input) - - if weight.numel() < input.numel(): #TODO: optimize - return torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias) - else: - return torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias) - - def convert_weight(self, weight, inplace=False, **kwargs): - if inplace: - weight *= self.scale_weight.to(device=weight.device, dtype=weight.dtype) - return weight - else: - return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype) - - def set_weight(self, weight, inplace_update=False, seed=None, **kwargs): - weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed) - if inplace_update: - self.weight.data.copy_(weight) - else: - self.weight = torch.nn.Parameter(weight, requires_grad=False) - - return scaled_fp8_op - CUBLAS_IS_AVAILABLE = False try: from cublas_ops import CublasLinear @@ -570,32 +395,30 @@ if CUBLAS_IS_AVAILABLE: # TODO check if this is actually faster I call BS def forward(self, *args, **kwargs): return super().forward(*args, **kwargs) -def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None): - fp8_compute = comfy.model_management.supports_fp8_compute(load_device) - # TODO consider the different support cases - # Potentially also allow for auto-quant with dynamic input quantizers - return disable_weight_init - # if scaled_fp8 is not None: - # return disable_weight_init - # # return scaled_fp8_ops(fp8_matrix_mult=fp8_compute and fp8_optimizations, scale_input=fp8_optimizations, override_dtype=scaled_fp8) - # - # if ( - # fp8_compute and - # (fp8_optimizations or PerformanceFeature.Fp8MatrixMultiplication in args.fast) and - # not disable_fast_fp8 - # ): - # return fp8_ops - # - # if ( - # PerformanceFeature.CublasOps in args.fast and - # CUBLAS_IS_AVAILABLE and - # weight_dtype == torch.float16 and - # (compute_dtype == torch.float16 or compute_dtype is None) - # ): - # logging.info("Using cublas ops") - # return cublas_ops - # - # if compute_dtype is None or weight_dtype == compute_dtype: - # return disable_weight_init - # - # return manual_cast +op_class_list = [ + cls for name, cls in inspect.getmembers(disable_weight_init, inspect.isclass) + if cls.__module__ == disable_weight_init.__module__ and name != "__class__" +] + +def operator_factory(**factory_kwargs): + class OpSet: + pass + op_set = OpSet() + + for k, v in factory_kwargs.items(): + assert hasattr(CastWeightBiasOp, k) + + for base_class in op_class_list: + new_class = type(base_class.__name__, (base_class,), factory_kwargs) + setattr(op_set, base_class.__name__, new_class) + + return op_set + +# 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): + fp8_compute = (comfy.model_management.supports_fp8_compute(load_device) and not disable_fast_fp8) + use_dynamic_quantizer = PerformanceFeature.DynamicQuantizer in args.fast + 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) + diff --git a/comfy/quant_tensor.py b/comfy/quant_tensor.py new file mode 100644 index 000000000..cb01349cc --- /dev/null +++ b/comfy/quant_tensor.py @@ -0,0 +1,91 @@ +import torch + +""" +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] + +def dynamic_tensor_quantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): + 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() + +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 + + +def tensor_dequantizer(x: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype): + x = (x.to(dtype=scale.dtype) * scale).to(dtype=dtype) + return x + +def woq_fwd(self, x): + dq_weight = self.dequantizer(self.weight, self.scale_weight, 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) + return torch.nn.functional.linear(x, dq_weight, bias) + +def quantized_fwd(self, input): + tensor_2d = False + if len(input.shape) == 2: + tensor_2d = True + input = input.unsqueeze(1) + + input_shape = input.shape + 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) + 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(), + bias=self.bias, out_dtype=input_dtype) + if isinstance(o, tuple): + o = o[0] + if tensor_2d: + return o.reshape(input_shape[0], -1) + return o.reshape((-1, input_shape[1], self.weight.shape[0])) diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py index f8a7c2a1b..0cfbb1da0 100644 --- a/comfy/sd1_clip.py +++ b/comfy/sd1_clip.py @@ -113,9 +113,11 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder): if operations is None: scaled_fp8 = model_options.get("scaled_fp8", None) if scaled_fp8 is not None: - operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8) + # operations = comfy.ops.scaled_fp8_ops(fp8_matrix_mult=False, override_dtype=scaled_fp8) # TODO + operations = comfy.ops.pick_operations(compute_dtype=torch.float8_e4m3fn, weight_dtype=torch.float8_e4m3fn) + else: - operations = comfy.ops.manual_cast + operations = comfy.ops.operator_factory(comfy_cast_weights=True) self.operations = operations self.transformer = model_class(config, dtype, device, self.operations)