| """Post-training quantization for KronLLM checkpoints. |
| |
| KronLLM does not contain any ``nn.Linear`` layer: every weight matrix is a |
| Kronecker composition of small factors, so ``torch.ao.quantization`` cannot see |
| it. Instead we quantize the *stored* tensors -- factors, embeddings and the |
| tied head -- with symmetric per-channel scales and dequantize them when the |
| checkpoint is loaded. The forward pass is therefore bit-identical in structure |
| and only sees float weights; the win is a 4x (int8) or 8x (int4) smaller file, |
| which is what makes the model practical on a phone or a microcontroller. |
| |
| Typical use:: |
| |
| from kronllm.quantization import quantize_checkpoint, load_any_checkpoint |
| |
| report = quantize_checkpoint("runs/Q2/best.pt", "runs/Q2/model_int8.pt", bits=8) |
| cfg, state = load_any_checkpoint("runs/Q2/model_int8.pt") |
| """ |
| from __future__ import annotations |
|
|
| import time |
| from pathlib import Path |
| from typing import Any, Dict, Mapping, Tuple |
|
|
| import torch |
|
|
| |
| |
| _MIN_NUMEL = 256 |
|
|
| |
| _SKIP_FRAGMENTS = ("norm", "tau", "gate", "theta", "anchor") |
|
|
|
|
| def _is_quantizable(name: str, tensor: torch.Tensor) -> bool: |
| """Return True if ``tensor`` should be stored in low precision. |
| |
| Args: |
| name: Parameter name inside the state dict. |
| tensor: Parameter tensor. |
| |
| Returns: |
| True for 2-D (or higher) float tensors of at least ``_MIN_NUMEL`` |
| elements whose name does not match a normalisation/scalar role. |
| """ |
| if not tensor.is_floating_point() or tensor.dim() < 2: |
| return False |
| if tensor.numel() < _MIN_NUMEL: |
| return False |
| return not any(fragment in name for fragment in _SKIP_FRAGMENTS) |
|
|
|
|
| def _pack_int4(codes: torch.Tensor) -> torch.Tensor: |
| """Pack signed int4 codes (range ``[-8, 7]``) two per byte. |
| |
| Args: |
| codes: Int8 tensor whose values are already clamped to ``[-8, 7]``. |
| |
| Returns: |
| Flat uint8 tensor with the low nibble holding the even elements. |
| """ |
| flat = codes.reshape(-1).to(torch.int16) |
| if flat.numel() % 2: |
| flat = torch.cat([flat, torch.zeros(1, dtype=flat.dtype)]) |
| shifted = flat + 8 |
| packed = shifted[0::2] | (shifted[1::2] << 4) |
| return packed.to(torch.uint8) |
|
|
|
|
| def _unpack_int4(packed: torch.Tensor, numel: int) -> torch.Tensor: |
| """Inverse of :func:`_pack_int4`. |
| |
| Args: |
| packed: Flat uint8 tensor produced by :func:`_pack_int4`. |
| numel: Number of original elements. |
| |
| Returns: |
| Int8 tensor of length ``numel`` in ``[-8, 7]``. |
| """ |
| low = (packed & 0x0F).to(torch.int16) - 8 |
| high = ((packed >> 4) & 0x0F).to(torch.int16) - 8 |
| out = torch.empty(packed.numel() * 2, dtype=torch.int16) |
| out[0::2] = low |
| out[1::2] = high |
| return out[:numel].to(torch.int8) |
|
|
|
|
| def _quantize_tensor(tensor: torch.Tensor, bits: int |
| ) -> Tuple[torch.Tensor, torch.Tensor]: |
| """Symmetric per-channel quantization of one tensor. |
| |
| Args: |
| tensor: Float tensor to quantize. |
| bits: Target bit width, either 8 or 4. |
| |
| Returns: |
| A pair ``(codes, scale)``. ``codes`` is int8 (values clamped to the |
| signed range) and ``scale`` broadcasts against ``tensor``. |
| """ |
| qmax = float(2 ** (bits - 1) - 1) |
| reduce_dims = tuple(range(1, tensor.dim())) |
| amax = tensor.detach().abs().amax(dim=reduce_dims, keepdim=True) |
| scale = (amax / qmax).clamp_min(1e-8) |
| codes = torch.round(tensor.detach() / scale).clamp(-qmax, qmax).to(torch.int8) |
| return codes, scale |
|
|
|
|
| def quantize_state_dict(state: Mapping[str, torch.Tensor], bits: int = 8 |
| ) -> Dict[str, Any]: |
| """Quantize every eligible tensor of a state dict. |
| |
| Args: |
| state: Float state dict as stored in a training checkpoint. |
| bits: 8 or 4. |
| |
| Returns: |
| A serialisable dict mapping parameter names to either a float tensor |
| (kept exact) or a ``{"q", "scale", "shape"}`` record. |
| """ |
| if bits not in (8, 4): |
| raise ValueError(f"bits must be 8 or 4, got {bits}") |
|
|
| out: Dict[str, Any] = {} |
| for name, tensor in state.items(): |
| if not _is_quantizable(name, tensor): |
| out[name] = tensor.detach().clone() |
| continue |
| codes, scale = _quantize_tensor(tensor, bits) |
| record: Dict[str, Any] = {"shape": tuple(tensor.shape), "scale": scale} |
| if bits == 8: |
| record["q"] = codes.reshape(tensor.shape) |
| else: |
| record["q"] = _pack_int4(codes) |
| record["numel"] = tensor.numel() |
| out[name] = record |
| return out |
|
|
|
|
| def dequantize_state_dict(quantized: Mapping[str, Any], bits: int |
| ) -> Dict[str, torch.Tensor]: |
| """Rebuild a float state dict from :func:`quantize_state_dict` output. |
| |
| Args: |
| quantized: Output of :func:`quantize_state_dict`. |
| bits: Bit width used during quantization. |
| |
| Returns: |
| Float state dict that can be passed to ``model.load_state_dict``. |
| """ |
| state: Dict[str, torch.Tensor] = {} |
| for name, value in quantized.items(): |
| if not isinstance(value, Mapping): |
| state[name] = value.float() if value.is_floating_point() else value |
| continue |
| scale = value["scale"].float() |
| if bits == 8: |
| codes = value["q"].float() |
| else: |
| codes = _unpack_int4(value["q"], int(value["numel"])).float() |
| codes = codes.reshape(value["shape"]) |
| state[name] = codes * scale |
| return state |
|
|
|
|
| def quantize_checkpoint(src: str | Path, dst: str | Path, bits: int = 8 |
| ) -> Dict[str, Any]: |
| """Quantize a training checkpoint and write a compact checkpoint. |
| |
| Args: |
| src: Path to a checkpoint written by ``kronllm.train`` (``best.pt``). |
| dst: Destination path for the quantized checkpoint. |
| bits: 8 or 4. |
| |
| Returns: |
| A report with the original and quantized sizes plus the parameter count. |
| |
| Raises: |
| FileNotFoundError: If ``src`` does not exist. |
| """ |
| src, dst = Path(src), Path(dst) |
| if not src.exists(): |
| raise FileNotFoundError(src) |
|
|
| checkpoint = torch.load(src, map_location="cpu", weights_only=False) |
| state = checkpoint.get("model", checkpoint) |
| quantized = quantize_state_dict(state, bits=bits) |
|
|
| payload = { |
| "quantized": True, |
| "bits": bits, |
| "cfg": checkpoint.get("cfg"), |
| "step": checkpoint.get("step"), |
| "val_ppl": checkpoint.get("val_ppl"), |
| "state": quantized, |
| } |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| torch.save(payload, dst) |
|
|
| return { |
| "src": str(src), |
| "dst": str(dst), |
| "bits": bits, |
| "params": int(sum(v.numel() for v in state.values())), |
| "size_before_mb": round(src.stat().st_size / 1e6, 3), |
| "size_after_mb": round(dst.stat().st_size / 1e6, 3), |
| "compression": round(src.stat().st_size / dst.stat().st_size, 2), |
| "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| } |
|
|
|
|
| def load_any_checkpoint(path: str | Path) -> Tuple[Dict[str, Any], Dict[str, torch.Tensor]]: |
| """Load a float or a quantized checkpoint. |
| |
| Args: |
| path: Checkpoint path. |
| |
| Returns: |
| A pair ``(cfg, state_dict)`` where ``state_dict`` is always float and |
| can be loaded into :class:`kronllm.model.KronTransformer`. |
| |
| Raises: |
| ValueError: If the checkpoint is quantized but carries no config. |
| """ |
| path = Path(path) |
| checkpoint = torch.load(path, map_location="cpu", weights_only=False) |
| if not checkpoint.get("quantized"): |
| cfg = checkpoint.get("cfg") |
| return cfg, checkpoint.get("model", checkpoint) |
|
|
| cfg = checkpoint.get("cfg") |
| if cfg is None: |
| raise ValueError(f"quantized checkpoint {path} has no config") |
| state = dequantize_state_dict(checkpoint["state"], int(checkpoint["bits"])) |
| return cfg, state |
|
|
