Repository navigation
GGUF DEQUANTIZE_TO_FP32 over-allocates: 1.1B Q4_K_M needs >12 GB heap transiently (~4.4 GB legit) #782
Description
Summary
Loading a 1.1B-parameter Q4_K_M GGUF via LlamaNetworkLoader.fromGguf(..., QuantPolicy.DEQUANTIZE_TO_FP32) over-allocates badly: the legitimate resident cost is ~4.4 GB (dense FP32), but the dequant path transiently needs >12 GB heap to get there. The model only loads cleanly with a 32 GB heap — unreasonable for a 1.1B model.
This is a transient-allocation / extra-copy problem in the dequant path, not a correctness bug — the produced tensors are correct.
Numbers
- Model: TinyLlama-1.1B-Chat v1.0, Q4_K_M (637 MB GGUF on disk).
- Parses correctly: 202 parameter tensors, ~1.1B params, correct shapes (
token_embd [32000, 2048], GQAk_proj [256, 2048]). - Legitimate dense FP32 footprint: ~4.4 GB (1.1e9 params × 4 bytes).
- Observed peak heap to complete the load: >12 GB transient (fails below; needs 32 GB to be safe) — roughly ~3× the dense floor.
Likely cause
The DEQUANTIZE_TO_FP32 path (QuantPolicy.kt → skainet-io-gguf dequant, Quants.kt) appears to materialize boxed Float / intermediate copies rather than unpacking each Q4_K/Q6_K block directly into a primitive FloatArray. With ~4.4 GB of final data, even one extra full-size intermediate (plus boxing overhead) blows past 12 GB.
Why it matters
This blocks adding TinyLlama-1.1B as a real-weights conformance/export reference on the IREE conformance side. The export use case (trace → StableHLO with weights baked as constants) fundamentally needs the ~4.4 GB FP32 resident, so the lazy RowDequantSource / ops.gather row-dequant path (a24f21d) does not help here — the issue is specifically the transient overshoot above the 4.4 GB floor during eager full materialization.
Suggested fix
- Stream the Q4_K/Q6_K unpack block-by-block straight into the destination
FloatArray(no boxedFloat, no full-size intermediate copy). - Target: peak heap ≈ dense FP32 footprint + a small per-block scratch (i.e. ~5–6 GB, not >12 GB, for this model).
Notes / related
- Do not work around with a giant
-Xmx; that masks the issue and is unusable in CI. - Related but distinct: [Research] Streaming / incremental tape API for lowering large models #740 (streaming/incremental tape for lowering large models) covers the next downstream step (trace/export memory at scale), not this load-time dequant.
🤖 Generated with Claude Code
Activity
michalharakal commentedon Jul 26, 2026
Second data point, from a much smaller model — the ~3× overshoot reproduces at 270M.
Model: FunctionGemma-270M, Q5_K_M, 248 MiB on disk. Loaded via
Gemma4WeightLoader(... QuantPolicy.DEQUANTIZE_TO_FP32).loadToMapStreaming<FP32, Float>inSKaiNET-transformers'FunctionGemmaExport.Measured from the emitted MLIR (324 weight globals), cross-checked against the on-disk archive:
832 MiB matches
gemma.safetensors(872,248,326 B) exactly, so the element count is ground truth rather than an estimate.Peak live heap ≈ 4.3 GiB against a 1.62 GiB dense floor — the same ~3× ratio this issue reports for TinyLlama-1.1B. The export OOMs at
-Xmx4gand passes at-Xmx12g.Three near-simultaneous FP32 copies are identifiable, only the first of which is this issue's scope:
TraceToGraphBuilder.extractFloatArray'sbuffer.copyOf()— closer to [Research] Streaming / incremental tape API for lowering large models #740BufferHandle.OwnedByteArrays behindmodule.externalParametersSo the ratio is not purely load-time here; but the load path is the first and unavoidable contributor, and it is the one that sets the floor everything else multiplies.
Two notes that may be useful for scoping:
VoidTensorOpsis@Backend(id = "void", displayName = "Shape-only")). So the FP32 materialization here is pure transport at double the needed width — argued as a separate lever in SKaiNET#884 (narrow dense storage), which is complementary to fixing the transient overshoot you describe here.t0andt2044, both262153x640) — tracked downstream as SKaiNET-transformers#260. Independent of this issue, but it means the absolute numbers above would drop by ~38% once deduplicated.Also confirming your note against a workaround: the downstream fix (SKaiNET-transformers#259) deliberately only raised a test heap and documented the arithmetic, precisely so the waste stays visible rather than masked.
🤖 Generated with Claude Code
Merge pull request #963 from SKaiNET-developers/skeep/003-storage-uni…
Merge pull request #965 from SKaiNET-developers/fix/782-dequant-overa…