"""
This file is part of ComfyUI.
Copyright (C) 2024 Stability AI
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
import torch
import logging
import contextlib
import dataclasses
import inspect
import comfy.model_management
from comfy.cli_args import args, PerformanceFeature
import comfy.float
import json
import comfy.memory_management
import comfy.pinned_memory
import comfy.rmsnorm
import comfy.utils
import comfy_aimdo.model_vbar
import comfy_aimdo.torch
def run_every_op():
if torch.compiler.is_compiling():
return
comfy.model_management.throw_exception_if_processing_interrupted()
def gqa_repeat_factor(query_heads, key_heads, value_heads):
if key_heads != value_heads:
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
if query_heads == key_heads:
return 1
if query_heads % key_heads != 0:
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
return query_heads // key_heads
def repeat_kv_for_gqa(k, v, query_heads, head_dim):
n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=head_dim)
v = v.repeat_interleave(n_rep, dim=head_dim)
return k, v
def scaled_dot_product_attention(q, k, v, *args, **kwargs):
attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
if kwargs.get("enable_gqa", False) and attn_mask is not None:
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
kwargs["enable_gqa"] = False
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
try:
if torch.cuda.is_available():
from torch.nn.attention import SDPBackend, sdpa_kernel
if "set_priority" in inspect.signature(sdpa_kernel).parameters:
SDPA_BACKEND_PRIORITY = [
SDPBackend.FLASH_ATTENTION,
SDPBackend.CUDNN_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
SDPA_CONTIGUOUS_KV_L2_SIZE = 0
if comfy.model_management.is_amd():
props = torch.cuda.get_device_properties(comfy.model_management.get_torch_device())
if props.gcnArchName.split(':')[0] == "gfx1151":
SDPA_CONTIGUOUS_KV_L2_SIZE = props.L2_cache_size
def scaled_dot_product_attention(q, k, v, *args, **kwargs):
if q.nelement() < 1024 * 128:
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
if kwargs.get("enable_gqa", False) and attn_mask is not None and not comfy.model_management.is_nvidia():
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
kwargs["enable_gqa"] = False
if SDPA_CONTIGUOUS_KV_L2_SIZE and sum(t.shape[-2] * t.shape[-1] for t in (k, v) if not t.is_contiguous()) * k.element_size() > SDPA_CONTIGUOUS_KV_L2_SIZE:
k, v = k.contiguous(), v.contiguous()
with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
if kwargs.get("enable_gqa", False) and attn_mask is not None and q.shape[-3] != k.shape[-3]:
dropout_p = args[1] if len(args) > 1 else kwargs.get("dropout_p", 0.0)
is_causal = args[2] if len(args) > 2 else kwargs.get("is_causal", False)
params = torch.backends.cuda.SDPAParams(q, k, v, attn_mask, dropout_p, is_causal, True)
supports_native_gqa = (
torch.backends.cuda.can_use_flash_attention(params)
or torch.backends.cuda.can_use_cudnn_attention(params)
or torch.backends.cuda.can_use_efficient_attention(params)
)
if not supports_native_gqa:
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
kwargs["enable_gqa"] = False
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
else:
logging.warning("Torch version too old to set sdpa backend priority.")
except (ModuleNotFoundError, TypeError):
logging.warning("Could not set sdpa backend priority.")
NVIDIA_MEMORY_CONV_BUG_WORKAROUND = False
try:
if comfy.model_management.is_nvidia():
cudnn_version = torch.backends.cudnn.version()
if (cudnn_version >= 91002 and cudnn_version < 91500) and comfy.model_management.torch_version_numeric >= (2, 9) and comfy.model_management.torch_version_numeric <= (2, 10):
NVIDIA_MEMORY_CONV_BUG_WORKAROUND = True
logging.info("working around nvidia conv3d memory bug.")
except:
pass
cast_to = comfy.model_management.cast_to
def cast_to_input(weight, input, non_blocking=False, copy=True):
return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
def materialize_meta_param(s, param_keys):
for param_key in param_keys:
param = getattr(s, param_key, None)
if param is not None and getattr(param, "is_meta", False):
setattr(s, param_key, torch.nn.Parameter(torch.zeros(param.shape, dtype=param.dtype), requires_grad=param.requires_grad))
def vbar_above_watermark(s):
if getattr(s, "_prefetch", None) is not None:
return False
vbar, alloc, size = s._v
return alloc - vbar.base_addr + size > vbar.get_watermark() * (32 << 20)
def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False):
offload_stream = None
cast_buffer = None
cast_buffer_offset = 0
if return_faulted:
fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias"))
def ensure_offload_stream(module, required_size, check_largest):
nonlocal offload_stream
nonlocal cast_buffer
if offload_stream is None:
offload_stream = comfy.model_management.get_offload_stream(device)
if offload_stream is None or not check_largest or len(comfy_modules) != 1:
return
current_size = 0 if cast_buffer is None else cast_buffer.size()
if current_size < required_size and module is comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[0]:
offload_stream = comfy.model_management.get_offload_stream(device)
cast_buffer = None
if required_size > comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT[1]:
comfy.model_management.LARGEST_AIMDO_CASTED_WEIGHT = (module, required_size)
def get_cast_buffer(buffer_size):
nonlocal offload_stream
nonlocal cast_buffer
nonlocal cast_buffer_offset
if buffer_size == 0:
return None
if offload_stream is None:
return torch.empty((buffer_size,), dtype=torch.uint8, device=device)
cast_buffer = comfy.model_management.get_aimdo_cast_buffer(offload_stream, device)
buffer = comfy_aimdo.torch.aimdo_to_tensor(cast_buffer.get(buffer_size, cast_buffer_offset), device)
cast_buffer_offset += buffer_size
return buffer
for s in comfy_modules:
signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
if return_faulted and (signature is None or not resident):
fully_faulted = False
prefetch = {
"signature": signature,
"resident": resident,
}
if resident:
s._prefetch = prefetch
continue
materialize_meta_param(s, ["weight", "bias"])
xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device) if signature is not None else None
cast_geometry = comfy.memory_management.tensors_to_geometries([ s.weight, s.bias ])
cast_dest = None
needs_cast = False
xfer_source = [ s.weight, s.bias ]
fast_disk = s._pin_state["fast_disk"]
subset = "weights-fast" if fast_disk else "weights"
pin = comfy.pinned_memory.get_pin(s, subset=subset)
if pin is None and not fast_disk:
loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded")
if loaded_pin is not None or signature is not None:
subset = "weights-loaded"
pin = loaded_pin
if pin is not None:
xfer_source = [ pin ]
for data, geometry in zip([ s.weight, s.bias ], cast_geometry):
if data is None:
continue
if data.dtype != geometry.dtype:
needs_cast = True
cast_dest = xfer_dest
xfer_dest = None
break
dest_size = comfy.memory_management.vram_aligned_size(xfer_source)
ensure_offload_stream(s, dest_size if xfer_dest is None else 0, True)
if xfer_dest is None:
xfer_dest = get_cast_buffer(dest_size)
def cast_maybe_lowvram_patch(xfer_source, xfer_dest, stream, xfer_dest2=None):
if xfer_source is not None:
if getattr(xfer_source, "is_lowvram_patch", False):
if xfer_dest is not None:
xfer_source.prepare(xfer_dest, stream, copy=True, commit=False)
xfer_source = [ xfer_dest ]
xfer_dest = xfer_dest2
xfer_dest2 = None
elif xfer_dest2 is not None:
xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False)
return
else:
return
comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2)
def handle_pin(m, pin, source, dest, subset="weights", size=None):
if pin is not None:
cast_maybe_lowvram_patch([pin], dest, offload_stream)
return
if signature is None or not fast_disk or args.high_ram:
comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
pin = comfy.pinned_memory.get_pin(m, subset=subset)
cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest)
handle_pin(s, pin, xfer_source, xfer_dest, subset=subset, size=dest_size)
for param_key in ("weight", "bias"):
lowvram_source = getattr(s, param_key + "_lowvram_function", None)
if lowvram_source is not None:
ensure_offload_stream(s, cast_buffer_offset, False)
lowvram_size = lowvram_source.memory_required()
lowvram_dest = get_cast_buffer(lowvram_size)
lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
subset = "patches-fast" if fast_disk else "patches"
pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset)
if pin is None and not fast_disk:
loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded")
if loaded_pin is not None:
subset = "patches-loaded"
pin = loaded_pin
elif signature is not None:
subset = "patches-loaded"
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size)
prefetch["xfer_dest"] = xfer_dest
prefetch["cast_dest"] = cast_dest
prefetch["cast_geometry"] = cast_geometry
prefetch["needs_cast"] = needs_cast
s._prefetch = prefetch
if return_faulted:
return offload_stream, fully_faulted
return offload_stream
def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True):
prefetch = getattr(s, "_prefetch", None)
if prefetch["resident"]:
weight = s._v_weight
bias = s._v_bias
else:
xfer_dest = prefetch["xfer_dest"]
if prefetch["needs_cast"]:
cast_dest = prefetch["cast_dest"] if prefetch["cast_dest"] is not None else torch.empty((comfy.memory_management.vram_aligned_size(prefetch["cast_geometry"]),), dtype=torch.uint8, device=device)
for pre_cast, post_cast in zip(comfy.memory_management.interpret_gathered_like([s.weight, s.bias ], xfer_dest),
comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], cast_dest)):
if post_cast is not None:
post_cast.copy_(pre_cast)
xfer_dest = cast_dest
params = comfy.memory_management.interpret_gathered_like(prefetch["cast_geometry"], xfer_dest)
weight = params[0]
bias = params[1]
if prefetch["signature"] is not None:
s._v_weight = weight
s._v_bias = bias
s._v_signature = prefetch["signature"]
def post_cast(s, param_key, x, dtype, resident, update_weight):
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
fns = getattr(s, param_key + "_function", [])
if x is None:
return None
orig = x
def to_dequant(tensor, dtype):
tensor = tensor.to(dtype=dtype)
if isinstance(tensor, QuantizedTensor):
tensor = tensor.dequantize()
return tensor
if (return_weights and orig.dtype != dtype) or len(fns) > 0:
x = to_dequant(x, dtype)
if not resident and lowvram_fn is not None:
x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
x = lowvram_fn(x)
if (want_requant and len(fns) == 0 or update_weight):
seed = comfy.utils.string_to_seed(s.seed_key)
if isinstance(orig, QuantizedTensor):
y = orig.requantize_from_float(x, scale="recalculate", stochastic_rounding=seed)
else:
y = comfy.float.stochastic_rounding(x, orig.dtype, seed=seed)
if want_requant and len(fns) == 0:
x = y
if update_weight:
orig.copy_(y)
for f in fns:
x = f(x)
return x
update_weight = prefetch["signature"] is not None
weight = post_cast(s, "weight", weight, dtype, prefetch["resident"], update_weight)
if bias is not None:
bias = post_cast(s, "bias", bias, bias_dtype, prefetch["resident"], update_weight)
if prefetch["signature"] is not None:
prefetch["resident"] = True
return (weight, bias) if return_weights else None
def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
if input is not None:
if dtype is None:
if isinstance(input, QuantizedTensor):
dtype = input.params.orig_dtype
else:
dtype = input.dtype
if bias_dtype is None:
bias_dtype = dtype
if device is None:
device = input.device
def format_return(result, offloadable):
weight, bias, offload_stream = result
return (weight, bias, offload_stream) if offloadable else (weight, bias)
non_blocking = comfy.model_management.device_supports_non_blocking(device)
if hasattr(s, "_v") and comfy.model_management.is_device_cpu(device):
materialize_meta_param(s, ["weight", "bias"])
weight = s.weight.to(dtype=dtype, copy=True)
if isinstance(weight, QuantizedTensor):
weight = weight.dequantize()
bias = s.bias.to(dtype=bias_dtype, copy=True) if s.bias is not None else None
return format_return((weight, bias, (None, None, None)), offloadable)
elif hasattr(s, "_v") and s.weight.device != device:
prefetched = hasattr(s, "_prefetch")
offload_stream = None
offload_device = None
if not prefetched:
offload_stream = cast_modules_with_vbar([s], dtype, device, bias_dtype, non_blocking)
comfy.model_management.sync_stream(device, offload_stream)
weight, bias = resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant)
if not prefetched:
if getattr(s, "_prefetch")["signature"] is not None:
offload_device = device
for param_key in ("weight", "bias"):
lowvram_fn = getattr(s, param_key + "_lowvram_function", None)
if lowvram_fn is not None:
lowvram_fn.clear_prepared()
delattr(s, "_prefetch")
return format_return((weight, bias, (offload_stream, offload_device, None)), offloadable)
if offloadable and (device != s.weight.device or
(s.bias is not None and device != s.bias.device)):
offload_stream = comfy.model_management.get_offload_stream(device)
else:
offload_stream = None
bias = None
weight = None
if offload_stream is not None and not args.cuda_malloc:
cast_buffer_size = comfy.memory_management.vram_aligned_size([ s.weight, s.bias ])
cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)
if cast_buffer is None:
offload_stream = comfy.model_management.get_offload_stream(device)
cast_buffer = comfy.model_management.get_cast_buffer(offload_stream, device, cast_buffer_size, s)
params = comfy.memory_management.interpret_gathered_like([ s.weight, s.bias ], cast_buffer)
weight = params[0]
bias = params[1]
weight_has_function = len(s.weight_function) > 0
bias_has_function = len(s.bias_function) > 0
weight = comfy.model_management.cast_to(s.weight, None, device, non_blocking=non_blocking, copy=weight_has_function, stream=offload_stream, r=weight)
if s.bias is not None:
bias = comfy.model_management.cast_to(s.bias, None, device, non_blocking=non_blocking, copy=bias_has_function, stream=offload_stream, r=bias)
comfy.model_management.sync_stream(device, offload_stream)
bias_a = bias
weight_a = weight
if s.bias is not None:
bias = bias.to(dtype=bias_dtype)
for f in s.bias_function:
bias = f(bias)
if weight_has_function or weight.dtype != dtype:
weight = weight.to(dtype=dtype)
if isinstance(weight, QuantizedTensor):
weight = weight.dequantize()
for f in s.weight_function:
weight = f(weight)
return format_return((weight, bias, (offload_stream, weight_a, bias_a)), offloadable)
def uncast_bias_weight(s, weight, bias, offload_stream):
if offload_stream is None:
return
os, weight_a, bias_a = offload_stream
device=None
if weight_a is not None and not isinstance(weight_a, torch.Tensor):
comfy_aimdo.model_vbar.vbar_unpin(s._v)
device = weight_a
if os is None:
return
if device is None:
if weight_a is not None:
device = weight_a.device
else:
if bias_a is None:
return
device = bias_a.device
os.wait_stream(comfy.model_management.current_stream(device))
class CastBiasWeightContext:
def __init__(self, *args, **kwargs):
self.slf = args[0] if len(args) else None
self.state = (None, None) if self.slf is None else cast_bias_weight(*args, **kwargs)
def __enter__(self):
result = self.state
if len(result) < 3 or result[2] is None:
self.state = self.slf = None
return result[:2]
def __exit__(self, *_args) -> None:
if self.slf is None:
return
slf, state = self.slf, self.state
self.state = self.slf = None
uncast_bias_weight(slf, *state)
class CastWeightBiasOp:
comfy_cast_weights = False
comfy_force_forward = False
weight_function = []
bias_function = []
class disable_weight_init:
@staticmethod
def _zero_init_parameter(module, name):
param = getattr(module, name)
device = None if getattr(param, "is_meta", False) else param.device
setattr(module, name, torch.nn.Parameter(torch.zeros(param.shape, device=device, dtype=param.dtype), requires_grad=False))
@staticmethod
def _lazy_load_from_state_dict(module, state_dict, prefix, local_metadata,
missing_keys, unexpected_keys, weight_shape,
bias_shape=None):
assign_to_params_buffers = local_metadata.get("assign_to_params_buffers", False)
prefix_len = len(prefix)
for k, v in state_dict.items():
key = k[prefix_len:]
if key == "weight":
if not assign_to_params_buffers:
v = v.clone()
module.weight = torch.nn.Parameter(v, requires_grad=False)
elif bias_shape is not None and key == "bias" and v is not None:
if not assign_to_params_buffers:
v = v.clone()
module.bias = torch.nn.Parameter(v, requires_grad=False)
else:
unexpected_keys.append(k)
if module.weight is None:
module.weight = torch.nn.Parameter(torch.zeros(weight_shape), requires_grad=False)
missing_keys.append(prefix + "weight")
if bias_shape is not None and module.bias is None and getattr(module, "comfy_need_lazy_init_bias", False):
module.bias = torch.nn.Parameter(torch.zeros(bias_shape), requires_grad=False)
missing_keys.append(prefix + "bias")
class Linear(torch.nn.Linear, CastWeightBiasOp):
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
super().__init__(in_features, out_features, bias, device, dtype)
return
torch.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
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
strict, missing_keys, unexpected_keys, error_msgs):
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
disable_weight_init._lazy_load_from_state_dict(
self,
state_dict,
prefix,
local_metadata,
missing_keys,
unexpected_keys,
weight_shape=(self.in_features, self.out_features),
bias_shape=(self.out_features,),
)
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return torch.nn.functional.linear(input, weight, bias)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class Conv1d(torch.nn.Conv1d, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return self._conv_forward(input, weight, bias)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class Conv2d(torch.nn.Conv2d, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return self._conv_forward(input, weight, bias)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class Conv3d(torch.nn.Conv3d, CastWeightBiasOp):
def reset_parameters(self):
return None
def _conv_forward(self, input, weight, bias, autopad=None, *args, **kwargs):
if autopad == "causal_zero":
weight = weight[:, :, -input.shape[2]:, :, :]
if NVIDIA_MEMORY_CONV_BUG_WORKAROUND and weight.dtype in (torch.float16, torch.bfloat16):
out = torch.cudnn_convolution(input, weight, self.padding, self.stride, self.dilation, self.groups, benchmark=False, deterministic=False, allow_tf32=True)
if bias is not None:
out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
return out
else:
return super()._conv_forward(input, weight, bias, *args, **kwargs)
def forward_comfy_cast_weights(self, input, autopad=None):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return self._conv_forward(input, weight, bias, autopad=autopad)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0 or "autopad" in kwargs:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class GroupNorm(torch.nn.GroupNorm, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class BatchNorm2d(torch.nn.BatchNorm2d, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class LayerNorm(torch.nn.LayerNorm, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class RMSNorm(torch.nn.RMSNorm, CastWeightBiasOp):
def reset_parameters(self):
self.bias = None
return None
def cast_weight(self, input):
return CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True)
def forward_comfy_cast_weights(self, input):
with self.cast_weight(input) as (weight, bias):
return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class ConvTranspose2d(torch.nn.ConvTranspose2d, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input, output_size=None):
num_spatial_dims = 2
output_padding = self._output_padding(
input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation)
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return torch.nn.functional.conv_transpose2d(
input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class ConvTranspose1d(torch.nn.ConvTranspose1d, CastWeightBiasOp):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input, output_size=None):
num_spatial_dims = 1
output_padding = self._output_padding(
input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation)
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return torch.nn.functional.conv_transpose1d(
input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
class Embedding(torch.nn.Embedding, CastWeightBiasOp):
def __init__(self, num_embeddings, embedding_dim, padding_idx=None, max_norm=None,
norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None,
_freeze=False, device=None, dtype=None):
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
super().__init__(num_embeddings, embedding_dim, padding_idx, max_norm,
norm_type, scale_grad_by_freq, sparse, _weight,
_freeze, device, dtype)
return
torch.nn.Module.__init__(self)
self.num_embeddings = num_embeddings
self.embedding_dim = embedding_dim
self.padding_idx = padding_idx
self.max_norm = max_norm
self.norm_type = norm_type
self.scale_grad_by_freq = scale_grad_by_freq
self.sparse = sparse
embedding_dtype = dtype if dtype is not None else torch.get_default_dtype()
self.weight = torch.nn.Parameter(
torch.empty((num_embeddings, embedding_dim), device="meta", dtype=embedding_dtype),
requires_grad=False,
)
self.bias = None
self.weight_comfy_model_dtype = dtype
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
strict, missing_keys, unexpected_keys, error_msgs):
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
disable_weight_init._lazy_load_from_state_dict(
self,
state_dict,
prefix,
local_metadata,
missing_keys,
unexpected_keys,
weight_shape=(self.num_embeddings, self.embedding_dim),
)
def reset_parameters(self):
self.bias = None
return None
def forward_comfy_cast_weights(self, input, out_dtype=None):
output_dtype = out_dtype
if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
out_dtype = None
with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):
return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
def host_rows(self, input, out_dtype=None):
x = self.weight[input.reshape(-1).cpu()].to(input.device)
x = x.view(*input.shape, x.shape[-1])
return x if out_dtype is None else x.to(dtype=out_dtype)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
if "out_dtype" in kwargs:
kwargs.pop("out_dtype")
return super().forward(*args, **kwargs)
@classmethod
def conv_nd(s, dims, *args, **kwargs):
if dims == 2:
return s.Conv2d(*args, **kwargs)
elif dims == 3:
return s.Conv3d(*args, **kwargs)
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 BatchNorm2d(disable_weight_init.BatchNorm2d):
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):
"""
Legacy FP8 linear function for backward compatibility.
Uses QuantizedTensor subclass for dispatch.
"""
dtype = self.weight.dtype
if dtype not in [torch.float8_e4m3fn]:
return None
input_dtype = input.dtype
input_shape = input.shape
tensor_3d = input.ndim == 3
if tensor_3d:
input = input.reshape(-1, input_shape[2])
if input.ndim != 2:
return None
lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
input = torch.clamp(input, min=-448, max=448, out=input)
input_fp8 = input.to(dtype).contiguous()
layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
w_shape = tuple(w.shape)
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=w_shape)
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
if tensor_3d:
o = o.reshape((input_shape[0], input_shape[1], w_shape[0]))
return o
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):
if len(self.weight_function) == 0 and len(self.bias_function) == 0:
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))
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return torch.nn.functional.linear(input, weight, bias)
CUBLAS_IS_AVAILABLE = False
try:
from cublas_ops import CublasLinear, cublas_half_matmul
CUBLAS_IS_AVAILABLE = True
except ImportError:
pass
if CUBLAS_IS_AVAILABLE:
class cublas_ops(manual_cast):
class Linear(CublasLinear, manual_cast.Linear):
def reset_parameters(self):
return None
def forward_comfy_cast_weights(self, input):
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
def forward(self, *args, **kwargs):
run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
return super().forward(*args, **kwargs)
from . import quant_ops
from .quant_ops import (
QuantizedTensor,
QUANT_ALGOS,
TensorCoreFP8Layout,
TensorWiseINT8Layout,
get_layout_class,
)
def _swiglu_eager(x):
gate, up = x.chunk(2, dim=-1)
return torch.nn.functional.silu(gate).mul_(up)
INPUT_ACT_EAGER = {
"gelu_tanh": lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
"swiglu": _swiglu_eager,
}
@contextlib.contextmanager
def _input_act_weight(x, input_act, act_weight, act_eps):
if input_act == "rms_norm" and isinstance(act_weight, torch.nn.Module):
with act_weight.cast_weight(x) as (weight, _):
eps = act_weight.eps
if eps is None:
eps = torch.finfo(torch.float64 if x.dtype == torch.float64 else torch.float32).eps
yield weight, eps
else:
yield act_weight, act_eps
def _eager_input_act(x, input_act, act_weight=None, act_eps=0.0):
if input_act is None:
return x
if input_act == "rms_norm":
with _input_act_weight(x, input_act, act_weight, act_eps) as (weight, eps):
return comfy.rmsnorm.rms_norm(x, weight, eps)
return INPUT_ACT_EAGER[input_act](x)
def _linear_residual(out, residual, residual_scale):
return out if residual is None else torch.addcmul(residual, out, residual_scale)
def _fp16_linear_wanted(x):
"""kitchen's fp16-accumulate GEMM replaces a plain linear when the user opted into
fp16 accumulation and the activation is fp16 on CUDA; weights come through cast_bias_weight."""
return (getattr(torch.backends.cuda.matmul, "allow_fp16_accumulation", False)
and x.dtype == torch.float16 and x.is_cuda)
def linear_input_act(linear, x, input_act, act_weight=None, act_eps=0.0,
residual=None, residual_scale=None):
run_every_op()
weight = linear.weight
quantized = isinstance(weight, QuantizedTensor)
if (comfy.model_management.in_training or getattr(linear, "comfy_force_forward", False) or getattr(linear, "_full_precision_mm", False)
or not ((quantized and weight._layout_cls == "TensorWiseINT8Layout"
and not getattr(weight._params, "transposed", False))
or (not quantized and _fp16_linear_wanted(x)))):
out = linear(_eager_input_act(x, input_act, act_weight, act_eps))
return _linear_residual(out, residual, residual_scale)
with CastBiasWeightContext(linear, x, offloadable=True,
compute_dtype=x.dtype if quantized else None,
want_requant=quantized) as (weight, bias):
return linear_input_act_(x, weight, bias, input_act, act_weight, act_eps,
residual, residual_scale, fp16_accumulation=not quantized)
def linear_input_act_(x, weight, bias, input_act=None, act_weight=None, act_eps=0.0,
residual=None, residual_scale=None, *, fp16_accumulation=True):
"""``linear(act(x))``, with ``act`` folded into an INT8 activation quantizer.
An INT8 linear quantizes its input anyway, so an elementwise activation can
ride along inside that kernel instead of writing a full-size intermediate to
HBM and reading it straight back. Worth it for an MLP's down-projection,
where the intermediate is several times the hidden size, and for a pre-norm
block's ``linear(rms_norm(x))`` ("rms_norm", which reads the norm's weight
and eps from act_weight/act_eps).
With ``residual``/``residual_scale`` the result is the pre-norm block's
addcmul, ``residual + residual_scale * linear(act(x))``, fused into the
INT8 GEMM epilogue where supported.
Linear weights are already prepared. An RMSNorm module in act_weight owns
its norm-weight cast context; a tensor uses the supplied act_eps directly.
"""
if input_act is None and residual is None:
return torch.nn.functional.linear(x, weight, bias)
if (comfy.model_management.in_training
or not isinstance(weight, QuantizedTensor)
or weight._layout_cls != "TensorWiseINT8Layout"
or getattr(weight._params, "transposed", False)):
x = _eager_input_act(x, input_act, act_weight, act_eps)
if (not comfy.model_management.in_training
and not isinstance(weight, QuantizedTensor)
and fp16_accumulation and _fp16_linear_wanted(x)):
return quant_ops.ck.fp16_linear(x, weight, bias, residual=residual, residual_scale=residual_scale)
return _linear_residual(torch.nn.functional.linear(x, weight, bias), residual, residual_scale)
qdata, scale = TensorWiseINT8Layout.get_plain_tensors(weight)
with _input_act_weight(x, input_act, act_weight, act_eps) as (act_weight, act_eps):
return quant_ops.ck.int8_linear(
x, qdata, scale, bias, x.dtype,
convrot=getattr(weight._params, "convrot", False),
convrot_groupsize=getattr(weight._params, "convrot_groupsize", 256),
input_act=input_act,
input_act_weight=act_weight,
input_act_eps=act_eps,
residual=residual,
residual_scale=residual_scale,
)
_GROUPED_INT8_FORMATS = {"asym_w4a8_int8": 4, "w6a8_int8": 6}
class QuantLinearFunc(torch.autograd.Function):
"""Custom autograd function for quantized linear: quantized forward, optionally FP8 backward.
When training_fp8_bwd is enabled:
- Forward: quantize input per layout (FP8/NVFP4), use quantized matmul
- Backward: all matmuls use FP8 tensor cores via torch.mm dispatch
- Cached input is FP8 (half the memory of bf16)
When training_fp8_bwd is disabled:
- Forward: quantize input per layout, use quantized matmul
- Backward: dequantize weight to compute_dtype, use standard matmul
"""
@staticmethod
def forward(ctx, input_float, weight, bias, layout_type, input_scale, compute_dtype):
input_shape = input_float.shape
inp = input_float.detach().flatten(0, -2)
if layout_type is not None:
q_input = QuantizedTensor.from_float(inp, layout_type, scale=input_scale)
else:
q_input = inp
w = weight.detach() if weight.requires_grad else weight
b = bias.detach() if bias is not None and bias.requires_grad else bias
output = torch.nn.functional.linear(q_input, w, b)
if len(input_shape) > 2:
output = output.unflatten(0, input_shape[:-1])
ctx.input_shape = input_shape
ctx.has_bias = bias is not None
ctx.compute_dtype = compute_dtype
ctx.weight_requires_grad = weight.requires_grad
ctx.fp8_bwd = comfy.model_management.training_fp8_bwd
if ctx.fp8_bwd:
if isinstance(q_input, QuantizedTensor) and layout_type.startswith('TensorCoreFP8'):
ctx.q_input = q_input
else:
ctx.q_input = QuantizedTensor.from_float(inp, "TensorCoreFP8E4M3Layout")
ctx.save_for_backward(weight)
else:
ctx.q_input = None
ctx.save_for_backward(input_float, weight)
return output
@staticmethod
@torch.autograd.function.once_differentiable
def backward(ctx, grad_output):
compute_dtype = ctx.compute_dtype
grad_2d = grad_output.flatten(0, -2).to(compute_dtype)
if ctx.fp8_bwd:
weight, = ctx.saved_tensors
grad_mm = QuantizedTensor.from_float(grad_2d, "TensorCoreFP8E5M2Layout")
if isinstance(weight, QuantizedTensor) and weight._layout_cls.startswith("TensorCoreFP8"):
weight_mm = weight
elif isinstance(weight, QuantizedTensor):
weight_mm = QuantizedTensor.from_float(weight.dequantize().to(compute_dtype), "TensorCoreFP8E4M3Layout")
else:
weight_mm = QuantizedTensor.from_float(weight.to(compute_dtype), "TensorCoreFP8E4M3Layout")
input_mm = ctx.q_input
else:
input_float, weight = ctx.saved_tensors
grad_mm = grad_2d
if isinstance(weight, QuantizedTensor):
weight_mm = weight.dequantize().to(compute_dtype)
else:
weight_mm = weight.to(compute_dtype)
input_mm = input_float.flatten(0, -2).to(compute_dtype) if ctx.weight_requires_grad else None
grad_input = torch.mm(grad_mm, weight_mm)
if len(ctx.input_shape) > 2:
grad_input = grad_input.unflatten(0, ctx.input_shape[:-1])
grad_weight = None
if ctx.weight_requires_grad:
grad_weight = torch.mm(grad_mm.t(), input_mm)
grad_bias = None
if ctx.has_bias:
grad_bias = grad_2d.sum(dim=0)
return grad_input, grad_weight, grad_bias, None, None, None
def _quantized_apply(module, fn, recurse=True):
"""Re-wrap Parameters after fn so .to()/.cuda() propagate through QuantizedTensor weights."""
if recurse:
for child in module.children():
child._apply(fn)
for key, param in module._parameters.items():
if param is None:
continue
p = fn(param)
if (not torch.is_inference_mode_enabled()) and p.is_inference():
p = p.clone()
module.register_parameter(key, torch.nn.Parameter(p, requires_grad=False))
for key, buf in module._buffers.items():
if buf is not None:
module._buffers[key] = fn(buf)
return module
def _load_quantized_module(module, super_load, state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs, load_extra_params=False):
"""Shared _load_from_state_dict body for quantized-weight modules.
Pops weight (+ scales, +/- extras), populates module.weight as a Parameter
or Parameter-wrapped QuantizedTensor, then calls super_load and strips
consumed keys from missing_keys. Reads compute_dtype from factory_kwargs
and disabled formats from module._disabled_formats.
"""
device = module.factory_kwargs["device"]
compute_dtype = module.factory_kwargs["dtype"]
disabled_formats = module._disabled_formats
layer_name = prefix.rstrip('.')
weight = state_dict.pop(f"{prefix}weight", None)
if weight is None:
logging.warning(f"Missing weight for layer {layer_name}")
module.weight = None
return
manually_loaded_keys = [f"{prefix}weight"]
def pop_scale(name, dtype=None):
key = f"{prefix}{name}"
v = state_dict.pop(key, None)
if v is not None:
v = v.to(device=device)
if dtype is not None:
v = v.view(dtype=dtype)
manually_loaded_keys.append(key)
return v
layer_conf = state_dict.pop(f"{prefix}comfy_quant", None)
if layer_conf is not None:
layer_conf = json.loads(layer_conf.numpy().tobytes())
if layer_conf is None:
module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False)
else:
module.quant_format = layer_conf.get("format", None)
module._full_precision_mm_config = layer_conf.get("full_precision_matrix_mult", False)
if not module._full_precision_mm:
module._full_precision_mm = module._full_precision_mm_config
if module.quant_format in disabled_formats:
module._full_precision_mm = True
if module.quant_format is None:
raise ValueError(f"Unknown quantization format for layer {layer_name}")
qconfig = QUANT_ALGOS[module.quant_format]
module.layout_type = qconfig["comfy_tensor_layout"]
layout_cls = get_layout_class(module.layout_type)
if module.quant_format in ("float8_e4m3fn", "float8_e5m2"):
scales = {"scale": pop_scale("weight_scale")}
elif module.quant_format == "mxfp8":
bs = pop_scale("weight_scale", torch.float8_e8m0fnu)
if bs is None:
raise ValueError(f"Missing MXFP8 block scales for layer {layer_name}")
scales = {"scale": bs}
elif module.quant_format == "nvfp4":
ts = pop_scale("weight_scale_2")
bs = pop_scale("weight_scale", torch.float8_e4m3fn)
if ts is None or bs is None:
raise ValueError(f"Missing NVFP4 scales for layer {layer_name}")
scales = {"scale": ts, "block_scale": bs}
elif module.quant_format == "int8_tensorwise":
scale = pop_scale("weight_scale")
if scale is None:
raise ValueError(f"Missing INT8 weight scale for layer {layer_name}")
scales = {"scale": scale}
params_conf = layer_conf.get("params", {})
if not isinstance(params_conf, dict):
params_conf = {}
if layer_conf.get("convrot", params_conf.get("convrot", False)):
scales["convrot"] = True
scales["convrot_groupsize"] = int(
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
)
elif module.quant_format == "convrot_w4a4":
scale = pop_scale("weight_scale")
if scale is None:
raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}")
params_conf = layer_conf.get("params", {})
if not isinstance(params_conf, dict):
params_conf = {}
scales = {
"scale": scale,
"convrot_groupsize": int(
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
),
"quant_group_size": 64,
"linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
}
elif module.quant_format in _GROUPED_INT8_FORMATS:
bits = _GROUPED_INT8_FORMATS[module.quant_format]
if weight.shape[1] * 8 != module._orig_shape[1] * bits:
raise ValueError(
f"{module.quant_format} layer {layer_name}: packed weight width {weight.shape[1]} "
f"does not match K={module._orig_shape[1]} at {bits} bits"
)
scale = pop_scale("weight_s_rel")
if scale is None:
raise ValueError(f"Missing {module.quant_format} group scale (weight_s_rel) for layer {layer_name}")
if scale.dtype == torch.uint8:
scale = scale.view(torch.float8_e4m3fn)
params_conf = layer_conf.get("params", {})
if not isinstance(params_conf, dict):
params_conf = {}
scales = {
"scale": scale,
"s_channel": pop_scale("weight_s_channel"),
"codebook": pop_scale("weight_codebook"),
"group_size": int(layer_conf.get("group_size", params_conf.get("group_size", 16))),
"convrot_groupsize": int(
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
),
}
else:
raise ValueError(f"Unsupported quantization format: {module.quant_format}")
params = layout_cls.Params(**scales, orig_dtype=compute_dtype, orig_shape=module._orig_shape)
module.weight = torch.nn.Parameter(
QuantizedTensor(weight.to(device=device, dtype=qconfig["storage_t"]), module.layout_type, params),
requires_grad=False,
)
if load_extra_params:
for param_name in qconfig["parameters"]:
if param_name in {"weight_scale", "weight_scale_2"}:
continue
param_key = f"{prefix}{param_name}"
_v = state_dict.pop(param_key, None)
if _v is None:
continue
module.register_parameter(param_name, torch.nn.Parameter(_v.to(device=device), requires_grad=False))
manually_loaded_keys.append(param_key)
super_load(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
for key in manually_loaded_keys:
if key in missing_keys:
missing_keys.remove(key)
def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extra_quant_params=()):
"""Shared state_dict body. extra_quant_conf merges into the comfy_quant JSON;
extra_quant_params names attributes written as additional top-level keys."""
if not hasattr(module, 'weight'):
logging.warning(f"Warning: state dict on uninitialized op {prefix}")
return sd
bias = getattr(module, 'bias', None)
if bias is not None:
sd[f"{prefix}bias"] = bias
if module.weight is None:
return sd
if isinstance(module.weight, QuantizedTensor):
sd.update(module.weight.state_dict(f"{prefix}weight"))
quant_conf = {"format": module.quant_format}
if getattr(module, '_full_precision_mm_config', False):
quant_conf["full_precision_matrix_mult"] = True
params = getattr(module.weight, "_params", None)
if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
quant_conf["convrot"] = True
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
elif module.quant_format == "convrot_w4a4":
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
linear_dtype = getattr(params, "linear_dtype", "int4")
if linear_dtype != "int4":
quant_conf["linear_dtype"] = linear_dtype
elif module.quant_format in _GROUPED_INT8_FORMATS:
quant_conf["group_size"] = getattr(params, "group_size", 16)
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
if extra_quant_conf:
quant_conf.update(extra_quant_conf)
sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
for name in extra_quant_params:
value = getattr(module, name, None)
if value is not None:
sd[f"{prefix}{name}"] = value
else:
sd[f"{prefix}weight"] = module.weight
return sd
class MixedPrecisionOp(CastWeightBiasOp):
quant_format = None
def can_use_quantized_matmul(self, disabled_formats):
return (self.quant_format in QUANT_ALGOS
and not self._full_precision_mm_config
and self.quant_format not in self._disabled_formats
and self.quant_format not in disabled_formats)
def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_precision_mm=False, disabled=[]):
class MixedPrecisionOps(manual_cast):
_quant_config = quant_config
_compute_dtype = compute_dtype
_full_precision_mm = full_precision_mm
_disabled = disabled
class Linear(torch.nn.Module, MixedPrecisionOp):
_disabled_formats = disabled
def __init__(self, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None):
super().__init__()
self.factory_kwargs = {"device": device, "dtype": MixedPrecisionOps._compute_dtype}
self.in_features = in_features
self.out_features = out_features
self._orig_shape = (out_features, in_features)
if bias:
self.bias = torch.nn.Parameter(torch.empty(out_features, **self.factory_kwargs))
else:
self.register_parameter("bias", None)
self.tensor_class = None
self._full_precision_mm = MixedPrecisionOps._full_precision_mm
self._full_precision_mm_config = False
def reset_parameters(self):
return None
def _load_from_state_dict(self, *args):
_load_quantized_module(self, super()._load_from_state_dict, *args, load_extra_params=True)
def state_dict(self, *args, destination=None, prefix="", **kwargs):
sd = destination if destination is not None else {}
return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale", "pre_quant_scale"))
def _forward(self, input, weight, bias):
return torch.nn.functional.linear(input, weight, bias)
def forward_comfy_cast_weights(
self,
input,
compute_dtype=None,
want_requant=False,
weight_only_quant=False,
):
if not weight_only_quant:
with CastBiasWeightContext(
self,
input,
offloadable=True,
compute_dtype=compute_dtype,
want_requant=want_requant,
) as (weight, bias):
if self._full_precision_mm and isinstance(weight, QuantizedTensor):
weight = weight.dequantize()
return self._forward(input, weight, bias)
with CastBiasWeightContext(
self,
input=None,
dtype=self.weight.dtype,
device=input.device,
bias_dtype=input.dtype,
offloadable=True,
compute_dtype=compute_dtype,
want_requant=True,
) as (weight, bias):
weight = weight.to(dtype=input.dtype)
return self._forward(input, weight, bias)
def forward(self, input, *args, **kwargs):
run_every_op()
pre_quant_scale = getattr(self, 'pre_quant_scale', None)
if pre_quant_scale is not None:
input = input * comfy.model_management.cast_to_device(pre_quant_scale, input.device, input.dtype)
input_shape = input.shape
reshaped_nd = False
compute_dtype = input.dtype
_use_quantized = (
getattr(self, 'layout_type', None) is not None and
not isinstance(input, QuantizedTensor) and not self._full_precision_mm and
not getattr(self, 'comfy_force_cast_weights', False) and
len(self.weight_function) == 0 and len(self.bias_function) == 0
)
quantize_input = QUANT_ALGOS.get(getattr(self, 'quant_format', None), {}).get("quantize_input", True)
if (input.requires_grad and _use_quantized and quantize_input):
with CastBiasWeightContext(
self,
input,
offloadable=True,
compute_dtype=compute_dtype,
want_requant=True
) as (weight, bias):
scale = getattr(self, 'input_scale', None)
if scale is not None:
scale = comfy.model_management.cast_to_device(scale, input.device, None)
return QuantLinearFunc.apply(
input, weight, bias, self.layout_type, scale, compute_dtype
)
if _use_quantized and quantize_input:
input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input
if input_reshaped.ndim == 2:
reshaped_nd = input.ndim >= 3
scale = getattr(self, 'input_scale', None)
if scale is not None:
scale = comfy.model_management.cast_to_device(scale, input.device, None)
input = QuantizedTensor.from_float(input_reshaped, self.layout_type, scale=scale)
weight_only_quant = _use_quantized and not quantize_input and isinstance(self.weight, QuantizedTensor)
output = self.forward_comfy_cast_weights(
input,
compute_dtype,
want_requant=isinstance(input, QuantizedTensor),
weight_only_quant=weight_only_quant,
)
if reshaped_nd:
output = output.reshape((*input_shape[:-1], self.weight.shape[0]))
return output
def convert_weight(self, weight, inplace=False, **kwargs):
if isinstance(weight, QuantizedTensor):
return weight.dequantize()
else:
return weight
def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
if getattr(self, 'layout_type', None) is not None:
weight = self.weight.requantize_from_float(weight, scale="recalculate", stochastic_rounding=seed, inplace_ops=True).to(self.weight.dtype)
else:
weight = weight.to(self.weight.dtype)
if return_weight:
return weight
assert inplace_update is False
self.weight = torch.nn.Parameter(weight, requires_grad=False)
def _apply(self, fn, recurse=True):
return _quantized_apply(self, fn, recurse)
class MoEExperts(torch.nn.Module, MixedPrecisionOp):
"""Container for E quantized expert weights, indexed via expert_weight(i).
The bank lives on self.weight as a single 3D tensor — either a
compute_dtype Parameter or a Parameter wrapping a QuantizedTensor
with leading expert dim.
State-dict layout matches mixed_precision_ops.Linear with a leading
expert dim:
{prefix}.weight quant data (storage_t), leading dim = E
{prefix}.weight_scale block / per-tensor scale
{prefix}.weight_scale_2 [E] or scalar NVFP4 only
{prefix}.bias [E, out_features] optional, compute_dtype
{prefix}.comfy_quant json -> {{"format": "...", "num_experts": E}}
Without comfy_quant the weight loads as a plain compute_dtype 3D Parameter [E, out, in].
"""
_disabled_formats = disabled
def __init__(self, num_experts: int, in_features: int, out_features: int, bias: bool = True, device=None, dtype=None):
super().__init__()
self.num_experts = num_experts
self.in_features = in_features
self.out_features = out_features
self._orig_shape = (num_experts, out_features, in_features)
self.factory_kwargs = {"device": device, "dtype": MixedPrecisionOps._compute_dtype}
if bias:
self.bias = torch.nn.Parameter(torch.empty(num_experts, out_features, **self.factory_kwargs))
else:
self.register_parameter("bias", None)
self.weight = None
self.quant_format = None
self.layout_type = None
self._full_precision_mm = MixedPrecisionOps._full_precision_mm
self._full_precision_mm_config = False
self._resident_bank = None
def reset_parameters(self):
return None
def _apply(self, fn, recurse=True):
return _quantized_apply(self, fn, recurse)
def _load_from_state_dict(self, state_dict, prefix, *args):
layer_conf = state_dict.get(f"{prefix}comfy_quant", None)
quant_format = json.loads(layer_conf.numpy().tobytes()).get("format") if layer_conf is not None else None
scale = state_dict.get(f"{prefix}weight_scale", None)
if quant_format in _GROUPED_INT8_FORMATS or (quant_format == "int8_tensorwise" and scale is not None and scale.ndim == 3):
for name in ("weight", "weight_scale", "weight_s_rel", "weight_s_channel"):
t = state_dict.get(f"{prefix}{name}", None)
if t is not None and t.ndim > 1 and t.shape[0] == self.num_experts:
state_dict[f"{prefix}{name}"] = t.reshape(self.num_experts * t.shape[1], *t.shape[2:])
self._orig_shape = (self.num_experts * self.out_features, self.in_features)
_load_quantized_module(self, super()._load_from_state_dict, state_dict, prefix, *args, load_extra_params=False)
def expert_weight(self, i: int):
"""Expert i's weight (Tensor or per-expert QuantizedTensor view)."""
if isinstance(self.weight, QuantizedTensor):
return self._expert_qt_from(self.weight, i)
return self.weight[i]
def _cast_bank(self, input):
if isinstance(self.weight, QuantizedTensor):
return CastBiasWeightContext(self, input=None, dtype=self.weight.dtype, device=input.device, bias_dtype=input.dtype, offloadable=True)
return CastBiasWeightContext(self, input, offloadable=True)
def _dequantize_bank(self, weight, dtype):
flat = QuantizedTensor(weight._qdata, weight._layout_cls, dataclasses.replace(weight._params, orig_dtype=dtype))
out = torch.empty((self.num_experts, self.out_features, self.in_features), dtype=dtype, device=weight.device)
step = max(1, (256 << 20) // (self.out_features * self.in_features * out.element_size()))
for first in range(0, self.num_experts, step):
last = min(first + step, self.num_experts)
out[first:last] = self._bank_rows(flat, first, last).dequantize().view(last - first, self.out_features, self.in_features)
return out
def _bank_rows(self, weight: QuantizedTensor, first: int, last: int) -> QuantizedTensor:
"""Experts [first, last) of a flat [E * out, in] bank as one QuantizedTensor."""
params = weight._params
rows = slice(first * self.out_features, last * self.out_features)
per_row = {f.name: getattr(params, f.name)[rows] for f in dataclasses.fields(params) if torch.is_tensor(getattr(params, f.name)) and getattr(params, f.name).ndim >= 1 and getattr(params, f.name).shape[0] == weight._qdata.shape[0]}
return QuantizedTensor(weight._qdata[rows], weight._layout_cls, dataclasses.replace(params, orig_shape=((last - first) * self.out_features, self.in_features), **per_row))
@contextlib.contextmanager
def bank_resident(self, input):
"""Cast the whole bank once; expert_linear inside reuses the cast.
Not re-entrant — do not nest calls on the same instance.
"""
with self._cast_bank(input) as (weight, bias):
if self._full_precision_mm and isinstance(weight, QuantizedTensor) and weight._qdata.ndim == 2:
weight = self._dequantize_bank(weight, input.dtype)
self._resident_bank = (weight, bias)
try:
yield self
finally:
self._resident_bank = None
def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
"""Linear against expert i's weight (with optional bias)."""
resident = getattr(self, "_resident_bank", None)
if resident is not None:
weight, bias = resident
return self._expert_linear_impl(input, weight, bias, i)
with self._cast_bank(input) as (weight, bias):
return self._expert_linear_impl(input, weight, bias, i)
def _expert_linear_impl(self, input, weight, bias, i):
if isinstance(weight, QuantizedTensor):
qw = self._expert_qt_from(weight, i)
else:
qw = cast_to_input(weight[i], input, copy=False)
b = cast_to_input(bias[i], input, copy=False) if bias is not None else None
if isinstance(qw, QuantizedTensor):
use_fast = (
not self._full_precision_mm
and qw.layout_cls.supports_fast_matmul()
and input.dim() == 2
)
if use_fast:
qin = QuantizedTensor.from_float(input, self.layout_type) if QUANT_ALGOS[self.quant_format].get("quantize_input", True) else input
return torch.nn.functional.linear(qin, qw, b)
qw = cast_to_input(qw.dequantize(), input, copy=False)
return torch.nn.functional.linear(input, qw, b)
def _expert_qt_from(self, weight: QuantizedTensor, i: int) -> QuantizedTensor:
"""Build a per-expert QuantizedTensor by indexing into a resident bank."""
if weight._qdata.ndim == 2:
return self._bank_rows(weight, i, i + 1)
params = weight._params
kwargs = {
"scale": params.scale[i] if params.scale.dim() else params.scale,
"orig_dtype": params.orig_dtype,
"orig_shape": (self.out_features, self.in_features),
}
if hasattr(params, "block_scale"):
kwargs["block_scale"] = params.block_scale[i]
if hasattr(params, "quant_group_size"):
kwargs["quant_group_size"] = params.quant_group_size
if hasattr(params, "convrot"):
kwargs["convrot"] = params.convrot
if hasattr(params, "convrot_groupsize"):
kwargs["convrot_groupsize"] = params.convrot_groupsize
if hasattr(params, "linear_dtype"):
kwargs["linear_dtype"] = params.linear_dtype
return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs))
def state_dict(self, *args, destination=None, prefix="", **kwargs):
sd = destination if destination is not None else {}
return _quantized_weight_state_dict(self, sd, prefix, extra_quant_conf={"num_experts": self.num_experts})
class Embedding(manual_cast.Embedding):
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
weight_key = f"{prefix}weight"
layer_conf = state_dict.pop(f"{prefix}comfy_quant", None)
if layer_conf is not None:
layer_conf = json.loads(layer_conf.numpy().tobytes())
quant_format = layer_conf.get("format") if layer_conf is not None else None
manually_loaded_keys = []
if quant_format in ("float8_e4m3fn", "float8_e5m2", "int8_tensorwise") and weight_key in state_dict:
self.quant_format = quant_format
qconfig = QUANT_ALGOS[quant_format]
self.layout_type = qconfig["comfy_tensor_layout"]
layout_cls = get_layout_class(self.layout_type)
weight = state_dict.pop(weight_key)
manually_loaded_keys.append(weight_key)
scale_key = f"{prefix}weight_scale"
scale = state_dict.pop(scale_key, None)
if scale is not None:
scale = scale.float()
manually_loaded_keys.append(scale_key)
extra = {}
if quant_format == "int8_tensorwise" and layer_conf.get("convrot", False):
extra["convrot"] = True
extra["convrot_groupsize"] = int(layer_conf.get("convrot_groupsize", 256))
params = layout_cls.Params(
scale=scale if scale is not None else torch.ones((), dtype=torch.float32),
orig_dtype=MixedPrecisionOps._compute_dtype,
orig_shape=(self.num_embeddings, self.embedding_dim),
**extra,
)
self.weight = torch.nn.Parameter(
QuantizedTensor(weight.to(dtype=qconfig["storage_t"]), qconfig["comfy_tensor_layout"], params),
requires_grad=False)
elif layer_conf is not None:
state_dict[f"{prefix}comfy_quant"] = torch.tensor(
list(json.dumps(layer_conf).encode('utf-8')), dtype=torch.uint8)
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
for k in manually_loaded_keys:
if k in missing_keys:
missing_keys.remove(k)
def state_dict(self, *args, destination=None, prefix="", **kwargs):
sd = destination if destination is not None else {}
return _quantized_weight_state_dict(self, sd, prefix)
def host_rows(self, input, out_dtype=None):
weight = self.weight
if not isinstance(weight, QuantizedTensor):
return super().host_rows(input, out_dtype=out_dtype)
if self.quant_format != "int8_tensorwise":
return self.forward_comfy_cast_weights(input, out_dtype=out_dtype)
idx = input.reshape(-1).cpu()
params = weight._params
scale = params.scale[idx] if params.scale.dim() >= 2 else params.scale
params = dataclasses.replace(params, scale=scale.to(input.device))
rows = weight._qdata[idx].to(input.device)
x = get_layout_class(self.layout_type).dequantize_embedding(rows, params, torch.arange(idx.numel(), device=input.device))
x = x.view(*input.shape, x.shape[-1])
return x if out_dtype is None else x.to(dtype=out_dtype)
def forward_comfy_cast_weights(self, input, out_dtype=None):
weight = self.weight
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
if isinstance(qdata, QuantizedTensor):
params = qdata._params
scale = params.scale
qdata = qdata._qdata
else:
params = weight._params
scale = None
if self.quant_format == "int8_tensorwise":
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
return x if out_dtype is None else x.to(dtype=out_dtype)
x = torch.nn.functional.embedding(
input, qdata, self.padding_idx, self.max_norm,
self.norm_type, self.scale_grad_by_freq, self.sparse)
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
x = x.to(dtype=target_dtype)
if scale is not None:
x = x * scale.to(dtype=target_dtype)
return x
return super().forward_comfy_cast_weights(input, out_dtype=out_dtype)
return MixedPrecisionOps
def get_disabled_quant_formats(device=None):
disabled = set()
if not comfy.model_management.supports_nvfp4_compute(device):
disabled.add("nvfp4")
if not comfy.model_management.supports_mxfp8_compute(device):
disabled.add("mxfp8")
if not comfy.model_management.supports_fp8_compute(device):
disabled.add("float8_e4m3fn")
disabled.add("float8_e5m2")
if not comfy.model_management.supports_int8_compute(device):
disabled.add("int8_tensorwise")
disabled.add("convrot_w4a4")
disabled.update(_GROUPED_INT8_FORMATS)
return disabled
@contextlib.contextmanager
def use_quantized_matmul(model, device):
disabled = get_disabled_quant_formats(device)
previous = []
try:
for module in model.modules():
if isinstance(module, MixedPrecisionOp) and module.can_use_quantized_matmul(disabled):
previous.append((module, module._full_precision_mm))
module._full_precision_mm = False
yield
finally:
for module, full_precision_mm in previous:
module._full_precision_mm = full_precision_mm
def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, model_config=None):
if model_config and hasattr(model_config, 'quant_config') and model_config.quant_config:
logging.info("Using mixed precision operations")
disabled = get_disabled_quant_formats(load_device)
logging.info("Native ops: {} {}".format(", ".join(QUANT_ALGOS.keys() - disabled), ", emulated ops: {}".format(", ".join(disabled)) if len(disabled) > 0 else ""))
return mixed_precision_ops(model_config.quant_config, compute_dtype, disabled=disabled)
fp8_compute = comfy.model_management.supports_fp8_compute(load_device)
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