{ "cells": [ { "cell_type": "markdown", "metadata": { "colab_type": "text", "id": "view-in-github" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# FABQ-VP: Variable Precision FABQ-RC\n", "\n", "**Author:** Zach Maronek · May 2026\n", "\n", "---\n", "\n", "## The Problem FABQ-VP Solves\n", "\n", "FABQ-RC operates at exactly 1-bit (binary ±1 + int4). FABQ-VP extends this to **variable precision across 2-8 bits per parameter**, targeting ~3-4 bpw for larger models.\n", "\n", "## Precision Pyramid\n", "\n", "| Precision | Fraction | Bits/Param |\n", "|-----------|----------|------------|\n", "| FP16 | 0.5% | 16.0 |\n", "| int8 | 4.5% | 8.0 |\n", "| int4 | 20% | 4.0 |\n", "| int2 | 25% | 2.0 |\n", "| binary | 50% | 1.0 |\n", "\n", "**Target:** ~3.0-4.0 bpw with <5% perplexity degradation\n", "\n", "---\n", "\n", "**Contents:** [1. Setup](#1) · [2. Method](#2) · [3. Implementation](#3) · [4. Evaluation](#4)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "## 1. Setup & Imports" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "setup_cell" }, "outputs": [], "source": [ "# Core dependencies\n", "!pip install -q torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121\n", "!pip install -q transformers accelerate bitsandbytes scikit-learn\n", "!pip install -q pandas numpy==1.26.4 tqdm matplotlib seaborn datasets\n", "\n", "import os, math, json, time, sys\n", "import numpy as np\n", "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "from torch.utils.data import DataLoader\n", "from transformers import AutoTokenizer, AutoModelForCausalLM\n", "from sklearn.cluster import MiniBatchKMeans\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "import seaborn as sns\n", "from tqdm.auto import tqdm\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "# Detect platform\n", "try:\n", " from google.colab import userdata\n", " COLAB = True\n", " print(\"Running on Google Colab\")\n", "except ImportError:\n", " COLAB = False\n", " print(\"Running locally or Kaggle\")\n", "\n", "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n", "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n", "\n", "print(f\"Device: {DEVICE}\")\n", "if torch.cuda.is_available():\n", " print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n", " print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "## 2. Method Overview\n", "\n", "FABQ-VP has four stages:\n", "\n", "```\n", "FP16 Weights\n", " │\n", " ▼\n", "Stage 1: Fisher-Weighted Channel Importance\n", " │\n", " ▼\n", "Stage 2: 5-Level Precision Allocation (fp16, int8, int4, int2, binary)\n", " │\n", " ▼\n", "Stage 3: Per-Precision Blocksize Selection\n", " │\n", " ▼\n", "Stage 4: Multi-Level Residual Codebooks\n", " │\n", " ▼\n", "FABQ-VP Quantized Model (~3.5 bpw)\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "## 3. Implementation" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "config_cell" }, "outputs": [], "source": [ "# Model configuration - using Qwen2.5-7B for 8B-class validation\n", "MODEL_NAME = \"Qwen/Qwen2.5-7B\"\n", "MAX_SEQ_LEN = 128\n", "CALIB_SIZE = 512\n", "\n", "# Precision pyramid fractions (must sum to 1.0)\n", "PRECISION_PYRAMID = {\n", " 'fp16': 0.005,\n", " 'int8': 0.045,\n", " 'int4': 0.200,\n", " 'int2': 0.250,\n", " 'binary': 0.500\n", "}\n", "\n", "# Blocksize candidates per precision level\n", "BS_CANDIDATES = {\n", " 'fp16': [1],\n", " 'int8': [1],\n", " 'int4': [16, 32, 64],\n", " 'int2': [32, 64, 128],\n", " 'binary': [64, 128, 256, 512]\n", "}\n", "\n", "# Get HF token for private models\n", "hf_token = None\n", "if COLAB:\n", " try:\n", " hf_token = userdata.get('HF_TOKEN')\n", " except Exception:\n", " pass\n", "else:\n", " hf_token = os.environ.get('HF_TOKEN')\n", "\n", "print(f\"Model: {MODEL_NAME}\")\n", "print(f\"Precision pyramid: {PRECISION_PYRAMID}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "tokenizer_cell" }, "outputs": [], "source": [ "print(f\"Loading tokenizer for {MODEL_NAME}...\")\n", "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True, token=hf_token)\n", "tokenizer.pad_token = tokenizer.eos_token\n", "print(\"Tokenizer loaded\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "model_load_cell" }, "outputs": [], "source": [ "print(f\"Loading {MODEL_NAME}...\")\n", "model = AutoModelForCausalLM.from_pretrained(\n", " MODEL_NAME,\n", " device_map='auto',\n", " torch_dtype=torch.float16,\n", " trust_remote_code=True,\n", " token=hf_token\n", ")\n", "model.eval()\n", "total_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Model loaded: {total_params / 1e9:.2f}B parameters\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3.1 Prepare Calibration Data" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from datasets import load_dataset\n", "\n", "print(\"Loading C4 calibration data...\")\n", "c4 = load_dataset(\n", " \"allenai/c4\",\n", " data_files={\"train\": \"en/c4-train.00000-of-01024.json.gz\"},\n", " split=f\"train[:{CALIB_SIZE}]\"\n", ")\n", "\n", "def tokenize_fn(batch):\n", " enc = tokenizer(\n", " batch['text'],\n", " truncation=True,\n", " max_length=MAX_SEQ_LEN,\n", " padding='max_length'\n", " )\n", " enc['labels'] = enc['input_ids'].copy()\n", " return enc\n", "\n", "cal_dataset = c4.map(tokenize_fn, batched=True, remove_columns=['text'])\n", "cal_dataset.set_format('torch', columns=['input_ids', 'labels'])\n", "cal_loader = DataLoader(cal_dataset, batch_size=1, shuffle=False)\n", "print(f\"{len(cal_loader)} calibration samples ready\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3.2 Stage 1 — Fisher-Weighted Channel Importance" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class FisherAccumulator:\n", " \"\"\"Accumulate Fisher Information per output channel.\"\"\"\n", " def __init__(self, model):\n", " self.model = model\n", " self.hooks = []\n", "\n", " def _hook_fn(self, module, grad_input, grad_output):\n", " if grad_output[0] is not None and hasattr(module, '_fisher_grad') and module.weight.grad is not None:\n", " grad_sq = module.weight.grad.data ** 2\n", " if grad_sq.dim() == 2:\n", " channel_fisher = grad_sq.sum(dim=1)\n", " else:\n", " channel_fisher = grad_sq.sum(dim=(1, 2, 3))\n", " module._fisher_grad += channel_fisher\n", "\n", " def __call__(self, cal_loader, device):\n", " for name, module in self.model.named_modules():\n", " if isinstance(module, nn.Linear):\n", " module._fisher_grad = torch.zeros_like(module.weight).sum(dim=1)\n", "\n", " for module in self.model.modules():\n", " if isinstance(module, nn.Linear):\n", " self.hooks.append(module.register_full_backward_hook(self._hook_fn))\n", "\n", " self.model.train()\n", " for batch in tqdm(cal_loader, desc=\"Computing Fisher\"):\n", " input_ids = batch['input_ids'].to(device)\n", " labels = batch['labels'].to(device)\n", " outputs = self.model(input_ids)\n", " loss = F.cross_entropy(outputs.logits.view(-1, outputs.logits.size(-1)), labels.view(-1))\n", " loss.backward()\n", " self.model.zero_grad()\n", "\n", " for h in self.hooks:\n", " h.remove()\n", " self.model.eval()\n", "\n", " fisher = {}\n", " for name, module in self.model.named_modules():\n", " if isinstance(module, nn.Linear) and hasattr(module, '_fisher_grad'):\n", " fisher[name] = module._fisher_grad.clone()\n", " del module._fisher_grad\n", " return fisher\n", "\n", "print(\"Computing Fisher importance...\")\n", "fisher_acc = FisherAccumulator(model)\n", "fisher_scores = fisher_acc(cal_loader, DEVICE)\n", "print(f\"Fisher computed for {len(fisher_scores)} layers\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3.3 Stage 2 — 5-Level Precision Allocation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def allocate_vp_precision(fisher_dict, pyramid=PRECISION_PYRAMID):\n", " \"\"\"\n", " Allocate precision levels per channel based on Fisher scores.\n", " \"\"\"\n", " allocation = {}\n", " for name, fisher in fisher_dict.items():\n", " out_channels = fisher.shape[0]\n", " order = torch.argsort(fisher, descending=True)\n", " \n", " alloc = {}\n", " cumulative = 0\n", " \n", " for precision, fraction in pyramid.items():\n", " threshold = int(out_channels * fraction)\n", " for i in range(cumulative, min(cumulative + threshold, out_channels)):\n", " alloc[int(order[i])] = precision\n", " cumulative += threshold\n", " \n", " allocation[name] = alloc\n", " return allocation\n", "\n", "allocation = allocate_vp_precision(fisher_scores, PRECISION_PYRAMID)\n", "\n", "# Summarize\n", "total_channels = sum(len(a) for a in allocation.values())\n", "precision_counts = {}\n", "for name, alloc in allocation.items():\n", " for prec in alloc.values():\n", " precision_counts[prec] = precision_counts.get(prec, 0) + 1\n", "\n", "print(\"Precision allocation summary:\")\n", "for prec, count in sorted(precision_counts.items(), key=lambda x: ['fp16', 'int8', 'int4', 'int2', 'binary'].index(x[0])):\n", " print(f\" {prec}: {count:,} channels ({100*count/total_channels:.1f}%)\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3.4 Stage 3 — Per-Precision Blocksize Selection" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def compute_reconstruction_error(weights, blocksize, fisher_channels):\n", " \"\"\"Compute Fisher-weighted reconstruction error.\"\"\"\n", " out_c, in_c = weights.shape\n", " total_err = 0.0\n", " for start in range(0, in_c, blocksize):\n", " end = min(start + blocksize, in_c)\n", " block = weights[:, start:end]\n", " scale = block.std() + 1e-8\n", " block_q = np.where(block > 0, 1.0, -1.0) * scale\n", " recon_err = ((block - block_q) ** 2).mean()\n", " block_fisher = fisher_channels.mean().item()\n", " total_err += block_fisher * recon_err\n", " return total_err\n", "\n", "\n", "def select_blocksize_for_precision(weights, fisher_channels, precision, candidates):\n", " \"\"\"Pick blocksize minimizing reconstruction error for given precision.\"\"\"\n", " if precision in ['fp16', 'int8']:\n", " return 1\n", " \n", " best_b, best_err = candidates[0], float('inf')\n", " for b in candidates:\n", " err = compute_reconstruction_error(weights, b, fisher_channels)\n", " if err < best_err:\n", " best_err = err\n", " best_b = b\n", " return best_b\n", "\n", "print(\"Selecting per-layer blocksize per precision level...\")\n", "blocksize_results = {}\n", "for name, module in tqdm(list(model.named_modules()), desc=\"Blocksize sweep\"):\n", " if not isinstance(module, nn.Linear):\n", " continue\n", " if name not in fisher_scores:\n", " continue\n", " \n", " weights = module.weight.data.cpu().numpy()\n", " fisher = fisher_scores[name]\n", " alloc = allocation[name]\n", " \n", " blocksize_results[name] = {}\n", " for prec, candidates in BS_CANDIDATES.items():\n", " chs = [ch for ch, p in alloc.items() if p == prec]\n", " if not chs:\n", " continue\n", " ch_weights = weights[chs, :]\n", " ch_fisher = fisher[chs]\n", " best_b = select_blocksize_for_precision(ch_weights, ch_fisher, prec, candidates)\n", " blocksize_results[name][prec] = best_b\n", "\n", "print(\"\\nBlocksize distribution by precision:\")\n", "for prec in ['int4', 'int2', 'binary']:\n", " bs_for_prec = [r.get(prec, 0) for r in blocksize_results.values() if prec in r]\n", " if bs_for_prec:\n", " print(f\" {prec}: {pd.Series(bs_for_prec).value_counts().to_dict()}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 3.5 Stage 4 — Multi-Level Residual Codebooks" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def build_multi_residual_codebooks(model, allocation, blocksize_results, cal_loader, device, n_clusters=256, max_samples=8192):\n", " \"\"\"\n", " Build residual codebooks for each precision transition:\n", " - int4->int8: clusters W_int8 - W_int4 residuals\n", " - int2->int4: clusters W_int4 - W_int2 residuals\n", " - binary->int2: clusters W_int2 - W_binary residuals\n", " \"\"\"\n", " codebooks = {}\n", " max_bs = max(BS_CANDIDATES['binary'])\n", " \n", " for (lower_prec, upper_prec) in [('int2', 'int4'), ('binary', 'int2')]:\n", " print(f\"\\nBuilding {lower_prec}->{upper_prec} codebook...\")\n", " all_residuals = []\n", " sample_count = 0\n", " \n", " for batch in tqdm(cal_loader, desc=f\"Collecting {lower_prec} residuals\", total=min(len(cal_loader), max_samples // 8)):\n", " if sample_count >= max_samples:\n", " break\n", " input_ids = batch['input_ids'].to(device)\n", " labels = batch['labels'].to(device)\n", " outputs = model(input_ids)\n", " model.zero_grad()\n", " \n", " for name, module in model.named_modules():\n", " if not isinstance(module, nn.Linear) or name not in allocation:\n", " continue\n", " \n", " weights = module.weight.detach().cpu().numpy()\n", " alloc = allocation[name]\n", " \n", " lower_chs = [ch for ch, p in alloc.items() if p == lower_prec]\n", " \n", " if not lower_chs:\n", " continue\n", " \n", " bs = blocksize_results.get(name, {}).get(lower_prec, 128)\n", " \n", " for ch in lower_chs:\n", " for start in range(0, weights.shape[1], bs):\n", " end = min(start + bs, weights.shape[1])\n", " \n", " # Skip padded blocks\n", " if end - start < bs:\n", " continue\n", " \n", " block = weights[ch, start:end]\n", " scale = block.std() + 1e-8\n", " block_q = np.where(block > 0, 1.0, -1.0) * scale\n", " residual = block - block_q\n", " \n", " res_flat = residual.flatten()\n", " pad_size = max_bs - len(res_flat)\n", " padded_res = np.pad(res_flat, (0, pad_size), mode='constant')\n", " all_residuals.append(padded_res)\n", " sample_count += 1\n", " \n", " if sample_count >= max_samples:\n", " break\n", " if sample_count >= max_samples:\n", " break\n", " if sample_count >= max_samples:\n", " break\n", " \n", " residuals_array = np.array(all_residuals, dtype=np.float32)\n", " print(f\" Collected {residuals_array.shape[0]} blocks\")\n", " \n", " # Cluster\n", " kmeans = MiniBatchKMeans(n_clusters=n_clusters, random_state=42, batch_size=1024, n_init=3)\n", " kmeans.fit(residuals_array)\n", " codebooks[(lower_prec, upper_prec)] = kmeans.cluster_centers_\n", " print(f\" Built {lower_prec}->{upper_prec} codebook with {n_clusters} centroids\")\n", " \n", " return codebooks\n", "\n", "codebooks = build_multi_residual_codebooks(model, allocation, blocksize_results, cal_loader, DEVICE)\n", "print(f\"\\nBuilt {len(codebooks)} residual codebooks\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "## 4. Evaluation" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def compute_bpw_vp(model, allocation, codebooks):\n", " \"\"\"Compute effective bits per parameter for FABQ-VP.\"\"\"\n", " precision_bits = {'fp16': 16, 'int8': 8, 'int4': 4, 'int2': 2, 'binary': 1}\n", " \n", " total_bits = 0\n", " total_params = 0\n", " \n", " for name, module in model.named_modules():\n", " if not isinstance(module, nn.Linear):\n", " continue\n", " if name not in allocation:\n", " continue\n", " \n", " shape = module.weight.shape\n", " n = shape[0] * shape[1]\n", " alloc = allocation[name]\n", " \n", " for prec, bits in precision_bits.items():\n", " ch_count = sum(1 for v in alloc.values() if v == prec)\n", " total_bits += ch_count * shape[1] * bits\n", " \n", " total_params += n\n", " \n", " # Codebook overhead\n", " for cb in codebooks.values():\n", " total_bits += cb.nbytes * 8\n", " \n", " bpw = total_bits / total_params\n", " return bpw\n", "\n", "bpw = compute_bpw_vp(model, allocation, codebooks)\n", "print(f\"FABQ-VP effective bits per parameter: {bpw:.4f}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import math\n", "\n", "def compute_perplexity(model, dataset, tokenizer, device, stride=128, max_samples=256):\n", " \"\"\"Compute perplexity on a dataset.\"\"\"\n", " model.eval()\n", " total_loss = 0.0\n", " total_tokens = 0\n", " \n", " texts = [dataset[i]['text'] for i in range(min(max_samples, len(dataset)))]\n", " \n", " for i in tqdm(range(len(texts)), desc=\"Computing perplexity\"):\n", " text = texts[i]\n", " enc = tokenizer(text, return_tensors='pt', truncation=True, max_length=512)\n", " input_ids = enc['input_ids'].to(device)\n", " seq_len = input_ids.size(1)\n", " \n", " for start in range(0, seq_len - 1, stride):\n", " end = min(start + stride, seq_len - 1)\n", " chunk = input_ids[:, start:end]\n", " labels = input_ids[:, start+1:end+1]\n", " \n", " with torch.no_grad():\n", " outputs = model(chunk)\n", " logits = outputs.logits\n", " \n", " shift_logits = logits[:, :-1, :].contiguous()\n", " shift_labels = labels.contiguous()\n", " \n", " loss = F.cross_entropy(\n", " shift_logits.view(-1, shift_logits.size(-1)),\n", " shift_labels.view(-1),\n", " reduction='sum'\n", " )\n", " total_loss += loss.item()\n", " total_tokens += shift_labels.numel()\n", " \n", " ppl = math.exp(total_loss / total_tokens)\n", " return ppl\n", "\n", "print(\"Computing FP16 baseline perplexity...\")\n", "fp16_ppl = compute_perplexity(model, c4, tokenizer, DEVICE)\n", "print(f\"FP16 perplexity: {fp16_ppl:.4f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 5. Results Summary" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(\"\\n\" + \"=\"*60)\n", "print(\"FABQ-VP PHASE 1 RESULTS\")\n", "print(\"=\"*60)\n", "print(f\"\\nModel: {MODEL_NAME}\")\n", "print(f\"Parameters: {sum(p.numel() for p in model.parameters()) / 1e9:.2f}B\")\n", "print(f\"\\nPrecision allocation:\")\n", "for prec, count in sorted(precision_counts.items(), key=lambda x: ['fp16', 'int8', 'int4', 'int2', 'binary'].index(x[0])):\n", " print(f\" {prec}: {count:,} channels ({100*count/total_channels:.1f}%)\")\n", "print(f\"\\nBits per parameter: {bpw:.4f}\")\n", "print(f\"FP16 perplexity: {fp16_ppl:.4f}\")\n", "print(\"\\nNote: FABQ-VP quantization not yet applied - this is calibration only\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "---\n", "\n", "## Next Steps\n", "\n", "1. **Apply FABQ-VP quantization** to model weights\n", "2. **Measure perplexity** after quantization\n", "3. **Compare** against uniform 3-bit and AWQ 4-bit baselines\n", "4. **Evaluate downstream** tasks (ARC, HellaSwag, etc.)\n", "\n", "---" ] } ], "metadata": { "colab": { "gpuType": "A100", "include_colab_link": true, "machine_shape": "hm", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 4 }