src/inference/quantize/quantize.py

"""
Frox AI Morph 1.1 — Quantization

Morph 1.0's `_load_4bit()` / `_load_8bit()` in the inference engine were
stubs: they printed a message and then returned a full-precision model
regardless. This module makes quantized loading actually work.

Two paths:
  - bitsandbytes NF4 (QLoRA-style) — best for training/fine-tuning
  - torchao / weight-only int8 — best for pure inference, no bnb dependency

Both operate on an already-constructed MorphForCausalLM by replacing
nn.Linear layers in-place, since Morph is a custom architecture (not a
HuggingFace AutoModel) and can't go through `from_pretrained(...,
quantization_config=...)` directly.
"""
from __future__ import annotations

from typing import List, Optional

import torch
import torch.nn as nn


# ── bitsandbytes NF4 (4-bit) ──────────────────────────────────────

def quantize_4bit(
    model: nn.Module,
    compute_dtype: torch.dtype = torch.float16,
    skip_modules: Optional[List[str]] = None,
) -> nn.Module:
    """
    Replace nn.Linear layers with bitsandbytes Linear4bit (NF4).
    Embedding and lm_head are skipped by default (quantizing the
    vocab projection tanks quality for a tiny VRAM saving).
    """
    try:
        import bitsandbytes as bnb
    except ImportError:
        raise ImportError(
            "bitsandbytes is required for 4-bit quantization. "
            "Install with: pip install bitsandbytes"
        )

    skip = set(skip_modules or ["lm_head", "embed_tokens"])
    replaced = 0

    def _replace(module: nn.Module, prefix: str = ""):
        nonlocal replaced
        for name, child in module.named_children():
            full_name = f"{prefix}.{name}" if prefix else name
            if any(s in full_name for s in skip):
                continue

            if isinstance(child, nn.Linear):
                new_layer = bnb.nn.Linear4bit(
                    child.in_features,
                    child.out_features,
                    bias=child.bias is not None,
                    compute_dtype=compute_dtype,
                    quant_type="nf4",
                )
                new_layer.weight = bnb.nn.Params4bit(
                    child.weight.data.clone(),
                    requires_grad=False,
                    quant_type="nf4",
                )
                if child.bias is not None:
                    new_layer.bias = nn.Parameter(child.bias.data.clone())
                setattr(module, name, new_layer)
                replaced += 1
            else:
                _replace(child, full_name)

    _replace(model)
    print(f"✓ 4-bit NF4 quantization applied to {replaced} linear layers")
    return model


# ── Weight-only int8 (inference-only, no bnb dependency) ─────────

class Int8Linear(nn.Module):
    """
    Weight-only int8 linear layer. Weights stored as int8 + per-channel
    scale; activations stay in the compute dtype. ~4x smaller weights
    than fp16 with a small quality cost — good for inference on cards
    without bitsandbytes support (e.g. some ARM / edge deployments).
    """

    def __init__(self, in_features: int, out_features: int, bias: bool = False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.register_buffer("weight_int8", torch.zeros(out_features, in_features, dtype=torch.int8))
        self.register_buffer("scale", torch.ones(out_features, dtype=torch.float32))
        self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None

    @classmethod
    def from_linear(cls, linear: nn.Linear) -> "Int8Linear":
        layer = cls(linear.in_features, linear.out_features, bias=linear.bias is not None)
        w = linear.weight.data.float()
        scale = w.abs().max(dim=1).values / 127.0
        scale = scale.clamp(min=1e-8)
        w_int8 = (w / scale.unsqueeze(1)).round().clamp(-127, 127).to(torch.int8)
        layer.weight_int8.copy_(w_int8)
        layer.scale.copy_(scale)
        if linear.bias is not None:
            layer.bias.data.copy_(linear.bias.data)
        return layer

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight_int8.to(x.dtype) * self.scale.unsqueeze(1).to(x.dtype)
        out = torch.nn.functional.linear(x, w, self.bias)
        return out


def quantize_8bit(
    model: nn.Module,
    skip_modules: Optional[List[str]] = None,
) -> nn.Module:
    """Replace nn.Linear layers with weight-only Int8Linear (no bnb needed)."""
    skip = set(skip_modules or ["lm_head", "embed_tokens"])
    replaced = 0

    def _replace(module: nn.Module, prefix: str = ""):
        nonlocal replaced
        for name, child in module.named_children():
            full_name = f"{prefix}.{name}" if prefix else name
            if any(s in full_name for s in skip):
                continue

            if isinstance(child, nn.Linear):
                setattr(module, name, Int8Linear.from_linear(child))
                replaced += 1
            else:
                _replace(child, full_name)

    _replace(model)
    print(f"✓ Weight-only int8 quantization applied to {replaced} linear layers")
    return model


# ── Size estimation ────────────────────────────────────────────────

def estimate_memory_footprint(
    num_params: int,
    dtype: str = "float16",
) -> dict:
    """Estimate model weight memory footprint at various precisions."""
    bytes_per_param = {
        "float32": 4, "float16": 2, "bfloat16": 2,
        "int8": 1, "nf4": 0.5, "int4": 0.5,
    }
    b = bytes_per_param.get(dtype, 2)
    total_bytes = num_params * b
    return {
        "dtype": dtype,
        "params_billions": round(num_params / 1e9, 3),
        "weights_gb": round(total_bytes / (1024 ** 3), 2),
        # Rule of thumb: inference needs weights + ~20% for activations/KV cache
        "estimated_inference_vram_gb": round(total_bytes * 1.2 / (1024 ** 3), 2),
    }


def print_quantization_report(model: nn.Module, dtype_label: str = "float16"):
    total_params = sum(p.numel() for p in model.parameters())
    footprint = estimate_memory_footprint(total_params, dtype_label)
    print("\n📦 Model Memory Footprint")
    print(f"   Parameters:       {footprint['params_billions']}B")
    print(f"   Precision:        {footprint['dtype']}")
    print(f"   Weight size:      {footprint['weights_gb']} GB")
    print(f"   Est. inference:   {footprint['estimated_inference_vram_gb']} GB "
          f"(weights + activations/KV headroom)\n")