flint-int4 does not load in mlx-lm (weights packed as uint8, mlx-lm expects uint32)

#1
by Nav8ha - opened

Thanks for releasing Flint with open weights.

RecursiveSI/flint-int4 loads fine in rec, but not with mlx-lm's Flint support (PR ml-explore/mlx-lm#1998, commit 810efca).

What happens

from mlx_lm import load
load("flint-int4")
# ValueError: [quantize] The requested number of bits 32 is not supported.

Why: the quantized weights in this checkpoint are stored as uint8 (for example model.layers.N.self_attn.q_proj.weight is uint8 (4096, 1024)). mlx-lm expects packed uint32. It works out the bit width from the packed shape, so it gets a wrong value. rec_mlx handles this in Model.sanitize (value.view(mx.uint32)). mlx-lm does not.

Workaround: view the packed weights as uint32 when loading. The bit layout is the same:

import mlx.core as mx
import mlx_lm.utils as u

raw_load = mx.load
u.mx.load = lambda f, *a, **k: {
    n: v.view(mx.uint32) if v.dtype == mx.uint8 else v for n, v in raw_load(f, *a, **k).items()
}

With this, the checkpoint loads in mlx-lm, and the 8-bit embeddings are detected from their shape.

Suggestion: store the packed weights as uint32 (or add a note to the model card). Then the INT4 checkpoint loads directly in mlx-lm and other MLX tools once Flint support lands.

We used this to build a 3-bit-experts variant for 24–32 GB Macs from your trained INT4 (only the routed experts are requantized). Method and numbers are here: https://huggingface.co/mlx-community/flint-mixed-3-4bit

Sign up or log in to comment