Kolibri-1 · 2.5-bit experts · German calibration (kolibri-1-2.5bit-de)

A research compression of Aleph-Alpha/Kolibri-1-BF16 (78.1B total / 3.46B active parameters, 50 layers, 384 routed experts, top-6 routing) down to 28.4 GB of weight tensors, so that the whole model stays resident on a single 32 GB-class GPU. All 384 experts per layer and the original routing are kept; nothing is pruned and no experts are merged. Calibration and evaluation are German only.

Status: experimental. This is a custom storage format with its own loader. It is not loadable with AutoModelForCausalLM.from_pretrained, vLLM, llama.cpp or other standard runtimes, and it is slow (≈6.7 tokens/s decode, see below). Quality was measured on small German subsets only. Please read Limitations before using it.

What was done

Reconstruction ("healing"). For every expert, the block scales and the rank-8 residual were fitted to reproduce the original expert's outputs on its own routed hidden states (teacher-forced through the original model, layer by layer). Initialization is a damped least-squares low-rank solution, followed by 64 Adam steps on the relative output error. A checkpoint is accepted only if it improves on a held-out selection set; the numbers below come from a third, disjoint report set. Codes and the router stay frozen. This is local, per-expert reconstruction — no logit distillation and no fine-tuning of the full model.

Calibration data. 256 × 512-token sequences from German Wikipedia (wikimedia/wikipedia, config 20231101.de, revision b04c8d1ceb2f5cd4588862100d08de323dccfbaa), split into train / validation / test at document level (SHA-256 of the article text). No benchmark questions were used for calibration.

Why no expert merging. The project started from PuzzleMoE-style pairwise expert merging (dual masks + 2-bit magnitudes) at the same byte budget. A pilot on layers 0, 6, 12 and 24 showed that, weighted by how often the router uses each expert on German text, plain per-expert quantization without merging had the lower error in layers 6–24 for 92–99 % of routed tokens (the merge only won in layer 0, where most experts are unused on German text). This model therefore uses the unmerged format.

Size and memory

Measured on an NVIDIA RTX PRO 6000 Blackwell (96 GB) with the PyTorch allocator capped at 30 GB and no CPU offload. This fits under 30 GiB with about 2 GiB headroom; it is not below 30,000,000,000 bytes (≈84 MB over). It has not been run on a physical 32 GB card.

Results

All evaluations use the identical pipeline for the original model and the compressed variants.

MMLU-ProX German (li-lab/MMLU-ProX, config de, revision 8e6106a6c6ce1c5027e66cc338143cf997b2aa09): a fixed, category-balanced subset of 140 questions (10 per category), zero-shot forced choice with thinking disabled (reasoning_effort="none"), scored from the answer-letter logits. Trailing padding options were removed. These numbers are not comparable to published chain-of-thought scores.

Held-out German text (64 × 512 tokens, German Wikipedia test split):

Reconstruction moves the model closer to the original's predictions (lower KL, higher agreement) while perplexity on raw Wikipedia text gets slightly worse. Note that the absolute perplexity of the original is itself high in this setup (512-token windows starting mid-article, no BOS); use it only for relative comparison.

Per-expert reconstruction error (normalized MSE of expert outputs, 19,200 experts, report set): mean 0.118 before → 0.097 after reconstruction (median 0.100, 95th percentile 0.146). The gain shrinks with depth (≈ −76 % in layer 0, −12 % in layer 24).

Speed (PyTorch decode path, batch 1, RTX PRO 6000):

Each active expert matrix is decoded on the fly; there is no fused/Triton kernel for this format yet.

Usage

Requirements: a CUDA GPU with ≥32 GB, PyTorch ≥ 2.8 and the Transformers commit that contains the official Kolibri-1 modeling code (see requirements.txt).

pip install -r requirements.txt
huggingface-cli download johannhartmann/kolibri-1-2.5bit-de --local-dir kolibri-1-2.5bit-de

import sys, torch
from pathlib import Path
from transformers import AutoTokenizer

root = Path("kolibri-1-2.5bit-de")
sys.path.insert(0, str(root))                      # bundled loader: kolibri_puzzle/
from kolibri_puzzle.model import build_resident
from kolibri_puzzle.pipeline import chat_input_ids

model = build_resident(root, backend="torch")      # everything on CUDA, no offload
tok = AutoTokenizer.from_pretrained(root / "tokenizer", local_files_only=True)
ids = torch.tensor([chat_input_ids(
    tok, [{"role": "user", "content": "Erkläre Mixture-of-Experts in drei Sätzen."}],
    reasoning_effort="none", enable_thinking=False)], device="cuda")
out = model.generate(ids, max_new_tokens=160, do_sample=False)
print(tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True))

The loader builds the official Kolibri1ForCausalLM from the pinned Transformers code and replaces only the expert modules and the INT8 linears, so attention (hybrid sliding/NoPE), routing and the KV cache are the upstream implementation.

Files

• checkpoints/healed/layer-XX.safetensors — one file per layer: INT8 core, router, norms and 384 VQ experts (codes, codebooks, scales, rank-8 residuals; the pre-reconstruction scales are stored alongside).
• checkpoints/top.safetensors — INT8 embeddings and LM head, final norm.
• model_config.json, generation_config.json, tokenizer/ — from the base model.
• kolibri_puzzle/ — loader and the complete compression/evaluation code (Apache-2.0).
• notebook/ — the Colab notebook used to produce this model (code embedded, hash-checked).
• reports/ — machine-readable results: summary, runtime/memory measurement, per-layer reconstruction metrics, preflight checks.
• config.json, runtime_environment.txt, deployment_manifest.json — exact run configuration and resolved Python environment.

Limitations

• Quality loss is real. −5 pp on a small German MMLU-ProX subset; about one in five tokens is predicted differently from the original. Errors grow with layer depth.
• Narrow evaluation. 140 multiple-choice questions and 32k tokens of Wikipedia are not enough to resolve differences of a few points. Long reasoning (thinking mode), instruction following, coding, tool use and English were not evaluated.
• German-only calibration. Behaviour on other languages and domains may degrade more.
• Custom format, slow runtime. No standard inference engine support; decode speed is limited by the pure-PyTorch decoder.
• Tested on one GPU type only (RTX PRO 6000 Blackwell, CUDA 13, PyTorch 2.11).

Reproducibility

• Base model: Aleph-Alpha/Kolibri-1-BF16 @ 07edef0bb225dcc22da2b6d055aced504eea7482
• Transformers: 55928cfcc4f1b8253cd7e6e06ed2fed19bd0d2b0
• Method vq2: 32 codewords, group 128, 12 k-means iterations, rank 8, 64 reconstruction steps, batch 64, lr 3e-3, seed 17; 256 train / 64 selection / 64 report tokens per expert.
• Run on Google Colab (G4, NVIDIA RTX PRO 6000 Blackwell 96 GB), ≈1 h compression for 50 layers.

License and acknowledgements

The base model weights are released by Aleph Alpha under Apache-2.0; this derivative is released under the same license, and Aleph Alpha's responsible-use guidance for Kolibri-1 continues to apply. The bundled code is Apache-2.0. Thanks to Aleph Alpha for Kolibri-1. The PuzzleMoE pilot used the dual-mask merging rule from PuzzleMoE (arXiv:2511.04805).