modeling_pmre_olmo.py

"""HuggingFace-compatible model class for PMRE factorized OLMo.

This module defines a custom model class that integrates with HuggingFace's
AutoModelForCausalLM via auto_map in config.json. Users load with:

    model = AutoModelForCausalLM.from_pretrained(
        "asystemoffields/OLMo-3-3.5B-Think-PMRE",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )

The factorized weights are stored in safetensors format. Each compressed Linear
becomes a FactorizedLinear with down/up factor matrices, preserving the 2x
storage reduction and ~2x inference speedup.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Optional, Tuple, Union

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast


class FactorizedLinear(nn.Module):
    """Low-rank replacement for nn.Linear: x -> up(down(x))."""

    def __init__(self, in_features: int, out_features: int, rank: int,
                 bias: bool = False, device=None, dtype=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
        self.up = nn.Linear(rank, out_features, bias=bias, device=device, dtype=dtype)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.up(self.down(x))

    def extra_repr(self) -> str:
        return (f"in_features={self.in_features}, out_features={self.out_features}, "
                f"rank={self.rank}, compression={self.in_features * self.out_features / (self.rank * (self.in_features + self.out_features)):.1f}x")


class FactorizedSparseLinear(nn.Module):
    """Low-rank linear plus a fixed sparse residual correction."""

    def __init__(self, in_features: int, out_features: int, rank: int,
                 sparse_nnz: int = 0, bias: bool = False, device=None, dtype=None):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank = rank
        self.down = nn.Linear(in_features, rank, bias=False, device=device, dtype=dtype)
        self.up = nn.Linear(rank, out_features, bias=bias, device=device, dtype=dtype)
        self.register_buffer("sparse_row_idx", torch.zeros(sparse_nnz, dtype=torch.long, device=device))
        self.register_buffer("sparse_col_idx", torch.zeros(sparse_nnz, dtype=torch.long, device=device))
        self.register_buffer("sparse_values", torch.zeros(sparse_nnz, dtype=dtype or torch.bfloat16, device=device))
        self._sparse_cache: Optional[torch.Tensor] = None

    @property
    def sparse_nnz(self) -> int:
        return int(self.sparse_values.numel())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        low_rank = self.up(self.down(x))
        if self.sparse_nnz == 0:
            return low_rank

        cached = self._sparse_cache
        if cached is None or cached.device != x.device:
            idx = torch.stack([self.sparse_row_idx, self.sparse_col_idx], dim=0)
            cached = torch.sparse_coo_tensor(
                idx, self.sparse_values.float(),
                size=(self.out_features, self.in_features), device=x.device,
            ).coalesce()
            self._sparse_cache = cached

        flat = x.reshape(-1, self.in_features)
        sparse_out = torch.sparse.mm(cached, flat.float().t()).t()
        sparse_out = sparse_out.reshape(*x.shape[:-1], self.out_features).to(dtype=x.dtype)
        return low_rank + sparse_out

    def extra_repr(self) -> str:
        return (f"in_features={self.in_features}, out_features={self.out_features}, "
                f"rank={self.rank}, sparse_nnz={self.sparse_nnz}")


def _parent_and_attr(model: nn.Module, path: str) -> Tuple[nn.Module, str]:
    parts = path.split(".")
    parent = model.get_submodule(".".join(parts[:-1]))
    return parent, parts[-1]


def apply_factorized_modules(model: nn.Module, metadata: list[dict[str, Any]],
                             device=None, dtype=None) -> nn.Module:
    """Replace target Linear modules with factorized versions based on metadata."""
    for entry in metadata:
        module_path = entry["module_path"]
        rank = int(entry["rank"])
        in_features = int(entry["in_features"])
        out_features = int(entry["out_features"])
        sparse_nnz = int(entry.get("sparse_nnz", 0) or 0)

        old = model.get_submodule(module_path)
        has_bias = isinstance(old, nn.Linear) and old.bias is not None

        if sparse_nnz > 0:
            replacement = FactorizedSparseLinear(
                in_features=in_features, out_features=out_features, rank=rank,
                sparse_nnz=sparse_nnz, bias=has_bias, device=device, dtype=dtype,
            )
        else:
            replacement = FactorizedLinear(
                in_features=in_features, out_features=out_features, rank=rank,
                bias=has_bias, device=device, dtype=dtype,
            )

        parent, attr = _parent_and_attr(model, module_path)
        setattr(parent, attr, replacement)
    return model


class PmreOlmoForCausalLM(PreTrainedModel):
    """HuggingFace-compatible wrapper for PMRE factorized OLMo models.

    This class handles:
    1. Building the base OLMo architecture
    2. Replacing compressed modules with FactorizedLinear
    3. Loading the factorized state dict from safetensors/bin files
    4. Standard generate() / forward() interface
    """

    config_class = None  # Will use the base model's config
    supports_gradient_checkpointing = True
    _no_split_modules = ["OLMo2DecoderLayer"]

    def __init__(self, config, **kwargs):
        super().__init__(config, **kwargs)

        # Build the base OLMo model
        base_config = AutoConfig.from_dict(config.to_dict())
        # Remove our custom fields before building base model
        for field in ("pmre_metadata", "auto_map", "model_type_override"):
            if hasattr(base_config, field):
                delattr(base_config, field)

        self.model_impl = AutoModelForCausalLM.from_config(
            base_config, trust_remote_code=True,
        )

        # Apply factorized module replacements
        metadata = getattr(config, "pmre_metadata", None)
        if metadata:
            apply_factorized_modules(self.model_impl, metadata, dtype=kwargs.get("torch_dtype"))

    def forward(self, **kwargs) -> CausalLMOutputWithPast:
        return self.model_impl(**kwargs)

    def generate(self, *args, **kwargs):
        return self.model_impl.generate(*args, **kwargs)

    def get_input_embeddings(self):
        return self.model_impl.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model_impl.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.model_impl.get_output_embeddings()

    def set_output_embeddings(self, value):
        self.model_impl.set_output_embeddings(value)

    def prepare_inputs_for_generation(self, *args, **kwargs):
        return self.model_impl.prepare_inputs_for_generation(*args, **kwargs)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        """Load a PMRE model from a HuggingFace repo or local directory.

        This overrides from_pretrained to:
        1. Load the config (which contains pmre_metadata)
        2. Instantiate with factorized modules
        3. Load the factorized state dict
        """
        model_dir = Path(pretrained_model_name_or_path)

        # Check for local factorized_modules.json
        metadata_path = model_dir / "factorized_modules.json" if model_dir.is_dir() else None

        config = kwargs.pop("config", None)
        if config is None:
            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, trust_remote_code=True)

        # Inject metadata into config if available from file
        if metadata_path and metadata_path.exists() and not getattr(config, "pmre_metadata", None):
            with metadata_path.open("r", encoding="utf-8") as f:
                config.pmre_metadata = json.load(f)

        # Build the model with factorized modules
        torch_dtype = kwargs.pop("torch_dtype", None)
        device_map = kwargs.pop("device_map", None)

        model = cls(config, torch_dtype=torch_dtype)

        # Load state dict
        from safetensors import safe_open
        from safetensors.torch import load_file

        safetensors_files = list(model_dir.glob("*.safetensors")) if model_dir.is_dir() else []

        if safetensors_files:
            state_dict = {}
            for sf_path in safetensors_files:
                state_dict.update(load_file(str(sf_path)))
            # Keys in safetensors are for model_impl
            missing, unexpected = model.model_impl.load_state_dict(state_dict, strict=False)
            if missing:
                print(f"  PMRE load: {len(missing)} missing keys (expected for non-factorized buffers)")
        else:
            # Fallback: try .pt/.bin files
            pt_files = list(model_dir.glob("pmre_model.pt")) + list(model_dir.glob("*.bin"))
            if pt_files:
                checkpoint = torch.load(pt_files[0], map_location="cpu", weights_only=False)
                sd = checkpoint.get("model_state_dict", checkpoint)
                model.model_impl.load_state_dict(sd, strict=False)

        if torch_dtype is not None:
            model = model.to(dtype=torch_dtype)
        if device_map:
            from accelerate import dispatch_model, infer_auto_device_map
            device_map_computed = infer_auto_device_map(model, max_memory=None)
            model = dispatch_model(model, device_map=device_map_computed)

        model.eval()
        return model