From ae337b5196919c02ad2549da73ad8deb80603ee6 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:34:28 +0300 Subject: [PATCH] fix(gguf): expand one tensor at a time instead of all of them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dequantizing produces float32 — four times the width the weight ends up at. Every tensor was collected before any was cast, so the peak was the whole checkpoint in float32 plus a copy of each: tens of GB of transient host allocation on a video model, enough to take the machine down rather than just the load. ThreadPoolExecutor.map made it worse by submitting everything at once and buffering results in order. Now cast-and-release per tensor, with a fixed in-flight window. Peak is the finished model plus a handful of tensors. Pool capped at 4. Also prints the path and tensor count before the wait, so a load killed part way leaves evidence of which backend it chose. --- gguf_support.py | 83 +++++++++++++++++++++------------ loaders.py | 8 +++- tests/test_gguf_support.py | 95 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 155 insertions(+), 31 deletions(-) diff --git a/gguf_support.py b/gguf_support.py index 9260633b..5196dfdb 100644 --- a/gguf_support.py +++ b/gguf_support.py @@ -27,6 +27,7 @@ from __future__ import annotations import importlib import importlib.util +import itertools import os import sys import types @@ -400,41 +401,44 @@ def _load_native(path: str) -> tuple[dict, dict, str]: Correct but not thrifty: the model ends up at its dequantized size, so this buys the ability to LOAD a .gguf, not GGUF's memory saving. The caller states that. + + Expansion is STREAMED. Dequantizing produces float32 — four times the size the weight + will occupy once it is cast — so collecting every tensor before converting any of them + peaks at several times the finished model in host RAM. On a video checkpoint that is + tens of gigabytes of transient allocation, which is enough to take the machine down + rather than merely the load. Each tensor is therefore cast and released as it is + produced, and only a small window is ever in flight. """ import gguf + import numpy as np import torch reader = gguf.GGUFReader(path) tensors = list(reader.tensors) - # Dequantizing serially is where the minute goes: gguf.quants.dequantize is numpy, and a - # video checkpoint is thousands of tensors. numpy releases the GIL for the array work, so - # 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)) + raw = (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16) + quantized = sum(1 for t in tensors if t.tensor_type not in raw) + # Bounded, and deliberately small: the window below keeps `workers * 2` float32 tensors + # alive at once, and the largest weight in a video model is around a gigabyte at that + # width. More threads would expand faster and spike higher; the durable speed win is the + # dequantized cache in `_load_native_cached`, not this pool. + workers = max(1, min(4, (os.cpu_count() or 2) - 1)) + print(f"[FunPack] gguf: expanding {quantized} of {len(tensors)} tensors from " + f"{os.path.basename(path)} across {workers} thread(s). This is the slow path — " + f"minutes on a video checkpoint, and the result is cached beside the file.", + flush=True) - def _dequant(tensor): + def _expand(tensor): + """One tensor, fully finished: nothing float32 outlives this call.""" qtype = tensor.tensor_type - if qtype in (gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16): - return tensor.data - # gguf.quants.dequantize is the reference implementation for every type the - # package knows — do not reimplement it per quant type. - return gguf.quants.dequantize(tensor.data, qtype) - - 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): + arr = tensor.data if qtype in raw else gguf.quants.dequantize(tensor.data, qtype) + if arr.dtype == np.float32: + # `.to` allocates the half-width copy and the float32 array is freed on return, + # so the wide form never accumulates. + t = torch.from_numpy(arr).to(torch.float16) + else: + # Detach from the memory-mapped file; the reader's mapping does not outlive us. + t = torch.from_numpy(arr.copy()) name = str(tensor.name) - t = torch.from_numpy(arr.copy()) # GGUF stores dimensions in the opposite order to torch. shape = tuple(int(d) for d in reversed(tensor.shape)) expected = 1 @@ -447,9 +451,30 @@ def _load_native(path: str) -> tuple[dict, dict, str]: f"{os.path.basename(path)}: tensor {name!r} dequantized to {t.numel()} " f"elements but its header declares {shape} ({expected}). This container is " f"not laid out the way the gguf package describes it.") - if shape: - t = t.reshape(shape) - sd[name] = t.to(torch.float16) if t.dtype == torch.float32 else t + return name, t.reshape(shape) if shape else t + + sd: dict = {} + if workers > 1 and len(tensors) > 1: + from collections import deque + from concurrent.futures import ThreadPoolExecutor + with ThreadPoolExecutor(max_workers=workers) as pool: + remaining = iter(tensors) + # A fixed window rather than `pool.map`, which submits every tensor at once and + # buffers each result until the one before it is consumed — the whole model in + # float32, which is the allocation this function exists to avoid. + flight = deque(pool.submit(_expand, t) + for t in itertools.islice(remaining, workers * 2)) + while flight: + name, t = flight.popleft().result() + sd[name] = t + nxt = next(remaining, None) + if nxt is not None: + flight.append(pool.submit(_expand, nxt)) + else: + for tensor in tensors: + name, t = _expand(tensor) + sd[name] = t + note = (f"gguf: dequantized at load ({quantized} quantized tensors expanded across " f"{workers} thread(s)) — the file loads, but it occupies its full size in VRAM, " f"and expanding it is why this is slower than a .safetensors. Install " diff --git a/loaders.py b/loaders.py index b42d7f78..71bfdc7c 100644 --- a/loaders.py +++ b/loaders.py @@ -130,8 +130,12 @@ def set_fp16_accumulation(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. +# file. Replaced whenever the key changes, so this never holds more than one model. +# +# The reference is strong on purpose: ComfyUI drops this node's cached output BEFORE it +# re-runs it, which is exactly the moment a weak one would be dead. The cost is that this +# entry outlives a "free model and node cache", so one checkpoint stays resident until a +# different file or dtype replaces it. Loading anything else releases it. _WEIGHT_CACHE: dict = {"key": None, "model": None} diff --git a/tests/test_gguf_support.py b/tests/test_gguf_support.py index 62a83e52..6baf717d 100644 --- a/tests/test_gguf_support.py +++ b/tests/test_gguf_support.py @@ -5,11 +5,14 @@ which backend gets picked, what happens when none is available, and whether .ggu are found at all (core's extension set excludes them, which is the whole reason this module exists). """ +import gc import os import sys import types +import weakref import pytest +import torch sys.path.insert(0, ".") import gguf_support # noqa: E402 @@ -370,3 +373,95 @@ def test_a_non_tensor_value_is_named(tmp_path): def test_a_real_state_dict_passes(): import torch gguf_support._assert_state_dict({"w": torch.zeros(2)}, "m.gguf") + + +def _fake_gguf(monkeypatch, tensors): + """A stand-in for the `gguf` package over `tensors` — (name, shape, values) triples. + + Every tensor is reported as quantized, so `_load_native` takes its dequantize path for + all of them, and each dequantize hands back a fresh float32 array whose lifetime the + caller can track. + """ + import numpy as np + + fake = types.ModuleType("gguf") + + class _QType: + F32 = "F32" + F16 = "F16" + Q8 = "Q8" + + class _Tensor: + def __init__(self, name, shape, values): + self.name = name + # GGUF reports dimensions in the opposite order to torch. + self.shape = tuple(reversed(shape)) + self.tensor_type = _QType.Q8 + self.data = np.asarray(values, dtype=np.float32) + + class _Reader: + def __init__(self, path): + self.tensors = [_Tensor(n, s, v) for n, s, v in tensors] + + produced = [] + live = [] + + def _dequantize(data, qtype): + # Sampled BEFORE this call's own array exists, so it counts what the loader is still + # holding on to. Peak is the whole point: every array is freed by the time the load + # returns no matter how it was written. + live.append(sum(1 for ref in produced if ref() is not None)) + arr = np.asarray(data, dtype=np.float32).copy() + produced.append(weakref.ref(arr)) + return arr + + quants = types.ModuleType("gguf.quants") + quants.dequantize = _dequantize + fake.GGMLQuantizationType = _QType + fake.GGUFReader = _Reader + fake.quants = quants + monkeypatch.setitem(sys.modules, "gguf", fake) + monkeypatch.setitem(sys.modules, "gguf.quants", quants) + return produced, live + + +def test_load_native_returns_half_width_tensors_in_torch_order(monkeypatch): + monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2) # forces the serial path + _fake_gguf(monkeypatch, [("a", (2, 3), [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])]) + + sd, options, note = gguf_support._load_native("/models/m.gguf") + + assert options == {} + assert list(sd) == ["a"] + assert tuple(sd["a"].shape) == (2, 3) + assert sd["a"].dtype is torch.float16 + assert sd["a"][1][2].item() == 6.0 + assert "1 quantized tensors" in note + + +def test_load_native_does_not_keep_every_expanded_tensor_alive(monkeypatch): + """The float32 form is four times the finished weight, so holding all of them at once + peaks at several times the model in host RAM — enough to take a machine down on a video + checkpoint. Each one must be cast and released as it is produced.""" + monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2) # forces the serial path + produced, live = _fake_gguf( + monkeypatch, [(f"t{i}", (4, 4), [[float(i)] * 4] * 4) for i in range(24)]) + + gguf_support._load_native("/models/m.gguf") + + assert len(produced) == 24 + # Collecting them all before converting any would climb 0, 1, 2, ... 23. + assert max(live) <= 2, f"{max(live)} expanded tensors were held at once" + gc.collect() + assert [ref for ref in produced if ref() is not None] == [] + + +def test_load_native_rejects_a_tensor_that_contradicts_its_header(monkeypatch): + monkeypatch.setattr(gguf_support.os, "cpu_count", lambda: 2) + _fake_gguf(monkeypatch, [("wrong", (5, 5), [1.0, 2.0, 3.0])]) + + with pytest.raises(RuntimeError) as excinfo: + gguf_support._load_native("/models/m.gguf") + + assert "'wrong'" in str(excinfo.value) + assert "(25)" in str(excinfo.value)