models/alignx.py

"""AlignX Model — full two-stage alignment framework.

Wraps a HuggingFace causal LM and surgically replaces the last transformer
FFN layer with the MoCaE AlignX layer (Stage 2 injection).
"""

import os
import torch
import torch.nn as nn
from typing import Optional, Dict, List, Union
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

from .mocae import MoCaE, AlignXLayer, ExpertFFN
from .task_feature_matrix import TaskFeatureMatrix


# ---------------------------------------------------------------------------
# Model-family helpers
# ---------------------------------------------------------------------------

def _get_layers(model) -> nn.ModuleList:
    """Return the transformer layer list for supported model families."""
    if hasattr(model, "model") and hasattr(model.model, "layers"):
        return model.model.layers
    if hasattr(model, "transformer") and hasattr(model.transformer, "h"):
        return model.transformer.h
    raise ValueError(f"Unsupported model architecture: {type(model)}")


def _get_mlp(layer) -> nn.Module:
    """Return the FFN/MLP sub-module from a transformer layer."""
    for attr in ("mlp", "feed_forward", "ffn", "ff"):
        if hasattr(layer, attr):
            return getattr(layer, attr)
    raise ValueError(f"Cannot locate MLP in layer {type(layer)}")


def _set_mlp(layer, new_module: nn.Module):
    """Replace the FFN/MLP sub-module in a transformer layer."""
    for attr in ("mlp", "feed_forward", "ffn", "ff"):
        if hasattr(layer, attr):
            setattr(layer, attr, new_module)
            return
    raise ValueError(f"Cannot locate MLP in layer {type(layer)}")


def _get_intermediate_dim(mlp_module) -> int:
    """Infer the FFN intermediate dimension from the gate/up projection."""
    for attr in ("gate_proj", "w1", "fc1", "dense_h_to_4h"):
        if hasattr(mlp_module, attr):
            proj = getattr(mlp_module, attr)
            return proj.out_features
    for attr in ("up_proj", "w2", "fc2"):
        if hasattr(mlp_module, attr):
            return getattr(mlp_module, attr).out_features
    raise ValueError(f"Cannot infer intermediate_dim from {type(mlp_module)}")


# ---------------------------------------------------------------------------
# AlignX model class
# ---------------------------------------------------------------------------

class AlignXModel(nn.Module):
    """AlignX-wrapped causal LM with MoCaE layer injected at the last FFN.

    Usage:
        model = AlignXModel(base_lm, hidden_dim=4096, intermediate_dim=11008)
        model.register_task_matrices(T_h, T_ha, T_ho)
        # Train only MoCaE params (base_lm is frozen)
        logits = model(input_ids, attention_mask)
    """

    def __init__(
        self,
        base_lm: nn.Module,
        hidden_dim: int,
        intermediate_dim: int,
        k: int = 256,
        lambda1: float = 0.6,
        lambda2: float = 0.4,
        epsilon: float = 0.05,
        n_clusters: int = 8,
        layer_idx: int = -1,
        freeze_base: bool = True,
    ):
        super().__init__()
        self.base_lm = base_lm
        self.hidden_dim = hidden_dim
        self.layer_idx = layer_idx
        self.config = base_lm.config

        # Build MoCaE
        self.mocae = MoCaE(
            hidden_dim=hidden_dim,
            intermediate_dim=intermediate_dim,
            k=k,
            lambda1=lambda1,
            lambda2=lambda2,
            epsilon=epsilon,
            n_clusters=n_clusters,
        )
        self.alignx_layer = AlignXLayer(self.mocae, hidden_dim)

        # Inject AlignX layer into the last transformer FFN
        layers = _get_layers(self.base_lm)
        target_layer = layers[layer_idx]
        original_mlp = _get_mlp(target_layer)

        # Initialise all experts from the original FFN weights (helpful default)
        self.mocae.init_experts_from_ffn(original_mlp)

        # Replace the FFN
        _set_mlp(target_layer, self.alignx_layer)

        # Move AlignXLayer (MoCaE) to the same device as the rest of the target layer.
        # input_layernorm is a plain fp32/bf16 LayerNorm — its device is reliable even
        # when the base model is 4-bit quantized.
        if hasattr(target_layer, "input_layernorm"):
            _target_device = target_layer.input_layernorm.weight.device
        elif hasattr(target_layer, "ln_1"):
            _target_device = target_layer.ln_1.weight.device
        else:
            _target_device = torch.device(
                f"cuda:{torch.cuda.device_count() - 1}" if torch.cuda.is_available() else "cpu"
            )
        _target_dtype = target_layer.input_layernorm.weight.dtype if hasattr(target_layer, "input_layernorm") else torch.float16
        self.alignx_layer.to(_target_device).to(_target_dtype)
        print(f"[AlignX] AlignXLayer (MoCaE) placed on {_target_device} dtype={_target_dtype}")

        # Freeze base model parameters (only MoCaE is trainable)
        if freeze_base:
            for name, param in self.base_lm.named_parameters():
                param.requires_grad_(False)
        for param in self.mocae.parameters():
            param.requires_grad_(True)

        total_base = sum(p.numel() for p in self.base_lm.parameters())
        total_mocae = sum(p.numel() for p in self.mocae.parameters())
        print(f"[AlignX] Base LM params: {total_base:,}  |  MoCaE params: {total_mocae:,}")

    # ------------------------------------------------------------------

    def register_task_matrices(self, T_helpful, T_harmless, T_honest):
        self.mocae.register_task_matrices(T_helpful, T_harmless, T_honest)

    def init_expert_from_finetuned(self, expert_idx: int, finetuned_lm: nn.Module):
        """Initialise expert `expert_idx` from the last FFN of a fine-tuned model."""
        layers = _get_layers(finetuned_lm)
        ffn = _get_mlp(layers[self.layer_idx])
        self.mocae.init_expert_from_ffn(expert_idx, ffn)

    # ------------------------------------------------------------------

    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
        return self.base_lm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
            **kwargs,
        )

    @torch.no_grad()
    def generate(self, input_ids, attention_mask=None, **kwargs):
        return self.base_lm.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            **kwargs,
        )

    def save_mocae(self, path: str):
        os.makedirs(path, exist_ok=True)
        torch.save(self.mocae.state_dict(), os.path.join(path, "mocae.pt"))
        print(f"[AlignX] Saved MoCaE weights to {path}/mocae.pt")

    def load_mocae(self, path: str):
        state = torch.load(os.path.join(path, "mocae.pt"), map_location="cpu")
        self.mocae.load_state_dict(state, strict=False)
        print(f"[AlignX] Loaded MoCaE weights from {path}/mocae.pt")


