modeling_yatnmn_gpt.py

"""HuggingFace-compatible CausalLM wrapper for YatNMN-Softplus GPT, with KV cache.

Mirrors modeling_gelu_gpt.py one-for-one. The only difference is the inner module
(`Yat_GPT` instead of `GELU_GPT`) — KV-cache, smear handling, generation glue, and
the `(kv_list, last_embed)` past_key_values format are identical.
"""
from __future__ import annotations

import math
from typing import Optional, Tuple, List

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation import GenerationMixin

try:
    from .configuration_yatnmn_gpt import YatGPTHfConfig
    from .yatnmn_gpt import Yat_GPT, YatGPTConfig
    from .torch_gpt import rms_norm, apply_rotary_emb
except ImportError:
    from torch_port.yatnmn.configuration_yatnmn_gpt import YatGPTHfConfig
    from torch_port.yatnmn.yatnmn_gpt import Yat_GPT, YatGPTConfig
    from torch_port.torch_gpt import rms_norm, apply_rotary_emb


def _kvcache_attn(
    attn_module: nn.Module,
    x_norm: torch.Tensor,
    ve: Optional[torch.Tensor],
    cos: torch.Tensor,
    sin: torch.Tensor,
    window_size: Tuple[int, int],
    past_k: Optional[torch.Tensor],
    past_v: Optional[torch.Tensor],
    input_raw_for_ve_gate: torch.Tensor,
):
    cfg = attn_module.config
    B, T_new, _ = x_norm.shape
    n_head, n_kv_head, head_dim = cfg.n_head, cfg.n_kv_head, cfg.head_dim

    q = attn_module.c_q(x_norm).reshape(B, T_new, n_head, head_dim)
    k = attn_module.c_k(x_norm).reshape(B, T_new, n_kv_head, head_dim)
    v = attn_module.c_v(x_norm).reshape(B, T_new, n_kv_head, head_dim)

    if attn_module._has_ve and ve is not None:
        ve_r = ve.reshape(B, T_new, n_kv_head, head_dim)
        gate = 3.0 * torch.sigmoid(attn_module.ve_gate(input_raw_for_ve_gate[..., :12]))
        v = v + gate.unsqueeze(-1) * ve_r

    q = apply_rotary_emb(q, cos, sin)
    k = apply_rotary_emb(k, cos, sin)
    q = rms_norm(q) * 1.2
    k = rms_norm(k) * 1.2

    k_bhtd = k.transpose(1, 2)
    v_bhtd = v.transpose(1, 2)
    q_bhtd = q.transpose(1, 2)

    if past_k is not None:
        k_bhtd = torch.cat([past_k, k_bhtd], dim=2)
        v_bhtd = torch.cat([past_v, v_bhtd], dim=2)
    new_k, new_v = k_bhtd, v_bhtd
    T_total = new_k.shape[2]

    if n_kv_head < n_head:
        repeats = n_head // n_kv_head
        k_bhtd = new_k.repeat_interleave(repeats, dim=1)
        v_bhtd = new_v.repeat_interleave(repeats, dim=1)
    else:
        k_bhtd, v_bhtd = new_k, new_v

    window_left = window_size[0]
    device = x_norm.device
    q_abs = torch.arange(T_total - T_new, T_total, device=device).unsqueeze(1)
    k_abs = torch.arange(T_total, device=device).unsqueeze(0)
    causal = k_abs <= q_abs
    if 0 < window_left < T_total:
        causal = causal & ((q_abs - k_abs) <= window_left)
    bias = torch.where(
        causal,
        torch.zeros((), dtype=x_norm.dtype, device=device),
        torch.full((), -1e9, dtype=x_norm.dtype, device=device),
    ).unsqueeze(0).unsqueeze(0)

    scale = 1.0 / math.sqrt(head_dim)
    att = torch.matmul(q_bhtd, k_bhtd.transpose(-2, -1)) * scale
    att = att + bias
    att = F.softmax(att, dim=-1)
    y = torch.matmul(att, v_bhtd)
    y = y.transpose(1, 2).contiguous().reshape(B, T_new, -1)
    return attn_module.c_proj(y), new_k, new_v


class YatGPTForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = YatGPTHfConfig
    base_model_prefix = "yatnmn_gpt"
    supports_gradient_checkpointing = False
    _no_split_modules = ["Block"]
    _supports_cache_class = False
    _supports_static_cache = False

    def _supports_default_dynamic_cache(self):
        return False

    def __init__(self, config: YatGPTHfConfig):
        super().__init__(config)
        inner = YatGPTConfig(
            sequence_len=config.sequence_len, vocab_size=config.vocab_size,
            n_layer=config.n_layer, n_head=config.n_head, n_kv_head=config.n_kv_head,
            n_embd=config.n_embd, window_pattern=config.window_pattern,
            tie_embeddings=config.tie_embeddings, rope_base=config.rope_base,
            pad_vocab_size_to=config.pad_vocab_size_to, mlp_type=config.mlp_type,
            scalar_bias=config.scalar_bias, softplus_bias=config.softplus_bias,
            learnable_epsilon=config.learnable_epsilon,
            epsilon_init=config.epsilon_init, constant_alpha=config.constant_alpha,
        )
        self.inner_config = inner
        self.model = Yat_GPT(inner)
        self.post_init()

    def get_input_embeddings(self): return self.model.wte
    def set_input_embeddings(self, v): self.model.wte = v
    def can_generate(self): return True

    def _forward_full(self, input_ids):
        return self.model(input_ids)

    def _forward_with_cache(self, input_ids_new, past_key_values, prev_token_embed=None):
        m = self.model
        cfg = m.config
        B, T_new = input_ids_new.shape
        past_len = 0 if past_key_values is None else past_key_values[0][0].shape[2]
        T_total = past_len + T_new

        cos_full, sin_full = m._get_rope(T_total, m.wte.weight.dtype, m.wte.weight.device)
        cos = cos_full[:, past_len:T_total]
        sin = sin_full[:, past_len:T_total]

        x_new = rms_norm(m.wte(input_ids_new))

        if past_len == 0:
            if T_new >= 2:
                gate = m.smear_lambda * torch.sigmoid(m.smear_gate(x_new[:, 1:, :24]))
                x_smeared = x_new[:, 1:] + gate * x_new[:, :-1]
                x = torch.cat([x_new[:, :1], x_smeared], dim=1)
            else:
                x = x_new
        else:
            assert prev_token_embed is not None, "prev_token_embed required for smear with past"
            x_cat = torch.cat([prev_token_embed, x_new], dim=1)
            gate = m.smear_lambda * torch.sigmoid(m.smear_gate(x_cat[:, 1:, :24]))
            x = x_cat[:, 1:] + gate * x_cat[:, :-1]

        x0 = x
        backout_layer = cfg.n_layer // 2
        x_backout = None

        new_past = []
        for i, block in enumerate(m.blocks):
            x = m.resid_lambdas[i] * x + m.x0_lambdas[i] * x0
            ve_key = str(i)
            ve = m.value_embeds[ve_key](input_ids_new).to(dtype=x.dtype) if ve_key in m.value_embeds else None
            past_k = past_key_values[i][0] if past_key_values is not None else None
            past_v = past_key_values[i][1] if past_key_values is not None else None
            x_norm = rms_norm(x)
            attn_out, new_k, new_v = _kvcache_attn(
                block.attn, x_norm, ve, cos, sin,
                m.window_sizes[i], past_k, past_v,
                input_raw_for_ve_gate=x_norm,
            )
            new_past.append((new_k, new_v))
            x = x + attn_out
            x = x + block.mlp(rms_norm(x))
            if i == backout_layer:
                x_backout = x

        if x_backout is not None:
            x = x - m.backout_lambda * x_backout
        x = rms_norm(x)

        softcap = 15.0
        logits = x @ m.wte.weight.t() if m.tie_embeddings else m.lm_head(x)
        logits = logits[..., : cfg.vocab_size].to(torch.float32)
        logits = softcap * torch.tanh(logits / softcap)
        last_embed = x_new[:, -1:, :]
        return logits, new_past, last_embed

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple] = None,
        labels: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.Tensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if input_ids is None:
            raise ValueError("input_ids required")
        use_cache = bool(use_cache) if use_cache is not None else (past_key_values is not None)

        kv_list = None
        prev_embed = None
        if past_key_values is not None:
            kv_list, prev_embed = past_key_values

        if use_cache:
            logits, new_past, new_last_embed = self._forward_with_cache(
                input_ids, kv_list, prev_token_embed=prev_embed,
            )
            pkv = (tuple(new_past), new_last_embed)
        else:
            logits = self._forward_full(input_ids)
            pkv = None

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return CausalLMOutputWithPast(
            loss=loss, logits=logits, past_key_values=pkv,
            hidden_states=None, attentions=None,
        )

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        return {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}

    def _get_cache_length(self, past_key_values):
        if past_key_values is None: return 0
        kv, _ = past_key_values
        if not kv: return 0
        return kv[0][0].shape[-2]


__all__ = ["YatGPTHfConfig", "YatGPTForCausalLM"]