modeling_bananamind2micro.py

"""BananaMind 2 Micro causal language model with an XSA refresh gate."""

from __future__ import annotations

import math
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation.utils import GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from .configuration_bananamind2micro import BananaMind2MicroConfig
except ImportError:  # Allows the standalone training script to import mounted code.
    from configuration_bananamind2micro import BananaMind2MicroConfig


class BananaMind2MicroRMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        states = hidden_states.float()
        states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.eps)
        return (states * self.weight.float()).to(hidden_states.dtype)


def _rope_cos_sin(
    head_dim: int,
    positions: torch.Tensor,
    theta: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (
        theta
        ** (
            torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device)
            / head_dim
        )
    )
    frequencies = torch.outer(positions.float(), inv_freq)
    return frequencies.cos(), frequencies.sin()


def _apply_rope(
    query: torch.Tensor,
    key: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    query_dtype = query.dtype
    key_dtype = key.dtype
    cosine = cosine[None, None, :, :]
    sine = sine[None, None, :, :]

    query_pairs = query.float().reshape(*query.shape[:-1], -1, 2)
    key_pairs = key.float().reshape(*key.shape[:-1], -1, 2)
    query_even, query_odd = query_pairs.unbind(-1)
    key_even, key_odd = key_pairs.unbind(-1)
    query = torch.stack(
        (query_even * cosine - query_odd * sine, query_even * sine + query_odd * cosine),
        dim=-1,
    ).flatten(-2)
    key = torch.stack(
        (key_even * cosine - key_odd * sine, key_even * sine + key_odd * cosine),
        dim=-1,
    ).flatten(-2)
    return query.to(query_dtype), key.to(key_dtype)


class BananaMind2MicroCache(DynamicCache):
    """Dynamic K/V cache plus per-layer refresh-convolution history."""

    def __init__(self, config: BananaMind2MicroConfig):
        try:
            super().__init__(config=config)
        except TypeError:  # Transformers releases before config-aware DynamicCache.
            super().__init__()
        self.refresh_states: list[torch.Tensor | None] = [
            None for _ in range(config.num_hidden_layers)
        ]

    def refresh_input(
        self,
        layer_idx: int,
        current: torch.Tensor,
        history_size: int,
    ) -> torch.Tensor:
        history = self.refresh_states[layer_idx]
        if history is None:
            previous = current.new_zeros(current.size(0), current.size(1), history_size)
        else:
            previous = history
            if previous.size(-1) < history_size:
                previous = F.pad(previous, (history_size - previous.size(-1), 0))
        convolution_input = torch.cat((previous, current), dim=-1)
        self.refresh_states[layer_idx] = convolution_input[..., -history_size:]
        return convolution_input

    def reorder_cache(self, beam_idx: torch.LongTensor):
        super().reorder_cache(beam_idx)
        self.refresh_states = [
            None if state is None else state.index_select(0, beam_idx.to(state.device))
            for state in self.refresh_states
        ]

    def batch_repeat_interleave(self, repeats: int):
        super().batch_repeat_interleave(repeats)
        self.refresh_states = [
            None if state is None else state.repeat_interleave(repeats, dim=0)
            for state in self.refresh_states
        ]

    def batch_select_indices(self, indices: torch.Tensor):
        super().batch_select_indices(indices)
        self.refresh_states = [
            None if state is None else state.index_select(0, indices.to(state.device))
            for state in self.refresh_states
        ]

    def crop(self, max_length: int):
        current_length = self.get_seq_length()
        target_length = current_length + max_length if max_length < 0 else max_length
        if target_length < current_length:
            raise NotImplementedError(
                "BananaMind2MicroCache cannot roll back its causal refresh state"
            )
        super().crop(max_length)


class BananaMind2MicroAttention(nn.Module):
    def __init__(self, config: BananaMind2MicroConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        self.o_proj.BANANAMIND_SCALE_INIT = True
        self.q_norm = BananaMind2MicroRMSNorm(self.head_dim, config.rms_norm_eps)
        self.k_norm = BananaMind2MicroRMSNorm(self.head_dim, config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        cosine: torch.Tensor,
        sine: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        past_key_values: Cache | None = None,
    ) -> torch.Tensor:
        batch_size, query_length, _ = hidden_states.shape
        query = self.q_proj(hidden_states).view(
            batch_size, query_length, self.num_heads, self.head_dim
        ).transpose(1, 2)
        key = self.k_proj(hidden_states).view(
            batch_size, query_length, self.num_kv_heads, self.head_dim
        ).transpose(1, 2)
        value = self.v_proj(hidden_states).view(
            batch_size, query_length, self.num_kv_heads, self.head_dim
        ).transpose(1, 2)

        query = self.q_norm(query)
        key = self.k_norm(key)
        query, key = _apply_rope(query, key, cosine, sine)

        past_length = 0
        if past_key_values is not None:
            past_length = past_key_values.get_seq_length(self.layer_idx)
            key, value = past_key_values.update(key, value, self.layer_idx)

        key_length = key.size(-2)
        key = key.repeat_interleave(self.num_kv_groups, dim=1)
        value = value.repeat_interleave(self.num_kv_groups, dim=1)

        is_causal = query_length > 1 and past_length == 0 and attention_mask is None
        sdpa_mask = None
        if not is_causal and query_length > 1:
            query_positions = past_length + torch.arange(query_length, device=query.device)
            key_positions = torch.arange(key_length, device=query.device)
            sdpa_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
        if attention_mask is not None:
            key_padding = attention_mask.to(torch.bool)
            if key_padding.size(-1) < key_length:
                key_padding = F.pad(key_padding, (key_length - key_padding.size(-1), 0), value=True)
            else:
                key_padding = key_padding[:, -key_length:]
            key_padding = key_padding[:, None, None, :]
            sdpa_mask = key_padding if sdpa_mask is None else sdpa_mask & key_padding
            is_causal = False

        output = F.scaled_dot_product_attention(
            query,
            key,
            value,
            attn_mask=sdpa_mask,
            is_causal=is_causal,
        )
        output = output.transpose(1, 2).contiguous().view(
            batch_size, query_length, self.num_heads * self.head_dim
        )
        return self.o_proj(output)


class BananaMind2MicroRefreshGate(nn.Module):
    def __init__(self, config: BananaMind2MicroConfig, layer_idx: int):
        super().__init__()
        hidden = config.hidden_size
        self.layer_idx = layer_idx
        self.kernel_size = config.refresh_kernel_size
        self.attention_norm = BananaMind2MicroRMSNorm(hidden, config.rms_norm_eps)
        self.embedding_norm = BananaMind2MicroRMSNorm(hidden, config.rms_norm_eps)
        self.output_norm = BananaMind2MicroRMSNorm(hidden, config.rms_norm_eps)
        self.gate_proj = nn.Linear(hidden, hidden, bias=False)
        self.value_proj = nn.Linear(hidden, hidden, bias=False)
        self.out_proj = nn.Linear(hidden, hidden, bias=False)
        self.out_proj.BANANAMIND_SCALE_INIT = True
        # Keeping [channels, kernel] as the actual Parameter lets stock Muon
        # optimize it directly; it is viewed as [channels, 1, kernel] for conv1d.
        self.depthwise_kernel = nn.Parameter(torch.empty(hidden, self.kernel_size))
        nn.init.normal_(self.depthwise_kernel, mean=0.0, std=config.initializer_range)
        self.alpha = nn.Parameter(torch.tensor(float(config.refresh_alpha_init)))

    def _causal_depthwise_conv(
        self,
        attention_signal: torch.Tensor,
        past_key_values: BananaMind2MicroCache | None,
    ) -> torch.Tensor:
        signal = attention_signal.transpose(1, 2)
        history_size = self.kernel_size - 1
        if past_key_values is None:
            convolution_input = F.pad(signal, (history_size, 0))
        else:
            convolution_input = past_key_values.refresh_input(
                self.layer_idx, signal, history_size
            )
        convolved = F.conv1d(
            convolution_input,
            self.depthwise_kernel.unsqueeze(1),
            groups=signal.size(1),
        )
        return convolved.transpose(1, 2)

    def forward(
        self,
        attention_output: torch.Tensor,
        original_embedding: torch.Tensor,
        past_key_values: BananaMind2MicroCache | None,
    ) -> torch.Tensor:
        attention_signal = self.attention_norm(attention_output.detach())
        embedding_value = self.embedding_norm(original_embedding)
        gate = self.gate_proj(attention_signal) + self._causal_depthwise_conv(
            attention_signal, past_key_values
        )
        value = self.value_proj(embedding_value)
        refreshed = self.out_proj(F.silu(gate) * value)
        return self.alpha * self.output_norm(refreshed)


class BananaMind2MicroMLP(nn.Module):
    def __init__(self, config: BananaMind2MicroConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.down_proj.BANANAMIND_SCALE_INIT = True

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states))


class BananaMind2MicroBlock(nn.Module):
    def __init__(self, config: BananaMind2MicroConfig, layer_idx: int):
        super().__init__()
        self.input_norm = BananaMind2MicroRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attention = BananaMind2MicroAttention(config, layer_idx)
        self.refresh = BananaMind2MicroRefreshGate(config, layer_idx)
        self.post_attention_norm = BananaMind2MicroRMSNorm(
            config.hidden_size, config.rms_norm_eps
        )
        self.mlp = BananaMind2MicroMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        original_embedding: torch.Tensor,
        cosine: torch.Tensor,
        sine: torch.Tensor,
        attention_mask: torch.Tensor | None,
        past_key_values: BananaMind2MicroCache | None,
    ) -> torch.Tensor:
        attention_output = self.attention(
            self.input_norm(hidden_states),
            cosine,
            sine,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
        )
        hidden_states = hidden_states + attention_output
        hidden_states = hidden_states + self.refresh(
            attention_output, original_embedding, past_key_values
        )
        return hidden_states + self.mlp(self.post_attention_norm(hidden_states))


