GGUF DEQUANTIZE_TO_FP32 over-allocates: 1.1B Q4_K_M needs >12 GB heap transiently (~4.4 GB legit)

[michalharakal]

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], GQA k_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 boxed Float, 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

[michalharakal]

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> in SKaiNET-transformers' FunctionGemmaExport.

Measured from the emitted MLIR (324 weight globals), cross-checked against the on-disk archive:

total elements: 436,111,680   = 1.62 GiB FP32   (832 MiB bf16)

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 -Xmx4g and passes at -Xmx12g.

Three near-simultaneous FP32 copies are identifiable, only the first of which is this issue's scope:

1. the loader's dequantized weights — this issue
2. TraceToGraphBuilder.extractFloatArray's buffer.copyOf() — closer to [Research] Streaming / incremental tape API for lowering large models #740
3. the BufferHandle.Owned ByteArrays behind module.externalParameters

So 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:

• The destination format is bf16, and the export trace does no arithmetic at all (VoidTensorOps is @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.
• 77% of that 436M element count is the vocab embedding emitted twice (t0 and t2044, both 262153x640) — 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