Initial upload: mixed-budget sub-4-bit artifacts + perplexity result

README.md

@@ -0,0 +1,60 @@
+# sub1quant — sub-4-bit quantization artifacts for gemma-4-E2B
+
+This repo holds the **mixed-budget sub-4-bit quantization artifacts** produced by the
+`sub1quant` project. The base model (`google/gemma-4-E2B`) is **not** mirrored here —
+pull it separately from Hugging Face.
+
+## Contents
+
+```
+quantized/
+  gemma_mixed_budget_full_g128_target4p0.pt    # 948 MB, 316 language_model weight tensors
+                                                # 301 groupwise_int4 + 14 int2_binary_residual
+                                                # + 1 int2_error_budget_residual
+                                                # avg BPW ≈ 4.0  (vs BF16 ≈ 16 BPW)
+eval_results/
+  mixed_budget_full_g128_target4p0_ppl_colab.json   # perplexity on wikitext test
+  mixed_budget_scan_full_g128_target4p0.json        # full-surface reconstruction scan
+  error_budget_residual_*.json                      # earlier int2+residual scan results
+src/                                                 # quantization/dequant primitives
+scripts/                                             # quantize_mixed_budget, eval_quantized, ...
+test_perplexity.py                                   # entry-point for perplexity eval
+data/wiki.test.txt                                   # wikitext-103 test, ~287k tokens
+```
+
+## Perplexity (latest)
+
+| format | BPW | perplexity | chunks | tokens | status |
+|--------|----:|-----------:|-------:|-------:|--------|
+| mixed_budget_full_g128_target4p0 | 4.00 | **107.2452** | 571 | 292282 | FAIL (>10.5) |
+
+Run on `NVIDIA L4` bf16→fp16, `gemma-4-E2B` from `google/gemma-4-E2B`, wikitext test
+(stride=512, max_length=512). Result file:
+`eval_results/mixed_budget_full_g128_target4p0_ppl_colab.json`.
+
+The 107 perplexity is materially higher than a working sub-4-bit quant on a 2B
+model — treat it as a measurement, not a quality claim. Reconstruction RMSE alone
+(see scan JSONs) does not predict this number.
+
+## Reproducing the perplexity eval
+
+```bash
+# install
+pip install "transformers>=5.5.0" torch accelerate safetensors
+
+# pull the base model (NOT in this repo)
+python -c "from huggingface_hub import snapshot_download; snapshot_download('google/gemma-4-E2B', local_dir='./models/gemma-4-E2B')"
+
+# run
+python test_perplexity.py \
+  --model models/gemma-4-E2B \
+  --quantized quantized/gemma_mixed_budget_full_g128_target4p0.pt \
+  --wikitext data/wiki.test.txt \
+  --device cuda \
+  --max-length 512 --stride 512
+```
+
+## License
+
+The base model is governed by Google's Gemma license. The quantization
+artifacts in this repo are released under Apache-2.0.

eval_results/mixed_budget_full_g128_target4p0_ppl_colab.json

@@ -0,0 +1,20 @@
+{
+  "timestamp_utc": 1782706594,
+  "method": "mixed_budget",
+  "group_size": 128,
+  "target_bpw": 4.0,
+  "quantized_pt": "/content/sub1quant/quantized/gemma_mixed_budget_full_g128_target4p0.pt",
+  "model_dir": "/content/sub1quant/models/gemma-4-E2B",
+  "wikitext": "/content/sub1quant/data/wiki.test.txt",
+  "device": "cuda",
+  "max_length": 512,
+  "stride": 512,
+  "seq_len_tokens": 292282,
+  "n_chunks": 571,
+  "weights_replaced": 276,
+  "weights_total_pt": 316,
+  "weights_skipped_shared_kv": 40,
+  "perplexity": 107.2452,
+  "target_ppl": 10.5,
+  "status": "FAIL"
+}

quantized/gemma_mixed_budget_full_g128_target4p0.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4337f598433a909e3895a3d3b47f2093dbdba2a1ed4e71502f5549e233f6326
+size 948181931

scripts/eval_quantized.py

