fix(gguf): expand one tensor at a time instead of all of them

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.
2026-08-23 05:34:28 +03:00
parent 993e5e4d04
commit ae337b5196
3 changed files with 155 additions and 31 deletions
@@ -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 "
@@ -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}


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