MiMo-V2.6-Distill-Qwen-9B — DCQ

Packed, servable checkpoints of XiaomiMiMo/MiMo-V2.6-Distill-Qwen-9B quantized with DCQ (Drift-Correcting Quantization): STE-learned integer rounding + recurrent-state anchoring for the Gated DeltaNet blocks + trust-region regularization + sensitivity-driven role-aware bit allocation. See tools/ for the quantizer, packer, loader and verifier; results/ holds the eval JSONs backing the table below.

Variants

Embeddings, untied lm_head, norms, conv1d, the vision tower and all biases stay bf16 in both variants. in_proj_a/b are also bf16 (the quantizer left ≥8-bit mappings untouched).

Usage — w4-gptq (standard path)

# pip install gptqmodel   (or serve the dir with vllm / sglang)
from huggingface_hub import snapshot_download
from gptqmodel import GPTQModel

d = snapshot_download("Akahsizrr/mimo-v2.6-distill-qwen-9b-dcq",
                      allow_patterns=["w4-gptq/*"])
model = GPTQModel.load(f"{d}/w4-gptq")

The layout is standard GPTQ v1, verified numerically per-tensor. Kernel-level serving (Marlin) is expected to work but was not tested on this pod (no compile toolchain) — validation used gptqmodel's torch dequant path.

Usage — w3mix-packed (custom path)

Mixed 3/4-bit tensors are not expressible in stock GPTQ kernels, so this variant ships its own loader (load_dcq.py, inside the dir and in tools/):

# pip install torch transformers safetensors
import sys
from huggingface_hub import snapshot_download
from transformers import AutoModelForImageTextToText, AutoTokenizer

d = snapshot_download("Akahsizrr/mimo-v2.6-distill-qwen-9b-dcq",
                      allow_patterns=["w3mix-packed/*"])
d = f"{d}/w3mix-packed"
sys.path.insert(0, d)                    # loader ships inside the dir
from load_dcq import attach_dcq

tok = AutoTokenizer.from_pretrained(d, trust_remote_code=True)
model = AutoModelForImageTextToText.from_pretrained(
    d, dtype="bfloat16", trust_remote_code=True, low_cpu_mem_usage=True)
attach_dcq(model, d, cache=False)        # fused-dequant linears
model = model.cuda()
# generate as usual — see serve_dcq.py for a ready example

Two inference modes:

• low-VRAM (cache=False) — weights stay packed; each forward dequantizes on the fly. Measured: 8.35 GiB peak for the whole model.
• cached (cache=True) — dequantizes once, keeps bf16 copies (~21 GiB peak). Faster decode; still a 7.5 GB download.

Verified on this artifact (L40S, torch 2.14+cu126)

• w3mix-packed: all 200 quantized tensors round-trip their stored bf16 lattice to ≤ 4.7e-3 relative error (bf16 snapping floor); all 560 passthrough tensors bit-identical — verify_packed.py.
• w4-gptq: all 248 quantized tensors dequantize within ≤ 3.9e-3 relative error under the standard GPTQ-v1 formula, and within ≤ 5.3e-3 through gptqmodel's own TorchLinear.dequantize_weight.
• End-to-end: serve_dcq.py greedy-decodes correctly (17*23 -> 391), logits match the bf16-lattice source at cosine ≥ 0.99994.
• Measured serving on w3mix-packed: 7.50 GiB resident after load; 8.35 GiB peak with fused-dequant (cache=False); ~21 GiB in cached-bf16 mode; ~16 tok/s decode with reference (non-FLA) GDN kernels. w4-gptq loads in gptqmodel at 8.0 GiB.
• Known limitation: gptqmodel's model integration path produced NaN logits on this pod (missing flash-linear-attention/causal-conv1d kernels on a custom multimodal arch). The quantized weights are verified; use vLLM/SGLang for kernel serving or the bundled loader.

Smoke benchmarks (n=20, thinking mode, directional only)

20-question smokes — one answer moves a score by 5 points. Directional, not a benchmark claim.

Method (what actually ran on these artifacts)

quant_learned_log.json in each source variant records the exact config. For both shipped runs:

• Learned STE rounding (Adam on pre-rounding weights + group log-scales)
• Recurrent-state anchoring: GDN final-state MSE in the rounding loss
• Trust-region anchor toward the original weights
• Role-aware bits from measured sensitivity (in_proj_qkv is the state-destroying outlier; in_proj_a/b are free)
• closed_loop: false — sequential layerwise calibration on the drifted prefix is not in these artifacts (the closed-loop variant exists in tools/quant_learned.py but isn't what shipped)

Prior-art neighbours: AdaRound (learned rounding), Qronos (drift correction), Quamba/MambaQuant (SSM-aware PTQ). The distinctive piece is putting the recurrent state itself inside the rounding loss.

Reproduce

tools/
  quant_learned.py      # DCQ quantizer (bf16-lattice output)
  pack_dcq.py           # lattice -> packed ints+scales (dcq-pack-1 | gptq)
  load_dcq.py           # packed loader + fused-dequant Linear
  serve_dcq.py          # reference generation script
  verify_packed.py      # round-trip + logit-parity gate
  sens_probe.py         # per-tensor-role sensitivity map
  eval_openai.py        # MMLU/GSM8K/AIME harness used for the table
results/                # the eval JSONs behind every number above