@@ -0,0 +1,383 @@
+import torch
+import gc
+import argparse
+import json
+import struct
+from pathlib import Path
+import sys
+from typing import Callable, Iterable, Optional
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from src.error_budget_residual import dequantize_binary_residual, dequantize_error_budget_residual
+from src.groupwise_int4 import dequantize_groupwise_int4
+from src.quantization import ternary_unpack, sigma_dequantize
+
+
+def _as_tensor(value, device: str, dtype=torch.float32) -> torch.Tensor:
+    if isinstance(value, torch.Tensor):
+        return value.to(device=device, dtype=dtype)
+    return torch.tensor(value, device=device, dtype=dtype)
+
+
+def _shape_tuple(shape) -> Optional[tuple]:
+    if shape is None:
+        return None
+    if isinstance(shape, torch.Tensor):
+        shape = shape.tolist()
+    return tuple(int(dim) for dim in shape)
+
+
+def quantized_entry_shape(q_entry: dict) -> Optional[tuple]:
+    shape = q_entry.get('original_shape', q_entry.get('orig_shape'))
+    if shape is not None:
+        return _shape_tuple(shape)
+    if 'U_shape' in q_entry and 'Vt_shape' in q_entry:
+        u_shape = _shape_tuple(q_entry['U_shape'])
+        vt_shape = _shape_tuple(q_entry['Vt_shape'])
+        return (u_shape[0], vt_shape[1])
+    return None
+
+
+def normalize_checkpoint_weight_keys(weight_keys: Optional[Iterable]) -> list[tuple[str, tuple]]:
+    if not weight_keys:
+        return []
+
+    normalized = []
+    for item in weight_keys:
+        if isinstance(item, dict):
+            key = item.get('key') or item.get('name')
+            shape = item.get('shape') or item.get('original_shape') or item.get('orig_shape')
+        else:
+            key, shape = item
+        if key:
+            normalized.append((key, _shape_tuple(shape)))
+    return normalized
+
+
+def load_model_weight_keys(model_dir: str | Path) -> list[tuple[str, tuple]]:
+    """Recover quantization key order from local safetensors metadata."""
+    model_dir = Path(model_dir)
+    safetensor_files = sorted(model_dir.glob("*.safetensors"))
+    weight_keys = []
+
+    for safetensor_path in safetensor_files:
+        with open(safetensor_path, 'rb') as f:
+            header_size = struct.unpack('<Q', f.read(8))[0]
+            header = json.loads(f.read(header_size))
+
+        for key, info in header.items():
+            if key == '__metadata__' or not isinstance(info, dict):
+                continue
+            shape = info.get('shape')
+            if 'weight' not in key or shape is None or len(shape) != 2:
+                continue
+            if any(x in key for x in [
+                'lm_head',
+                'embed_tokens',
+                'norm',
+                'audio_tower',
+                'vision_tower',
+                'embed_vision',
+            ]):
+                continue
+            weight_keys.append((key, _shape_tuple(shape)))
+
+    language_model_keys = [item for item in weight_keys if 'language_model' in item[0]]
+    return language_model_keys or weight_keys
+
+
+def resolve_quantized_weight_key(
+    entry_index,
+    q_entry: dict,
+    fallback_weight_keys: Optional[list[tuple[str, tuple]]] = None,
+) -> str:
+    key = q_entry.get('key')
+    if key:
+        return key
+
+    if fallback_weight_keys:
+        try:
+            fallback_index = int(entry_index)
+        except (TypeError, ValueError) as exc:
+            raise KeyError(
+                f"Quantized entry {entry_index!r} has no key and cannot be mapped by index"
+            ) from exc
+
+        if fallback_index >= len(fallback_weight_keys):
+            raise KeyError(
+                f"Quantized entry {entry_index!r} has no key and exceeds "
+                f"{len(fallback_weight_keys)} fallback model weights"
+            )
+
+        key, expected_shape = fallback_weight_keys[fallback_index]
+        entry_shape = quantized_entry_shape(q_entry)
+        if entry_shape is not None and expected_shape is not None and entry_shape != expected_shape:
+            raise ValueError(
+                f"Legacy quantized entry {entry_index!r} maps to {key}, but "
+                f"checkpoint shape {entry_shape} != model shape {expected_shape}"
+            )
+        return key
+
+    raise KeyError(
+        f"Quantized entry {entry_index!r} has no source key. Pass a local "
+        "model directory with safetensors metadata or use a checkpoint that stores keys."
+    )
+
+
+def reconstruct_weight(q_entry: dict, device: str = 'cpu') -> torch.Tensor:
+    U_scale = _as_tensor(q_entry['U_scale'], device)
+    Vt_scale = _as_tensor(q_entry['Vt_scale'], device)
+    S_scale = _as_tensor(q_entry['S_scale'], device)
+    U = ternary_unpack(q_entry['U_packed'].to(device), q_entry['U_shape']).float() * U_scale
+    Vt = ternary_unpack(q_entry['Vt_packed'].to(device), q_entry['Vt_shape']).float() * Vt_scale
+    S = sigma_dequantize(q_entry['S'].to(device), S_scale)
+    return torch.matmul(U * S.unsqueeze(0), Vt)
+
+
+def reconstruct_quantized_entry(q_entry: dict, device: str = 'cpu') -> torch.Tensor:
+    if q_entry.get('format') == 'groupwise_int4' or 'packed_int4' in q_entry:
+        return dequantize_groupwise_int4(q_entry, device)
+
+    if q_entry.get('format') == 'int2_error_budget_residual':
+        return dequantize_error_budget_residual(q_entry, device)
+
+    if q_entry.get('format') == 'int2_base':
+        return dequantize_binary_residual(q_entry, device, include_residual=False)
+
+    if q_entry.get('format') == 'int2_binary_residual' or 'base_packed' in q_entry:
+        return dequantize_binary_residual(q_entry, device)
+
+    if {'U_packed', 'Vt_packed', 'S'}.issubset(q_entry):
+        return reconstruct_weight(q_entry, device)
+
+    if 'packed' in q_entry:
+        shape = q_entry['orig_shape']
+        scale = _as_tensor(q_entry['scale'], device)
+        return ternary_unpack(q_entry['packed'].to(device), shape).float() * scale
+
+    if 'q' in q_entry:
+        num_bits = int(q_entry.get('num_bits', 0))
+        qmax = 2 ** (num_bits - 1) - 1
+        if qmax <= 0:
+            raise ValueError(f"Unsupported num_bits for magnitude checkpoint: {num_bits}")
+
+        q = q_entry['q'].to(device).float()
+        shape = quantized_entry_shape(q_entry)
+        if shape is not None and q.numel() == shape[0] * shape[1]:
+            q = q.reshape(shape)
+
+        scale = q_entry['scale']
+        if q_entry.get('per_channel'):
+            scale = _as_tensor(scale, device).reshape(-1, 1)
+        else:
+            scale = _as_tensor(scale, device)
+        return q * scale / qmax
+
+    raise ValueError(f"Unknown quantized entry format: {sorted(q_entry.keys())}")
+
+
+def build_model_weight_map(model) -> dict[str, object]:
+    weight_map = {}
+    for name, module in model.named_modules():
+        if getattr(module, 'weight', None) is None:
+            continue
+        weight_key = f"{name}.weight" if name else "weight"
+        weight_map[weight_key] = module
+    return weight_map
+
+
+def is_expected_missing_shared_kv_weight(key: str, weight_map: dict[str, object]) -> bool:
+    if key.endswith('.self_attn.k_proj.weight'):
+        q_key = key.replace('.k_proj.weight', '.q_proj.weight')
+        o_key = key.replace('.k_proj.weight', '.o_proj.weight')
+        return q_key in weight_map or o_key in weight_map
+    if key.endswith('.self_attn.v_proj.weight'):
+        q_key = key.replace('.v_proj.weight', '.q_proj.weight')
+        o_key = key.replace('.v_proj.weight', '.o_proj.weight')
+        return q_key in weight_map or o_key in weight_map
+    return False
+
+
+def apply_quantized_weights(
+    model,
+    quantized: dict,
+    device: str = 'cpu',
+    model_dir: str | Path | None = None,
+    checkpoint_weight_keys: Optional[Iterable] = None,
+    reconstruct_fn: Callable[[dict, str], torch.Tensor] = reconstruct_quantized_entry,
+    strict: bool = True,
+) -> dict:
+    fallback_weight_keys = normalize_checkpoint_weight_keys(checkpoint_weight_keys)
+    if not fallback_weight_keys and model_dir is not None:
+        fallback_weight_keys = load_model_weight_keys(model_dir)
+
+    weight_map = build_model_weight_map(model)
+    stats = {'replaced': 0, 'skipped': [], 'missing': [], 'shape_mismatches': []}
+
+    for entry_index, q_entry in quantized.items():
+        key = resolve_quantized_weight_key(entry_index, q_entry, fallback_weight_keys)
+        module = weight_map.get(key)
+        if module is None:
+            if is_expected_missing_shared_kv_weight(key, weight_map):
+                stats['skipped'].append(key)
+                continue
+            message = f"No model module found for quantized weight {key}"
+            if strict:
+                raise KeyError(message)
+            stats['missing'].append(message)
+            continue
+
+        reconstructed = reconstruct_fn(q_entry, device)
+        target_shape = tuple(module.weight.shape)
+        if tuple(reconstructed.shape) != target_shape:
+            message = (
+                f"Shape mismatch for {key}: reconstructed {tuple(reconstructed.shape)} "
+                f"!= model {target_shape}"
+            )
+            if strict:
+                raise ValueError(message)
+            stats['shape_mismatches'].append(message)
+            continue
+
+        with torch.no_grad():
+            module.weight.data = reconstructed.to(
+                dtype=module.weight.dtype,
+                device=module.weight.device,
+            )
+        stats['replaced'] += 1
+
+    return stats
+
+
+def eval_perplexity(model, tokenizer, wikitext_path: str, device: str,
+                    max_length: int = 512, stride: int = 512):
+    with open(wikitext_path, 'r', encoding='utf-8') as f:
+        text = f.read()
+
+    print("Tokenizing...")
+    encodings = tokenizer(text, return_tensors='pt')
+    encodings = {k: v.to(device) for k, v in encodings.items()}
+    seq_len = encodings['input_ids'].shape[1]
+    print(f"  Sequence length: {seq_len} tokens")
+
+    nlls = []
+    prev_end_loc = 0
+
+    for begin_loc in range(0, seq_len, stride):
+        end_loc = min(begin_loc + max_length, seq_len)
+        trg_len = end_loc - prev_end_loc
+
+        input_ids = encodings['input_ids'][:, begin_loc:end_loc]
+        target_ids = input_ids.clone()
+        target_ids[:, :-trg_len] = -100
+
+        with torch.no_grad():
+            outputs = model(input_ids, labels=target_ids)
+            neg_log_likelihood = outputs.loss * trg_len
+
+        nlls.append(neg_log_likelihood)
+        prev_end_loc = end_loc
+
+        if end_loc >= seq_len:
+            break
+
+    avg_nll = torch.stack(nlls).sum() / seq_len
+    perplexity = torch.exp(avg_nll).item()
+    return perplexity, {'n_chunks': len(nlls), 'seq_len': seq_len}
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate quantized model perplexity")
+    parser.add_argument('--quantized-pt', default='quantized/gemma-4-E2B-sub1bit.pt',
+                        help='Path to quantized .pt checkpoint')
+    parser.add_argument('--model-dir', default='models/gemma-4-E2B',
+                        help='Path to base model directory')
+    parser.add_argument('--wikitext', default='data/wiki.test.txt',
+                        help='Path to WikiText test file')
+    parser.add_argument('--device', default=None,
+                        help='Device (auto-detect if not set)')
+    parser.add_argument('--max-length', type=int, default=512)
+    parser.add_argument('--stride', type=int, default=512)
+    args = parser.parse_args()
+
+    if args.device:
+        device = args.device
+    else:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    wikitext_path = Path(args.wikitext)
+    if not wikitext_path.exists():
+        print(f"WikiText not found: {wikitext_path}")
+        return
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    print("=" * 60)
+    print("QUANTIZED MODEL PERPLEXITY EVALUATION")
+    print("=" * 60)
+    print(f"Device: {device}")
+    print(f"Quantized: {args.quantized_pt}")
+    print(f"Base model: {args.model_dir}")
+    print(f"WikiText: {args.wikitext}")
+    print()
+
+    # 1. Load quantized checkpoint
+    print("[1] Loading quantized checkpoint...")
+    q_data = torch.load(args.quantized_pt, map_location='cpu', weights_only=True)
+    quantized = q_data['quantized']
+    print(f"  {len(quantized)} quantized entries")
+    print()
+
+    # 2. Load base model
+    print("[2] Loading base model...")
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_dir,
+        device_map=device,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True
+    )
+    model.eval()
+    print()
+
+    # 3. Reconstruct weights
+    print("[3] Applying quantized weights...")
+    apply_stats = apply_quantized_weights(
+        model,
+        quantized,
+        device=device,
+        model_dir=args.model_dir,
+        checkpoint_weight_keys=q_data.get('weight_keys'),
+    )
+    print(f"  Replaced {apply_stats['replaced']}/{len(quantized)} weights")
+    if apply_stats['skipped']:
+        print(f"  Skipped {len(apply_stats['skipped'])} shared-KV checkpoint entries")
+    print()
+
+    # 4. Evaluate perplexity
+    print("[4] Evaluating perplexity...")
+    ppl, stats = eval_perplexity(
+        model, tokenizer, str(wikitext_path), device,
+        max_length=args.max_length, stride=args.stride
+    )
+    print()
+
+    print("=" * 60)
+    print("RESULTS")
+    print("=" * 60)
+    print(f"  Perplexity: {ppl:.4f}")
+    print(f"  Chunks: {stats['n_chunks']}")
+    print(f"  Target: <= 10.5")
+    status = "PASS" if ppl <= 10.5 else "FAIL"
+    print(f"  Status: {status}")
+    print("=" * 60)
+
+    del model
+    gc.collect()
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

scripts/limited_ppl_bench.py

