perf(loaders): stop re-reading the checkpoint for inference settings
Attention and SLA are applied to a clone of a one-entry weight cache keyed on file identity + dtypes, so a sparsity change skips the disk. GGUF dequantization now runs across threads instead of serially. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
6f0464202d
commit
6bf533e808
@@ -3,6 +3,15 @@
|
|||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
- **Changing an attention or SLA setting no longer re-reads the checkpoint.** They live on
|
||||||
|
the loader node, so ComfyUI invalidated the whole node and the file was loaded again — a
|
||||||
|
minute of dequantizing a GGUF to alter a sparsity ratio. The weights depend only on the
|
||||||
|
file (identity, size, mtime) and the dtypes; everything else is applied to a clone. One
|
||||||
|
entry, replaced whenever the key changes.
|
||||||
|
- **GGUF dequantizes across threads instead of one tensor at a time.** The native path
|
||||||
|
expands every quantized tensor at load with numpy, serially, which is where the wait went.
|
||||||
|
numpy releases the GIL for that work, so it now runs on up to 8 threads — bounded rather
|
||||||
|
than one per core, because each worker holds a dequantized tensor and its copy.
|
||||||
- **The updater installs only what is MISSING, and never upgrades anything.** It ran
|
- **The updater installs only what is MISSING, and never upgrades anything.** It ran
|
||||||
`pip install -r requirements.txt`, which upgrades any package below its version floor —
|
`pip install -r requirements.txt`, which upgrades any package below its version floor —
|
||||||
and on an install full of compiled extensions (torch, comfy-kitchen, comfy-aimdo,
|
and on an install full of compiled extensions (torch, comfy-kitchen, comfy-aimdo,
|
||||||
|
|||||||
+31
-13
@@ -262,18 +262,35 @@ def _load_native(path: str) -> tuple[dict, dict, str]:
|
|||||||
import torch
|
import torch
|
||||||
|
|
||||||
reader = gguf.GGUFReader(path)
|
reader = gguf.GGUFReader(path)
|
||||||
sd: dict = {}
|
tensors = list(reader.tensors)
|
||||||
quantized = 0
|
# Dequantizing serially is where the minute goes: gguf.quants.dequantize is numpy, and a
|
||||||
for tensor in reader.tensors:
|
# video checkpoint is thousands of tensors. numpy releases the GIL for the array work, so
|
||||||
name = str(tensor.name)
|
# threads are real parallelism here. Bounded, not one per core: each worker holds a
|
||||||
|
# dequantized tensor AND its copy, so an unbounded pool trades load time for a memory
|
||||||
|
# spike on the largest weights.
|
||||||
|
workers = max(1, min(8, (os.cpu_count() or 2) - 1))
|
||||||
|
quantized = sum(
|
||||||
|
1 for t in tensors
|
||||||
|
if t.tensor_type not in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16))
|
||||||
|
|
||||||
|
def _dequant(tensor):
|
||||||
qtype = tensor.tensor_type
|
qtype = tensor.tensor_type
|
||||||
if qtype in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
|
if qtype in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
|
||||||
arr = tensor.data
|
return tensor.data
|
||||||
else:
|
# gguf.quants.dequantize is the reference implementation for every type the
|
||||||
# gguf.quants.dequantize is the reference implementation for every type the
|
# package knows — do not reimplement it per quant type.
|
||||||
# package knows — do not reimplement it per quant type.
|
return gguf.quants.dequantize(tensor.data, qtype)
|
||||||
arr = gguf.quants.dequantize(tensor.data, qtype)
|
|
||||||
quantized += 1
|
if workers > 1 and len(tensors) > 1:
|
||||||
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
|
with ThreadPoolExecutor(max_workers=workers) as pool:
|
||||||
|
arrays = list(pool.map(_dequant, tensors))
|
||||||
|
else:
|
||||||
|
arrays = [_dequant(t) for t in tensors]
|
||||||
|
|
||||||
|
sd: dict = {}
|
||||||
|
for tensor, arr in zip(tensors, arrays):
|
||||||
|
name = str(tensor.name)
|
||||||
t = torch.from_numpy(arr.copy())
|
t = torch.from_numpy(arr.copy())
|
||||||
# GGUF stores dimensions in the opposite order to torch.
|
# GGUF stores dimensions in the opposite order to torch.
|
||||||
shape = tuple(int(d) for d in reversed(tensor.shape))
|
shape = tuple(int(d) for d in reversed(tensor.shape))
|
||||||
@@ -290,7 +307,8 @@ def _load_native(path: str) -> tuple[dict, dict, str]:
|
|||||||
if shape:
|
if shape:
|
||||||
t = t.reshape(shape)
|
t = t.reshape(shape)
|
||||||
sd[name] = t.to(torch.float16) if t.dtype == torch.float32 else t
|
sd[name] = t.to(torch.float16) if t.dtype == torch.float32 else t
|
||||||
note = (f"gguf: dequantized at load ({quantized} quantized tensors expanded) — the file "
|
note = (f"gguf: dequantized at load ({quantized} quantized tensors expanded across "
|
||||||
f"loads, but it occupies its full size in VRAM. Install ComfyUI-GGUF to keep it "
|
f"{workers} thread(s)) — the file loads, but it occupies its full size in VRAM, "
|
||||||
f"quantized.")
|
f"and expanding it is why this is slower than a .safetensors. Install "
|
||||||
|
f"ComfyUI-GGUF for a lazy quantized load.")
|
||||||
return sd, {}, note
|
return sd, {}, note
|
||||||
|
|||||||
+52
@@ -11,6 +11,7 @@ SageAttention is installed, flash when flash-attn is, xformers when enabled —
|
|||||||
lies about the rest.
|
lies about the rest.
|
||||||
"""
|
"""
|
||||||
import logging
|
import logging
|
||||||
|
import os
|
||||||
|
|
||||||
import comfy.sd
|
import comfy.sd
|
||||||
import comfy.utils
|
import comfy.utils
|
||||||
@@ -128,6 +129,31 @@ def set_fp16_accumulation(enabled):
|
|||||||
return bool(enabled)
|
return bool(enabled)
|
||||||
|
|
||||||
|
|
||||||
|
# One loaded checkpoint, so changing an inference setting on the loader does not re-read the
|
||||||
|
# file. Replaced whenever the key changes, so this never holds more than one model — the same
|
||||||
|
# lifetime ComfyUI's node-output cache already gave it.
|
||||||
|
_WEIGHT_CACHE: dict = {"key": None, "model": None}
|
||||||
|
|
||||||
|
|
||||||
|
def _weight_cache_key(model_name, weight_dtype, compute_dtype):
|
||||||
|
"""Everything that changes the WEIGHTS: the file (by identity, size and mtime) and the
|
||||||
|
dtypes. Not the attention backend, not SLA — those are applied to a clone.
|
||||||
|
|
||||||
|
Size and mtime are in the key so a file replaced under the same name is re-read. An
|
||||||
|
unreadable stat returns None, which never matches and so always reloads: failing to
|
||||||
|
notice a changed file is worse than an occasional extra read.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
path = gguf_support.gguf_path("diffusion_models", model_name) if gguf_support.is_gguf(model_name) \
|
||||||
|
else folder_paths.get_full_path("diffusion_models", model_name)
|
||||||
|
if not path:
|
||||||
|
return None
|
||||||
|
st = os.stat(path)
|
||||||
|
return (path, st.st_size, st.st_mtime_ns, str(weight_dtype), str(compute_dtype))
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
class FunPackDiffusionModelLoader:
|
class FunPackDiffusionModelLoader:
|
||||||
"""Loads a diffusion model with quantization and an attention backend chosen per node."""
|
"""Loads a diffusion model with quantization and an attention backend chosen per node."""
|
||||||
|
|
||||||
@@ -240,6 +266,20 @@ class FunPackDiffusionModelLoader:
|
|||||||
model_options = weight_model_options(weight_dtype)
|
model_options = weight_model_options(weight_dtype)
|
||||||
notes.append(f"weight dtype: {weight_dtype}")
|
notes.append(f"weight dtype: {weight_dtype}")
|
||||||
|
|
||||||
|
# Attention and SLA are INFERENCE settings, but they live on this node, so ComfyUI
|
||||||
|
# invalidates the whole node when one changes and the checkpoint is read from disk
|
||||||
|
# again — a minute of dequantizing a GGUF to alter a sparsity ratio. The weights only
|
||||||
|
# depend on the file and the dtypes, so those are the cache key and everything else
|
||||||
|
# is applied to a clone. One entry, replaced whenever the key changes: this holds no
|
||||||
|
# more than ComfyUI's own node-output cache already did.
|
||||||
|
cache_key = _weight_cache_key(model_name, weight_dtype, compute_dtype)
|
||||||
|
cached = _WEIGHT_CACHE.get("model") if _WEIGHT_CACHE.get("key") == cache_key else None
|
||||||
|
if cached is not None:
|
||||||
|
notes.append("weights: reused (unchanged file and dtypes)")
|
||||||
|
model = cached.clone()
|
||||||
|
return self._finish(model, notes, attention, sla, sla_sparsity, sla_block_size,
|
||||||
|
sla_min_seq_len, sla_dense_last_steps, sla_protect_audio)
|
||||||
|
|
||||||
path = None
|
path = None
|
||||||
misnamed = False
|
misnamed = False
|
||||||
if gguf_support.is_gguf(model_name):
|
if gguf_support.is_gguf(model_name):
|
||||||
@@ -280,6 +320,18 @@ class FunPackDiffusionModelLoader:
|
|||||||
elif dtype is not None:
|
elif dtype is not None:
|
||||||
notes.append("compute dtype: unsupported by this ComfyUI, ignored")
|
notes.append("compute dtype: unsupported by this ComfyUI, ignored")
|
||||||
|
|
||||||
|
# Kept BEFORE the attention work, so what is cached is the plain weights. Everything
|
||||||
|
# below is applied to a clone, which is what lets a sparsity change skip the disk.
|
||||||
|
_WEIGHT_CACHE["key"] = cache_key
|
||||||
|
_WEIGHT_CACHE["model"] = model
|
||||||
|
return self._finish(model.clone(), notes, attention, sla, sla_sparsity,
|
||||||
|
sla_block_size, sla_min_seq_len, sla_dense_last_steps,
|
||||||
|
sla_protect_audio)
|
||||||
|
|
||||||
|
def _finish(self, model, notes, attention, sla, sla_sparsity, sla_block_size,
|
||||||
|
sla_min_seq_len, sla_dense_last_steps, sla_protect_audio):
|
||||||
|
"""Attention + SLA onto an already-loaded model. Separate from the read so it can run
|
||||||
|
against a cached one."""
|
||||||
override = attention_override(attention)
|
override = attention_override(attention)
|
||||||
|
|
||||||
# SLA wraps the chosen backend rather than replacing it: there is one override
|
# SLA wraps the chosen backend rather than replacing it: there is one override
|
||||||
|
|||||||
@@ -231,3 +231,73 @@ def test_sla_is_a_toggle_not_a_rival_backend():
|
|||||||
assert sla_attention.SLA_NAME not in loaders.attention_choices()
|
assert sla_attention.SLA_NAME not in loaders.attention_choices()
|
||||||
spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES()
|
spec = loaders.FunPackDiffusionModelLoader.INPUT_TYPES()
|
||||||
assert spec["optional"]["sla"][1]["default"] is False
|
assert spec["optional"]["sla"][1]["default"] is False
|
||||||
|
|
||||||
|
|
||||||
|
# ── changing an inference setting must not re-read the checkpoint ─────────────
|
||||||
|
# Attention and SLA live on the loader node, so ComfyUI invalidates the whole node when one
|
||||||
|
# changes and the file is read again — a minute of dequantizing a GGUF to alter a sparsity
|
||||||
|
# ratio. Only the file and the dtypes decide the weights.
|
||||||
|
|
||||||
|
|
||||||
|
class _CacheModel:
|
||||||
|
def __init__(self):
|
||||||
|
self.model_options = {}
|
||||||
|
self.clones = 0
|
||||||
|
|
||||||
|
def clone(self):
|
||||||
|
c = _CacheModel()
|
||||||
|
c.clones = self.clones + 1
|
||||||
|
return c
|
||||||
|
|
||||||
|
|
||||||
|
def _key_env(tmp_path, monkeypatch):
|
||||||
|
import loaders
|
||||||
|
f = tmp_path / "m.safetensors"
|
||||||
|
f.write_bytes(b"x" * 16)
|
||||||
|
monkeypatch.setattr(loaders.folder_paths, "get_full_path", lambda folder, n: str(f),
|
||||||
|
raising=False)
|
||||||
|
loaders._WEIGHT_CACHE.update({"key": None, "model": None})
|
||||||
|
return loaders, f
|
||||||
|
|
||||||
|
|
||||||
|
def test_the_key_ignores_attention_and_sla(tmp_path, monkeypatch):
|
||||||
|
loaders, _ = _key_env(tmp_path, monkeypatch)
|
||||||
|
a = loaders._weight_cache_key("m.safetensors", "default", "bf16")
|
||||||
|
b = loaders._weight_cache_key("m.safetensors", "default", "bf16")
|
||||||
|
assert a == b and a is not None
|
||||||
|
|
||||||
|
|
||||||
|
def test_a_different_dtype_is_a_different_key(tmp_path, monkeypatch):
|
||||||
|
loaders, _ = _key_env(tmp_path, monkeypatch)
|
||||||
|
assert (loaders._weight_cache_key("m.safetensors", "default", "bf16")
|
||||||
|
!= loaders._weight_cache_key("m.safetensors", "fp8_e4m3fn", "bf16"))
|
||||||
|
|
||||||
|
|
||||||
|
def test_a_replaced_file_is_re_read(tmp_path, monkeypatch):
|
||||||
|
"""Same name, different bytes: size and mtime are in the key so the cache misses."""
|
||||||
|
loaders, f = _key_env(tmp_path, monkeypatch)
|
||||||
|
before = loaders._weight_cache_key("m.safetensors", "default", "bf16")
|
||||||
|
f.write_bytes(b"y" * 32)
|
||||||
|
assert loaders._weight_cache_key("m.safetensors", "default", "bf16") != before
|
||||||
|
|
||||||
|
|
||||||
|
def test_an_unreadable_file_never_matches(tmp_path, monkeypatch):
|
||||||
|
"""Failing to notice a changed file is worse than an occasional extra read."""
|
||||||
|
import loaders
|
||||||
|
monkeypatch.setattr(loaders.folder_paths, "get_full_path", lambda folder, n: None,
|
||||||
|
raising=False)
|
||||||
|
assert loaders._weight_cache_key("gone.safetensors", "default", "bf16") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_the_cached_model_is_cloned_never_handed_out(tmp_path, monkeypatch):
|
||||||
|
"""The cached object must not collect the model_options of every run that reused it."""
|
||||||
|
loaders, _ = _key_env(tmp_path, monkeypatch)
|
||||||
|
base = _CacheModel()
|
||||||
|
key = loaders._weight_cache_key("m.safetensors", "default", "bf16")
|
||||||
|
loaders._WEIGHT_CACHE.update({"key": key, "model": base})
|
||||||
|
node = loaders.FunPackDiffusionModelLoader()
|
||||||
|
out, _status = node._finish(loaders._WEIGHT_CACHE["model"].clone(), [], "default",
|
||||||
|
False, 0.9, "64", 8192, 0, True)
|
||||||
|
out.model_options.setdefault("transformer_options", {})["marker"] = 1
|
||||||
|
assert base.model_options == {} # untouched
|
||||||
|
assert out.clones == 1
|
||||||
|
|||||||
Reference in New Issue
Block a user