transformer: weight arena undersized for Ternary-Bonsai-2-27B on Metal (GEIST_WEIGHT_MMAP=0)

[geisten]

Problem

Loading Ternary-Bonsai-2-27B PQ2_0 on Metal with GEIST_WEIGHT_MMAP=0 (arena mode) fails with weight arena exhausted at blk.63.attn_output. The arena comes up about 7.7 MB short. The mmap-alias mode loads, and its logits match cpu_neon.

Reproduced on main a99ab7c, M1 Max. Found while checking #464 against ternary models; main and the #464 branch behave the same.

Likely cause (inferred, not verified)

compute_weight_arena_capacity (src/archs/transformer/weight_load/tensor_views.c) adds a second staging copy plus the F32 copy only for small F16/BF16 matrices, plus 64 MB of headroom. Bonsai's 96 BF16 DeltaNet alpha/beta tensors, together with other widened or derived buffers on Metal, appear to exceed that headroom on the 27B.

Expected

The arena capacity counts every buffer the loader derives on this backend, so arena mode loads the model.

[geisten]

Cause verified

compute_weight_arena_capacity adds the widen surcharge only when the backend sets caps.weights_need_backend_arena:

if (be->desc->caps.weights_need_backend_arena && t->n_dims == 2 &&
    (t->dtype == GGUF_TYPE_F16 || t->dtype == GGUF_TYPE_BF16) && elems <= (4u << 20)) {
    total += aligned + <F32 bytes>;
}

Only Vulkan sets that cap. Metal does not, which is why it defaults to mmap-alias. But load_layer_proj widens small F16/BF16 matrices when resolve_weight returns GEIST_E_UNSUPPORTED, whatever the caps say. So on Metal with GEIST_WEIGHT_MMAP=0, each widened matrix takes its file bytes plus a second staged copy plus the F32 copy from the arena, and only the file bytes were counted. The 64 MB headroom absorbs part of that, so a model fails only once it has enough such matrices.

Second reproducer (small and fast)

llama-3.2-3B Q4_K_M with only attn_k requantized to F16 (llama-quantize --allow-requantize --tensor-type attn_k=f16 … Q4_K_M). attn_k is 1024×3072, which is under the 4M-element widen cap, so each of the 28 layers widens one matrix: 6 MB staged again + 12 MB F32 that the capacity never counted, about 504 MB in total.

• main a99ab7c, GEIST_DUMP_BACKEND=metal GEIST_WEIGHT_MMAP=0 dump_geist_logits: weight arena exhausted at 'blk.22.ffn_up.weight' (used 2201610240, capacity 2205264128, need 14155776).
• I instrumented arena_alloc to log every allocation. No tensor is staged twice except the widened attn_k, and no allocation lacks a tensor name. Capacity before headroom is 2138155264, which is about the file's tensor bytes.

Fix checked locally

Dropping be->desc->caps.weights_need_backend_arena && from that condition makes the model load. The Metal arena-mode logits are byte-identical to Metal mmap-alias (top1 " Paris"), on main and on the #464 branch (79c73ad).

Cost: on a CPU backend, where a small F16 matrix resolves natively and is never widened, the arena is oversized by 1.5× those tensors' F16 bytes plus their F32 size. A tighter version would ask the backend whether it has a half-precision linear. Neither variant is a PR yet.

[geisten]

This also reproduces on a small model, on main a99ab7c. It is smollm2-360m Q8_0 rewritten so that attn_q/attn_k are F16 or BF16, run on Metal with GEIST_WEIGHT_MMAP=0:

weight arena exhausted at 'blk.23.attn_q.weight' (used 488467200, capacity 488591104, need 1843200)

The half-precision q/k are widened to F32 at load, so they take twice their file size in the arena. The arena capacity seems to be sized from the file bytes and does not account for that.

---

Generated by Claude Code