gemma4 GGUF: embedding quant type hardcoded to Q6_K, so a Q4_0-embedding GGUF fails to load with a bare AssertionError

[salekseev]

Summary

convert_gemma4_to_gguf (python/freetoken/models/gemma4/gguf.py) hardcodes GGML_Q6_K
for the token embedding table and the tied LM head:

embed = GGUFEmbedding(
    num_embeddings=config.vocab_size,
    embedding_dim=config.hidden_size,
    quant_type=GGML_Q6_K,          # <-- assumed, not read from the file
    embed_scale=config.embedding_scale,
)
...
if config.tie_word_embeddings:
    model.lm_head = GGUFTiedLMHead(embed, GGML_Q6_K)

token_embd.weight is not always Q6_K. google's gemma-4-26B-A4B-it-qat-q4_0-gguf keeps it
at Q6_K, but unsloth's gemma-4-26B-A4B-it-qat-GGUF quantizes it to Q4_0, and row stride
differs by block type — 2310 bytes/row for Q6_K (2816/256 * 210) against 1584 for Q4_0
(2816/32 * 18) at hidden_size = 2816. GGUFEmbedding.__init__ sizes qweight with
row_bytes(embedding_dim, quant_type), so the buffer is allocated 2310 wide and the file
delivers 1584.

The load then dies in the shared loader on an assert with no message:

File "freetoken/layers/base.py", line 45, in load_state_dict
    assert param.shape == item.shape and param.dtype == item.dtype
AssertionError

layers/base.py:45 names neither the parameter nor the shapes, so the failure gives no
indication of which tensor or why. Instrumenting that traversal shows it:

MISMATCH model.embed_tokens.qweight: model wants (262144, 2310) torch.uint8 | file has (262144, 1584) torch.uint8

Nothing else in the checkpoint is affected — the projections and routed experts are Q4_0
and load fine, so the failure is purely this one assumption.

Environment

• FreeToken 0.1.2 (PyPI wheel + freetoken_kernel_cache-0.1.2+cu130)
• RTX 4080 SUPER (sm_89, 16 GB), driver 610.57.04, CUDA UMD 13.3
• Model: unsloth/gemma-4-26B-A4B-it-qat-GGUF / gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf
  (despite the filename this file is general.file_type = 2 (MOSTLY_Q4_0):
  266 Q4_0 + 392 F32 tensors, general.name = "Gemma-4 26B-A4B IT (smart Q4_0, QAT-lossless)")
• ft serve --model .../gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf

Reproduction

Any gemma4 GGUF whose token_embd.weight is not Q6_K. Reading the type straight out of the
file is enough to show the mismatch without a load:

from freetoken.models.gguf.reader import iter_gguf_tensors
from freetoken.models.gguf.dequant import GGML_NAME, GGML_Q6_K, row_bytes

path = ".../gemma-4-26B-A4B-it-qat-UD-Q4_K_XL.gguf"
for t in iter_gguf_tensors(path):
    if t.name == "token_embd.weight":
        print(f"file: {GGML_NAME[t.ggml_type]}, {row_bytes(2816, t.ggml_type)} bytes/row")
        print(f"code: {GGML_NAME[GGML_Q6_K]}, {row_bytes(2816, GGML_Q6_K)} bytes/row")

file: Q4_0, 1584 bytes/row
code: Q6_K, 2310 bytes/row

Suggested fix