# ---------------------------------------------------------------------------
# Builder helper
# ---------------------------------------------------------------------------

def build_alignx_model(
    base_model_name_or_path: str,
    finetuned_paths: Optional[Dict[str, str]] = None,
    task_matrix_paths: Optional[Dict[str, str]] = None,
    load_in_4bit: bool = True,
    device_map: str = "auto",
    k: int = 256,
    lambda1: float = 0.6,
    lambda2: float = 0.4,
    freeze_base: bool = True,
    layer_idx: int = -1,
) -> AlignXModel:
    """Build a full AlignX model from a base checkpoint.

    Optionally initialises each expert from a per-axis fine-tuned checkpoint
    and loads precomputed task-feature matrices.
    """
    bnb_config = None
    if load_in_4bit:
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
        )

    base_lm = AutoModelForCausalLM.from_pretrained(
        base_model_name_or_path,
        quantization_config=bnb_config,
        device_map=device_map,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
    )

    hidden_dim = base_lm.config.hidden_size
    layers = _get_layers(base_lm)
    mlp = _get_mlp(layers[layer_idx])
    intermediate_dim = _get_intermediate_dim(mlp)

    alignx = AlignXModel(
        base_lm=base_lm,
        hidden_dim=hidden_dim,
        intermediate_dim=intermediate_dim,
        k=k,
        lambda1=lambda1,
        lambda2=lambda2,
        layer_idx=layer_idx,
        freeze_base=freeze_base,
    )

    # Initialise each expert from axis-specific fine-tuned model
    if finetuned_paths:
        axis_to_idx = {"helpful": 0, "harmless": 1, "honest": 2}
        for axis, ckpt_path in finetuned_paths.items():
            if axis in axis_to_idx and os.path.exists(ckpt_path):
                print(f"[AlignX] Loading expert {axis} from {ckpt_path}")
                ft_lm = AutoModelForCausalLM.from_pretrained(
                    ckpt_path,
                    torch_dtype=torch.float16,
                    device_map="cpu",
                    trust_remote_code=True,
                )
                alignx.init_expert_from_finetuned(axis_to_idx[axis], ft_lm)
                del ft_lm

    # Load task-feature matrices
    if task_matrix_paths:
        T = {}
        for axis in ("helpful", "harmless", "honest"):
            path = task_matrix_paths.get(axis, "")
            if path and os.path.exists(path):
                T[axis] = torch.load(path, map_location="cpu")
        if len(T) == 3:
            alignx.register_task_matrices(T["helpful"], T["harmless"], T["honest"])

    return alignx