{ "nbformat": 4, "nbformat_minor": 4, "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "# Install dependencies\n", "!pip install -q transformers torch accelerate scikit-learn\n", "!pip install -q pandas numpy tqdm matplotlib seaborn datasets\n", "\n", "import os, sys, math, json, time, gc\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, BitsAndBytesConfig\n", "from sklearn.cluster import MiniBatchKMeans\n", "import pandas as pd\n", "import matplotlib.pyplot as plt\n", "from tqdm.auto import tqdm\n", "import warnings\n", "warnings.filterwarnings('ignore')\n", "\n", "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n", "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n", "os.environ['PYTORCH_ALLOC_CONF'] = 'expandable_segments:True'\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\")\n", "\n", "# Config\n", "MODEL_NAME = \"deepseek-ai/DeepSeek-V4-Flash\"\n", "CALIB_SIZE = 2048\n", "MAX_SEQ_LEN = 32\n", "INT4_FRACTION = 0.05\n", "BS_CANDIDATES = [64, 128, 256, 512]\n", "N_CLUSTERS = 64\n", "HF_TOKEN = os.environ.get('HF_TOKEN', None)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 1.1 Load Tokenizer" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "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 or tokenizer.pad_token\n", "print(\"Tokenizer loaded. Vocab size:\", tokenizer.vocab_size)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "os.makedirs(\"offload\", exist_ok=True)\n", "\n", "print(f\"Loading {MODEL_NAME} with CPU offloading...\")\n", "t0 = time.time()\n", "\n", "model = AutoModelForCausalLM.from_pretrained(\n", " MODEL_NAME,\n", " device_map=\"auto\",\n", " offload_folder=\"offload\",\n", " offload_state_dict=True,\n", " torch_dtype=\"auto\",\n", " low_cpu_mem_usage=True,\n", " trust_remote_code=True,\n", " token=HF_TOKEN\n", ")\n", "\n", "t1 = time.time()\n", "print(f\"Model loaded in {t1-t0:.1f}s\")\n", "if torch.cuda.is_available():\n", " print(f\"VRAM used: {torch.cuda.memory_allocated() / 1e9:.2f} GB\")\n", "\n", "# Inspect model architecture to find MoE layers\n", "print(\"\\nModel architecture inspection:\")\n", "total_params = 0\n", "moe_layers = 0\n", "dense_layers = 0\n", "gate_layers = 0\n", "for name, module in model.named_modules():\n", " if isinstance(module, nn.Linear):\n", " total_params += module.weight.numel()\n", " if 'gate' in name.lower() or 'router' in name.lower():\n", " gate_layers += 1\n", " elif 'expert' in name.lower() or 'mlp' in name.lower():\n", " moe_layers += 1\n", " else:\n", " dense_layers += 1\n", "\n", "print(f\" Total Linear layers: {dense_layers + moe_layers + gate_layers}\")\n", "print(f\" Dense layers: {dense_layers}\")\n", "print(f\" MoE expert layers: {moe_layers}\")\n", "print(f\" Gate/router layers: {gate_layers}\")\n", "print(f\" Total parameters: {total_params/1e9:.1f}B\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "class MoEFisherAccumulator:\n", " def __init__(self, model):\n", " self.model = model\n", " self.hooks = []\n", "\n", " def _hook_fn(self, name, module, grad_input, grad_output):\n", " if grad_output[0] is None:\n", " return\n", " grad = grad_output[0].detach().clone().to(torch.float32).cpu()\n", " if grad.dim() >= 2:\n", " sum_dims = list(range(grad.dim() - 1))\n", " channel_fisher = (grad ** 2).sum(dim=sum_dims)\n", " else:\n", " channel_fisher = (grad ** 2)\n", " if hasattr(module, '_fisher_buf'):\n", " if module._fisher_buf.device.type != 'cpu':\n", " module._fisher_buf = module._fisher_buf.cpu()\n", " if channel_fisher.shape[0] == module._fisher_buf.shape[0]:\n", " module._fisher_buf.add_(channel_fisher)\n", " else:\n", " module._fisher_buf.add_(channel_fisher.sum())\n", " del grad, channel_fisher\n", "\n", " def compute(self, cal_loader, max_batches=16):\n", " for name, module in self.model.named_modules():\n", " if not isinstance(module, nn.Linear):\n", " continue\n", " if 'gate' in name.lower() or 'router' in name.lower():\n", " continue\n", " buf = torch.zeros(module.out_features, device='cpu', dtype=torch.float32)\n", " module.register_buffer('_fisher_buf', buf)\n", " h = module.register_full_backward_hook(\n", " lambda mod, gi, go, n=name: self._hook_fn(n, mod, gi, go)\n", " )\n", " self.hooks.append(h)\n", "\n", " self.model.train()\n", " if hasattr(self.model, 'gradient_checkpointing_enable'):\n", " self.model.gradient_checkpointing_enable()\n", "\n", " pbar = tqdm(cal_loader, desc=\"MoE Fisher\", total=max_batches)\n", " for batch_idx, batch in enumerate(pbar):\n", " if batch_idx >= max_batches:\n", " break\n", " input_ids = batch['input_ids'].to(DEVICE)\n", " labels = batch['labels'].to(DEVICE)\n", " try:\n", " with torch.amp.autocast(device_type=\"cuda\", dtype=torch.float16):\n", " outputs = self.model(input_ids, labels=labels)\n", " loss = outputs.loss\n", " if loss is not None:\n", " loss.backward()\n", " self.model.zero_grad(set_to_none=True)\n", " except RuntimeError as e:\n", " print(f\" Batch {batch_idx} error: {e}\")\n", " self.model.zero_grad(set_to_none=True)\n", " torch.cuda.empty_cache()\n", " continue\n", " del outputs, loss, input_ids, labels\n", " torch.cuda.empty_cache()\n", " gc.collect()\n", "\n", " self.model.eval()\n", " if hasattr(self.model, 'gradient_checkpointing_disable'):\n", " self.model.gradient_checkpointing_disable()\n", " for h in self.hooks:\n", " h.remove()\n", " result = {}\n", " for name, module in self.model.named_modules():\n", " if hasattr(module, '_fisher_buf'):\n", " result[name] = module._fisher_buf.clone()\n", " return result\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 2.1 Prepare Calibration Data" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "from datasets import load_dataset\n", "\n", "print(\"Loading calibration dataset (C4 subset)...\")\n", "pile = 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 = pile.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\"Loaded {len(cal_loader)} calibration samples (seq_len={MAX_SEQ_LEN})\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 2.2 Compute MoE Fisher" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "print(\"Computing MoE Fisher Information...\")\n", "fisher = MoEFisherAccumulator(model)\n", "fisher_scores = fisher.compute(cal_loader, max_batches=16)\n", "print(f\"Fisher computed for {len(fisher_scores)} layers/modules\")\n", "\n", "sorted_fisher = sorted(fisher_scores.items(), key=lambda x: x[1].max().item(), reverse=True)\n", "print(\"\\nTop 10 most Fisher-sensitive layers:\")\n", "for name, scores in sorted_fisher[:10]:\n", " print(f\" {name}: max={scores.max().item():.6f}, mean={scores.mean().item():.6f}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "def is_gate_layer(name):\n", " return 'gate' in name.lower() or 'router' in name.lower() or 'score' in name.lower()\n", "\n", "def allocate_precision_moe(fisher_dict, int4_fraction=0.05):\n", " allocation = {}\n", " for name, fisher in fisher_dict.items():\n", " if fisher.dim() == 0:\n", " fisher = fisher.unsqueeze(0)\n", " out_ch = fisher.shape[0]\n", " n_int4 = max(1, int(out_ch * int4_fraction))\n", " if out_ch <= 1:\n", " n_int4 = 1\n", " order = torch.argsort(fisher, descending=True)\n", " alloc = {}\n", " for rank, ch in enumerate(order):\n", " alloc[int(ch)] = 'int4' if rank < n_int4 else 'binary'\n", " allocation[name] = alloc\n", " return allocation\n", "\n", "allocation = allocate_precision_moe(fisher_scores, INT4_FRACTION)\n", "\n", "total = sum(len(a) for a in allocation.values())\n", "int4_c = sum(sum(1 for v in a.values() if v == 'int4') for a in allocation.values())\n", "binary_c = total - int4_c\n", "print(f\"Precision allocation: {len(allocation)} layers\")\n", "print(f\" int4 channels: {int4_c:>10,} ({100*int4_c/max(1,total):.1f}%)\")\n", "print(f\" binary channels: {binary_c:>10,} ({100*binary_c/max(1,total):.1f}%)\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "BS_PENALTIES = {64: 1.1, 128: 1.0, 256: 0.9, 512: 0.85}\n", "\n", "def blocksize_recon_error(weights, blocksize, fisher_channels):\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 = float(block.std()) + 1e-8\n", " block_q = np.where(block > 0, 1.0, -1.0) * scale\n", " recon_err = float(((block - block_q) ** 2).mean())\n", " block_fisher = float(fisher_channels[start:end].mean()) if len(fisher_channels) > 0 else 1.0\n", " total_err += block_fisher * recon_err\n", " return total_err * BS_PENALTIES.get(blocksize, 1.0)\n", "\n", "def select_best_blocksize(weights, fisher_channels, candidates=BS_CANDIDATES):\n", " best_b, best_err = candidates[0], float('inf')\n", " for b in candidates:\n", " err = blocksize_recon_error(weights, b, fisher_channels)\n", " if err < best_err:\n", " best_err = err\n", " best_b = b\n", " return best_b\n", "\n", "gate_layer_names = {name for name, _ in model.named_modules() if is_gate_layer(name)}\n", "\n", "blocksize_results = {}\n", "for name, module in tqdm(model.named_modules(), desc=\"Adaptive BS sweep\"):\n", " if name in gate_layer_names:\n", " continue\n", " if not isinstance(module, nn.Linear) or not hasattr(module, 'weight'):\n", " continue\n", " if name not in allocation:\n", " continue\n", " weights = module.weight.data.float().cpu().numpy()\n", " if name in fisher_scores:\n", " fisher = fisher_scores[name].float().cpu().numpy()\n", " else:\n", " fisher = np.ones(module.out_features)\n", " best_b = select_best_blocksize(weights, fisher)\n", " blocksize_results[name] = best_b\n", "\n", "print(f\"Blocksize sweep complete: {len(blocksize_results)} layers\")\n", "bs_counts = pd.Series(list(blocksize_results.values())).value_counts().sort_index()\n", "print(\"Distribution:\")\n", "for bs, cnt in bs_counts.items():\n", " print(f\" bs={bs:3d}: {cnt:3d} layers\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "def build_codebook_moe(model, allocation, blocksize_results, n_clusters=64, max_samples=16384):\n", " model.eval()\n", " all_residuals = []\n", " sample_count = 0\n", " max_bs = max(BS_CANDIDATES)\n", " skip_names = {name for name, _ in model.named_modules() if is_gate_layer(name)}\n", "\n", " for name, module in tqdm(model.named_modules(), desc=\"Building codebook\"):\n", " if name in skip_names:\n", " continue\n", " if not isinstance(module, nn.Linear) or name not in allocation:\n", " continue\n", " if sample_count >= max_samples:\n", " break\n", " weights = module.weight.detach().cpu().numpy()\n", " bs = blocksize_results.get(name, 128)\n", " alloc = allocation[name]\n", " binary_chs = [ch for ch, prec in alloc.items() if prec == 'binary']\n", " if not binary_chs:\n", " continue\n", " step = max(1, len(binary_chs) // 20)\n", " for ch in binary_chs[::step]:\n", " for start in range(0, weights.shape[1], bs):\n", " end = min(start + bs, weights.shape[1])\n", " if end - start < bs:\n", " continue\n", " block = weights[ch, start:end]\n", " scale = float(block.std()) + 1e-8\n", " block_q = np.where(block > 0, 1.0, -1.0) * scale\n", " residual = block - block_q\n", " res_flat = residual.flatten()\n", " padded = np.pad(res_flat, (0, max_bs - len(res_flat)), mode='constant')\n", " all_residuals.append(padded)\n", " sample_count += 1\n", " if sample_count >= max_samples:\n", " break\n", " if sample_count >= max_samples:\n", " break\n", "\n", " if len(all_residuals) == 0:\n", " print(\"WARNING: No residuals collected!\")\n", " return np.zeros((n_clusters, max_bs), dtype=np.float32)\n", "\n", " residuals_array = np.array(all_residuals, dtype=np.float32)\n", " mask = ~np.any(np.isnan(residuals_array) | np.isinf(residuals_array), axis=1)\n", " residuals_array = residuals_array[mask]\n", " print(f\"Collected {len(residuals_array)} residual blocks, shape={residuals_array.shape}\")\n", "\n", " kmeans = MiniBatchKMeans(n_clusters=n_clusters, random_state=42, batch_size=1024, n_init=3)\n", " kmeans.fit(residuals_array)\n", " print(f\"Built codebook: {n_clusters} centroids x {max_bs} dims\")\n", " return kmeans.cluster_centers_\n", "\n", "print(\"Building residual codebook...\")\n", "codebook = build_codebook_moe(model, allocation, blocksize_results, n_clusters=N_CLUSTERS)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "class QuantizedLinear(nn.Module):\n", " 'FABQ-RC quantized linear layer with MoE support.'\n", " def __init__(self, in_features, out_features, bias,\n", " int8_channels, binary_channels,\n", " int8_weights, int8_scales,\n", " binary_reconstructed_weights, blocksize, codebook, codebook_idx):\n", " super().__init__()\n", " self.in_features = in_features\n", " self.out_features = out_features\n", " self.blocksize = blocksize\n", " self.register_buffer('int8_channels', int8_channels.cpu())\n", " self.register_buffer('binary_channels', binary_channels.cpu())\n", " self.register_buffer('int8_weights', int8_weights.cpu())\n", " self.register_buffer('int8_scales', int8_scales.cpu())\n", " self.register_buffer('binary_reconstructed', binary_reconstructed_weights.cpu())\n", " self.register_buffer('codebook_idx', torch.tensor(codebook_idx, dtype=torch.long))\n", " if bias is not None:\n", " self.register_buffer('bias', bias.cpu())\n", " else:\n", " self.bias = None\n", "\n", " def forward(self, x):\n", " w = self.binary_reconstructed.to(x.dtype).to(x.device)\n", " if self.int8_channels.numel() > 0:\n", " int8_w = self.int8_weights.to(x.dtype).to(x.device) * self.int8_scales.to(x.dtype).to(x.device)\n", " w[self.int8_channels] = int8_w\n", " return F.linear(x, w, self.bias)\n", "\n", "\n", "def quantize_layer(name, module, allocation, blocksize_results, codebook, codebook_idx=0):\n", " raw_w = module.weight.data.float()\n", " out_c, in_c = raw_w.shape\n", " alloc = allocation.get(name, {})\n", " int8_chs = sorted([ch for ch, prec in alloc.items() if prec == 'int4'])\n", " binary_chs = sorted([ch for ch, prec in alloc.items() if prec == 'binary'])\n", " bs = blocksize_results.get(name, 128)\n", "\n", " if not int8_chs:\n", " binary_recon = torch.zeros_like(raw_w)\n", " for start in range(0, in_c, bs):\n", " end = min(start + bs, in_c)\n", " block = raw_w[:, start:end]\n", " scale = block.std(dim=1, keepdim=True) + 1e-8\n", " block_q = torch.where(block > 0, 1.0, -1.0) * scale\n", " binary_recon[:, start:end] = block_q\n", " return QuantizedLinear(\n", " in_c, out_c, module.bias,\n", " int8_channels=torch.tensor([], dtype=torch.long),\n", " binary_channels=torch.tensor(list(range(out_c)), dtype=torch.long),\n", " int8_weights=torch.tensor([], dtype=torch.int8),\n", " int8_scales=torch.tensor([], dtype=torch.float16),\n", " binary_reconstructed_weights=binary_recon.half(),\n", " blocksize=bs,\n", " codebook=codebook,\n", " codebook_idx=codebook_idx\n", " )\n", "\n", " int8_w = raw_w[int8_chs]\n", " int8_scale = int8_w.abs().max(dim=1, keepdim=True)[0] / 127.0 + 1e-8\n", " int8_quant = torch.round(int8_w / int8_scale).clamp(-127, 127).to(torch.int8)\n", "\n", " binary_recon = torch.zeros_like(raw_w)\n", " for start in range(0, in_c, bs):\n", " end = min(start + bs, in_c)\n", " block = raw_w[:, start:end]\n", " scale = block.std(dim=1, keepdim=True) + 1e-8\n", " block_q = torch.where(block > 0, 1.0, -1.0) * scale\n", " binary_recon[:, start:end] = block_q\n", "\n", " return QuantizedLinear(\n", " in_c, out_c, module.bias,\n", " int8_channels=torch.tensor(int8_chs, dtype=torch.long),\n", " binary_channels=torch.tensor(binary_chs, dtype=torch.long),\n", " int8_weights=int8_quant,\n", " int8_scales=int8_scale.half().squeeze(-1),\n", " binary_reconstructed_weights=binary_recon.half(),\n", " blocksize=bs,\n", " codebook=codebook,\n", " codebook_idx=codebook_idx\n", " )\n", "\n", "\n", "def get_parent_module(model, name):\n", " parts = name.split('.')\n", " child_name = parts[-1]\n", " parent = model\n", " for p in parts[:-1]:\n", " parent = getattr(parent, p)\n", " return parent, child_name\n", "\n", "\n", "def quantize_fabq_rc_moe(model, allocation, blocksize_results, codebook):\n", " quantized_info = {'dense': 0, 'expert': 0, 'gate_skipped': 0}\n", " skip_names = {name for name, _ in model.named_modules() if is_gate_layer(name)}\n", "\n", " for name, module in tqdm(list(model.named_modules()), desc=\"Quantizing\"):\n", " if not isinstance(module, nn.Linear):\n", " continue\n", " if name in skip_names or name not in allocation:\n", " quantized_info['gate_skipped'] += 1\n", " continue\n", " q_layer = quantize_layer(name, module, allocation, blocksize_results, codebook)\n", " parent, child_name = get_parent_module(model, name)\n", " setattr(parent, child_name, q_layer)\n", " if 'expert' in name.lower():\n", " quantized_info['expert'] += 1\n", " else:\n", " quantized_info['dense'] += 1\n", " del module\n", " torch.cuda.empty_cache()\n", " gc.collect()\n", " return quantized_info\n", "\n", "\n", "print(\"Applying FABQ-RC quantization...\")\n", "t0 = time.time()\n", "stats = quantize_fabq_rc_moe(model, allocation, blocksize_results, codebook)\n", "t1 = time.time()\n", "print(f\"Quantization complete in {t1-t0:.1f}s\")\n", "print(f\" Dense layers quantized: {stats['dense']}\")\n", "print(f\" Expert layers quantized: {stats['expert']}\")\n", "print(f\" Gate/router skipped: {stats['gate_skipped']}\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "def compute_bpw(model, allocation, blocksize_results, codebook):\n", " total_bits = 0\n", " total_weights = 0\n", " skip_names = {name for name, _ in model.named_modules() if is_gate_layer(name)}\n", "\n", " for name, module in model.named_modules():\n", " if isinstance(module, QuantizedLinear):\n", " out_c, in_c = module.out_features, module.in_features\n", " n_int4 = module.int8_channels.numel()\n", " n_binary = module.binary_channels.numel()\n", " bs = module.blocksize\n", " total_bits += n_int4 * in_c * 4\n", " total_bits += n_int4 * 16\n", " total_bits += n_binary * in_c * 1\n", " n_blocks = (in_c + bs - 1) // bs\n", " total_bits += n_binary * n_blocks * 16\n", " total_bits += out_c * n_blocks * 4\n", " total_weights += out_c * in_c\n", " elif isinstance(module, nn.Linear) and hasattr(module, 'weight'):\n", " total_bits += module.weight.numel() * 16\n", " total_weights += module.weight.numel()\n", "\n", " if codebook is not None:\n", " total_bits += codebook.nbytes * 8\n", " bpw = total_bits / max(1, total_weights)\n", " return bpw, total_bits, total_weights\n", "\n", "bpw, total_bits, total_weights = compute_bpw(model, allocation, blocksize_results, codebook)\n", "print(f\"BPW Analysis:\")\n", "print(f\" Total weights: {total_weights/1e9:.2f}B\")\n", "print(f\" Total bits: {total_bits/1e9:.2f}B\")\n", "print(f\" Effective bpw: {bpw:.4f}\")\n", "print(f\" Estimated size: {total_bits / 8 / 1e9:.2f} GB\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "def compute_perplexity(model, tokenizer, max_samples=128, stride=512):\n", " from datasets import load_dataset\n", " print(\"Loading evaluation dataset...\")\n", " try:\n", " wiki = load_dataset(\"wikitext\", \"wikitext-2-raw-v1\", split=\"test\")\n", " except:\n", " print(\"WikiText-2 not available, using C4 validation instead\")\n", " wiki = load_dataset(\n", " \"allenai/c4\",\n", " data_files={\"train\": \"en/c4-train.00000-of-01024.json.gz\"},\n", " split=\"train[:128]\"\n", " )\n", " model.eval()\n", " total_loss = 0.0\n", " total_tokens = 0\n", " n_samples = 0\n", " for i, example in enumerate(tqdm(wiki, desc=\"Evaluating\", total=max_samples)):\n", " if i >= max_samples:\n", " break\n", " text = example.get('text', example.get('page', ''))\n", " if not text or len(text) < 10:\n", " continue\n", " enc = tokenizer(text, return_tensors='pt', truncation=True, max_length=stride)\n", " input_ids = enc['input_ids'].to(DEVICE)\n", " if input_ids.shape[1] < 10:\n", " continue\n", " with torch.no_grad():\n", " try:\n", " outputs = model(input_ids, labels=input_ids)\n", " loss = outputs.loss\n", " if loss is not None:\n", " n_tokens = input_ids.shape[1]\n", " total_loss += loss.item() * n_tokens\n", " total_tokens += n_tokens\n", " n_samples += 1\n", " except RuntimeError as e:\n", " print(f\" Sample {i}: {e}\")\n", " torch.cuda.empty_cache()\n", " continue\n", " if total_tokens == 0:\n", " print(\"WARNING: No valid evaluation samples!\")\n", " return float('inf')\n", " ppl = math.exp(total_loss / total_tokens)\n", " print(f\"\\nPerplexity: {ppl:.2f} (over {n_samples} samples, {total_tokens} tokens)\")\n", " return ppl\n", "\n", "print(\"\\n\" + \"=\"*60)\n", "print(\"EVALUATING QUANTIZED MODEL\")\n", "print(\"=\"*60)\n", "ppl = compute_perplexity(model, tokenizer, max_samples=64)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "print(\"Memory Analysis:\")\n", "print(f\" Estimated quantized model size: {total_bits / 8 / 1e9:.2f} GB\")\n", "print(f\" Original FP16 size: ~{total_weights * 16 / 8 / 1e9:.2f} GB\")\n", "print(f\" Compression ratio: {total_weights * 16 / max(1, total_bits):.2f}x\")\n", "if torch.cuda.is_available():\n", " vram_used = torch.cuda.memory_allocated() / 1e9\n", " vram_total = torch.cuda.get_device_properties(0).total_memory / 1e9\n", " print(f\" VRAM used: {vram_used:.2f} / {vram_total:.2f} GB\")\n", " print(f\" VRAM free: {vram_total - vram_used:.2f} GB\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "import pickle\n", "\n", "SAVE_PATH = \"fabq_rc_v4_flash.pt\"\n", "META_PATH = \"fabq_rc_v4_flash_meta.pkl\"\n", "\n", "print(f\"Saving quantized model to {SAVE_PATH}...\")\n", "state_dict = model.state_dict()\n", "torch.save(state_dict, SAVE_PATH)\n", "\n", "meta = {\n", " 'model_name': MODEL_NAME,\n", " 'architecture': 'DeepSeek-V4-Flash',\n", " 'quantization': 'FABQ-RC',\n", " 'bpw': bpw,\n", " 'total_weights': total_weights,\n", " 'allocation': {k: {str(ck): v for ck, v in av.items()} for k, av in allocation.items()},\n", " 'blocksize_results': blocksize_results,\n", " 'codebook': codebook,\n", " 'perplexity': ppl,\n", "}\n", "with open(META_PATH, 'wb') as f:\n", " pickle.dump(meta, f)\n", "\n", "model_size = os.path.getsize(SAVE_PATH) / 1e9\n", "meta_size = os.path.getsize(META_PATH) / 1e9\n", "print(f\" Model weights: {model_size:.2f} GB\")\n", "print(f\" Metadata: {meta_size:.2f} MB\")\n", "print(f\" Total: {model_size + meta_size:.2f} GB\")\n", "print(\"\\nFABQ-RC quantization of DeepSeek V4-Flash complete!\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "from huggingface_hub import HfApi, create_repo\n", "import os\n", "\n", "REPO_ID = \"toxzak/DeepSeek-V4-Flash-FABQ-RC\"\n", "\n", "if HF_TOKEN:\n", " api = HfApi(token=HF_TOKEN)\n", " create_repo(repo_id=REPO_ID, token=HF_TOKEN, exist_ok=True, repo_type='model')\n", " print(f\"Uploading to https://huggingface.co/{REPO_ID}\")\n", "\n", " files_to_upload = [SAVE_PATH, META_PATH]\n", " for fpath in files_to_upload:\n", " if os.path.exists(fpath):\n", " fname = os.path.basename(fpath)\n", " size = os.path.getsize(fpath) / 1e9\n", " print(f\" Uploading {fname} ({size:.2f} GB)...\")\n", " api.upload_file(path_or_fileobj=fpath, path_in_repo=fname, repo_id=REPO_ID, repo_type='model')\n", " print(f\" Done\")\n", " print(\"\\nAll quantized checkpoint files uploaded to Hugging Face.\")\n", "else:\n", " print(\"HF_TOKEN not set. Set os.environ['HF_TOKEN'] to upload.\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "import struct\n", "import numpy as np\n", "import pickle\n", "\n", "GGUF_PATH = \"fabq_rc_v4_flash.gguf\"\n", "GGUF_MAGIC = 0x46554747\n", "GGUF_VERSION = 3\n", "ALIGN = 32\n", "GGML_TYPE_F16 = 1\n", "GGUF_TYPE_UINT32 = 4\n", "GGUF_TYPE_FLOAT32 = 6\n", "GGUF_TYPE_STRING = 8\n", "\n", "print(\"Reconstructing FP16 weights from quantized state for GGUF export...\")\n", "\n", "state_dict = torch.load(SAVE_PATH, map_location='cpu')\n", "with open(META_PATH, 'rb') as f:\n", " meta = pickle.load(f)\n", "\n", "fp16_tensors = []\n", "for name, param in state_dict.items():\n", " fp16_tensors.append((name, param.to(torch.float16)))\n", "\n", "print(f\"Collected {len(fp16_tensors)} tensors for GGUF\")\n", "\n", "def align_to(x, a):\n", " return ((x + a - 1) // a) * a\n", "\n", "def write_str(f, s):\n", " enc = s.encode('utf-8')\n", " f.write(struct.pack('