modeling_lean_mixtral.py

"""
LeanMixtral: Mixtral with adaptive KV-cache compression.

Load with:
    AutoModelForCausalLM.from_pretrained("LeanMixtral-8x7B", trust_remote_code=True)
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
from torch import nn
from transformers import MixtralConfig, MixtralForCausalLM
from transformers.cache_utils import Cache
from transformers.modeling_outputs import CausalLMOutputWithPast


# ---------------------------------------------------------------------------
# Quantization-safe linear (not nn.Linear, so bitsandbytes won't quantize)
# ---------------------------------------------------------------------------

class _SafeLinear(nn.Module):
    """Drop-in replacement for nn.Linear that bitsandbytes will not quantize."""

    def __init__(self, in_features: int, out_features: int, bias: bool = True) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.empty(out_features)) if bias else None
        nn.init.kaiming_uniform_(self.weight, a=5**0.5)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return nn.functional.linear(x, self.weight, self.bias)


# ---------------------------------------------------------------------------
# Fixed-rate KV compressor
# ---------------------------------------------------------------------------

class KVCompressorModule(nn.Module):
    """Autoencoder for per-token KV vector compression."""

    def __init__(self, input_dim: int, bottleneck_dim: int,
                 decoder_hidden_dim: int = 128, decoder_depth: int = 2,
                 fixed_alpha: float | None = 1.0) -> None:
        super().__init__()
        self.encoder = _SafeLinear(input_dim, bottleneck_dim)
        if decoder_depth == 2:
            self.decoder = nn.Sequential(
                _SafeLinear(bottleneck_dim, decoder_hidden_dim),
                nn.GELU(),
                _SafeLinear(decoder_hidden_dim, input_dim),
            )
        else:
            self.decoder = _SafeLinear(bottleneck_dim, input_dim)
        self._alpha_const = nn.Parameter(
            torch.tensor(fixed_alpha if fixed_alpha is not None else 1.0),
            requires_grad=False,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.decoder(self.encoder(x))


# ---------------------------------------------------------------------------
# Main model
# ---------------------------------------------------------------------------

class LeanMixtralForCausalLM(MixtralForCausalLM):
    config_class = MixtralConfig
    # _SafeLinear already prevents bitsandbytes from quantizing compressor modules

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        """Load LeanMixtral: resolves base model, loads compressor weights automatically."""
        model_dir = Path(pretrained_model_name_or_path)

        config_path = model_dir / "config.json" if model_dir.is_dir() else None
        base_model = None
        if config_path and config_path.exists():
            import json
            with config_path.open() as f:
                cfg = json.load(f)
            base_model = cfg.get("leanmix_base_model")

        if base_model:
            from transformers import AutoConfig
            config = AutoConfig.from_pretrained(str(model_dir), trust_remote_code=True)
            kwargs["config"] = config
            model = super().from_pretrained(base_model, *args, **kwargs)
        else:
            model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)

        # Auto-load compressor weights if present
        comp_path = model_dir / "leanmix_compressors.pt" if model_dir.is_dir() else None
        if comp_path and comp_path.exists():
            state = torch.load(comp_path, map_location="cpu", weights_only=False)
            missing, unexpected = model.load_state_dict(state, strict=False)
            real_missing = [k for k in missing if "leanmix" in k]
            if real_missing:
                print(f"Warning: missing LeanMixtral keys: {real_missing}")
            print(f"[LeanMixtral] Loaded {len(state)} compressor params for {len(model._kv_layers)} layers")

        return model

    def __init__(self, config: MixtralConfig) -> None:
        super().__init__(config)

        # Read compression config
        self._kv_layers: list[int] = [
            int(x) for x in getattr(config, "leanmix_compressed_kv_layers", [])
        ]

        kv_input_dims: dict[str, int] = dict(getattr(config, "leanmix_kv_input_dims", {}))
        kv_key_dims: dict[str, int] = dict(getattr(config, "leanmix_kv_key_dims", {}))
        kv_value_dims: dict[str, int] = dict(getattr(config, "leanmix_kv_value_dims", {}))
        kv_decoder_depth: dict[str, int] = dict(getattr(config, "leanmix_kv_decoder_depth", {}))
        kv_decoder_hidden: dict[str, int] = dict(getattr(config, "leanmix_kv_decoder_hidden_dim", {}))
        kv_fixed_alpha: dict[str, float | None] = dict(getattr(config, "leanmix_kv_fixed_alpha", {}))

        for li in self._kv_layers:
            ls = str(li)
            in_dim = kv_input_dims.get(ls, 1024)
            k_dim = kv_key_dims.get(ls, 4)
            v_dim = kv_value_dims.get(ls, 4)
            depth = kv_decoder_depth.get(ls, 2)
            hidden = kv_decoder_hidden.get(ls, 128)
            alpha = kv_fixed_alpha.get(ls, 1.0)

            attn = self.model.layers[li].self_attn
            attn.leanmix_k_compressor = KVCompressorModule(
                in_dim, k_dim, hidden, depth, alpha,
            )
            attn.leanmix_v_compressor = KVCompressorModule(
                in_dim, v_dim, hidden, depth, alpha,
            )

        self._compressed_up_to: dict[int, int] = {li: 0 for li in self._kv_layers}

    # ------------------------------------------------------------------
    # Compression helpers
    # ------------------------------------------------------------------

    def _compress_tensor(
        self, t: torch.Tensor, compressor: KVCompressorModule, start: int,
    ) -> torch.Tensor:
        """Compress new tokens in a [B, n_heads, T, head_dim] KV tensor."""
        if start >= t.shape[2]:
            return t
        B, H, T, D = t.shape
        old = t[:, :, :start, :]
        new = t[:, :, start:, :]
        flat = new.permute(0, 2, 1, 3).reshape(-1, H * D)
        compressed = compressor(flat.to(compressor.encoder.weight.dtype))
        compressed = compressed.to(t.dtype)
        result = compressed.view(B, T - start, H, D).permute(0, 2, 1, 3)
        if start > 0:
            return torch.cat([old, result], dim=2)
        return result

    def _compress_past(self, past_key_values: Any) -> Any:
        """Compress KV cache entries for all configured layers."""
        if past_key_values is None or not self._kv_layers:
            return past_key_values

        if isinstance(past_key_values, Cache):
            for layer_idx in self._kv_layers:
                start = self._compressed_up_to[layer_idx]
                attn = self.model.layers[layer_idx].self_attn

                kv = past_key_values[layer_idx]
                k, v = kv[0], kv[1]
                total_tokens = k.shape[2]

                k_new = self._compress_tensor(k, attn.leanmix_k_compressor, start)
                v_new = self._compress_tensor(v, attn.leanmix_v_compressor, start)

                # Write back to cache
                cache_layer = past_key_values.layers[layer_idx]
                cache_layer.keys = k_new
                cache_layer.values = v_new
                self._compressed_up_to[layer_idx] = total_tokens
            return past_key_values

        if isinstance(past_key_values, tuple):
            past_list = list(past_key_values)
            for layer_idx in self._kv_layers:
                k, v = past_list[layer_idx]
                start = self._compressed_up_to[layer_idx]
                total_tokens = k.shape[2]
                attn = self.model.layers[layer_idx].self_attn
                k_new = self._compress_tensor(k, attn.leanmix_k_compressor, start)
                v_new = self._compress_tensor(v, attn.leanmix_v_compressor, start)
                past_list[layer_idx] = (k_new, v_new)
                self._compressed_up_to[layer_idx] = total_tokens
            return tuple(past_list)

        return past_key_values

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(self, *args: Any, **kwargs: Any) -> CausalLMOutputWithPast:
        # Reset compression tracking on new sequence (no past_key_values)
        past = kwargs.get("past_key_values", None)
        if past is None and len(args) < 5:
            for layer_idx in self._kv_layers:
                self._compressed_up_to[layer_idx] = 0

        outputs = super().forward(*args, **kwargs)

        if hasattr(outputs, "past_key_values") and outputs.past_key_values is not None:
            outputs.past_key_values = self._compress_past(outputs.past_key_values)

        return outputs