Repository navigation
[bug]: Quantized-embedding RAM reservation ignores the packed weights already held #9979
Description
What happened
materialize_quantized_embedding (invokeai/backend/model_manager/load/quantized_embedding.py:63) dequantizes a GGUF embedding in one piece. It calls ram_cache.make_room(_PEAK_COPIES * numel * itemsize) with a fixed _PEAK_COPIES = 6.
That reservation is absolute: it frees room in the cache's tracked budget. At that point the loader already holds the whole packed state dict, which the cache doesn't track. So the real peak is the packed weights plus the decode transient, and the flat 6× is right only by coincidence.
Measured on CPU (private bytes) for large-vocabulary encoders. The reservation is the same 11.74 GiB for every UMT5 build, because it depends only on the embedding size:
| Model / file | Real peak | Reservation | Off by |
|---|---|---|---|
| UMT5-XXL Q8_0 (Q8_0 embedding, 256k × 4096) | 9.55 GiB | 11.74 GiB | over by 2.2 GiB |
| UMT5-XXL Q4_K_M (Q6_K embedding) | 11.53 GiB | 11.74 GiB | about right |
| UMT5-XXL Q6_K / F16 / F32 | 11.74 GiB | under by 0.7–2 GiB (estimated) | |
| Qwen2.5-VL-7B Q4_K_M (#9969) | under by ~1.1 GB | ||
| Qwen2.5-VL-7B Q8_0 (#9969) | under by ~2.2 GB | ||
| Qwen3-VL-8B Q8_0 | under by ~2.5 GB |
When it's short, the load can briefly exceed the cache budget and push a RAM-tight machine into swap.
The whole-tensor decode transient itself is large:
| Embedding type | Transient (result included) |
|---|---|
| Q8_0 | 2.0× the bf16 result |
| Q6_K | 4.1× |
What you expected to happen
The reservation matches the real peak, and the peak itself is lower.
Proposed fix
Decode GGML embeddings in row slices into a preallocated output, as unpack_ggml_at_load already does on #9932. GGUF blocks are row-aligned, so this is bit-exact. Then reserve absolutely, counting the packed bytes the caller holds.
Measured with 16384-row slices, bit-exact (torch.equal):
| Embedding type | Transient | Change | Real peak (UMT5 file) |
|---|---|---|---|
| Q8_0 | 2.28 GiB (1.16×) | same speed | 9.55 → 7.9 GiB |
| Q6_K | 2.69 GiB (1.37×) | slower at 2 threads; bigger slices cut that | 11.5 → 6.1 GiB (Q4_K_M file) |
SDNQ keeps the whole-tensor path. This affects all callers: Krea-2, Mistral, Z-Image, Wan and, with the open PRs, Qwen-Image.
Activity