| """T2.7 · Export packed ternary students to llama.cpp GGUF. |
| |
| Ternary weights go out as mainline llama.cpp's Q2_0: blocks of 64 values, an fp16 scale and 2-bit codes, where |
| the stored value q+1 means q·d (00 = −1, 01 = 0, 10 = +1). Our packing (`quant.ternary.pack_2bit`) uses the same |
| bit layout, and our groups are 128 wide, so each block takes 16 of our code bytes unchanged and its group's scale, |
| written twice per group. The export is lossless as long as every scale is an fp16 normal (≥6.1e-5; checked). |
| |
| D-51's N-bit experts go out as the GGUF type whose grid they were trained on (`quant.ternary.GRIDS`): int2 as |
| Q2_0 (the same 2-bit packing, groups of 64, so one block per group with its signed scale), int4 as Q4_0, int3 as |
| Q3_K with its down projection as Q4_0 (rows of 384 cannot be Q3_K). Their packers are below; gguf-py has |
| Q4_0's quantiser and both types' dequantisers, and `check_round_trip` dequantises with gguf-py. |
| |
| Everything else is written as the student holds it: bf16 matrices as BF16 (their bytes copied), and 1-D norms |
| and the router as F32, as llama.cpp's converter writes them. Metadata (tokenizer, hyperparameters, chat template) |
| comes from a template GGUF of the same architecture made by llama.cpp's `convert_hf_to_gguf.py`, whose tensor list |
| and order the export follows. |
| """ |
| from __future__ import annotations |
|
|
| import re |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| Q2_0_BLOCK, Q2_0_BYTES = 64, 18 |
| Q4_0_BLOCK, Q4_0_BYTES = 32, 18 |
| Q3_K_BLOCK, Q3_K_BYTES = 256, 110 |
| IQ4_XS_BLOCK, IQ4_XS_BYTES = 256, 136 |
| FP16_MIN_NORMAL = 6.103515625e-05 |
|
|
|
|
| def gguf_module(llama_cpp: str | Path): |
| sys.path.insert(0, str(Path(llama_cpp) / "gguf-py")) |
| import gguf |
| return gguf |
|
|
|
|
| def q2_0_blocks(codes: torch.Tensor, scales: torch.Tensor, group: int) -> np.ndarray: |
| """Packed codes [..., C/4] (uint8) and scales [..., C/group] (fp16) → Q2_0 bytes [..., C/64 · 18].""" |
| if group % Q2_0_BLOCK: |
| raise ValueError(f"group size {group} is not a multiple of Q2_0's {Q2_0_BLOCK}") |
| c = codes.contiguous().numpy() |
| s = scales.to(torch.float16).contiguous().numpy() |
| lead, n_bytes = c.shape[:-1], c.shape[-1] |
| nb = n_bytes * 4 // Q2_0_BLOCK |
| d = np.repeat(s.reshape(-1, s.shape[-1]), group // Q2_0_BLOCK, axis=1) |
| out = np.empty((d.shape[0], nb, Q2_0_BYTES), dtype=np.uint8) |
| out[:, :, :2] = d.view(np.uint8).reshape(d.shape[0], nb, 2) |
| out[:, :, 2:] = c.reshape(-1, nb, Q2_0_BLOCK // 4) |
| return out.reshape(*lead, nb * Q2_0_BYTES) |
|
|
|
|
| def dequant_q2_0(raw: np.ndarray) -> np.ndarray: |
| """Q2_0 bytes [..., nb · 18] → float32 [..., nb · 64], as ggml's dequantize_row_q2_0.""" |
| lead = raw.shape[:-1] |
| b = raw.reshape(-1, Q2_0_BYTES) |
| d = b[:, :2].copy().view(np.float16).astype(np.float32) |
| q = b[:, 2:] |
| vals = np.stack([(q >> s) & 3 for s in (0, 2, 4, 6)], axis=-1).reshape(-1, Q2_0_BLOCK).astype(np.float32) - 1.0 |
| return (vals * d).reshape(*lead, -1) |
|
|
|
|
| def q4_0_blocks(codes: torch.Tensor, scales: torch.Tensor) -> np.ndarray: |
| """int4 packed codes [..., C/2] (`pack_4bit` of q+8) and scales [..., C/32] (fp16) → Q4_0 bytes [..., C/32 · 18]: |
| per block the fp16 d, then byte j holds element j (low nibble) and element j+16 (high), stored as q+8.""" |
| from lowbit.quant.ternary import unpack_4bit |
| u = unpack_4bit(codes).contiguous().numpy() |
| d = scales.to(torch.float16).contiguous().numpy() |
| lead, nb = u.shape[:-1], u.shape[-1] // Q4_0_BLOCK |
| u = u.reshape(-1, nb, 2, Q4_0_BLOCK // 2) |
| out = np.empty((u.shape[0], nb, Q4_0_BYTES), dtype=np.uint8) |
| out[:, :, :2] = d.reshape(-1, nb, 1).view(np.uint8) |
| out[:, :, 2:] = u[:, :, 0] | (u[:, :, 1] << 4) |
| return out.reshape(*lead, nb * Q4_0_BYTES) |
|
|
|
|
| def q3_k_blocks(codes: torch.Tensor, d: torch.Tensor, m: torch.Tensor) -> np.ndarray: |
| """int3 packed codes [..., C/2] (`pack_4bit` of L = q+4), super-block scales d [..., C/256] (fp16) and sub-block |
| scales m [..., C/16] (int8, -32..31) → Q3_K bytes [..., C/256 · 110], the layout of ggml's |
| quantize_row_q3_K_ref: hmask[32] (bit b of byte i: L[32b+i] > 3), qs[64] (low two bits of L; per 128-half, |
| byte l holds elements l, l+32, l+64, l+96 at bits 0, 2, 4, 6), scales[12] (m+32 in 6 bits: low nibbles of |
| sub-blocks 0-7 then 8-15 in bytes 0-7, the high two bits of sub-block j in byte 8 + j%4 at bit 2·(j//4)), d.""" |
| from lowbit.quant.ternary import unpack_4bit |
| L = unpack_4bit(codes).contiguous().numpy() |
| lead, nb = L.shape[:-1], L.shape[-1] // Q3_K_BLOCK |
| L = L.reshape(-1, nb, Q3_K_BLOCK) |
| rows = L.shape[0] |
| hb = ((L >> 2) & 1).reshape(rows, nb, 8, 32) |
| hmask = np.zeros((rows, nb, 32), dtype=np.uint8) |
| for b in range(8): |
| hmask |= hb[:, :, b] << b |
| low = (L & 3).reshape(rows, nb, 2, 4, 32) |
| qs = np.zeros((rows, nb, 2, 32), dtype=np.uint8) |
| for k in range(4): |
| qs |= low[:, :, :, k] << (2 * k) |
| sc = (m.to(torch.int16) + 32).to(torch.uint8).contiguous().numpy().reshape(rows, nb, 16) |
| lo4 = sc & 15 |
| hi2 = ((sc >> 4) & 3).reshape(rows, nb, 4, 4) |
| scales = np.zeros((rows, nb, 12), dtype=np.uint8) |
| scales[:, :, :8] = lo4[:, :, :8] | (lo4[:, :, 8:] << 4) |
| for r in range(4): |
| scales[:, :, 8:] |= hi2[:, :, r] << (2 * r) |
| dd = d.to(torch.float16).contiguous().numpy().reshape(rows, nb, 1).view(np.uint8) |
| out = np.concatenate([hmask, qs.reshape(rows, nb, 64), scales, dd], axis=-1) |
| return out.reshape(*lead, nb * Q3_K_BYTES) |
|
|
|
|
| def iq4_xs_blocks(codes: torch.Tensor, d: torch.Tensor, m: torch.Tensor) -> np.ndarray: |
| """D-57 iq4xs packed codes [..., C/2] (`pack_4bit` of table indices 0..15), super-block scales d [..., C/256] |
| (fp16) and sub-block scales m [..., C/32] (int8, -32..31) → IQ4_XS bytes [..., C/256 · 136], ggml's block_iq4_xs: |
| d, scales_h (uint16: bits 2j, 2j+1 = the high two bits of m_j + 32), scales_l[4] (low nibbles of m_j + 32, sub-block |
| 2i in the low half of byte i), then per sub-block of 32 the 16 bytes Q4_0 and IQ4_NL use (j low, j+16 high).""" |
| from lowbit.quant.ternary import unpack_4bit |
| u = unpack_4bit(codes).contiguous().numpy() |
| lead, nb = u.shape[:-1], u.shape[-1] // IQ4_XS_BLOCK |
| u = u.reshape(-1, nb, 8, 2, 16) |
| rows = u.shape[0] |
| qs = (u[:, :, :, 0] | (u[:, :, :, 1] << 4)).reshape(rows, nb, 128) |
| ls = (m.to(torch.int16) + 32).contiguous().numpy().astype(np.uint16).reshape(rows, nb, 8) |
| scales_l = ((ls[:, :, 0::2] & 15) | ((ls[:, :, 1::2] & 15) << 4)).astype(np.uint8) |
| scales_h = np.zeros((rows, nb), dtype=np.uint16) |
| for j in range(8): |
| scales_h |= ((ls[:, :, j] >> 4) & 3) << (2 * j) |
| dd = d.to(torch.float16).contiguous().numpy().reshape(rows, nb, 1).view(np.uint8) |
| out = np.concatenate([dd, scales_h.astype("<u2").reshape(rows, nb, 1).view(np.uint8), scales_l, qs], axis=-1) |
| return out.reshape(*lead, nb * IQ4_XS_BYTES) |
|
|
|
|
| def packed_blocks(p: dict, gguf) -> tuple: |
| """(GGUF type, bytes) of one packed tensor in its native type: ternary and int2 → Q2_0, int4 → Q4_0, int3 → Q3_K.""" |
| T = gguf.GGMLQuantizationType |
| if "blocks" in p: |
| b = p["blocks"].numpy() if torch.is_tensor(p["blocks"]) else p["blocks"] |
| return getattr(T, p["format"]), np.ascontiguousarray(b.reshape(*b.shape[:-2], -1)) |
| method = p.get("method") |
| if method == "int4": |
| return T.Q4_0, q4_0_blocks(p["codes"], p["scales"]) |
| if method in ("int3", "int3nl"): |
| return T.Q3_K, q3_k_blocks(p["codes"], p["scales"], p["sub_scales"]) |
| if method == "iq4nl": |
| return T.IQ4_NL, q4_0_blocks(p["codes"], p["scales"]) |
| if method == "iq4xs": |
| return T.IQ4_XS, iq4_xs_blocks(p["codes"], p["scales"], p["sub_scales"]) |
| return T.Q2_0, q2_0_blocks(p["codes"], p["scales"], p["group_size"]) |
|
|
|
|
| def dequant_native(raw: np.ndarray, qtype, gguf) -> np.ndarray: |
| """A native tensor's bytes → float32: Q2_0 by `dequant_q2_0` (gguf-py has no Q2_0), the rest by gguf-py.""" |
| if qtype == gguf.GGMLQuantizationType.Q2_0: |
| return dequant_q2_0(raw) |
| return gguf.quants.dequantize(raw, qtype) |
|
|
|
|
| def bf16_bytes(t: torch.Tensor) -> np.ndarray: |
| t = t.to(torch.bfloat16).contiguous() |
| return t.view(torch.int16).numpy().view(np.uint8).reshape(*t.shape[:-1], t.shape[-1] * 2) |
|
|
|
|
| _LAYER = {"input_layernorm.weight": "attn_norm.weight", "post_attention_layernorm.weight": "ffn_norm.weight", |
| "self_attn.q_proj.weight": "attn_q.weight", "self_attn.k_proj.weight": "attn_k.weight", |
| "self_attn.v_proj.weight": "attn_v.weight", "self_attn.o_proj.weight": "attn_output.weight", |
| "self_attn.q_norm.weight": "attn_q_norm.weight", "self_attn.k_norm.weight": "attn_k_norm.weight", |
| "mlp.gate.weight": "ffn_gate_inp.weight", "mlp.gate_proj.weight": "ffn_gate.weight", |
| "mlp.up_proj.weight": "ffn_up.weight", "mlp.down_proj.weight": "ffn_down.weight"} |
| _TOP = {"model.embed_tokens.weight": "token_embd.weight", "model.norm.weight": "output_norm.weight", |
| "lm_head.weight": "output.weight"} |
|
|
|
|
| def gguf_name(hf: str) -> str: |
| if hf in _TOP: |
| return _TOP[hf] |
| m = re.fullmatch(r"model\.layers\.(\d+)\.(.+)", hf) |
| if not m or m.group(2) not in _LAYER: |
| raise KeyError(f"no GGUF name for {hf}") |
| return f"blk.{m.group(1)}.{_LAYER[m.group(2)]}" |
|
|
|
|
| def student_tensors(dense: dict[str, torch.Tensor], packed: dict[str, dict], gguf, |
| ternary_as: str = "q2_0", dense_as: str = "bf16") -> dict[str, tuple]: |
| """GGUF name → (type, bytes) for a student: `dense` holds plain tensors by HF name; `packed` holds ternary |
| weights by HF name (without `.weight`) as {codes, scales, group_size}. MoE experts (`...mlp.experts.gate_up_proj`, |
| `...down_proj`) split the fused gate/up rows the way llama.cpp's converter does (gate first). |
| |
| `ternary_as` "q2_0" writes each packed tensor in its native type (`packed_blocks`; D-51's N-bit formats too), |
| "bf16" the same weights dequantised. `dense_as` sets the other matrices (attention, embedding): "bf16" writes them exactly as trained; "q8_0" |
| quantises them post hoc, as the Q8 teacher GGUF does. On CPU the bf16 attention and embedding are ~85% of the |
| bytes a token reads (the ternary experts ~10%), so Q8_0 there makes the student faster than the Q8 teacher. |
| Norms and the router stay F32 either way.""" |
| T = gguf.GGMLQuantizationType |
| out = {} |
| for hf, t in dense.items(): |
| name = gguf_name(hf) |
| if t.dim() == 1 or name.endswith("ffn_gate_inp.weight"): |
| out[name] = (T.F32, t.float().contiguous().numpy()) |
| elif dense_as == "q8_0": |
| out[name] = (T.Q8_0, gguf.quants.quantize(t.float().contiguous().numpy(), T.Q8_0)) |
| else: |
| out[name] = (T.BF16, bf16_bytes(t)) |
| for hf, p in packed.items(): |
| if ternary_as == "bf16": |
| from lowbit.quant.ternary import dequant_packed |
| w = dequant_packed(p, torch.bfloat16) |
| blocks, qt = bf16_bytes(w), T.BF16 |
| else: |
| qt, blocks = packed_blocks(p, gguf) |
| m = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf) |
| if m: |
| i = m.group(1) |
| if m.group(2) == "gate_up": |
| half = blocks.shape[1] // 2 |
| out[f"blk.{i}.ffn_gate_exps.weight"] = (qt, np.ascontiguousarray(blocks[:, :half])) |
| out[f"blk.{i}.ffn_up_exps.weight"] = (qt, np.ascontiguousarray(blocks[:, half:])) |
| else: |
| out[f"blk.{i}.ffn_down_exps.weight"] = (qt, blocks) |
| else: |
| out[gguf_name(hf + ".weight")] = (qt, blocks) |
| return out |
|
|
|
|
| def write_gguf(template: str | Path, out: str | Path, tensors: dict[str, tuple], gguf, |
| metadata: dict | None = None) -> dict: |
| """Write `tensors` (GGUF name → (type, bytes or float32 array)) in the template's tensor order, with the |
| template's metadata. Every template tensor must be supplied, and nothing else. `metadata` (key → str or int) |
| replaces or adds general.* keys, e.g. a release's name and its base model.""" |
| metadata = metadata or {} |
| reader = gguf.GGUFReader(str(template)) |
| arch = reader.fields["general.architecture"].contents() |
| names = [t.name for t in reader.tensors] |
| missing, extra = set(names) - set(tensors), set(tensors) - set(names) |
| if missing or extra: |
| raise ValueError(f"tensor sets differ: missing {sorted(missing)[:5]}, extra {sorted(extra)[:5]}") |
| writer = gguf.GGUFWriter(str(out), arch=arch) |
| for field in reader.fields.values(): |
| if field.name == gguf.Keys.General.ARCHITECTURE or field.name.startswith("GGUF.") or field.name in metadata: |
| continue |
| val_type = field.types[0] |
| sub_type = field.types[-1] if val_type == gguf.GGUFValueType.ARRAY else None |
| value = field.contents() |
| if field.name == gguf.Keys.General.FILE_TYPE: |
| T, kinds = gguf.GGMLQuantizationType, {q for q, _ in tensors.values()} |
| cbk = sorted(k.name for k in kinds if k.name in ("IQ2_XXS", "IQ2_XS", "IQ2_S", "IQ3_XXS", "IQ3_S")) |
| if cbk: |
| value = int(getattr(gguf.LlamaFileType, f"MOSTLY_{cbk[0]}")) |
| elif T.Q3_K in kinds: |
| value = int(gguf.LlamaFileType.MOSTLY_Q3_K_M) |
| elif T.IQ4_XS in kinds: |
| value = int(gguf.LlamaFileType.MOSTLY_IQ4_XS) |
| elif T.Q2_0 in kinds: |
| value = int(gguf.LlamaFileType.MOSTLY_Q2_0) |
| elif T.Q4_0 in kinds: |
| value = int(gguf.LlamaFileType.MOSTLY_Q4_0) |
| elif T.IQ4_NL in kinds: |
| value = int(gguf.LlamaFileType.MOSTLY_IQ4_NL) |
| writer.add_key_value(field.name, value, val_type, sub_type=sub_type) |
| for key, value in metadata.items(): |
| vt = gguf.GGUFValueType.UINT32 if isinstance(value, int) else gguf.GGUFValueType.STRING |
| writer.add_key_value(key, value, vt) |
| for name in names: |
| qtype, data = tensors[name] |
| writer.add_tensor_info(name, data.shape, data.dtype, data.nbytes, raw_dtype=qtype) |
| writer.write_header_to_file() |
| writer.write_kv_data_to_file() |
| writer.write_ti_data_to_file() |
| for name in names: |
| writer.write_tensor_data(tensors[name][1]) |
| writer.close() |
| return {"tensors": len(names), "bytes": Path(out).stat().st_size} |
|
|
|
|
| def check_round_trip(out: str | Path, packed: dict[str, dict], gguf) -> dict: |
| """Acceptance: every scale an fp16 normal, and the GGUF's packed tensors (Q2_0; D-51: also Q4_0, Q3_K) dequantise |
| bit-exactly to the packed student's weights (`dequant_packed`).""" |
| from lowbit.quant.ternary import dequant_packed |
| reader = gguf.GGUFReader(str(out)) |
| by = {t.name: t for t in reader.tensors} |
| cbp = {k: p for k, p in packed.items() if "blocks" in p} |
| if cbp: |
| bad = [] |
| for hf, p in cbp.items(): |
| i, name = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf).groups() |
| b = p["blocks"].numpy() |
| parts = {"gate": b[:, :b.shape[1] // 2], "up": b[:, b.shape[1] // 2:]} if name == "gate_up" else {"down": b} |
| for k, want in parts.items(): |
| got = np.asarray(by[f"blk.{i}.ffn_{k}_exps.weight"].data) |
| if got.tobytes() != np.ascontiguousarray(want).tobytes(): |
| bad.append(f"blk.{i}.ffn_{k}_exps") |
| packed = {k: p for k, p in packed.items() if "blocks" not in p} |
| if not packed: |
| return {"codebook_tensors": 2 * len(cbp) // 2, "mismatched_tensors": bad, "bit_exact": not bad} |
| scales = torch.cat([p["scales"].float().flatten() for p in packed.values()]) |
| nz = scales[scales != 0] |
| res = {"scale_min_nonzero": float(nz.abs().min()) if nz.numel() else None, "zero_scales": int((scales == 0).sum()), |
| "scales_below_fp16_normal": int(((nz.abs() < FP16_MIN_NORMAL)).sum()), "mismatched_tensors": []} |
|
|
| def deq(name: str) -> np.ndarray: |
| return dequant_native(np.asarray(by[name].data), by[name].tensor_type, gguf) |
|
|
| for hf, p in packed.items(): |
| ref = dequant_packed(p, torch.float32).numpy() |
| m = re.fullmatch(r"model\.layers\.(\d+)\.mlp\.experts\.(gate_up|down)_proj", hf) |
| if m and m.group(2) == "gate_up": |
| got = np.concatenate([deq(f"blk.{m.group(1)}.ffn_{k}_exps.weight") for k in ("gate", "up")], axis=1) |
| elif m: |
| got = deq(f"blk.{m.group(1)}.ffn_down_exps.weight") |
| else: |
| got = deq(gguf_name(hf + ".weight")) |
| if got.shape != ref.shape or not np.array_equal(got, ref): |
| res["mismatched_tensors"].append(hf) |
| res["bit_exact"] = not res["mismatched_tensors"] and res["scales_below_fp16_normal"] == 0 |
| return res |
|
|