@@ -0,0 +1,134 @@
+"""Run a short WikiText perplexity smoke benchmark for base/quantized checkpoints."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from scripts.eval_quantized import apply_quantized_weights
+
+
+def eval_limited_ppl(
+    model,
+    tokenizer,
+    text: str,
+    device: str,
+    tokens: int,
+    max_length: int,
+    stride: int,
+) -> dict:
+    encodings = tokenizer(text, return_tensors="pt")
+    input_ids = encodings["input_ids"][:, :tokens].to(device)
+    seq_len = input_ids.shape[1]
+    nlls = []
+    prev_end_loc = 0
+
+    for begin_loc in range(0, seq_len, stride):
+        end_loc = min(begin_loc + max_length, seq_len)
+        trg_len = end_loc - prev_end_loc
+
+        batch = input_ids[:, begin_loc:end_loc]
+        target = batch.clone()
+        target[:, :-trg_len] = -100
+
+        with torch.no_grad():
+            outputs = model(batch, labels=target)
+            nlls.append(outputs.loss.detach() * trg_len)
+
+        prev_end_loc = end_loc
+        if end_loc >= seq_len:
+            break
+
+    ppl = torch.exp(torch.stack(nlls).sum() / seq_len).item()
+    return {"ppl": ppl, "seq_len": seq_len, "chunks": len(nlls)}
+
+
+def run(args: argparse.Namespace) -> dict:
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    model_dir = Path(args.model_dir)
+    device = args.device or ("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+    tokenizer = AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        str(model_dir),
+        dtype=dtype,
+        device_map=device,
+        trust_remote_code=True,
+    )
+    model.eval()
+
+    apply_stats = {"replaced": 0, "skipped": []}
+    checkpoint_stats = None
+    if args.quantized_pt:
+        q_data = torch.load(args.quantized_pt, map_location="cpu", weights_only=True)
+        checkpoint_stats = q_data.get("stats")
+        apply_stats = apply_quantized_weights(
+            model,
+            q_data["quantized"],
+            device=device,
+            model_dir=model_dir,
+            checkpoint_weight_keys=q_data.get("weight_keys"),
+            strict=False,
+        )
+        del q_data
+        gc.collect()
+
+    text = Path(args.wikitext).read_text(encoding="utf-8")
+    metrics = eval_limited_ppl(
+        model,
+        tokenizer,
+        text,
+        device,
+        tokens=args.tokens,
+        max_length=args.max_length,
+        stride=args.stride,
+    )
+    metrics.update(
+        {
+            "label": args.label,
+            "mode": "quantized" if args.quantized_pt else "base",
+            "quantized_pt": args.quantized_pt,
+            "apply_stats": apply_stats,
+            "checkpoint_stats": checkpoint_stats,
+            "device": device,
+        }
+    )
+    return metrics
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--model-dir", default="models/gemma-4-E2B")
+    parser.add_argument("--wikitext", default="data/wiki.test.txt")
+    parser.add_argument("--quantized-pt", default=None)
+    parser.add_argument("--tokens", type=int, default=4096)
+    parser.add_argument("--max-length", type=int, default=512)
+    parser.add_argument("--stride", type=int, default=512)
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--output", required=True)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    start = time.time()
+    result = run(args)
+    result["elapsed_s"] = round(time.time() - start, 1)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
+    print("RESULT=" + json.dumps(result, indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/quantize_mixed_budget.py

@@ -0,0 +1,214 @@
+"""Build a mixed-budget checkpoint from a scan_mixed_budget allocation JSON."""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import re
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+from scripts.eval_quantized import reconstruct_quantized_entry
+from scripts.quantize_groupwise_int4 import iter_model_weights
+from src.error_budget_residual import estimate_binary_residual_bpw, quantize_binary_residual, quantize_error_budget_residual
+from src.groupwise_int4 import quantize_groupwise_int4
+
+
+ERROR_BUDGET_RE = re.compile(r"^int2_error_budget_k(\d+)$")
+
+
+def build_selection_map(scan: dict) -> dict[str, str]:
+    selected = scan.get("mixed_allocation", {}).get("selected_layers", [])
+    if not selected:
+        raise ValueError("scan JSON does not contain mixed_allocation.selected_layers")
+    selection = {}
+    for item in selected:
+        key = item.get("key")
+        method = item.get("method")
+        if not key or not method:
+            raise ValueError(f"invalid selected layer entry: {item!r}")
+        selection[str(key)] = str(method)
+    return selection
+
+
+def quantize_weight_for_method(weight: torch.Tensor, method: str, group_size: int) -> dict:
+    if method == "groupwise_int4":
+        return quantize_groupwise_int4(weight, group_size=group_size)
+    if method == "int2_binary_residual":
+        return quantize_binary_residual(weight, group_size=group_size)
+    if method == "int2_base":
+        entry = quantize_binary_residual(weight, group_size=group_size)
+        entry["format"] = "int2_base"
+        entry["bpw"] = estimate_binary_residual_bpw(weight.shape, group_size=group_size) - 1.0 - (16.0 / group_size)
+        return entry
+
+    match = ERROR_BUDGET_RE.match(method)
+    if match:
+        return quantize_error_budget_residual(
+            weight,
+            group_size=group_size,
+            outliers_per_group=int(match.group(1)),
+        )
+
+    raise ValueError(f"Unsupported mixed-budget method {method!r}")
+
+
+def build_checkpoint(
+    scan: dict,
+    model_dir: Path,
+    group_size: int,
+    max_layers: int | None,
+    skip_reconstruction_metrics: bool,
+) -> dict:
+    selection = build_selection_map(scan)
+    quantized: dict[int, dict] = {}
+    layer_stats: list[dict] = []
+    total_params = 0
+    total_bits = 0.0
+    total_weighted_mse = 0.0
+
+    for idx, key, weight in iter_model_weights(model_dir, max_layers=max_layers):
+        method = selection.get(key)
+        if method is None:
+            raise KeyError(f"No mixed-budget selection for {key}")
+
+        entry = quantize_weight_for_method(weight, method, group_size=group_size)
+        entry["key"] = key
+        quantized[idx] = entry
+
+        params = weight.numel()
+        total_params += params
+        total_bits += float(entry["bpw"]) * params
+
+        mse = None
+        rmse = None
+        if not skip_reconstruction_metrics:
+            restored = reconstruct_quantized_entry(entry)
+            error = restored - weight
+            mse = error.pow(2).mean().item()
+            rmse = mse**0.5
+            total_weighted_mse += mse * params
+            del restored, error
+
+        layer_stats.append(
+            {
+                "idx": idx,
+                "key": key,
+                "method": method,
+                "shape": list(weight.shape),
+                "params": params,
+                "bpw": entry["bpw"],
+                "mse": mse,
+                "rmse": rmse,
+            }
+        )
+        print(
+            f"Layer {idx:3d}: {method}, bpw={entry['bpw']:.4f}, shape={list(weight.shape)}"
+            + ("" if mse is None else f", mse={mse:.6f}, rmse={rmse:.6f}"),
+            flush=True,
+        )
+
+        del weight
+        gc.collect()
+
+    if not quantized:
+        raise RuntimeError(f"No quantizable language_model 2D weights found in {model_dir}")
+
+    selected_keys = {item["key"] for item in layer_stats}
+    extra_keys = sorted(set(selection) - selected_keys)
+    if extra_keys:
+        raise KeyError(f"Scan allocation had selections not emitted by this build: {extra_keys[:5]}")
+
+    avg_bpw = total_bits / total_params
+    stats = {
+        "method": "mixed_budget",
+        "format_version": 1,
+        "group_size": group_size,
+        "layers": len(quantized),
+        "total_params": total_params,
+        "avg_bpw": avg_bpw,
+        "compression_vs_bf16": 16.0 / avg_bpw,
+        "weighted_mse": None if skip_reconstruction_metrics else total_weighted_mse / total_params,
+        "weighted_rmse": None if skip_reconstruction_metrics else (total_weighted_mse / total_params) ** 0.5,
+        "method_counts": _method_counts(layer_stats),
+        "layer_stats": layer_stats,
+        "source_allocation": scan.get("mixed_allocation"),
+        "uniform_int4": scan.get("uniform_int4"),
+    }
+    return {
+        "quantized": quantized,
+        "stats": stats,
+        "method": "mixed_budget",
+        "config": {
+            "model_dir": str(model_dir),
+            "group_size": group_size,
+            "max_layers": max_layers,
+            "source_scan": scan.get("source_scan"),
+        },
+        "weight_keys": [
+            {"key": item["key"], "shape": item["shape"], "method": item["method"]}
+            for item in layer_stats
+        ],
+    }
+
+
+def _method_counts(layer_stats: list[dict]) -> dict[str, int]:
+    counts: dict[str, int] = {}
+    for item in layer_stats:
+        method = str(item["method"])
+        counts[method] = counts.get(method, 0) + 1
+    return dict(sorted(counts.items()))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--scan-json", required=True)
+    parser.add_argument("--model-dir", default="models/gemma-4-E2B")
+    parser.add_argument("--output", default="quantized/gemma_mixed_budget.pt")
+    parser.add_argument("--group-size", type=int, default=None)
+    parser.add_argument("--max-layers", type=int, default=None)
+    parser.add_argument("--skip-reconstruction-metrics", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    scan_path = Path(args.scan_json)
+    scan = json.loads(scan_path.read_text(encoding="utf-8"))
+    scan["source_scan"] = str(scan_path)
+    group_size = args.group_size if args.group_size is not None else int(scan.get("group_size", 128))
+    max_layers = args.max_layers if args.max_layers is not None else scan.get("max_layers")
+
+    checkpoint = build_checkpoint(
+        scan=scan,
+        model_dir=Path(args.model_dir),
+        group_size=group_size,
+        max_layers=max_layers,
+        skip_reconstruction_metrics=args.skip_reconstruction_metrics,
+    )
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(checkpoint, output)
+
+    stats = checkpoint["stats"]
+    print()
+    print("=" * 72)
+    print("MIXED BUDGET RESULTS")
+    print("=" * 72)
+    print(f"Layers: {stats['layers']}")
+    print(f"Average BPW: {stats['avg_bpw']:.4f}")
+    print(f"Compression vs BF16: {stats['compression_vs_bf16']:.2f}x")
+    if stats["weighted_mse"] is not None:
+        print(f"Weighted MSE: {stats['weighted_mse']:.8f}")
+        print(f"Weighted RMSE: {stats['weighted_rmse']:.6f}")
+    print("Method counts:", json.dumps(stats["method_counts"], sort_keys=True))
+    print("Saved:", output)
+
+
+if __name__ == "__main__":
+    main()

src/Sub1BitLLM.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+from dataclasses import dataclass
+from typing import Dict, Optional, List, Iterator, Tuple, Any
+from pathlib import Path
+import os
+import numpy as np
+
+from .gguf_writer import GGUFWriter, GGML_TYPES
+from .quantization import dequantize_factor
+
+
+@dataclass
+class Sub1BitConfig:
+    codebook_dim: int = 128
+    energy_threshold: float = 0.95
+    rank: int = 16
+    U_bits: float = 0.5
+    S_bits: float = 2.0
+    Vt_bits: float = 0.5
+    model_name: str = "llama-2-7b-sub1bit"
+    architecture: str = "llama"
+
+
+class LowRankFactor(torch.nn.Module):
+    def __init__(self, U: torch.Tensor, S: torch.Tensor, Vt: torch.Tensor):
+        super().__init__()
+        self.register_buffer('U', U.half())
+        self.register_buffer('S', S.half())
+        self.register_buffer('Vt', Vt.half())
+        self.rank = U.shape[1]
+
+    def forward(self) -> torch.Tensor:
+        return torch.matmul(self.U * self.S.unsqueeze(0), self.Vt)
+
+    def forward_lowrank(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.matmul(torch.matmul(x, self.Vt.T) * self.S.unsqueeze(0), self.U.T)
+
+
+class TernaryQuantizedFactor(torch.nn.Module):
+    def __init__(self, data: torch.Tensor, scale: torch.Tensor, rank: int):
+        super().__init__()
+        self.register_buffer('data', data)
+        self.register_buffer('scale', scale)
+        self.rank = rank
+
+    def forward(self) -> torch.Tensor:
+        return self.data.float() * self.scale
+
+
+class Sub1BitLLM(torch.nn.Module):
+    def __init__(
+        self,
+        model_path: str,
+        config: Optional[Sub1BitConfig] = None,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    ):
+        super().__init__()
+        self.model_path = model_path
+        self.config = config or Sub1BitConfig()
+        self.device = device
+        self.layers: Dict[int, LowRankFactor] = {}
+        self.metadata: Dict = {}
+
+    @classmethod
+    def from_fp16(
+        cls,
+        model_path: str,
+        config: Optional[Sub1BitConfig] = None,
+        checkpoint_dir: Optional[str] = None,
+        device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    ) -> "Sub1BitLLM":
+        instance = cls(model_path, config, device)
+        if checkpoint_dir is None:
+            checkpoint_dir = Path(model_path).parent / "checkpoints"
+        else:
+            checkpoint_dir = Path(checkpoint_dir)
+        if not checkpoint_dir.exists():
+            raise FileNotFoundError(f"Checkpoint directory not found: {checkpoint_dir}")
+        for ckpt_file in sorted(checkpoint_dir.glob("layer_*.pt")):
+            layer_idx = int(ckpt_file.stem.split("_")[1])
+            factor = torch.load(ckpt_file, weights_only=False, map_location=device)
+            U = torch.from_numpy(factor['U']).to(device)
+            S = torch.from_numpy(factor['S']).to(device)
+            Vt = torch.from_numpy(factor['Vt']).to(device)
+            instance.layers[layer_idx] = LowRankFactor(U, S, Vt)
+        instance.metadata = {
+            'num_layers': len(instance.layers),
+            'rank': instance.config.rank,
+            'energy_threshold': instance.config.energy_threshold
+        }
+        return instance
+
+    def load_checkpoint(self, checkpoint_path: str) -> "Sub1BitLLM":
+        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+        if 'layers' in checkpoint:
+            for layer_idx, factor_data in checkpoint['layers'].items():
+                self.layers[int(layer_idx)] = LowRankFactor(
+                    torch.from_numpy(factor_data['U']).to(self.device),
+                    torch.from_numpy(factor_data['S']).to(self.device),
+                    torch.from_numpy(factor_data['Vt']).to(self.device),
+                )
+        return self
+
+    def state_dict(self) -> Dict[str, torch.Tensor]:
+        state = {}
+        for layer_idx, layer in self.layers.items():
+            state[f'layers.{layer_idx}.U'] = layer.U
+            state[f'layers.{layer_idx}.S'] = layer.S
+            state[f'layers.{layer_idx}.Vt'] = layer.Vt
+        return state
+
+    def forward(self, x: torch.Tensor, layer_indices: Optional[List[int]] = None) -> Dict[int, torch.Tensor]:
+        outputs = {}
+        indices = layer_indices if layer_indices is not None else list(self.layers.keys())
+        for idx in indices:
+            if idx in self.layers:
+                outputs[idx] = self.layers[idx].forward_lowrank(x)
+        return outputs
+
+    def get_weight(self, layer_idx: int) -> torch.Tensor:
+        if layer_idx not in self.layers:
+            raise KeyError(f"Layer {layer_idx} not found")
+        return self.layers[layer_idx]()
+
+    def iter_layers(self) -> Iterator[Tuple[int, LowRankFactor]]:
+        for idx in sorted(self.layers.keys()):
+            yield idx, self.layers[idx]
+
+    def compression_stats(self) -> Dict[str, float]:
+        total_original = 0
+        total_factor = 0
+        for _, layer in self.iter_layers():
+            orig_size = layer.U.shape[0] * layer.Vt.shape[1]
+            factor_size = layer.U.numel() + layer.S.numel() + layer.Vt.numel()
+            total_original += orig_size
+            total_factor += factor_size
+        return {
+            'compression_ratio': total_original / total_factor if total_factor > 0 else 0,
+            'avg_rank': sum(l.rank for _, l in self.iter_layers()) / max(len(self.layers), 1)
+        }
+
+    def to_gguf(self, output_path: str, metadata: Optional[Dict] = None):
+        writer = GGUFWriter(output_path)
+        writer.add_key_value("general.architecture", self.config.architecture)
+        writer.add_key_value("general.name", self.config.model_name)
+        writer.add_key_value("quantization.type", "sub1bit_lowrank")
+        writer.add_key_value("quantization.U_bits", self.config.U_bits)
+        writer.add_key_value("quantization.S_bits", self.config.S_bits)
+        writer.add_key_value("quantization.Vt_bits", self.config.Vt_bits)
+        if metadata:
+            for key, value in metadata.items():
+                writer.add_key_value(key, value)
+        for layer_idx, layer in self.iter_layers():
+            writer.add_tensor(
+                f"model.layers.{layer_idx}.U",
+                layer.U.cpu().numpy().astype(np.float16),
+                GGML_TYPES['float16']
+            )
+            writer.add_tensor(
+                f"model.layers.{layer_idx}.S",
+                layer.S.cpu().numpy().astype(np.float16),
+                GGML_TYPES['float16']
+            )
+            writer.add_tensor(
+                f"model.layers.{layer_idx}.Vt",
+                layer.Vt.cpu().numpy().astype(np.float16),
+                GGML_TYPES['float16']
+            )
+            writer.add_tensor(
+                f"model.layers.{layer_idx}.rank",
+                np.array([layer.rank], dtype=np.int32),
+                GGML_TYPES['int32']
+            )
+        writer.write()
+        return os.path.getsize(output_path)
+
+
+def from_fp16(
+    model_path: str,
+    config: Optional[Sub1BitConfig] = None,
+    checkpoint_dir: Optional[str] = None,
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+) -> Sub1BitLLM:
+    return Sub1BitLLM.from_fp16(model_path, config, checkpoint_dir, device)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="Sub1BitLLM API Demo")
+    parser.add_argument("--model", type=str, required=True, help="Path to model weights")
+    parser.add_argument("--checkpoint-dir", type=str, default=None, help="Path to checkpoint directory")
+    parser.add_argument("--device", type=str, default="cuda", help="Device")
+    args = parser.parse_args()
+    config = Sub1BitConfig(
+        codebook_dim=128,
+        energy_threshold=0.95,
+        model_name="llama-2-7b-sub1bit"
+    )
+    model = from_fp16(args.model, config=config, checkpoint_dir=args.checkpoint_dir, device=args.device)
+    print(f"Loaded Sub1BitLLM with {len(model.layers)} layers")
+    print(f"Compression stats: {model.compression_stats()}")

src/__init__.py

@@ -0,0 +1,63 @@
+import torch
+from typing import Dict, Optional, List
+from .Sub1BitLLM import Sub1BitLLM, Sub1BitConfig, from_fp16
+from .lowrank_factorization import low_rank_factorize, factorize_model_weights, compute_optimal_rank
+from .quantization import (
+    ternary_quantize, ternary_pack, ternary_unpack,
+    sigma_quantize, sigma_dequantize,
+    quantize_factor, pack_factor, unpack_factor, dequantize_factor,
+)
+from .groupwise_int4 import (
+    dequantize_groupwise_int4,
+    estimate_groupwise_int4_bpw,
+    pack_signed_int4,
+    quantize_groupwise_int4,
+    unpack_signed_int4,
+)
+from .mixed_budget import allocate_mixed_budget, summarize_allocation
+from .error_budget_residual import (
+    dequantize_binary_residual,
+    dequantize_error_budget_residual,
+    estimate_binary_residual_bpw,
+    estimate_error_budget_residual_bpw,
+    quantize_binary_residual,
+    quantize_error_budget_residual,
+)
+from .gguf_writer import GGUFWriter, GGML_TYPES, GGUF_TYPES
+from .pack_gguf import pack_sub1bit_model, QuantizedLayer
+
+__all__ = [
+    "Sub1BitLLM",
+    "Sub1BitConfig",
+    "from_fp16",
+    "low_rank_factorize",
+    "factorize_model_weights",
+    "compute_optimal_rank",
+    "ternary_quantize",
+    "ternary_pack",
+    "ternary_unpack",
+    "sigma_quantize",
+    "sigma_dequantize",
+    "quantize_factor",
+    "pack_factor",
+    "unpack_factor",
+    "dequantize_factor",
+    "dequantize_groupwise_int4",
+    "estimate_groupwise_int4_bpw",
+    "pack_signed_int4",
+    "quantize_groupwise_int4",
+    "unpack_signed_int4",
+    "allocate_mixed_budget",
+    "summarize_allocation",
+    "dequantize_binary_residual",
+    "dequantize_error_budget_residual",
+    "estimate_binary_residual_bpw",
+    "estimate_error_budget_residual_bpw",
+    "quantize_binary_residual",
+    "quantize_error_budget_residual",
+    "GGUFWriter",
+    "GGML_TYPES",
+    "GGUF_TYPES",
+    "pack_sub1bit_model",
+    "QuantizedLayer",
+]

src/error_budget_residual.py

@@ -0,0 +1,397 @@
+from __future__ import annotations
+
+from math import ceil, log2
+from typing import Sequence
+
+import torch
+
+from .groupwise_int4 import pack_signed_int4, unpack_signed_int4
+
+
+def pack_signed_int2(values: torch.Tensor) -> torch.Tensor:
+    """Pack signed INT2 values in {-1, 0, 1}, four values per byte."""
+    flat = values.to(torch.int8).flatten()
+    if flat.numel() and (flat.min().item() < -1 or flat.max().item() > 1):
+        raise ValueError("signed int2 values must be in {-1, 0, 1}")
+    encoded = (flat + 1).to(torch.uint8)
+    pad = (-encoded.numel()) % 4
+    if pad:
+        encoded = torch.cat([encoded, torch.ones(pad, dtype=torch.uint8, device=encoded.device)])
+    encoded = encoded.reshape(-1, 4)
+    return (
+        encoded[:, 0]
+        | torch.bitwise_left_shift(encoded[:, 1], 2)
+        | torch.bitwise_left_shift(encoded[:, 2], 4)
+        | torch.bitwise_left_shift(encoded[:, 3], 6)
+    ).contiguous()
+
+
+def unpack_signed_int2(packed: torch.Tensor, count: int) -> torch.Tensor:
+    """Unpack bytes produced by pack_signed_int2."""
+    if count < 0:
+        raise ValueError("count must be non-negative")
+    packed = packed.to(torch.uint8).flatten()
+    values = torch.stack(
+        (
+            torch.bitwise_and(packed, 0x03),
+            torch.bitwise_and(torch.bitwise_right_shift(packed, 2), 0x03),
+            torch.bitwise_and(torch.bitwise_right_shift(packed, 4), 0x03),
+            torch.bitwise_and(torch.bitwise_right_shift(packed, 6), 0x03),
+        ),
+        dim=1,
+    ).flatten()[:count]
+    return values.to(torch.int16).sub(1).to(torch.int8)
+
+
+def pack_binary_sign(values: torch.Tensor) -> torch.Tensor:
+    """Pack residual signs {-1, +1}, eight signs per byte."""
+    flat = values.to(torch.int8).flatten()
+    if flat.numel() and not torch.all((flat == -1) | (flat == 1) | (flat == 0)):
+        raise ValueError("binary residual signs must be -1, 0, or 1")
+    encoded = torch.where(flat > 0, torch.ones_like(flat, dtype=torch.uint8), torch.zeros_like(flat, dtype=torch.uint8))
+    pad = (-encoded.numel()) % 8
+    if pad:
+        encoded = torch.cat([encoded, torch.zeros(pad, dtype=torch.uint8, device=encoded.device)])
+    encoded = encoded.reshape(-1, 8)
+    shifts = torch.arange(8, dtype=torch.uint8, device=encoded.device)
+    return torch.sum(torch.bitwise_left_shift(encoded, shifts), dim=1).to(torch.uint8).contiguous()
+
+
+def unpack_binary_sign(packed: torch.Tensor, count: int) -> torch.Tensor:
+    """Unpack bytes produced by pack_binary_sign into {-1, +1} signs."""
+    if count < 0:
+        raise ValueError("count must be non-negative")
+    packed = packed.to(torch.uint8).flatten()
+    shifts = torch.arange(8, dtype=torch.uint8, device=packed.device)
+    bits = torch.bitwise_and(torch.bitwise_right_shift(packed.unsqueeze(1), shifts), 1).flatten()[:count]
+    return torch.where(bits > 0, torch.ones_like(bits, dtype=torch.int8), -torch.ones_like(bits, dtype=torch.int8))
+
+
+def _validate_2d_shape(shape: Sequence[int]) -> tuple[int, int]:
+    if len(shape) != 2:
+        raise ValueError("shape must be a 2D weight matrix shape")
+    rows, cols = int(shape[0]), int(shape[1])
+    if rows <= 0 or cols <= 0:
+        raise ValueError("shape dimensions must be positive")
+    return rows, cols
+
+
+def estimate_binary_residual_bpw(
+    shape: Sequence[int],
+    group_size: int = 128,
+    scale_bits: int = 16,
+) -> float:
+    """Estimate BPW for INT2 base + 1-bit residual signs + two group scales."""
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+    rows, cols = _validate_2d_shape(shape)
+    groups_per_row = ceil(cols / group_size)
+    weight_bits = rows * cols * 3
+    scale_bits_total = rows * groups_per_row * scale_bits * 2
+    return (weight_bits + scale_bits_total) / (rows * cols)
+
+
+def estimate_error_budget_residual_bpw(
+    shape: Sequence[int],
+    group_size: int = 128,
+    outliers_per_group: int = 8,
+    scale_bits: int = 16,
+    correction_bits: int = 4,
+) -> float:
+    """Estimate BPW for INT2 base plus sparse residual corrections.
+
+    Each group stores a base scale, a residual correction scale, and up to
+    outliers_per_group index+signed-correction pairs.
+    """
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+    if outliers_per_group < 0:
+        raise ValueError("outliers_per_group must be non-negative")
+    rows, cols = _validate_2d_shape(shape)
+    groups_per_row = ceil(cols / group_size)
+    groups = rows * groups_per_row
+    index_bits = max(1, ceil(log2(group_size)))
+    base_and_binary_bits = rows * cols * 3
+    base_and_binary_scale_bits = groups * scale_bits * 2
+    outlier_scale_bits = groups * scale_bits if outliers_per_group else 0
+    outlier_bits = groups * outliers_per_group * (index_bits + correction_bits)
+    return (base_and_binary_bits + base_and_binary_scale_bits + outlier_scale_bits + outlier_bits) / (rows * cols)
+
+
+def _reshape_blocks(weight: torch.Tensor, group_size: int) -> tuple[torch.Tensor, int]:
+    if weight.ndim != 2:
+        raise ValueError("weight must be a 2D tensor")
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+
+    source = weight.detach().to(torch.float32)
+    rows, cols = source.shape
+    groups_per_row = ceil(cols / group_size)
+    padded_cols = groups_per_row * group_size
+    if padded_cols != cols:
+        padded = torch.zeros((rows, padded_cols), dtype=source.dtype, device=source.device)
+        padded[:, :cols] = source
+        source = padded
+    return source.reshape(rows, groups_per_row, group_size), padded_cols
+
+
+def quantize_binary_residual(
+    weight: torch.Tensor,
+    group_size: int = 128,
+    scale_dtype: torch.dtype = torch.float16,
+) -> dict:
+    """Quantize with an INT2 base and per-group 1-bit residual correction.
+
+    The residual code stores only the sign of the base reconstruction error plus
+    one learned magnitude per group. This is a compact candidate below 4 BPW.
+    """
+    blocks, padded_cols = _reshape_blocks(weight, group_size)
+    rows, cols = weight.shape
+
+    base_max = blocks.abs().amax(dim=2)
+    safe_max = torch.where(base_max > 0, base_max, torch.ones_like(base_max))
+    multipliers = torch.tensor([0.375, 0.5, 0.625, 0.75, 0.875, 1.0], dtype=torch.float32, device=blocks.device)
+
+    best_sse = torch.full_like(base_max, float("inf"))
+    best_q = torch.zeros_like(blocks, dtype=torch.int8)
+    best_base_scales = torch.ones_like(base_max)
+    best_residual_sign = torch.zeros_like(blocks, dtype=torch.int8)
+    best_residual_scales = torch.zeros_like(base_max)
+
+    for multiplier in multipliers:
+        candidate_scales = (safe_max * multiplier).to(scale_dtype).to(torch.float32)
+        q = torch.round(blocks / candidate_scales.unsqueeze(-1)).clamp(-1, 1)
+        base_rec = q * candidate_scales.unsqueeze(-1)
+        residual = blocks - base_rec
+        residual_scales = residual.abs().mean(dim=2).to(scale_dtype).to(torch.float32)
+        residual_sign = torch.where(residual >= 0, 1.0, -1.0)
+        residual_rec = residual_sign * residual_scales.unsqueeze(-1)
+        sse = (residual - residual_rec).pow(2).sum(dim=2)
+
+        better = sse < best_sse
+        best_sse = torch.where(better, sse, best_sse)
+        best_base_scales = torch.where(better, candidate_scales, best_base_scales)
+        best_residual_scales = torch.where(better, residual_scales, best_residual_scales)
+        best_q = torch.where(better.unsqueeze(-1), q.to(torch.int8), best_q)
+        best_residual_sign = torch.where(better.unsqueeze(-1), residual_sign.to(torch.int8), best_residual_sign)
+
+    zero_groups = base_max <= 0
+    best_residual_sign = torch.where(
+        zero_groups.unsqueeze(-1),
+        torch.zeros_like(best_residual_sign),
+        best_residual_sign,
+    )
+    best_residual_scales = torch.where(zero_groups, torch.zeros_like(best_residual_scales), best_residual_scales)
+
+    return {
+        "format": "int2_binary_residual",
+        "base_bits": 2,
+        "residual_bits": 1,
+        "group_size": int(group_size),
+        "orig_shape": [int(rows), int(cols)],
+        "padded_shape": [int(rows), int(padded_cols)],
+        "base_packed": pack_signed_int2(best_q.cpu()),
+        "base_count": int(best_q.numel()),
+        "base_scales": best_base_scales.to(scale_dtype).cpu(),
+        "residual_packed": pack_binary_sign(best_residual_sign.cpu()),
+        "residual_count": int(best_residual_sign.numel()),
+        "residual_scales": best_residual_scales.to(scale_dtype).cpu(),
+        "bpw": estimate_binary_residual_bpw((rows, cols), group_size=group_size, scale_bits=16),
+    }
+
+
+def _valid_block_mask(rows: int, cols: int, padded_cols: int, group_size: int, device: torch.device) -> torch.Tensor:
+    valid = torch.zeros((rows, padded_cols), dtype=torch.bool, device=device)
+    valid[:, :cols] = True
+    return valid.reshape(rows, padded_cols // group_size, group_size)
+
+
+def quantize_error_budget_residual(
+    weight: torch.Tensor,
+    group_size: int = 128,
+    outliers_per_group: int = 8,
+    scale_dtype: torch.dtype = torch.float16,
+) -> dict:
+    """Quantize with an INT2 base plus sparse top-error residual corrections."""
+    if outliers_per_group < 0:
+        raise ValueError("outliers_per_group must be non-negative")
+    if group_size > 256:
+        raise ValueError("group_size must be <= 256 so residual indices fit in uint8")
+
+    blocks, padded_cols = _reshape_blocks(weight, group_size)
+    rows, cols = weight.shape
+    groups_per_row = padded_cols // group_size
+    valid_blocks = _valid_block_mask(rows, cols, padded_cols, group_size, blocks.device)
+    k = min(int(outliers_per_group), group_size)
+
+    base_max = blocks.abs().amax(dim=2)
+    safe_max = torch.where(base_max > 0, base_max, torch.ones_like(base_max))
+    multipliers = torch.tensor([0.375, 0.5, 0.625, 0.75, 0.875, 1.0], dtype=torch.float32, device=blocks.device)
+
+    best_sse = torch.full_like(base_max, float("inf"))
+    best_q = torch.zeros_like(blocks, dtype=torch.int8)
+    best_base_scales = torch.ones_like(base_max)
+    best_binary_sign = torch.zeros_like(blocks, dtype=torch.int8)
+    best_binary_scales = torch.zeros_like(base_max)
+    best_indices = torch.zeros((rows, groups_per_row, k), dtype=torch.long, device=blocks.device)
+    best_counts = torch.zeros((rows, groups_per_row), dtype=torch.long, device=blocks.device)
+    best_outlier_q = torch.zeros((rows, groups_per_row, k), dtype=torch.int8, device=blocks.device)
+    best_outlier_scales = torch.zeros_like(base_max)
+
+    for multiplier in multipliers:
+        candidate_scales = (safe_max * multiplier).to(scale_dtype).to(torch.float32)
+        q = torch.round(blocks / candidate_scales.unsqueeze(-1)).clamp(-1, 1)
+        base_rec = q * candidate_scales.unsqueeze(-1)
+        residual = blocks - base_rec
+        binary_scales = residual.abs().mean(dim=2).to(scale_dtype).to(torch.float32)
+        binary_sign = torch.where(residual >= 0, 1.0, -1.0)
+        binary_rec = binary_sign * binary_scales.unsqueeze(-1)
+        remaining = residual - binary_rec
+
+        outlier_rec = torch.zeros_like(blocks)
+        top_idx = torch.zeros((rows, groups_per_row, k), dtype=torch.long, device=blocks.device)
+        selected_counts = torch.zeros((rows, groups_per_row), dtype=torch.long, device=blocks.device)
+        outlier_q = torch.zeros((rows, groups_per_row, k), dtype=torch.int8, device=blocks.device)
+        outlier_scales = torch.zeros_like(base_max)
+
+        if k:
+            masked_abs = torch.where(valid_blocks, remaining.abs(), torch.full_like(remaining, -1.0))
+            top_abs, top_idx = torch.topk(masked_abs, k=k, dim=2)
+            selected = top_abs >= 0
+            selected_counts = selected.sum(dim=2)
+            gathered = torch.gather(remaining, 2, top_idx)
+            max_abs = torch.where(top_abs[:, :, 0] > 0, top_abs[:, :, 0], torch.zeros_like(top_abs[:, :, 0]))
+            outlier_scales = (max_abs / 7.0).to(scale_dtype).to(torch.float32)
+            denom = torch.where(outlier_scales > 0, outlier_scales, torch.ones_like(outlier_scales))
+            outlier_q = torch.round(gathered / denom.unsqueeze(-1)).clamp(-7, 7).to(torch.int8)
+            outlier_q = torch.where(selected, outlier_q, torch.zeros_like(outlier_q))
+            correction = outlier_q.to(torch.float32) * outlier_scales.unsqueeze(-1)
+            correction = torch.where(selected, correction, torch.zeros_like(correction))
+            outlier_rec.scatter_add_(2, top_idx, correction)
+
+        error = torch.where(valid_blocks, base_rec + binary_rec + outlier_rec - blocks, torch.zeros_like(blocks))
+        sse = error.pow(2).sum(dim=2)
+        better = sse < best_sse
+        best_sse = torch.where(better, sse, best_sse)
+        best_base_scales = torch.where(better, candidate_scales, best_base_scales)
+        best_q = torch.where(better.unsqueeze(-1), q.to(torch.int8), best_q)
+        best_binary_sign = torch.where(better.unsqueeze(-1), binary_sign.to(torch.int8), best_binary_sign)
+        best_binary_scales = torch.where(better, binary_scales, best_binary_scales)
+        best_indices = torch.where(better.unsqueeze(-1), top_idx, best_indices)
+        best_counts = torch.where(better, selected_counts, best_counts)
+        best_outlier_q = torch.where(better.unsqueeze(-1), outlier_q, best_outlier_q)
+        best_outlier_scales = torch.where(better, outlier_scales, best_outlier_scales)
+
+    return {
+        "format": "int2_error_budget_residual",
+        "base_bits": 2,
+        "correction_bits": 4,
+        "group_size": int(group_size),
+        "outliers_per_group": int(outliers_per_group),
+        "orig_shape": [int(rows), int(cols)],
+        "padded_shape": [int(rows), int(padded_cols)],
+        "base_packed": pack_signed_int2(best_q.cpu()),
+        "base_count": int(best_q.numel()),
+        "base_scales": best_base_scales.to(scale_dtype).cpu(),
+        "binary_residual_packed": pack_binary_sign(best_binary_sign.cpu()),
+        "binary_residual_count": int(best_binary_sign.numel()),
+        "binary_residual_scales": best_binary_scales.to(scale_dtype).cpu(),
+        "outlier_indices": best_indices.to(torch.uint8).cpu(),
+        "outlier_counts": best_counts.to(torch.uint8).cpu(),
+        "outlier_packed": pack_signed_int4(best_outlier_q.cpu()),
+        "outlier_count": int(best_outlier_q.numel()),
+        "outlier_scales": best_outlier_scales.to(scale_dtype).cpu(),
+        "bpw": estimate_error_budget_residual_bpw(
+            (rows, cols),
+            group_size=group_size,
+            outliers_per_group=outliers_per_group,
+            scale_bits=16,
+            correction_bits=4,
+        ),
+    }
+
+
+def dequantize_binary_residual(
+    entry: dict,
+    device: str | torch.device = "cpu",
+    include_residual: bool = True,
+) -> torch.Tensor:
+    """Dequantize an INT2 base plus optional 1-bit residual correction."""
+    rows, cols = [int(dim) for dim in entry["orig_shape"]]
+    padded_rows, padded_cols = [int(dim) for dim in entry.get("padded_shape", entry["orig_shape"])]
+    group_size = int(entry["group_size"])
+    if padded_rows != rows:
+        raise ValueError("padded row count must match original row count")
+    if padded_cols % group_size != 0:
+        raise ValueError("padded columns must be divisible by group_size")
+
+    groups_per_row = padded_cols // group_size
+    if "base_packed" in entry:
+        base_q = unpack_signed_int2(entry["base_packed"].to(device), int(entry["base_count"]))
+    else:
+        base_q = entry["base_q"].to(device=device, dtype=torch.int8).flatten()
+    base_q = base_q.to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, group_size)
+    base_scales = entry["base_scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+    restored = base_q * base_scales
+    if include_residual:
+        if "residual_packed" in entry:
+            residual_sign = unpack_binary_sign(entry["residual_packed"].to(device), int(entry["residual_count"]))
+        else:
+            residual_sign = entry["residual_sign"].to(device=device, dtype=torch.int8).flatten()
+        residual_sign = residual_sign.to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, group_size)
+        residual_scales = entry["residual_scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+        restored = restored + residual_sign * residual_scales
+
+    return restored.reshape(rows, padded_cols)[:, :cols].contiguous()
+
+
+def dequantize_error_budget_residual(
+    entry: dict,
+    device: str | torch.device = "cpu",
+    include_residual: bool = True,
+) -> torch.Tensor:
+    """Dequantize an INT2 base plus sparse residual correction side channel."""
+    rows, cols = [int(dim) for dim in entry["orig_shape"]]
+    padded_rows, padded_cols = [int(dim) for dim in entry.get("padded_shape", entry["orig_shape"])]
+    group_size = int(entry["group_size"])
+    if padded_rows != rows:
+        raise ValueError("padded row count must match original row count")
+    if padded_cols % group_size != 0:
+        raise ValueError("padded columns must be divisible by group_size")
+
+    groups_per_row = padded_cols // group_size
+    base_q = unpack_signed_int2(entry["base_packed"].to(device), int(entry["base_count"]))
+    base_q = base_q.to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, group_size)
+    base_scales = entry["base_scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+    restored = base_q * base_scales
+
+    if include_residual:
+        binary_sign = unpack_binary_sign(
+            entry["binary_residual_packed"].to(device),
+            int(entry["binary_residual_count"]),
+        )
+        binary_sign = binary_sign.to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, group_size)
+        binary_scales = entry["binary_residual_scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+        restored = restored + binary_sign * binary_scales
+
+        k = int(entry["outliers_per_group"])
+        if k <= 0:
+            return restored.reshape(rows, padded_cols)[:, :cols].contiguous()
+
+        outlier_q = unpack_signed_int4(entry["outlier_packed"].to(device), int(entry["outlier_count"]))
+        outlier_q = outlier_q.to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, k)
+        indices = entry["outlier_indices"].to(device=device, dtype=torch.long).reshape(rows, groups_per_row, k)
+        counts = entry["outlier_counts"].to(device=device, dtype=torch.long).reshape(rows, groups_per_row)
+        outlier_scales = entry["outlier_scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+
+        valid = torch.arange(k, device=device).reshape(1, 1, k) < counts.unsqueeze(-1)
+        correction = torch.where(valid, outlier_q * outlier_scales, torch.zeros_like(outlier_q))
+        restored_flat = restored.reshape(rows * groups_per_row, group_size)
+        group_ids = torch.arange(rows * groups_per_row, device=device).reshape(rows, groups_per_row, 1).expand_as(indices)
+        restored_flat.index_put_(
+            (group_ids.reshape(-1)[valid.reshape(-1)], indices.reshape(-1)[valid.reshape(-1)]),
+            correction.reshape(-1)[valid.reshape(-1)],
+            accumulate=True,
+        )
+
+    return restored.reshape(rows, padded_cols)[:, :cols].contiguous()

src/gguf_writer.py

@@ -0,0 +1,95 @@
+import struct
+import os
+import numpy as np
+from typing import List, Tuple, Any
+
+
+GGUF_MAGIC = b'GGUF'
+GGUF_VERSION = 3
+
+
+GGUF_TYPES = {
+    'uint8': 0, 'int8': 1,
+    'uint16': 2, 'int16': 3,
+    'uint32': 4, 'int32': 5,
+    'float32': 6, 'bool': 7,
+    'string': 8, 'array': 9,
+    'uint64': 10, 'int64': 11, 'float64': 12,
+}
+
+
+GGML_TYPES = {
+    'float32': 0, 'float16': 1,
+    'q8_0': 2, 'q4_0': 3, 'q4_1': 4,
+    'int8': 5, 'int16': 6, 'int32': 7,
+}
+
+
+class GGUFWriter:
+    def __init__(self, filepath: str):
+        self.filepath = filepath
+        self.kv_pairs: List[Tuple[str, Any]] = []
+        self.tensors: List[Tuple[str, np.ndarray, int]] = []
+
+    def add_key_value(self, key: str, value: Any):
+        self.kv_pairs.append((key, value))
+
+    def add_tensor(self, name: str, data: np.ndarray, ggml_type: int = 0):
+        self.tensors.append((name, data, ggml_type))
+
+    def write(self):
+        os.makedirs(os.path.dirname(self.filepath) or ".", exist_ok=True)
+        with open(self.filepath, 'wb') as f:
+            f.write(GGUF_MAGIC)
+            f.write(struct.pack('<I', GGUF_VERSION))
+            f.write(struct.pack('<Q', len(self.tensors)))
+            f.write(struct.pack('<Q', len(self.kv_pairs)))
+            for key, value in self.kv_pairs:
+                self._write_string(f, key)
+                self._write_value(f, value)
+            sizes = [data.nbytes for _, data, _ in self.tensors]
+            data_offset = f.tell() + sum(
+                8 + len(name.encode('utf-8')) + 4 + 8 * data.ndim + 4 + 8
+                for name, data, _ in self.tensors
+            )
+            for i, (name, data, ggml_type) in enumerate(self.tensors):
+                self._write_string(f, name)
+                f.write(struct.pack('<I', data.ndim))
+                for dim in data.shape:
+                    f.write(struct.pack('<q', dim))
+                f.write(struct.pack('<I', ggml_type))
+                f.write(struct.pack('<Q', data_offset))
+                data_offset += sizes[i]
+            for _, data, _ in self.tensors:
+                f.write(data.tobytes())
+
+    def _write_string(self, f, s: str):
+        encoded = s.encode('utf-8')
+        f.write(struct.pack('<Q', len(encoded)))
+        f.write(encoded)
+
+    def _write_value(self, f, value: Any):
+        if isinstance(value, str):
+            f.write(struct.pack('<i', GGUF_TYPES['string']))
+            self._write_string(f, value)
+        elif isinstance(value, bool):
+            f.write(struct.pack('<i', GGUF_TYPES['bool']))
+            f.write(struct.pack('<b', 1 if value else 0))
+        elif isinstance(value, int):
+            f.write(struct.pack('<i', GGUF_TYPES['int32']))
+            f.write(struct.pack('<i', value))
+        elif isinstance(value, float):
+            f.write(struct.pack('<i', GGUF_TYPES['float32']))
+            f.write(struct.pack('<f', value))
+        elif isinstance(value, list):
+            f.write(struct.pack('<i', GGUF_TYPES['array']))
+            if value:
+                f.write(struct.pack('<i', GGUF_TYPES.get(type(value[0]).__name__, 8)))
+            f.write(struct.pack('<Q', len(value)))
+            for item in value:
+                if isinstance(item, str):
+                    self._write_string(f, item)
+                else:
+                    self._write_value(f, item)
+        else:
+            raise ValueError(f"Unsupported type: {type(value)}")

src/groupwise_int4.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from math import ceil
+from typing import Sequence
+
+import torch
+
+
+INT4_QMAX = 7
+
+
+def _validate_signed_int4(values: torch.Tensor) -> torch.Tensor:
+    flat = values.to(torch.int8).flatten()
+    if flat.numel() == 0:
+        return flat
+    if flat.min().item() < -8 or flat.max().item() > 7:
+        raise ValueError("signed int4 values must be in [-8, 7]")
+    return flat
+
+
+def pack_signed_int4(values: torch.Tensor) -> torch.Tensor:
+    """Pack signed int4 values into uint8 bytes, two values per byte."""
+    flat = _validate_signed_int4(values)
+    encoded = (flat + 8).to(torch.uint8)
+    if encoded.numel() % 2:
+        encoded = torch.cat([encoded, torch.full((1,), 8, dtype=torch.uint8, device=encoded.device)])
+
+    low = encoded[0::2]
+    high = torch.bitwise_left_shift(encoded[1::2], 4)
+    return torch.bitwise_or(low, high).contiguous()
+
+
+def unpack_signed_int4(packed: torch.Tensor, count: int) -> torch.Tensor:
+    """Unpack uint8 bytes produced by pack_signed_int4."""
+    if count < 0:
+        raise ValueError("count must be non-negative")
+    packed = packed.to(torch.uint8).flatten()
+    low = torch.bitwise_and(packed, 0x0F)
+    high = torch.bitwise_right_shift(packed, 4)
+    encoded = torch.stack((low, high), dim=1).flatten()[:count]
+    return encoded.to(torch.int16).sub(8).to(torch.int8)
+
+
+def estimate_groupwise_int4_bpw(
+    shape: Sequence[int],
+    group_size: int = 128,
+    scale_bits: int = 16,
+) -> float:
+    """Estimate bits per weight for row-wise group INT4 plus per-group scales."""
+    if len(shape) != 2:
+        raise ValueError("shape must be a 2D weight matrix shape")
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+
+    rows, cols = int(shape[0]), int(shape[1])
+    if rows <= 0 or cols <= 0:
+        raise ValueError("shape dimensions must be positive")
+
+    groups_per_row = ceil(cols / group_size)
+    weight_bits = rows * cols * 4
+    scale_overhead_bits = rows * groups_per_row * scale_bits
+    return (weight_bits + scale_overhead_bits) / (rows * cols)
+
+
+def quantize_groupwise_int4(
+    weight: torch.Tensor,
+    group_size: int = 128,
+    scale_dtype: torch.dtype = torch.float16,
+) -> dict:
+    """Symmetric row-wise group INT4 quantization for 2D weight matrices.
+
+    This is an inference-oriented storage format: the packed nibbles and group
+    scales can be consumed by an INT4 kernel later, while still being easy to
+    dequantize for correctness/perplexity evaluation today.
+    """
+    if weight.ndim != 2:
+        raise ValueError("weight must be a 2D tensor")
+    if group_size <= 0:
+        raise ValueError("group_size must be positive")
+
+    source = weight.detach().to(torch.float32)
+    rows, cols = source.shape
+    groups_per_row = ceil(cols / group_size)
+    padded_cols = groups_per_row * group_size
+    if padded_cols != cols:
+        padded = torch.zeros((rows, padded_cols), dtype=source.dtype, device=source.device)
+        padded[:, :cols] = source
+        source = padded
+
+    blocks = source.reshape(rows, groups_per_row, group_size)
+    max_abs = blocks.abs().amax(dim=2)
+    scales = torch.where(
+        max_abs > 0,
+        max_abs / INT4_QMAX,
+        torch.ones_like(max_abs),
+    )
+    stored_scales = scales.to(scale_dtype).to(torch.float32)
+    q = torch.round(blocks / stored_scales.unsqueeze(-1)).clamp(-INT4_QMAX, INT4_QMAX).to(torch.int8)
+
+    return {
+        "format": "groupwise_int4",
+        "bits": 4,
+        "group_size": int(group_size),
+        "orig_shape": [int(rows), int(cols)],
+        "padded_shape": [int(rows), int(padded_cols)],
+        "scales": stored_scales.to(scale_dtype).cpu(),
+        "packed_int4": pack_signed_int4(q.cpu()),
+        "bpw": estimate_groupwise_int4_bpw((rows, cols), group_size=group_size, scale_bits=16),
+    }
+
+
+def dequantize_groupwise_int4(entry: dict, device: str | torch.device = "cpu") -> torch.Tensor:
+    """Dequantize an entry produced by quantize_groupwise_int4."""
+    rows, cols = [int(dim) for dim in entry["orig_shape"]]
+    padded_rows, padded_cols = [int(dim) for dim in entry.get("padded_shape", entry["orig_shape"])]
+    group_size = int(entry["group_size"])
+    if padded_rows != rows:
+        raise ValueError("padded row count must match original row count")
+    if padded_cols % group_size != 0:
+        raise ValueError("padded columns must be divisible by group_size")
+
+    groups_per_row = padded_cols // group_size
+    count = rows * groups_per_row * group_size
+    q = unpack_signed_int4(entry["packed_int4"].to(device), count).to(torch.float32)
+    q = q.reshape(rows, groups_per_row, group_size)
+    scales = entry["scales"].to(device=device, dtype=torch.float32).reshape(rows, groups_per_row, 1)
+    restored = (q * scales).reshape(rows, padded_cols)
+    return restored[:, :cols].contiguous()

src/lowrank_factorization.py

@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+from typing import Dict, Tuple, List
+
+
+def compute_optimal_rank(singular_values: torch.Tensor, energy_threshold: float = 0.95) -> int:
+    squared_sv = singular_values ** 2
+    cumulative = torch.cumsum(squared_sv, dim=0)
+    total = torch.sum(squared_sv)
+    normalized = cumulative / total
+
+    r = torch.searchsorted(normalized, energy_threshold).item() + 1
+    return min(r, len(singular_values))
+
+
+def low_rank_factorize(weight: torch.Tensor, energy_threshold: float = 0.95) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    original_shape = weight.shape
+
+    if weight.dim() > 2:
+        weight = weight.reshape(-1, weight.shape[-1])
+
+    U, S, Vt = torch.linalg.svd(weight, full_matrices=False)
+
+    r = compute_optimal_rank(S, energy_threshold)
+
+    U_r = U[:, :r]
+    S_r = S[:r]
+    Vt_r = Vt[:r, :]
+
+    U_r = U_r.reshape(original_shape[0], r)
+    Vt_r = Vt_r.reshape(r, original_shape[-1])
+
+    return U_r, S_r, Vt_r, r
+
+
+def factorize_model_weights(weights: Dict[int, torch.Tensor], energy_threshold: float = 0.95) -> Dict:
+    factors = {}
+
+    for layer_idx, W in weights.items():
+        U, S, Vt, r = low_rank_factorize(W, energy_threshold)
+        factors[layer_idx] = {
+            'U': U,
+            'S': S,
+            'Vt': Vt,
+            'rank': r,
+            'original_shape': W.shape,
+            'energy_captured': (S ** 2).sum().item() / ((W ** 2).sum().item() + 1e-8)
+        }
+
+    return factors
+
+
+def reconstruct_weight(factors: Dict) -> torch.Tensor:
+    U = factors['U']
+    S = factors['S']
+    Vt = factors['Vt']
+
+    return torch.matmul(U * S.unsqueeze(0), Vt)
+
+
+def compute_compression_ratio(original_size: int, factors: Dict) -> float:
+    U_size = sum(f['U'].numel() for f in factors.values())
+    S_size = sum(f['S'].numel() for f in factors.values())
+    Vt_size = sum(f['Vt'].numel() for f in factors.values())
+
+    low_rank_size = U_size + S_size + Vt_size
+    return original_size / low_rank_size
+
+
+if __name__ == "__main__":
+    W = torch.randn(4096, 4096)
+
+    U, S, Vt, r = low_rank_factorize(W, energy_threshold=0.95)
+    print(f"Original shape: {W.shape}")
+    print(f"Rank: {r}")
+    print(f"U shape: {U.shape}, S shape: {S.shape}, Vt shape: {Vt.shape}")
+
+    reconstructed = torch.matmul(U * S.unsqueeze(0), Vt)
+    error = torch.norm(W - reconstructed) / torch.norm(W)
+    print(f"Reconstruction error: {error:.6f}")
+
+    weights = {0: torch.randn(4096, 4096), 1: torch.randn(4096, 4096)}
+    factors = factorize_model_weights(weights)
+    for layer_idx, f in factors.items():
+        print(f"Layer {layer_idx}: rank={f['rank']}, energy={f['energy_captured']:.4f}")

src/mixed_budget.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from math import sqrt
+from typing import Iterable
+
+
+def _candidate_score(layer: dict, candidate: dict) -> dict:
+    params = int(layer["params"])
+    activation_weight = float(layer.get("activation_weight", 1.0))
+    scored = deepcopy(candidate)
+    scored["idx"] = layer.get("idx")
+    scored["key"] = layer.get("key")
+    scored["params"] = params
+    scored["activation_weight"] = activation_weight
+    scored["weighted_mse"] = float(candidate["mse"]) * activation_weight
+    scored["total_bits"] = float(candidate["bpw"]) * params
+    scored["weighted_sse"] = scored["weighted_mse"] * params
+    return scored
+
+
+def summarize_allocation(selected_layers: Iterable[dict]) -> dict:
+    selected = list(selected_layers)
+    total_params = sum(int(item["params"]) for item in selected)
+    if total_params <= 0:
+        raise ValueError("allocation must contain at least one parameter")
+
+    total_bits = sum(float(item["bpw"]) * int(item["params"]) for item in selected)
+    weighted_sse = sum(float(item.get("weighted_mse", item["mse"])) * int(item["params"]) for item in selected)
+    weighted_mse = weighted_sse / total_params
+    return {
+        "layers": len(selected),
+        "total_params": total_params,
+        "avg_bpw": total_bits / total_params,
+        "weighted_mse": weighted_mse,
+        "weighted_rmse": sqrt(weighted_mse),
+        "compression_vs_bf16": 16.0 / (total_bits / total_params),
+    }
+
+
+def allocate_mixed_budget(layers: list[dict], target_avg_bpw: float) -> dict:
+    if target_avg_bpw <= 0:
+        raise ValueError("target_avg_bpw must be positive")
+    if not layers:
+        raise ValueError("layers must not be empty")
+
+    candidate_layers: list[list[dict]] = []
+    for layer in layers:
+        candidates = [_candidate_score(layer, item) for item in layer.get("candidates", [])]
+        if not candidates:
+            raise ValueError(f"layer {layer.get('key', layer.get('idx'))} has no candidates")
+        candidates.sort(key=lambda item: (item["bpw"], item["weighted_mse"], item["method"]))
+        candidate_layers.append(candidates)
+
+    selected_indices = [0 for _ in candidate_layers]
+    selected = [candidates[0] for candidates in candidate_layers]
+    total_params = sum(item["params"] for item in selected)
+    budget_bits = target_avg_bpw * total_params
+    current_bits = sum(item["total_bits"] for item in selected)
+    if current_bits > budget_bits + 1e-9:
+        raise ValueError("cheapest candidates exceed target_avg_bpw")
+
+    while True:
+        best_upgrade = None
+        for layer_idx, candidates in enumerate(candidate_layers):
+            current = selected[layer_idx]
+            for candidate_idx, candidate in enumerate(candidates):
+                if candidate_idx == selected_indices[layer_idx]:
+                    continue
+                extra_bits = candidate["total_bits"] - current["total_bits"]
+                if extra_bits <= 0:
+                    continue
+                if current_bits + extra_bits > budget_bits + 1e-9:
+                    continue
+                error_reduction = current["weighted_sse"] - candidate["weighted_sse"]
+                if error_reduction <= 0:
+                    continue
+                score = error_reduction / extra_bits
+                contender = (score, error_reduction, -extra_bits, layer_idx, candidate_idx, candidate)
+                if best_upgrade is None or contender > best_upgrade:
+                    best_upgrade = contender
+
+        if best_upgrade is None:
+            break
+
+        _, _, _, layer_idx, candidate_idx, candidate = best_upgrade
+        current_bits += candidate["total_bits"] - selected[layer_idx]["total_bits"]
+        selected_indices[layer_idx] = candidate_idx
+        selected[layer_idx] = candidate
+
+    summary = summarize_allocation(selected)
+    summary.update(
+        {
+            "target_avg_bpw": float(target_avg_bpw),
+            "selected_layers": selected,
+            "method_counts": _method_counts(selected),
+        }
+    )
+    return summary
+
+
+def _method_counts(selected_layers: Iterable[dict]) -> dict:
+    counts: dict[str, int] = {}
+    for item in selected_layers:
+        method = str(item["method"])
+        counts[method] = counts.get(method, 0) + 1
+    return dict(sorted(counts.items()))

src/pack_gguf.py

@@ -0,0 +1,98 @@
+import torch
+import os
+import numpy as np
+from typing import Dict, Optional, List, Tuple
+from dataclasses import dataclass
+
+from .gguf_writer import GGUFWriter, GGML_TYPES
+
+
+@dataclass
+class QuantizedLayer:
+    U: torch.Tensor
+    S: torch.Tensor
+    Vt: torch.Tensor
+    U_scale: torch.Tensor
+    S_scale: torch.Tensor
+    bit_allocations: Tuple[int, int, int]
+
+
+def create_lowrank_type_defs() -> str:
+    return """
+enum ggml_type {
+    GGML_TYPE_F32     = 0,
+    GGML_TYPE_F16     = 1,
+    GGML_TYPE_Q8_0    = 2,
+    GGML_TYPE_Q4_0    = 3,
+    GGML_TYPE_Q4_1    = 4,
+    GGML_TYPE_LOWRANK_UV_0BIT  = 100,
+    GGML_TYPE_LOWRANK_UV_1BIT  = 101,
+    GGML_TYPE_LOWRANK_SIGMA_2BIT = 102,
+};
+"""
+
+
+def pack_sub1bit_model(
+    factors: Dict[int, Dict],
+    output_path: str,
+    model_name: str = "llama-2-7b-sub1bit",
+    metadata: Optional[Dict] = None
+):
+    writer = GGUFWriter(output_path)
+    writer.add_key_value("general.architecture", "llama")
+    writer.add_key_value("general.name", model_name)
+    writer.add_key_value("general.file_type", "sub1bit")
+    if metadata:
+        for key, value in metadata.items():
+            writer.add_key_value(key, value)
+    for layer_idx, layer_data in factors.items():
+        writer.add_tensor(
+            f"layer.{layer_idx}.U",
+            layer_data.get('U_packed', layer_data['U']).cpu().numpy(),
+            GGML_TYPES['int8']
+        )
+        writer.add_tensor(
+            f"layer.{layer_idx}.S",
+            layer_data['S'].cpu().numpy().astype(np.float16),
+            GGML_TYPES['float16']
+        )
+        writer.add_tensor(
+            f"layer.{layer_idx}.Vt",
+            layer_data.get('Vt_packed', layer_data['Vt']).cpu().numpy(),
+            GGML_TYPES['int8']
+        )
+        if 'U_scale' in layer_data:
+            writer.add_tensor(
+                f"layer.{layer_idx}.U_scale",
+                np.array([layer_data['U_scale'].item()], dtype=np.float32),
+                GGML_TYPES['float32']
+            )
+        if 'Vt_scale' in layer_data:
+            writer.add_tensor(
+                f"layer.{layer_idx}.Vt_scale",
+                np.array([layer_data['Vt_scale'].item()], dtype=np.float32),
+                GGML_TYPES['float32']
+            )
+        if 'S_scale' in layer_data:
+            writer.add_tensor(
+                f"layer.{layer_idx}.S_scale",
+                np.array([layer_data['S_scale'].item()], dtype=np.float32),
+                GGML_TYPES['float32']
+            )
+    writer.write()
+    return os.path.getsize(output_path)
+
+
+if __name__ == "__main__":
+    dummy_factors = {
+        0: {
+            'U': torch.randn(4096, 16),
+            'S': torch.randn(16),
+            'Vt': torch.randn(16, 4096),
+            'rank': 16
+        }
+    }
+    output_path = "C:/Users/Zwmar/projects/sub1quant/quantized/test.gguf"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    size = pack_sub1bit_model(dummy_factors, output_path)
+    print(f"GGUF file created: {size} bytes")

src/quantization.py

@@ -0,0 +1,96 @@
+import torch
+import numpy as np
+
+
+def ternary_quantize(x: torch.Tensor):
+    scale = x.abs().max()
+    if scale == 0:
+        scale = 1.0
+    normalized = x / scale
+    ternary = torch.sign(normalized)
+    ternary[ternary == 0] = 1
+    return ternary.to(torch.int8), scale
+
+
+def ternary_pack(t: torch.Tensor) -> torch.Tensor:
+    encoded = t.to(torch.int8) + 1
+    n = encoded.numel()
+    pad = (5 - n % 5) % 5
+    if pad:
+        encoded = torch.cat([encoded.flatten(), torch.zeros(pad, dtype=torch.int8, device=encoded.device)])
+    weights = torch.tensor([81, 27, 9, 3, 1], dtype=torch.int32, device=encoded.device)
+    packed = (encoded.reshape(-1, 5).to(torch.int32) * weights).sum(dim=1)
+    return packed.to(torch.uint8)
+
+
+def ternary_unpack(packed: torch.Tensor, original_shape: tuple) -> torch.Tensor:
+    weights = torch.tensor([81, 27, 9, 3, 1], dtype=torch.int32, device=packed.device)
+    expanded = packed.to(torch.int32).unsqueeze(-1) // weights % 3
+    flat = (expanded - 1).flatten()
+    return flat[:np.prod(original_shape)].reshape(original_shape).to(torch.int8)
+
+
+def sigma_quantize(s: torch.Tensor, num_bits: int = 2):
+    max_val = s.abs().max()
+    if max_val == 0:
+        max_val = 1.0
+    qmax = 2 ** (num_bits - 1) - 1
+    scale = max_val / qmax
+    quantized = (s / scale).round().clamp(-qmax, qmax)
+    return quantized.to(torch.int8), scale
+
+
+def sigma_dequantize(q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    return q.float() * scale
+
+
+def quantize_factor(U: torch.Tensor, S: torch.Tensor, Vt: torch.Tensor, sigma_bits: int = 2):
+    U_q, U_scale = ternary_quantize(U)
+    Vt_q, Vt_scale = ternary_quantize(Vt)
+    S_q, S_scale = sigma_quantize(S, sigma_bits)
+    return {
+        'U': U_q, 'U_scale': U_scale,
+        'Vt': Vt_q, 'Vt_scale': Vt_scale,
+        'S': S_q, 'S_scale': S_scale,
+    }
+
+
+def pack_factor(data: dict):
+    return {
+        'U_packed': ternary_pack(data['U']),
+        'U_scale': data['U_scale'],
+        'Vt_packed': ternary_pack(data['Vt']),
+        'Vt_scale': data['Vt_scale'],
+        'S': data['S'],
+        'S_scale': data['S_scale'],
+        'U_shape': data['U'].shape,
+        'Vt_shape': data['Vt'].shape,
+    }
+
+
+def unpack_factor(data: dict):
+    if 'U_packed' in data:
+        U = ternary_unpack(data['U_packed'], data['U_shape'])
+        Vt = ternary_unpack(data['Vt_packed'], data['Vt_shape'])
+    else:
+        U = data['U']
+        Vt = data['Vt']
+    return {
+        'U': U,
+        'U_scale': data['U_scale'],
+        'Vt': Vt,
+        'Vt_scale': data['Vt_scale'],
+        'S': data['S'],
+        'S_scale': data['S_scale'],
+    }
+
+
+def dequantize_factor(data: dict) -> torch.Tensor:
+    if 'U_packed' in data:
+        U = ternary_unpack(data['U_packed'], data['U_shape']).float()
+        Vt = ternary_unpack(data['Vt_packed'], data['Vt_shape']).float()
+    else:
+        U = data['U'].float()
+        Vt = data['Vt'].float()
+    S = sigma_dequantize(data['S'], data['S_scale'])
+    return torch.matmul(U * S.unsqueeze(0), Vt)

test_perplexity.py

@@ -0,0 +1,115 @@
+"""
+Perplexity evaluation for quantized .pt checkpoints.
+
+This runs a real transformer forward pass after applying quantized weights.
+For GGUF/GEMV latency checks, use run_gemv_chain.py instead.
+"""
+
+import argparse
+import gc
+from pathlib import Path
+
+import torch
+
+from scripts.eval_quantized import apply_quantized_weights, eval_perplexity
+
+
+def evaluate_quantized_perplexity(
+    model_name="models/gemma-4-E2B",
+    quantized_path="quantized/gemma-4-E2B-sub1bit.pt",
+    wikitext_path="data/wiki.test.txt",
+    device=None,
+    max_length=512,
+    stride=512,
+):
+    """Evaluate WikiText perplexity after applying quantized weights."""
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if not Path(wikitext_path).exists():
+        raise FileNotFoundError(f"WikiText not found: {wikitext_path}")
+
+    print("=" * 60)
+    print("QUANTIZED MODEL PERPLEXITY EVALUATION")
+    print("=" * 60)
+    print(f"Device: {device}")
+    print(f"Model: {model_name}")
+    print(f"Quantized: {quantized_path}")
+    print(f"WikiText: {wikitext_path}")
+
+    print("\n[1] Loading quantized checkpoint...")
+    q_data = torch.load(quantized_path, map_location="cpu", weights_only=True)
+    quantized = q_data["quantized"]
+    print(f"  {len(quantized)} quantized entries")
+
+    print("\n[2] Loading base model...")
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map=device,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True,
+    )
+    model.eval()
+
+    print("\n[3] Applying quantized weights...")
+    apply_stats = apply_quantized_weights(
+        model,
+        quantized,
+        device=device,
+        model_dir=model_name if Path(model_name).is_dir() else None,
+        checkpoint_weight_keys=q_data.get("weight_keys"),
+    )
+    print(f"  Replaced {apply_stats['replaced']}/{len(quantized)} weights")
+    if apply_stats["skipped"]:
+        print(f"  Skipped {len(apply_stats['skipped'])} shared-KV checkpoint entries")
+
+    print("\n[4] Evaluating perplexity...")
+    ppl, stats = eval_perplexity(
+        model,
+        tokenizer,
+        wikitext_path,
+        device,
+        max_length=max_length,
+        stride=stride,
+    )
+
+    print()
+    print("=" * 60)
+    print("RESULTS")
+    print("=" * 60)
+    print(f"  Perplexity: {ppl:.4f}")
+    print(f"  Chunks: {stats['n_chunks']}")
+    print(f"  Target: <= 10.5")
+    print(f"  Status: {'PASS' if ppl <= 10.5 else 'FAIL'}")
+    print("=" * 60)
+
+    del model
+    gc.collect()
+    if device == "cuda":
+        torch.cuda.empty_cache()
+
+    return ppl, stats
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Evaluate quantized model perplexity")
+    parser.add_argument("--model", default="models/gemma-4-E2B")
+    parser.add_argument("--quantized", default="quantized/gemma-4-E2B-sub1bit.pt")
+    parser.add_argument("--wikitext", default="data/wiki.test.txt")
+    parser.add_argument("--device", default=None)
+    parser.add_argument("--max-length", type=int, default=512)
+    parser.add_argument("--stride", type=int, default=512)
+    args = parser.parse_args()
+
+    evaluate_quantized_perplexity(
+        model_name=args.model,
+        quantized_path=args.quantized,
+        wikitext_path=args.wikitext,
+        device=args.device,
+        max_length=args.max_length,
+        stride=args.stride,
+    )