class BananaMind2MicroPreTrainedModel(PreTrainedModel):
    config_class = BananaMind2MicroConfig
    base_model_prefix = "transformer"
    supports_gradient_checkpointing = False
    _no_split_modules = ["BananaMind2MicroBlock"]
    _supports_sdpa = True
    _supports_cache_class = True

    def _init_weights(self, module: nn.Module):
        std = self.config.initializer_range
        if hasattr(module, "BANANAMIND_SCALE_INIT"):
            std *= (2 * self.config.num_hidden_layers) ** -0.5
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)


class BananaMind2MicroForCausalLM(BananaMind2MicroPreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        # The refresh gate carries convolution history in addition to K/V state,
        # so generation must let forward() create BananaMind2MicroCache.
        return False

    def __init__(self, config: BananaMind2MicroConfig):
        super().__init__(config)
        self.transformer = nn.ModuleDict(
            {
                "wte": nn.Embedding(config.vocab_size, config.hidden_size),
                "h": nn.ModuleList(
                    [BananaMind2MicroBlock(config, index) for index in range(config.num_hidden_layers)]
                ),
                "ln_f": BananaMind2MicroRMSNorm(config.hidden_size, config.rms_norm_eps),
            }
        )
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.embedding_scale = math.sqrt(config.hidden_size)
        self.post_init()
        if config.tie_word_embeddings:
            self.tie_weights()

    def get_input_embeddings(self):
        return self.transformer["wte"]

    def set_input_embeddings(self, value):
        self.transformer["wte"] = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if use_cache is None:
            use_cache = self.config.use_cache and labels is None
        if use_cache and past_key_values is None:
            past_key_values = BananaMind2MicroCache(self.config)
        if use_cache and not isinstance(past_key_values, BananaMind2MicroCache):
            raise TypeError("BananaMind 2 Micro requires BananaMind2MicroCache for refresh state")
        if not use_cache:
            past_key_values = None

        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
        sequence_length = input_ids.size(1)
        total_length = past_length + sequence_length
        if total_length > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {total_length} exceeds {self.config.max_position_embeddings}"
            )

        original_embedding = self.transformer["wte"](input_ids) * self.embedding_scale
        hidden_states = original_embedding
        positions = torch.arange(
            past_length, total_length, dtype=torch.float32, device=input_ids.device
        )
        cosine, sine = _rope_cos_sin(
            self.config.head_dim, positions, self.config.rope_theta
        )
        for block in self.transformer["h"]:
            hidden_states = block(
                hidden_states,
                original_embedding,
                cosine,
                sine,
                attention_mask,
                past_key_values,
            )
        hidden_states = self.transformer["ln_f"](hidden_states)
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[..., :-1, :].float().reshape(-1, logits.size(-1)),
                labels[..., 1:].reshape(-1),
            )
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values,
        )


BananaMind2MicroForCausalLM.register_for_auto_class("AutoModelForCausalLM")

__all__ = [
    "BananaMind2MicroCache",
    "BananaMind2MicroForCausalLM",
    "BananaMind2MicroPreTrainedModel",
]