xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-08-14 23:23:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

Clean up and weight_loading dtype fix

Browse Source
This commit is contained in:
lspindler 2025-09-01 09:45:39 +02:00
parent 22261e1049
commit c265708d47
4 changed files with 64 additions and 69 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/model_base.py
View File

@ -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)

35
comfy/ops.py
View File

@ -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)

92
comfy/quant_tensor.py
View File

@ -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]

4
comfy/supported_models_base.py
View File

@ -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