""" 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