GGUFEmbedding and GGUFTiedLMHead are already type-generic — they thread quant_type
into row_bytes() and ggml_dequantize(), and Q4_0 is in _MMVQ / _MMQ / _DEQUANT in
layers/gguf.py — so only the plumbing is missing. Reading the type from the tensor table
in parse_gguf_config and using it in convert_gemma4_to_gguf is sufficient; a
metadata-only GGUF (an FTW dir's source_metadata.gguf, which lists no tensors) has no
entry to read, so the Q6_K default still applies there.

With that change the model loads and serves correctly on this checkpoint: 179 tok/s decode
at 1200 tokens, --moe-backend offload, sane output, tool calls and reasoning parsed.

Happy to open a PR with the patch and a test if the approach looks right.

One adjacent thing, possibly worth its own issue: layers/base.py:45 should name the
tensor. Every shape/dtype mismatch in every model surfaces as a bare AssertionError,
and the traversal has the key and both shapes in hand at that point. A one-line message
there would have turned this from an afternoon of instrumenting the loader into a minute.

[erichstuntebeck]

Second data point, and it widens the problem: the quant type varies in more than one
direction, so a Q4_0-aware fix would still be incomplete.

ggml-org/gemma-4-26B-A4B-it-GGUF → gemma-4-26B-A4B-it-Q4_0.gguf stores
token_embd.weight as Q8_0 (ggml type 8), not Q6_K and not Q4_0:

source	token_embd.weight	bytes/row at hidden=2816
google ...-qat-q4_0-gguf (this issue)	Q6_K	2816/256*210 = 2310
unsloth ...-qat-GGUF (this issue)	Q4_0	2816/32*18 = 1584
ggml-org gemma-4-26B-A4B-it-GGUF	Q8_0	2816/32*34 = 2992

Same failure, opposite direction — the file is wider than the assumed buffer here:

MISMATCH model.embed_tokens.qweight: param [262144, 2310] vs item [262144, 2992]

Confirmed the type independently of FreeToken with a standalone GGUF header parser
(token_embd.weight, dims (2816, 262144), ggml_type 8 = Q8_0; whole-file histogram
F32 392 / Q8_0 1 / Q4_0 265), and confirmed provenance by matching the local file size
14618145824 against HuggingFace's x-linked-size for that repo. Reproduced on a clean
checkout of main @ 2757bb5: 512 state_dict entries, exactly one shape mismatch, bare
AssertionError at layers/base.py:45.

So the fix has to read the type from the tensor table rather than special-case any value.

Three things worth folding into the fix

1. Both hardcode sites need it, not just the one that raises. models/gemma4/gguf.py:363
   (GGUFEmbedding(..., quant_type=GGML_Q6_K)) is where the load dies, but :375
   (GGUFTiedLMHead(embed, GGML_Q6_K)) passes the same assumed type into
   fused_mul_mat_gguf. Fixing only 363 would leave the tied head decoding Q8_0 bytes as
   Q6_K — wrong output rather than a loud failure.

2. An equivalent fix already exists, unmerged, inside an unrelated feature PR. feat(laguna): native GGUF support for poolside Laguna (S/XS) #102
   ("feat(laguna): native GGUF support for poolside Laguna") adds exactly this:

   def _embed_quant(shim: "GgufConfigShim") -> int:
       return gguf_tensor_type(shim.model_path, "token_embd.weight") or GGML_Q6_K

   and threads it through as quant_type=embed_quant. Its own docstring notes the same
   Q6_K/Q4_0 split this issue reports. Might be worth splitting that hunk out so the loader
   fix can land independently of the Laguna work.

3. Add Gemma-4 E-series (E2B/E4B) support: PLE, KV-layer sharing, double-wide MLP, and a dense GGUF parser #59 adds a third hardcode site. "Add Gemma-4 E-series (E2B/E4B) support" introduces
   another quant_type=GGML_Q6_K, in the E-series path, so whatever lands here should cover
   it or it will regress on merge.

Environment

FreeToken 0.1.2 from source at 2757bb5 · Tesla T4 16 GiB sm_75 · driver 580.178.04 ·
CUDA 13.0.88 · torch 2.11.0+cu130 · Ubuntu 24.04.4, kernel 6.17.0-1022-azure · AMD EPYC 7V12
(4 vCPU), 27 GiB RAM · no other GPUs, no WSL, no proxy.