diff --git a/README.md b/README.md index 3cdfdb5..56f9d19 100644 --- a/README.md +++ b/README.md @@ -17,12 +17,12 @@ | Project | Venue | Public Release | | --- | --- | --- | -| [CAT-Q](projects/cat-q) | ICML 2026 Oral | Model checkpoints, inference, and evaluation code | +| [CAT-Q](projects/cat-q) | ICML 2026 Oral | Model checkpoints, inference, evaluation, and packed ternary deployment code | ## Latest News - `[04/08/2026]` πŸ”₯ [The technical report of ScaleQ-1.58](https://arxiv.org/abs/2608.01078) "**Attend to Your Own Thoughts: Breaking the Barrier for Post-Training Quantization of Reasoning LLMs through the Lens of 1.58-Bit Quantization**" is now available on arXiv. -- `[22/07/2026]` πŸš€ [The CAT-Q model checkpoints](projects/cat-q) (including Qwen3-1.7B/4B/8B/14B/32B, Llama2-7B, Qwen3-30B-A3B and Qwen3-235B-A22B), inference, and evaluation code are now available. +- `[22/07/2026]` πŸš€ [The CAT-Q model checkpoints](projects/cat-q) (including Qwen3-1.7B/4B/8B/14B/32B, Llama2-7B, Qwen3-30B-A3B and Qwen3-235B-A22B), inference, evaluation, and **real ternary deployment** code are now available. - `[25/06/2026]` πŸ”₯ [The CAT-Q paper](https://arxiv.org/abs/2606.26650) is now available on arXiv. - `[01/05/2026]` πŸŽ‰Our paper "**CAT-Q: Cost-efficient and Accurate Ternary Quantization for LLMs**" is accepted to **ICML 2026 as an oral**. The project page for our sliding-layer reconstruction framework used in CAT-Q is available at [SliderQuant (ICLR 2026)](https://github.com/deep-optimization/SliderQuant). diff --git a/projects/cat-q/README.md b/projects/cat-q/README.md index 6534607..2f04040 100644 --- a/projects/cat-q/README.md +++ b/projects/cat-q/README.md @@ -15,10 +15,11 @@ This repository contains the official implementation of **CAT-Q (ICML 2026 Oral) ## Latest News -- `[Stay tuned]` We are preparing to release the CAT-Q training code. -- `[22/07/2026]` We release the CAT-Q model checkpoints, inference, and evaluation code. -- `[25/06/2026]` The [CAT-Q paper](https://arxiv.org/abs/2606.26650) is available on arXiv. -- `[01/05/2026]` πŸŽ‰πŸŽ‰πŸŽ‰**CAT-Q: Cost-efficient and Accurate Ternary Quantization for LLMs** is accepted to ICML 2026 as an oral. +- `[Stay tuned]` We are preparing to release the CAT-Q training code, etc. +- `[04/08/2026]` [The technical report](https://arxiv.org/abs/2608.01078) "**Attend to Your Own Thoughts: Breaking the Barrier for Post-Training Quantization of Reasoning LLMs through the Lens of 1.58-Bit Quantization**" is now available on arXiv. +- `[22/07/2026]` πŸŽ‰πŸŽ‰πŸŽ‰The CAT-Q model checkpoints, **scaling from Qwen3-1.7B all the way to Qwen3-235B-A22B** (Qwen3-1.7B/4B/8B/14B/32B, Llama2-7B, Qwen3-30B-A3B and Qwen3-235B-A22B), inference, evaluation, and **real ternary deployment** code are now available. +- `[25/06/2026]` [The CAT-Q paper](https://arxiv.org/abs/2606.26650) is now available on arXiv. +- `[01/05/2026]` πŸŽ‰πŸŽ‰πŸŽ‰Our paper "**CAT-Q: Cost-efficient and Accurate Ternary Quantization for LLMs**" is accepted to **ICML 2026 as an oral**. The project page for our sliding-layer reconstruction framework used in CAT-Q is available at [SliderQuant (ICLR 2026)](https://github.com/deep-optimization/SliderQuant). ## Overview @@ -49,6 +50,7 @@ Using only 512 calibration samples, CAT-Q scales W1.58 quantization from 1.7B de - [Installation](#installation) - [Evaluation](#evaluation) - [Hugging Face Export](#hugging-face-export) +- [Packed Ternary Deployment](#packed-ternary-deployment) - [Citation](#citation) - [Acknowledgement](#acknowledgement) - [License](#license) @@ -87,6 +89,9 @@ The following W1.58A16 checkpoints are available on Hugging Face: | Qwen3-235B-A22B | MoE | [qwen3-moe-235B-A22B](https://huggingface.co/IntelLabsChina/CAT-Q/tree/main/qwen3-moe-235B-A22B) | All checkpoints are hosted under [IntelLabsChina/CAT-Q](https://huggingface.co/IntelLabsChina/CAT-Q). +Every folder holds the learnable CAT-Q parameters (`parameters.pth`) with the config that +produced them, plus a ready-to-run packed ternary `*-catq-q2_0.gguf`; see +[deployment/README.md](deployment/README.md) for how to serve it. > **Note:** The code was refactored for open-source release, so checkpoint accuracy may differ slightly from the paper results (typically within Β±0.2 percentage points). @@ -137,6 +142,16 @@ Export the restored model as a fake-quantized Hugging Face model: The exported model retains the original Hugging Face architecture and stores merged fake-quantized floating-point weights; it is not a packed ternary checkpoint. +## Packed Ternary Deployment + +Export the restored model as a packed ternary GGUF, where each quantized weight occupies 2 bits next to one fp16 scale per group of 128: + +```bash +./export_gguf.sh +``` + +Conversion is self-contained: it reads the checkpoint and writes the GGUF, with no intermediate model and no inference runtime involved. The result runs on the ternary kernels of the [Bonsai](https://github.com/PrismML-Eng/Bonsai-demo) runtime, with real weight compression rather than fake quantization. See [`deployment/README.md`](deployment/README.md) for the full export-and-serve walkthrough. + ## Citation If CAT-Q is useful in your research, please cite: @@ -152,7 +167,7 @@ If CAT-Q is useful in your research, please cite: ## Acknowledgement -CAT-Q is implemented based on [SliderQuant](https://github.com/deep-optimization/SliderQuant). +CAT-Q is implemented based on [SliderQuant](https://github.com/deep-optimization/SliderQuant). Packed ternary deployment builds on the group-128 ternary kernels of [Bonsai](https://github.com/PrismML-Eng/Bonsai-demo) and its [llama.cpp fork](https://github.com/PrismML-Eng/llama.cpp). ## License diff --git a/projects/cat-q/deployment/README.md b/projects/cat-q/deployment/README.md new file mode 100644 index 0000000..a709894 --- /dev/null +++ b/projects/cat-q/deployment/README.md @@ -0,0 +1,157 @@ +# Deploying CAT-Q Models with Packed Ternary Weights + +The [Hugging Face export](../README.md#hugging-face-export) writes a *fake-quantized* +model: the weights are ternary, but each one still occupies 16 bits and runs through +ordinary floating-point matrix multiplication. + +This document covers the other export path, which stores the quantized projections as +**group-128 packed ternary weights** and runs them through the ternary kernels of +[llama.cpp](https://github.com/PrismML-Eng/llama.cpp) as used by the +[Bonsai demo](https://github.com/PrismML-Eng/Bonsai-demo). + +CAT-Q quantizes a weight group as + +``` +W_g = s_g * T_g, T_g in {-1, 0, +1}, |g| = 128 +``` + +which is exactly the `Q2_0` block type of that runtime (`block_q2_0`: one fp16 scale +plus 128 two-bit codes, 34 bytes, 2.125 bits per weight). The exporter takes the codes +and scales straight out of the CAT-Q quantizer and packs them, so the deployed weights +are bit-for-bit the ones the fake-quantized model uses. Embeddings, norms, the LM head +and MoE routers stay in floating point, as they are outside the quantized set. + +The implementation lives with the rest of the quantization code: +`quantize/ternary_export.py` recovers the codes and scales, `quantize/q2_0.py` packs +them into `Q2_0` blocks, and `quantize/gguf_export.py` writes the GGUF. + +## 1. Export the model + +Exporting is a pure CAT-Q step: it reads the checkpoint and writes the GGUF directly. +No inference runtime is involved, and no intermediate model is produced. Beyond the +project requirements it only needs the `gguf` package (`pip install gguf`, already in +`pyproject.toml`). + +```bash +cd BitTern/projects/cat-q + +# select the checkpoint, exactly as for evaluation +# result_dir=configs/qwen3-4b in task_list.conf +./export_gguf.sh +``` + +This writes `configs//export-gguf/-catq-q2_0.gguf`. + +The launcher is a thin wrapper; the same thing can be run directly: + +```bash +python main.py \ + --config configs/qwen3-4b/config.yaml \ + --checkpoint configs/qwen3-4b/parameters.pth \ + --output_dir configs/qwen3-4b/export-gguf \ + --export_gguf_path configs/qwen3-4b/export-gguf/Qwen3-4B-catq-q2_0.gguf +``` + +`--export_gguf_path` takes either a `.gguf` file or a directory, in which case the file +is named `-catq-q2_0.gguf`. `--gguf_float_type {f16,bf16,f32}` selects the dtype of +the tensors CAT-Q leaves in floating point and defaults to `f16`; norms and MoE routers +are always `F32`, as in a stock llama.cpp conversion. + +Dense (Qwen3, LLaMA) and MoE (Qwen3-MoE) checkpoints are both supported. For MoE models +the per-expert `gate_proj`/`up_proj`/`down_proj` weights are packed and stacked into the +`ffn_*_exps` tensors the runtime expects, while the router stays in `F32`. + +### Very large checkpoints + +By default the exporter keeps the whole 16-bit model resident while it packs, which needs +roughly the size of the original model plus the size of the GGUF. Add `--gguf_low_memory` +when that does not fit - Qwen3-235B-A22B, for instance, needs it on a 512 GiB host: + +```bash +python main.py \ + --config configs/qwen3-moe-235B-A22B/config.yaml \ + --checkpoint configs/qwen3-moe-235B-A22B/parameters.pth \ + --output_dir configs/qwen3-moe-235B-A22B/export-gguf \ + --export_gguf_path configs/qwen3-moe-235B-A22B/export-gguf \ + --gguf_low_memory +``` + +Each weight is then dropped as soon as it has been packed and the GGUF is assembled +through a temporary file (set `TMPDIR` to a filesystem with room for the result), which +holds host memory to about the size of the packed model. The file is byte-for-byte the +same as without the flag. Because the 16-bit weights are gone by the time packing ends, +the flag cannot be combined with `--tasks` or `--export_model_path`. + +## 2. Get a runtime with ternary kernels + +The GGUF needs a llama.cpp build with group-128 ternary kernels: + +```bash +git clone -b prism https://github.com/PrismML-Eng/llama.cpp.git +export LLAMA_CPP_DIR=$PWD/llama.cpp +``` + +The easiest way to build it is with the Bonsai demo's own scripts, which also fetch the +runtime for you if it is missing: + +```bash +git clone https://github.com/PrismML-Eng/Bonsai-demo.git +cd Bonsai-demo +./scripts/build_cuda_linux.sh "$LLAMA_CPP_DIR" # CUDA; build_cpu_linux.sh / build_mac.sh also exist +``` + +Binaries land in `Bonsai-demo/bin//`. Prebuilt binaries and other backends +(Metal, Vulkan, ROCm) are described in the Bonsai demo README. + +## 3. Run it + +The result is a standard GGUF file, so any tool from the runtime works with it: + +```bash +BIN=/path/to/Bonsai-demo/bin/cuda +export LD_LIBRARY_PATH="$BIN:$LD_LIBRARY_PATH" + +# one-off generation +"$BIN/llama-cli" -m Qwen3-4B-catq-q2_0.gguf -ngl 99 -p "Explain ternary quantization." + +# OpenAI-compatible server + web UI on http://localhost:8080 +"$BIN/llama-server" -m Qwen3-4B-catq-q2_0.gguf -ngl 99 -c 8192 -fa on + +# throughput and memory +"$BIN/llama-bench" -m Qwen3-4B-catq-q2_0.gguf -ngl 99 +``` + +Notes for Qwen3 checkpoints, which are thinking models: + +- `--reasoning-format deepseek` keeps `` blocks out of `message.content` and puts + them in `message.reasoning_content`. +- `--chat-template-kwargs '{"enable_thinking": false}'` turns thinking off. +- `-fit off` stops the server from growing the KV cache to fill the device memory, which + is worth setting when measuring the memory footprint. + +For model management, the web UI, tool calling, speculative decoding and non-Linux +platforms, follow the [Bonsai demo](https://github.com/PrismML-Eng/Bonsai-demo) +documentation; a CAT-Q GGUF can be used wherever it expects a Bonsai ternary model. + +## Reference numbers + +Qwen3-4B, single NVIDIA L40, context 2048, measured with `llama-bench`: + +| | packed ternary | fake-quantized `F16` | ratio | +| --- | ---: | ---: | ---: | +| file size | 1.63 GiB | 7.50 GiB | 4.60x | +| device memory | 2093 MiB | 8745 MiB | 4.18x | +| decode (tg128) | 285.1 t/s | 93.5 t/s | 3.05x | + +252 of the 398 tensors are packed ternary and hold 3.63 B of the weights at 2.125 bits +each; the remaining floating-point tensors (mostly the token embedding) account for most +of what is left, which is why the whole-file ratio is below the 7.53x of the quantized +part alone. Task accuracy matches the fake-quantized model to within run-to-run noise. + +## Acknowledgement + +The packed ternary format and the kernels used here come from +[Bonsai](https://github.com/PrismML-Eng/Bonsai-demo) by PrismML and its +[llama.cpp fork](https://github.com/PrismML-Eng/llama.cpp) (branch `prism`), built on +[llama.cpp](https://github.com/ggml-org/llama.cpp). We thank their authors for making +efficient ternary inference available to the community. diff --git a/projects/cat-q/export_gguf.sh b/projects/cat-q/export_gguf.sh new file mode 100755 index 0000000..2ad8e33 --- /dev/null +++ b/projects/cat-q/export_gguf.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Export the selected CAT-Q checkpoint as a packed ternary GGUF model. +# +# The GGUF is written directly from the checkpoint; no llama.cpp checkout and no +# intermediate model are involved. See deployment/README.md for how to run the +# result. +set -euo pipefail + +source ./task_list.conf +source ./scripts/gpu_lock.sh + +config_path="${result_dir}/config.yaml" +checkpoint_path="${result_dir}/parameters.pth" +export_path="${result_dir}/export-gguf" +if [[ ! -f "${config_path}" ]]; then + echo "Config file not found: ${config_path}" >&2 + exit 2 +fi +if [[ ! -f "${checkpoint_path}" ]]; then + echo "Checkpoint not found: ${checkpoint_path}" >&2 + exit 2 +fi + +mkdir -p "${export_path}" +catq_acquire_gpus 1 "${THRESHOLD}" "${WAIT_MODE}" "${WAIT_INTERVAL}" +trap catq_release_gpus EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +python main.py \ + --config "${config_path}" \ + --output_dir "${export_path}" \ + --export_gguf_path "${export_path}" \ + --gguf_float_type f16 \ + --checkpoint "${checkpoint_path}" diff --git a/projects/cat-q/main.py b/projects/cat-q/main.py index 6dd79eb..f48902f 100644 --- a/projects/cat-q/main.py +++ b/projects/cat-q/main.py @@ -7,7 +7,10 @@ def build_parser(): parser = argparse.ArgumentParser( - description="Load released CAT-Q parameters for evaluation or fake-quantized Hugging Face export." + description=( + "Load released CAT-Q parameters for evaluation, fake-quantized Hugging Face " + "export, or packed ternary GGUF export." + ) ) parser.add_argument( "--config", @@ -29,6 +32,30 @@ def build_parser(): default=None, help="Directory for the fake-quantized HF model", ) + parser.add_argument( + "--export_gguf_path", + type=str, + default=None, + help=( + "Where to write the packed ternary GGUF: a .gguf file, or a directory " + "in which -catq-q2_0.gguf is created" + ), + ) + parser.add_argument( + "--gguf_float_type", + choices=["f16", "bf16", "f32"], + default="f16", + help="Dtype for the GGUF tensors CAT-Q keeps in floating point", + ) + parser.add_argument( + "--gguf_low_memory", + action="store_true", + help=( + "Free each 16-bit weight as soon as it has been packed and spool the " + "GGUF through a temporary file; needed for models that do not fit in " + "host memory twice. Cannot be combined with --tasks or --export_model_path" + ), + ) parser.add_argument("--net", type=str, default=None) parser.add_argument( "--quant_layer_list", @@ -112,12 +139,36 @@ def parse_arguments(argv=None): parser.error("--model is required (directly or through --config)") if not args.checkpoint: parser.error("--checkpoint is required") - if not args.tasks and not args.export_model_path: - parser.error("select at least one action: --tasks or --export_model_path") + if not args.tasks and not args.export_model_path and not args.export_gguf_path: + parser.error("select at least one action: --tasks, --export_model_path, or --export_gguf_path") + if args.gguf_low_memory: + if not args.export_gguf_path: + parser.error("--gguf_low_memory requires --export_gguf_path") + if args.tasks or args.export_model_path: + parser.error( + "--gguf_low_memory discards the 16-bit weights while packing, so it " + "cannot be combined with --tasks or --export_model_path" + ) args.ignored_config_keys = ignored_config_keys return args +def _save_hf_model(lm, directory, logger): + directory.mkdir(parents=True, exist_ok=True) + lm.model.save_pretrained(directory) + lm.tokenizer.save_pretrained(directory) + logger.info("Saved fake-quantized Hugging Face model to %s", directory) + + +def _gguf_outfile(args): + path = Path(args.export_gguf_path) + if path.suffix == ".gguf": + path.parent.mkdir(parents=True, exist_ok=True) + return path + path.mkdir(parents=True, exist_ok=True) + return path / f"{args.net}-catq-q2_0.gguf" + + def main(argv=None): args = parse_arguments(argv) @@ -146,15 +197,31 @@ def main(argv=None): args.net = args.model.rstrip("/").split("/")[-1] args.quant_rate = 1.0 + exporter = None + if args.export_gguf_path: + from quantize.gguf_export import TernaryGGUFExporter + + exporter = TernaryGGUFExporter( + args.model, + _gguf_outfile(args), + float_type=args.gguf_float_type, + low_memory=args.gguf_low_memory, + ) + lm = LMClass(args) - merge_catq_checkpoint(lm, args, logger) + merge_catq_checkpoint( + lm, + args, + logger, + ternary_sink=exporter.capture if exporter else None, + release_packed_weights=args.gguf_low_memory, + ) if args.export_model_path: - export_dir = Path(args.export_model_path) - export_dir.mkdir(parents=True, exist_ok=True) - lm.model.save_pretrained(export_dir) - lm.tokenizer.save_pretrained(export_dir) - logger.info("Saved fake-quantized Hugging Face model to %s", export_dir) + _save_hf_model(lm, Path(args.export_model_path), logger) + + if exporter is not None: + logger.info("Saved packed ternary model to %s", exporter.write(lm.model)) if args.tasks: evaluate(lm, args, logger) diff --git a/projects/cat-q/pyproject.toml b/projects/cat-q/pyproject.toml index ed537d0..e828dfd 100644 --- a/projects/cat-q/pyproject.toml +++ b/projects/cat-q/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "catq" version = "0.1.0" -description = "Evaluation and fake-quantized Hugging Face export for released CAT-Q checkpoints." +description = "Evaluation, fake-quantized Hugging Face export and packed ternary GGUF export for released CAT-Q checkpoints." readme = "README.md" license = {file = "LICENSE"} requires-python = ">=3.10,<3.11" @@ -19,6 +19,7 @@ classifiers = [ dependencies = [ "accelerate>=0.30.0,<2.0.0", "datasets==2.18.0", + "gguf>=0.17.0", "huggingface-hub==0.32.4", "lm-eval==0.4.7", "numpy>=1.26.0", diff --git a/projects/cat-q/quantize/gguf_export.py b/projects/cat-q/quantize/gguf_export.py new file mode 100644 index 0000000..20e7be6 --- /dev/null +++ b/projects/cat-q/quantize/gguf_export.py @@ -0,0 +1,468 @@ +"""Write a CAT-Q checkpoint out as a packed ternary GGUF file. + +The quantized projections are stored as `Q2_0`, the group-128 ternary weight +type of the Bonsai llama.cpp fork (https://github.com/PrismML-Eng/llama.cpp, +branch `prism`); see `quantize/q2_0.py` for the block layout. Everything the +CAT-Q recipe leaves in floating point - token embeddings, the LM head, every +norm and the MoE router - is written as F16 (or `--gguf_float_type`) and F32 +exactly as a stock conversion would. + +Only the `gguf` PyPI package is required. Producing the file has no dependency +on a llama.cpp checkout; llama.cpp is needed to *run* the result, not to build +it. + +Usage is through `main.py --export_gguf_path`, which feeds `capture()` from the +checkpoint restoration in `quantize.merge` and then calls `write()`; see +`deployment/README.md` for the export-and-serve walkthrough. +""" + +import json +import logging +from pathlib import Path + +import gguf +import numpy as np +import torch + +from .q2_0 import LLAMA_FTYPE_MOSTLY_Q2_0, pack_q2_0, register_q2_0 +from .ternary_export import iter_layer_ternary + +logger = logging.getLogger(__name__) + +FLOAT_TYPES = { + "f16": gguf.GGMLQuantizationType.F16, + "bf16": gguf.GGMLQuantizationType.BF16, + "f32": gguf.GGMLQuantizationType.F32, +} + +# Tokenizer pre-tokenizer identifiers used by llama.cpp. Guessing one silently +# changes tokenization, so only architectures that have been checked are listed. +BPE_PRE_TOKENIZERS = { + "qwen3": "qwen2", + "qwen3moe": "qwen2", +} + + +class _Arch: + """Per-architecture conversion rules.""" + + def __init__(self, name, arch, vocab, permute_qk=False, rope_dim=False, vocab_size=False): + self.name = name + self.arch = arch + self.vocab = vocab + self.permute_qk = permute_qk + self.rope_dim = rope_dim + self.vocab_size = vocab_size + + +ARCHITECTURES = { + "Qwen3ForCausalLM": _Arch("qwen3", gguf.MODEL_ARCH.QWEN3, "bpe"), + "Qwen3MoeForCausalLM": _Arch("qwen3moe", gguf.MODEL_ARCH.QWEN3MOE, "bpe"), + "LlamaForCausalLM": _Arch( + "llama", gguf.MODEL_ARCH.LLAMA, "spm", permute_qk=True, rope_dim=True, vocab_size=True + ), +} + + +def _permute(tensor, n_head, n_head_kv): + """The q/k row interleaving llama.cpp undoes for LLaMA-style RoPE.""" + if n_head_kv is not None and n_head != n_head_kv: + n_head = n_head_kv + return ( + tensor.reshape(n_head, 2, tensor.shape[0] // n_head // 2, *tensor.shape[1:]) + .swapaxes(1, 2) + .reshape(tensor.shape) + ) + + +def _pin_malloc_thresholds(): + """Keep glibc handing large blocks back to the kernel. + + The exporter frees a whole decoder layer worth of float tensors on every + iteration. glibc grows its dynamic mmap threshold whenever such a block is + released, after which same-sized allocations come from the heap and their + memory is never returned, so a long export slowly accumulates hundreds of + gigabytes of freed-but-resident memory. Pinning the threshold disables that + heuristic. + """ + import ctypes + + M_MMAP_THRESHOLD = -3 + M_TRIM_THRESHOLD = -1 + try: + libc = ctypes.CDLL("libc.so.6", use_errno=True) + libc.mallopt(ctypes.c_int(M_MMAP_THRESHOLD), ctypes.c_int(128 * 1024)) + libc.mallopt(ctypes.c_int(M_TRIM_THRESHOLD), ctypes.c_int(128 * 1024)) + except OSError: # pragma: no cover - non-glibc platforms + logger.debug("Could not pin the malloc thresholds; this is glibc-only") + + +class TernaryGGUFExporter: + """Turn a restored CAT-Q model into a single `Q2_0` GGUF file. + + `capture` is handed to `merge_catq_checkpoint`, which calls it once per + decoder layer while the ternary codes and scales still exist. `write` then + walks the merged model and swaps the packed bytes in for the corresponding + float weights. + """ + + def __init__(self, model_dir, outfile, float_type="f16", low_memory=False): + if float_type not in FLOAT_TYPES: + raise ValueError(f"Unknown float type {float_type!r}; expected one of {sorted(FLOAT_TYPES)}") + self.model_dir = Path(model_dir) + self.outfile = Path(outfile) + self.float_type = FLOAT_TYPES[float_type] + self.low_memory = low_memory + if low_memory: + _pin_malloc_thresholds() + # AutoConfig rather than raw config.json: llama.cpp reads the same + # normalised view, so defaults such as `head_dim` and `rope_theta` are + # filled in for older checkpoints that omit them. + from transformers import AutoConfig + + self.hparams = AutoConfig.from_pretrained(self.model_dir).to_dict() + + architectures = self.hparams.get("architectures") or [] + if not architectures or architectures[0] not in ARCHITECTURES: + raise ValueError( + f"Packed ternary export does not know architecture {architectures}; " + f"supported: {sorted(ARCHITECTURES)}" + ) + self.arch = ARCHITECTURES[architectures[0]] + self.block_count = self.hparams["num_hidden_layers"] + self.n_head = self.hparams["num_attention_heads"] + self.n_head_kv = self.hparams.get("num_key_value_heads", self.n_head) + + self.qtype_q2_0 = register_q2_0() + self.tensor_map = gguf.get_tensor_name_map(self.arch.arch, self.block_count) + # In low-memory mode the packed bytes are spooled to a temporary file as + # soon as they are handed to the writer, so only one tensor at a time has + # to be resident while the GGUF is being assembled. + self.writer = gguf.GGUFWriter(path=None, arch=self.arch.name, use_temp_file=low_memory) + self.packed = {} # hf weight name -> packed uint8 array + self.n_packed = 0 + + # ------------------------------------------------------------------ + # collection + # ------------------------------------------------------------------ + @torch.no_grad() + def capture(self, layer_id, qlayer): + """Pack every ternary projection of one restored decoder layer. + + Returns the Hugging Face weight names that were packed, so that the + caller can drop the corresponding floating-point weights (see + `quantize.merge`). + """ + names = [] + for name, ternary in iter_layer_ternary(layer_id, qlayer): + codes = ternary.codes.cpu().numpy() + scales = ternary.scales.cpu().numpy().reshape(codes.shape[:-1] + (-1,)) + self.packed[name] = self._pack(name, codes, scales) + names.append(name) + self.n_packed += len(names) + logger.info("Packed %d ternary tensors from layer %d", len(names), layer_id) + return names + + def _pack(self, name, codes, scales): + """Pack one weight, undoing the LLaMA q/k row interleaving if needed. + + The scales are laid out as one column per group, so the row permutation + applies to them unchanged. + """ + if self.arch.permute_qk and name.endswith("q_proj.weight"): + codes = _permute(codes, self.n_head, self.n_head) + scales = _permute(scales, self.n_head, self.n_head) + elif self.arch.permute_qk and name.endswith("k_proj.weight"): + codes = _permute(codes, self.n_head, self.n_head_kv) + scales = _permute(scales, self.n_head, self.n_head_kv) + return pack_q2_0(codes, scales) + + # ------------------------------------------------------------------ + # metadata + # ------------------------------------------------------------------ + def _set_parameters(self): + writer, hparams = self.writer, self.hparams + writer.add_block_count(self.block_count) + writer.add_context_length(hparams["max_position_embeddings"]) + writer.add_embedding_length(hparams["hidden_size"]) + writer.add_feed_forward_length(hparams["intermediate_size"]) + writer.add_head_count(self.n_head) + writer.add_head_count_kv(self.n_head_kv) + + rope = hparams.get("rope_parameters") or hparams.get("rope_scaling") or {} + rope_type = rope.get("rope_type", rope.get("type")) + factor = rope.get("factor") + if rope_type == "linear" and factor is not None: + writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + writer.add_rope_scaling_factor(factor) + elif rope_type == "yarn" and factor is not None: + writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + writer.add_rope_scaling_factor(factor) + writer.add_rope_scaling_orig_ctx_len(rope["original_max_position_embeddings"]) + elif rope_type not in (None, "default"): + raise ValueError(f"Unsupported rope_scaling type {rope_type!r} for packed ternary export") + + writer.add_rope_freq_base(hparams.get("rope_theta") or rope["rope_theta"]) + writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"]) + + if (n_experts := hparams.get("num_experts")) is not None: + writer.add_expert_count(n_experts) + if (n_used := hparams.get("num_experts_per_tok")) is not None: + writer.add_expert_used_count(n_used) + if (moe_ff := hparams.get("moe_intermediate_size")) is not None: + writer.add_expert_feed_forward_length(moe_ff) + + head_dim = hparams.get("head_dim") + if head_dim is not None: + writer.add_key_length(head_dim) + writer.add_value_length(head_dim) + if self.arch.vocab_size: + writer.add_vocab_size(hparams["vocab_size"]) + if self.arch.rope_dim: + writer.add_rope_dimension_count(head_dim or hparams["hidden_size"] // self.n_head) + + writer.add_file_type(LLAMA_FTYPE_MOSTLY_Q2_0) + + def _set_vocab(self): + if self.arch.vocab == "bpe": + self._set_vocab_bpe() + else: + self._set_vocab_spm() + gguf.SpecialVocab(self.model_dir, load_merges=self.arch.vocab == "bpe").add_to_gguf(self.writer) + + def _looks_special(self, token): + """Added tokens that llama.cpp treats as control tokens even when the + Hugging Face tokenizer does not flag them as special.""" + if isinstance(token, (bytes, bytearray)): + token = token.decode("utf-8") + return ( + token in ("", "", "<2mass>", "[@BOS@]") + or (token.startswith("<|") and token.endswith("|>")) + or (token.startswith("<\uff5c") and token.endswith("\uff5c>")) + or (token.startswith("")) + ) + + def _set_vocab_bpe(self): + from transformers import AutoTokenizer + + pre = BPE_PRE_TOKENIZERS.get(self.arch.name) + if pre is None: + raise ValueError(f"No known llama.cpp pre-tokenizer for architecture {self.arch.name!r}") + + tokenizer = AutoTokenizer.from_pretrained(self.model_dir) + vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) + reverse_vocab = {index: token for token, index in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() + added_decoder = tokenizer.added_tokens_decoder + + tokens, toktypes = [], [] + for index in range(vocab_size): + if index not in reverse_vocab: + tokens.append(f"[PAD{index}]") + toktypes.append(gguf.TokenType.UNUSED) + continue + token = reverse_vocab[index] + if token in added_vocab: + if not added_decoder[index].normalized: + # llama.cpp expects added tokens in their normalized form. + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) + if added_decoder[index].special or self._looks_special(token): + toktypes.append(gguf.TokenType.CONTROL) + else: + token = token.replace("\u2581", " ") + toktypes.append(gguf.TokenType.USER_DEFINED) + else: + toktypes.append(gguf.TokenType.NORMAL) + tokens.append(token) + + self.writer.add_tokenizer_model("gpt2") + self.writer.add_tokenizer_pre(pre) + self.writer.add_token_list(tokens) + self.writer.add_token_types(toktypes) + + def _set_vocab_spm(self): + from sentencepiece import SentencePieceProcessor + + model_file = self.model_dir / "tokenizer.model" + if not model_file.is_file(): + raise FileNotFoundError( + f"{model_file} is required for a SentencePiece vocabulary; " + "BPE-only LLaMA variants are not supported by this exporter" + ) + sp = SentencePieceProcessor() + sp.LoadFromFile(str(model_file)) + + vocab_size = self.hparams.get("vocab_size", sp.vocab_size()) + tokens = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores = [-10000.0] * vocab_size + toktypes = [gguf.TokenType.UNUSED] * vocab_size + + for index in range(min(sp.vocab_size(), vocab_size)): + tokens[index] = sp.IdToPiece(index).encode("utf-8") + scores[index] = sp.GetScore(index) + if sp.IsUnknown(index): + toktypes[index] = gguf.TokenType.UNKNOWN + elif sp.IsControl(index): + toktypes[index] = gguf.TokenType.CONTROL + elif sp.IsUnused(index): + toktypes[index] = gguf.TokenType.UNUSED + elif sp.IsByte(index): + toktypes[index] = gguf.TokenType.BYTE + else: + toktypes[index] = gguf.TokenType.NORMAL + + for index, token in self._added_tokens().items(): + if index >= vocab_size: + continue + content = token["content"] + if token.get("special") or self._looks_special(content): + toktypes[index] = gguf.TokenType.CONTROL + else: + content = content.replace("\u2581", " ") + toktypes[index] = gguf.TokenType.USER_DEFINED + tokens[index] = content.encode("utf-8") + scores[index] = -1000.0 + + self.writer.add_tokenizer_model("llama") + self.writer.add_tokenizer_pre("default") + self.writer.add_token_list(tokens) + self.writer.add_token_scores(scores) + self.writer.add_token_types(toktypes) + + def _added_tokens(self): + added = {} + path = self.model_dir / "added_tokens.json" + if path.is_file(): + for content, index in json.loads(path.read_text(encoding="utf-8")).items(): + added[int(index)] = {"content": content} + path = self.model_dir / "tokenizer_config.json" + if path.is_file(): + config = json.loads(path.read_text(encoding="utf-8")) + for index, token in config.get("added_tokens_decoder", {}).items(): + added[int(index)] = token + return added + + # ------------------------------------------------------------------ + # tensors + # ------------------------------------------------------------------ + def _gguf_name(self, name): + new_name = self.tensor_map.get_name(name, try_suffixes=(".weight", ".bias")) + if new_name is None: + raise ValueError(f"No GGUF name for tensor {name!r}") + return new_name + + def _is_router(self, new_name): + for key in (gguf.MODEL_TENSOR.FFN_GATE_INP, gguf.MODEL_TENSOR.FFN_GATE_INP_SHEXP): + if key not in gguf.MODEL_TENSORS[self.arch.arch]: + continue + template = gguf.TENSOR_NAMES[key] + for bid in range(self.block_count): + if new_name == template.format(bid=bid) + ".weight": + return True + return False + + def _float_dtype(self, new_name, n_dims): + # 1-D tensors, norms and the MoE router stay in F32, as in a stock + # llama.cpp conversion; everything else follows --gguf_float_type. + if n_dims <= 1 or new_name.endswith("_norm.weight") or self._is_router(new_name): + return gguf.GGMLQuantizationType.F32 + return self.float_type + + def _add_float_tensor(self, new_name, tensor): + dtype = self._float_dtype(new_name, tensor.ndim) + if dtype == gguf.GGMLQuantizationType.F32: + data = tensor.to(torch.float32).numpy() + elif dtype == gguf.GGMLQuantizationType.F16: + data = tensor.to(torch.float16).numpy() + else: + data = tensor.to(torch.bfloat16).view(torch.uint16).numpy() + self.writer.add_tensor(new_name, data, raw_shape=tuple(tensor.shape), raw_dtype=dtype) + + def _expert_slot(self, name): + """`...experts...weight` -> (merged key, slot), else None.""" + parts = name.split(".") + if "experts" not in parts: + return None + index = parts.index("experts") + if index + 1 >= len(parts) or not parts[index + 1].isdigit(): + return None + slot = int(parts[index + 1]) + merged = ".".join(parts[:index + 1] + parts[index + 2:]) + return merged, slot + + @torch.no_grad() + def write(self, model): + if not self.packed: + raise RuntimeError("No ternary tensors were captured; nothing to pack") + + state_dict = model.state_dict() + tied = self.hparams.get("tie_word_embeddings", False) + experts = {} + n_ternary = 0 + + for name, tensor in state_dict.items(): + if name.endswith(".inv_freq") or (tied and name == "lm_head.weight"): + continue + + slot = self._expert_slot(name) + if slot is not None: + merged, index = slot + bucket = experts.setdefault(merged, {}) + bucket[index] = self.packed.pop(name, None) + if bucket[index] is None: + bucket[index] = tensor + if len(bucket) < self.hparams["num_experts"]: + continue + stacked = [bucket[i] for i in range(self.hparams["num_experts"])] + new_name = self._gguf_name(merged) + if isinstance(stacked[0], np.ndarray): + n_ternary += len(stacked) + stacked = np.stack(stacked) + del experts[merged], bucket + self.writer.add_tensor(new_name, stacked, raw_dtype=self.qtype_q2_0) + del stacked + else: + self._add_float_tensor(new_name, torch.stack(stacked)) + del experts[merged] + continue + + new_name = self._gguf_name(name) + if name in self.packed: + self.writer.add_tensor(new_name, self.packed.pop(name), raw_dtype=self.qtype_q2_0) + n_ternary += 1 + else: + self._add_float_tensor(new_name, tensor) + + if experts: + raise RuntimeError(f"Incomplete expert groups: {sorted(experts)}") + if n_ternary != self.n_packed: + raise RuntimeError( + f"{self.n_packed - n_ternary} captured ternary tensors are absent from the model: " + f"{sorted(self.packed)[:4]}" + ) + + self._write_metadata() + self.writer.write_header_to_file(path=self.outfile) + self.writer.write_kv_data_to_file() + self.writer.write_tensors_to_file(progress=True) + self.writer.close() + logger.info( + "Wrote %s (%d tensors, %d ternary)", + self.outfile, + len(self.writer.tensors[0]), + n_ternary, + ) + return self.outfile + + def _write_metadata(self): + total_params, shared_params, expert_params, expert_count = self.writer.get_total_parameter_count() + metadata = gguf.Metadata.load(None, self.model_dir, None, total_params) + if metadata.name is None: + metadata.name = self.model_dir.name + if metadata.size_label is None and total_params > 0: + metadata.size_label = gguf.size_label(total_params, shared_params, expert_params, expert_count) + metadata.set_gguf_meta_model(self.writer) + self.writer.add_type(gguf.GGUFType.MODEL) + self._set_parameters() + self._set_vocab() + self.writer.add_quantization_version(gguf.GGML_QUANT_VERSION) diff --git a/projects/cat-q/quantize/merge.py b/projects/cat-q/quantize/merge.py index 6348dea..578a7b9 100644 --- a/projects/cat-q/quantize/merge.py +++ b/projects/cat-q/quantize/merge.py @@ -1,5 +1,7 @@ """Checkpoint restoration and one-way merging for released CAT-Q parameters.""" +import gc + import torch from quantize.checkpoint import load_catq_parameters @@ -99,6 +101,38 @@ def _native_layer_state(qlayer, native_layer): return {name: value for name, value in merged.items() if name in expected} +def _resident_gib(): + """Anonymous resident memory of this process, for the low-memory export log. + + File-backed pages are left out on purpose: the memory-mapped shards the + loader hands out are reclaimable, so only anonymous memory can push the + export into the out-of-memory killer. + """ + try: + with open("/proc/self/status", "r", encoding="utf-8") as handle: + for line in handle: + if line.startswith("RssAnon:"): + return int(line.split()[1]) / 1024 ** 2 + except OSError: + pass + return float("nan") + + +def _release_weights(model, names, dtype): + """Drop the references to weights that are already packed. + + The names stay in `state_dict()` so the exporter can still map them onto + GGUF tensor names, but the model no longer pins the memory-mapped shards + they were read from. + """ + for name in names: + module_path, _, attribute = name.rpartition(".") + module = model.get_submodule(module_path) + module._parameters[attribute] = torch.nn.Parameter( + torch.empty(0, dtype=dtype), requires_grad=False + ) + + def _filter_inert_quantizer_bounds(layer_parameters, qlayer): expected = qlayer.state_dict() filtered = {} @@ -120,11 +154,21 @@ def _filter_inert_quantizer_bounds(layer_parameters, qlayer): @torch.no_grad() -def merge_catq_checkpoint(lm, args, logger): - """Restore released per-layer parameters and merge them into native HF weights.""" - if args.use_scaling and args.export_model_path: +def merge_catq_checkpoint(lm, args, logger, ternary_sink=None, release_packed_weights=False): + """Restore released per-layer parameters and merge them into native HF weights. + + `ternary_sink(layer_id, qlayer)` is called for every restored layer just + before its weights are folded into floating point, which is the only point + where the ternary codes and per-group scales are still available (see + `quantize.ternary_export`). It may return the weight names it consumed; + with `release_packed_weights` those floating-point weights are then dropped + from the model, which keeps peak host memory close to the size of the packed + result rather than the size of the 16-bit model. The merged model is then + only good for the packed export. + """ + if args.use_scaling and (args.export_model_path or ternary_sink is not None): raise ValueError( - "Native Hugging Face export currently requires use_scaling: false; " + "Model export currently requires use_scaling: false; " "equivalent-scaling checkpoints remain supported for evaluation." ) _set_quantizer_options(args) @@ -177,15 +221,37 @@ def merge_catq_checkpoint(lm, args, logger): ) qlayer.float() + captured = ternary_sink(layer_id, qlayer) if ternary_sink is not None else None qlayer.update_quant_mode("weight_merge", args=args) if args.use_scaling: layers[layer_id] = qlayer.to(dtype=dtype) else: - native_layer.load_state_dict( - _native_layer_state(qlayer, native_layer), - strict=False, - ) + merged_state = _native_layer_state(qlayer, native_layer) + if release_packed_weights and captured: + # The packed bytes already carry these weights, and writing the + # merged floats back would fault in a private copy of every + # memory-mapped page the loader handed out. + prefix = f"model.layers.{layer_id}." + packed_names = {name[len(prefix):] for name in captured} + merged_state = { + name: value + for name, value in merged_state.items() + if name not in packed_names + } + native_layer.load_state_dict(merged_state, strict=False) layers[layer_id] = native_layer.to(dtype=dtype) + if release_packed_weights and captured: + del qlayer + _release_weights(lm.model, captured, dtype) + del parameters[layer_id] + gc.collect() + logger.info( + "Merged CAT-Q parameters into layer %d (packed %d weights, host memory %.1f GiB)", + layer_id, + len(captured), + _resident_gib(), + ) + continue logger.info("Merged CAT-Q parameters into layer %d", layer_id) lm.model.to(dtype=dtype) diff --git a/projects/cat-q/quantize/q2_0.py b/projects/cat-q/quantize/q2_0.py new file mode 100644 index 0000000..943bbb6 --- /dev/null +++ b/projects/cat-q/quantize/q2_0.py @@ -0,0 +1,88 @@ +"""Pack CAT-Q ternary codes and scales into GGML `Q2_0` blocks. + +`Q2_0` is the group-128 ternary weight type used by the Bonsai llama.cpp fork +(https://github.com/PrismML-Eng/llama.cpp, branch `prism`). The layout below +is taken from that source tree, not from documentation: + + ggml/src/ggml-common.h + #define QK2_0 128 + typedef struct { ggml_half d; uint8_t qs[QK2_0 / 4]; } block_q2_0; + + ggml/src/ggml-quants.c :: dequantize_row_q2_0 + byte_index = j / 4; bit_offset = (j % 4) * 2 + q = (qs[byte_index] >> bit_offset) & 0x03 + y[j] = (q - 1) * d -> 00: -1 01: 0 10: +1 11: +2 + +A block is 34 bytes per 128 weights (2.125 bits per weight). A CAT-Q group +therefore maps with `code -> code + 1` and `d = fp16(scale)`; the `+2` level is +never produced. GGML rows run along `ne[0] = in_features`, the same direction +CAT-Q groups over, so no transpose or regrouping is needed. +""" + +import numpy as np + +QK2_0 = 128 +BLOCK_BYTES = 2 + QK2_0 // 4 # fp16 scale + 32 packed bytes + +# ggml/include/ggml.h :: GGML_TYPE_Q2_0 = 42 +GGML_TYPE_Q2_0 = 42 +# src/llama-model.cpp / include/llama.h :: LLAMA_FTYPE_MOSTLY_Q2_0 = 41 +LLAMA_FTYPE_MOSTLY_Q2_0 = 41 + + +def register_q2_0(): + """Teach the `gguf` package about `Q2_0`. + + The upstream PyPI package stops at type 41, so type 42 is free. Only the + block geometry is needed: the writer uses it to recover the logical tensor + shape from the packed byte shape. + """ + from gguf.constants import GGML_QUANT_SIZES + + GGML_QUANT_SIZES.setdefault(GGML_TYPE_Q2_0, (QK2_0, BLOCK_BYTES)) + return GGML_TYPE_Q2_0 + + +def pack_q2_0(codes, scales): + """codes: (..., out, in) int8 in {-1, 0, +1}; scales: one float per group. + + Returns a uint8 array of shape (..., out, in // 128 * 34). Leading + dimensions (used by stacked MoE expert tensors, shape (n_expert, out, in)) + pass through untouched: GGML stores those as plain row-major rows as well. + """ + codes = np.asarray(codes) + if codes.ndim < 2: + raise ValueError(f"expected at least a 2-D weight, got shape {codes.shape}") + lead_shape = tuple(codes.shape[:-2]) + out_features, in_features = codes.shape[-2:] + if in_features % QK2_0 != 0: + raise ValueError(f"in_features={in_features} is not a multiple of {QK2_0}") + + n_blocks = int(np.prod(codes.shape)) // QK2_0 + scales = np.asarray(scales, dtype=np.float32).reshape(-1) + if scales.size != n_blocks: + raise ValueError(f"expected {n_blocks} scales, got {scales.size}") + + quants = codes.astype(np.int8).reshape(n_blocks, QK2_0) + if quants.min() < -1 or quants.max() > 1: + raise ValueError("ternary codes must lie in {-1, 0, +1}") + + levels = (quants + np.int8(1)).astype(np.uint8).reshape(n_blocks, QK2_0 // 4, 4) + shifted = levels << np.array([0, 2, 4, 6], dtype=np.uint8).reshape(1, 1, 4) + qs = shifted[..., 0] | shifted[..., 1] | shifted[..., 2] | shifted[..., 3] + + d = scales.astype(np.float16).reshape(n_blocks, 1).view(np.uint8) + + blocks = np.concatenate([d, qs], axis=-1) + assert blocks.shape == (n_blocks, BLOCK_BYTES) + return blocks.reshape(*lead_shape, out_features, in_features // QK2_0 * BLOCK_BYTES) + + +def unpack_q2_0(blocks, shape): + """Inverse of `pack_q2_0`, returning float32 weights of `shape`.""" + raw = np.ascontiguousarray(blocks).reshape(-1, BLOCK_BYTES) + d = raw[:, :2].copy().view(np.float16).astype(np.float32) + qs = raw[:, 2:] + shifted = qs.reshape(qs.shape[0], -1, 1) >> np.array([0, 2, 4, 6], dtype=np.uint8).reshape(1, 1, 4) + codes = (shifted & 0x03).reshape(qs.shape[0], QK2_0).astype(np.int8) - np.int8(1) + return (codes.astype(np.float32) * d).reshape(shape) diff --git a/projects/cat-q/quantize/ternary_export.py b/projects/cat-q/quantize/ternary_export.py new file mode 100644 index 0000000..950de7b --- /dev/null +++ b/projects/cat-q/quantize/ternary_export.py @@ -0,0 +1,131 @@ +"""Packed ternary export for released CAT-Q checkpoints. + +`quantize.merge` restores a checkpoint and folds the result back into ordinary +floating-point Hugging Face weights, which is a *fake-quantized* model: the +values are ternary but every weight still occupies 16 bits. + +This module taps the same restoration path one step earlier and keeps the two +factors the quantizer works with, + + W_g = s_g * T_g, T_g in {-1, 0, +1}, |g| = group_size + +so that a deployment backend can store `T` at 2 bits per weight next to one +scale per group. `quantize.gguf_export` consumes the tensors yielded here +and packs them straight into a GGUF file. + +Nothing here re-estimates a scale from the merged float weights: the numbers +come from the quantizer itself, and every tensor is checked against the stock +fake-quantized output before it is handed on. +""" + +from dataclasses import dataclass + +import torch + +from quantize.int_linear import QuantLinear +from quantize.int_linear_lora import LoRAQuantLinear + + +@dataclass +class TernaryTensor: + """Ternary codes plus one scale per group, as stored by the exporter.""" + + codes: torch.Tensor # int8, weight shape, values in {-1, 0, +1} + scales: torch.Tensor # float32, (numel // group_size, 1) + group_size: int + + def dequantize(self): + return ( + self.codes.to(torch.float32).reshape(-1, self.group_size) * self.scales + ).reshape(self.codes.shape) + + +def extract_ternary(quantizer, weight): + """Re-run `TernaryQuantizer.forward` and keep its ternary codes and scales. + + The returned factors are verified to reproduce `quantizer(weight)` exactly, + so the packed model is numerically identical to the fake-quantized one. + """ + group_size = quantizer.group_size + if not group_size: + raise ValueError("Packed ternary export requires a grouped ternary quantizer") + original_shape = weight.shape + if original_shape[-1] % group_size != 0: + raise ValueError( + f"in_features={original_shape[-1]} is not a multiple of group_size={group_size}; " + "the packed ternary format requires whole groups" + ) + + grouped = weight.reshape(-1, group_size) + if quantizer.shift_mu: + mean = grouped.mean(dim=-1, keepdim=True) + else: + mean = grouped.new_zeros((grouped.shape[0], 1)) + + if quantizer.ter_scale_type == "absmean": + absmean_values = grouped if quantizer.init_scale_from_raw_weights else grouped - mean + scale = absmean_values.abs().mean(dim=-1, keepdim=True) + 1e-6 + elif quantizer.ter_scale_type == "variance": + scale = grouped.std(dim=-1, keepdim=True, unbiased=False) + 1e-6 + else: + raise ValueError(f"Unsupported ternary scale type: {quantizer.ter_scale_type}") + + if quantizer.learnable_mu: + mean = mean + quantizer.generate_mu_factor() * scale + if quantizer.learnable_scale: + scale = quantizer.generate_scale_factor() * scale + threshold = quantizer.init_round_thd + if quantizer.learnable_round: + threshold = threshold * quantizer.generate_round_factor() + + if quantizer.shift_mu and not quantizer.drop_quant_mu: + raise ValueError( + "drop_quant_mu is disabled: the merged weight keeps a per-group mean offset " + "and is therefore not representable as scale * ternary" + ) + + codes = torch.clamp(torch.round(((grouped - mean) / scale) * 0.5 / threshold), -1, 1) + reconstructed = (codes * scale).reshape(original_shape) + if not torch.equal(reconstructed, quantizer(weight)): + raise AssertionError( + "Extracted ternary codes and scales do not reproduce the fake-quantized weight" + ) + + return TernaryTensor( + codes=codes.reshape(original_shape).to(torch.int8), + scales=scale.to(torch.float32).contiguous(), + group_size=group_size, + ) + + +def _merged_weight(module): + """Weight seen by the quantizer, i.e. with the LoRA update already applied.""" + weight = module.weight + if isinstance(module, LoRAQuantLinear) and module.r > 0: + for index in range(module.lora_iter_num): + weight = weight + module.lora_B[index] @ module.lora_A[index] * module.scaling + return weight + + +@torch.no_grad() +def iter_layer_ternary(layer_id, qlayer): + """Yield `(weight_name, TernaryTensor)` for one restored decoder layer. + + Called from `merge_catq_checkpoint` while the layer still carries its + quantizer, i.e. before the ternary weights are folded into float tensors. + """ + found = False + for name, module in qlayer.named_modules(): + if not isinstance(module, (QuantLinear, LoRAQuantLinear)): + continue + quantizer = getattr(module, "weight_quantizer", None) + if type(quantizer).__name__ != "TernaryQuantizer": + continue + found = True + yield ( + f"model.layers.{layer_id}.{name}.weight", + extract_ternary(quantizer, _merged_weight(module)), + ) + + if not found: + raise RuntimeError(f"Layer {layer_id} produced no ternary tensors") diff --git a/projects/cat-q/tests/test_ternary_export.py b/projects/cat-q/tests/test_ternary_export.py new file mode 100644 index 0000000..6c100f2 --- /dev/null +++ b/projects/cat-q/tests/test_ternary_export.py @@ -0,0 +1,220 @@ +import unittest + +try: + import numpy as np + import torch + from torch import nn +except ImportError: # pragma: no cover - minimal documentation environments + np = None + torch = None + nn = None + + +def _ternary_params(**overrides): + params = { + "n_bits": 1, + "group_size": 128, + "shift_mu": False, + "drop_quant_mu": True, + "ter_scale_type": "absmean", + "init_scale_from_raw_weights": True, + "learnable_scale": False, + "learnable_mu": False, + "learnable_round": False, + "learnable_factor_act": "sigmoid", + "init_round_thd": 0.5, + "per_channel_axes": [0], + "symmetric": False, + "dynamic_method": "per_channel", + "disable_zero_point": False, + } + params.update(overrides) + return params + + +@unittest.skipIf(torch is None, "PyTorch is not installed") +class ExtractTernaryTest(unittest.TestCase): + def _quantizer(self, shape, **overrides): + from quantize.quantizer import TernaryQuantizer + + return TernaryQuantizer(weight_quant_params=_ternary_params(**overrides), shape=shape) + + def test_codes_are_ternary_and_grouped(self): + from quantize.ternary_export import extract_ternary + + torch.manual_seed(0) + weight = torch.randn(64, 256) + ternary = extract_ternary(self._quantizer(weight.shape), weight) + + self.assertEqual(ternary.group_size, 128) + self.assertEqual(ternary.codes.dtype, torch.int8) + self.assertEqual(tuple(ternary.codes.shape), (64, 256)) + self.assertEqual(ternary.scales.numel(), weight.numel() // 128) + self.assertTrue(bool(torch.isin(ternary.codes, torch.tensor([-1, 0, 1], dtype=torch.int8)).all())) + + def test_codes_and_scales_reproduce_the_fake_quantized_weight(self): + from quantize.ternary_export import extract_ternary + + torch.manual_seed(1) + weight = torch.randn(32, 128) + for overrides in ({}, {"learnable_scale": True}, {"learnable_round": True}): + with self.subTest(**overrides): + quantizer = self._quantizer(weight.shape, **overrides) + ternary = extract_ternary(quantizer, weight) + self.assertTrue(torch.equal(ternary.dequantize(), quantizer(weight))) + + def test_partial_groups_are_rejected(self): + from quantize.ternary_export import extract_ternary + + weight = torch.randn(8, 192) + with self.assertRaises(ValueError): + extract_ternary(self._quantizer(weight.shape), weight) + + def test_a_retained_group_mean_is_rejected(self): + from quantize.ternary_export import extract_ternary + + weight = torch.randn(8, 128) + quantizer = self._quantizer(weight.shape, shift_mu=True, drop_quant_mu=False) + with self.assertRaises(ValueError): + extract_ternary(quantizer, weight) + + def test_lora_update_is_applied_before_ternarization(self): + from quantize.int_linear_lora import LoRAQuantLinear + from quantize.ternary_export import _merged_weight, extract_ternary + + torch.manual_seed(2) + linear = nn.Linear(128, 32, bias=False) + module = LoRAQuantLinear( + org_module=linear, + weight_quant_params=_ternary_params(), + act_quant_params={"n_bits": 16}, + r=4, + lora_alpha=4, + ) + with torch.no_grad(): + module.lora_A[0].normal_() + module.lora_B[0].normal_() + + merged = _merged_weight(module) + self.assertFalse(torch.equal(merged, module.weight)) + ternary = extract_ternary(module.weight_quantizer, merged) + self.assertTrue(torch.equal(ternary.dequantize(), module.weight_quantizer(merged))) + + +@unittest.skipIf(np is None, "NumPy is not installed") +class PackQ2_0Test(unittest.TestCase): + @staticmethod + def _pack_module(): + from quantize import q2_0 + + return q2_0 + + def test_block_layout_matches_the_runtime_struct(self): + q2_0 = self._pack_module() + codes = np.zeros((1, 128), dtype=np.int8) + codes[0, :4] = [-1, 0, 1, 0] + packed = q2_0.pack_q2_0(codes, np.array([0.5], dtype=np.float32)) + + self.assertEqual(packed.shape, (1, 34)) + self.assertEqual(packed.dtype, np.uint8) + self.assertEqual(np.float16(packed[0, :2].copy().view(np.float16)[0]), np.float16(0.5)) + # 00 | 01 | 10 | 01 packed low bits first -> 0b01100100 + self.assertEqual(packed[0, 2], 0b01100100) + # the remaining weights are zero, i.e. level 01 in every slot + self.assertTrue((packed[0, 3:] == 0b01010101).all()) + + def test_round_trip(self): + q2_0 = self._pack_module() + rng = np.random.default_rng(0) + codes = rng.integers(-1, 2, size=(6, 256)).astype(np.int8) + scales = rng.random(codes.size // 128).astype(np.float32) + + packed = q2_0.pack_q2_0(codes, scales) + self.assertEqual(packed.shape, (6, 2 * 34)) + restored = q2_0.unpack_q2_0(packed, codes.shape) + expected = codes.reshape(-1, 128) * scales.astype(np.float16).astype(np.float32).reshape(-1, 1) + self.assertTrue(np.array_equal(restored, expected.reshape(codes.shape))) + + def test_expert_stacks_keep_their_leading_dimension(self): + q2_0 = self._pack_module() + codes = np.zeros((3, 4, 128), dtype=np.int8) + packed = q2_0.pack_q2_0(codes, np.ones(12, dtype=np.float32)) + self.assertEqual(packed.shape, (3, 4, 34)) + + def test_non_ternary_codes_are_rejected(self): + q2_0 = self._pack_module() + codes = np.full((1, 128), 2, dtype=np.int8) + with self.assertRaises(ValueError): + q2_0.pack_q2_0(codes, np.ones(1, dtype=np.float32)) + + + + +@unittest.skipIf(np is None, "NumPy is not installed") +class GGUFExportRulesTest(unittest.TestCase): + """Conversion rules that decide what ends up in the GGUF and how.""" + + @staticmethod + def _exporter(arch="qwen3moe"): + import gguf + + from quantize.gguf_export import ARCHITECTURES, TernaryGGUFExporter + + exporter = TernaryGGUFExporter.__new__(TernaryGGUFExporter) + exporter.arch = ARCHITECTURES["Qwen3MoeForCausalLM" if arch == "qwen3moe" else "LlamaForCausalLM"] + exporter.block_count = 2 + exporter.n_head = 4 + exporter.n_head_kv = 2 + exporter.float_type = gguf.GGMLQuantizationType.F16 + exporter.tensor_map = gguf.get_tensor_name_map(exporter.arch.arch, exporter.block_count) + return exporter + + def test_tensor_names_follow_the_gguf_convention(self): + exporter = self._exporter() + self.assertEqual(exporter._gguf_name("model.embed_tokens.weight"), "token_embd.weight") + self.assertEqual( + exporter._gguf_name("model.layers.1.self_attn.q_proj.weight"), "blk.1.attn_q.weight" + ) + self.assertEqual( + exporter._gguf_name("model.layers.1.mlp.experts.down_proj.weight"), + "blk.1.ffn_down_exps.weight", + ) + with self.assertRaises(ValueError): + exporter._gguf_name("model.layers.1.mystery.weight") + + def test_norms_and_the_router_stay_in_f32(self): + import gguf + + exporter = self._exporter() + f32 = gguf.GGMLQuantizationType.F32 + self.assertEqual(exporter._float_dtype("blk.0.attn_norm.weight", 1), f32) + self.assertEqual(exporter._float_dtype("blk.0.attn_q_norm.weight", 2), f32) + self.assertEqual(exporter._float_dtype("blk.0.ffn_gate_inp.weight", 2), f32) + self.assertEqual( + exporter._float_dtype("token_embd.weight", 2), gguf.GGMLQuantizationType.F16 + ) + + def test_expert_slots_are_grouped_by_projection(self): + exporter = self._exporter() + self.assertEqual( + exporter._expert_slot("model.layers.3.mlp.experts.7.up_proj.weight"), + ("model.layers.3.mlp.experts.up_proj.weight", 7), + ) + self.assertIsNone(exporter._expert_slot("model.layers.3.mlp.up_proj.weight")) + + def test_llama_permutation_applies_to_codes_and_scales(self): + from quantize.gguf_export import _permute + from quantize.q2_0 import unpack_q2_0 + + rng = np.random.default_rng(0) + codes = rng.integers(-1, 2, size=(8, 256)).astype(np.int8) + scales = rng.random((8, 2)).astype(np.float16).astype(np.float32) + weight = codes.astype(np.float32) * np.repeat(scales, 128, axis=1) + + exporter = self._exporter(arch="llama") + packed = exporter._pack("model.layers.0.self_attn.q_proj.weight", codes, scales) + restored = unpack_q2_0(packed, codes.shape) + np.testing.assert_allclose(restored, _permute(weight, 4, 4), rtol=0, atol=0) + +if __name__ == "__main__": + unittest.main()