modeling_cali.py

from typing import Optional, Tuple, Union

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

from transformers import PreTrainedModel
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.utils import logging

from .configuration_cali import CALIConfig

logger = logging.get_logger(__name__)


class RMSNorm(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, x: torch.Tensor) -> torch.Tensor:
        return x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() * self.weight


def build_rope_cache(seq_len, head_dim, theta=10000.0, device=None):
    half = head_dim // 2
    freqs = 1.0 / (theta ** (torch.arange(0, half, device=device).float() / half))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(t, freqs)
    cos = torch.cat([freqs.cos(), freqs.cos()], dim=-1)[None, None, :, :]
    sin = torch.cat([freqs.sin(), freqs.sin()], dim=-1)[None, None, :, :]
    return cos, sin


def apply_rope(x, cos, sin):
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return x * cos + torch.cat([-x2, x1], dim=-1) * sin


class GroupedQueryAttention(nn.Module):
    def __init__(self, config: CALIConfig):
        super().__init__()
        self.num_heads = config.num_heads
        self.num_kv_heads = config.num_kv_heads
        self.head_dim = config.head_dim
        self.groups = config.num_heads // config.num_kv_heads
        self.scale = config.head_dim ** -0.5

        self.q_proj = nn.Linear(config.hidden_dim, config.num_heads * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_dim, config.num_kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_dim, config.num_kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_heads * config.head_dim, config.hidden_dim, bias=False)

    def forward(self, x, cos, sin, attention_mask=None, past_key_value=None, use_cache=False):
        B, T, _ = x.shape

        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)

        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)

        if past_key_value is not None:
            k = torch.cat([past_key_value[0], k], dim=2)
            v = torch.cat([past_key_value[1], v], dim=2)

        present = (k, v) if use_cache else None

        k = k.repeat_interleave(self.groups, dim=1)
        v = v.repeat_interleave(self.groups, dim=1)

        full_T = k.shape[2]
        att = torch.matmul(q, k.transpose(-2, -1)) * self.scale

        causal = torch.triu(
            torch.ones(T, full_T, device=x.device, dtype=torch.bool),
            diagonal=full_T - T + 1
        )
        att = att.masked_fill(causal[None, None], float("-inf"))

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                padding_mask = attention_mask[:, None, None, :full_T].to(dtype=att.dtype)
                padding_mask = (1.0 - padding_mask) * torch.finfo(att.dtype).min
                att = att + padding_mask
            elif attention_mask.dim() == 4:
                att = att + attention_mask

        att = F.softmax(att.float(), dim=-1).to(q.dtype)
        out = torch.matmul(att, v).transpose(1, 2).contiguous().view(B, T, self.num_heads * self.head_dim)
        return self.o_proj(out), present


class GatedFFN(nn.Module):
    def __init__(self, config: CALIConfig):
        super().__init__()
        ffn_dim = (int(config.hidden_dim * config.ffn_multiplier) + 255) // 256 * 256
        self.gate_proj = nn.Linear(config.hidden_dim, ffn_dim, bias=False)
        self.up_proj   = nn.Linear(config.hidden_dim, ffn_dim, bias=False)
        self.down_proj = nn.Linear(ffn_dim, config.hidden_dim, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class CALIBlock(nn.Module):
    def __init__(self, config: CALIConfig):
        super().__init__()
        self.norm1 = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.attn  = GroupedQueryAttention(config)
        self.norm2 = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.ffn   = GatedFFN(config)

    def forward(self, x, cos, sin, attention_mask=None, past_key_value=None, use_cache=False):
        attn_out, present = self.attn(
            self.norm1(x), cos, sin,
            attention_mask=attention_mask,
            past_key_value=past_key_value,
            use_cache=use_cache,
        )
        x = x + attn_out
        x = x + self.ffn(self.norm2(x))
        return x, present


class CALIPreTrainedModel(PreTrainedModel):
    config_class = CALIConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["CALIBlock"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=std)

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, CALIModel):
            module.gradient_checkpointing = value


class CALIModel(CALIPreTrainedModel):
    def __init__(self, config: CALIConfig):
        super().__init__(config)
        self.gradient_checkpointing = False
        self.embed = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.layers = nn.ModuleList([CALIBlock(config) for _ in range(config.num_layers)])
        self.norm = RMSNorm(config.hidden_dim, eps=config.rms_norm_eps)
        self.post_init()

    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        past_key_values=None,
        inputs_embeds=None,
        use_cache=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is None:
            inputs_embeds = self.embed(input_ids)

        B, T, _ = inputs_embeds.shape
        device = inputs_embeds.device

        past_len = past_key_values[0][0].shape[2] if past_key_values else 0
        cos, sin = build_rope_cache(T + past_len, self.config.head_dim, self.config.rope_theta, device)
        cos = cos[:, :, past_len:past_len + T, :]
        sin = sin[:, :, past_len:past_len + T, :]

        hidden_states = inputs_embeds
        all_hidden_states = () if output_hidden_states else None
        present_key_values = () if use_cache else None

        for i, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_kv = past_key_values[i] if past_key_values else None

            if self.gradient_checkpointing and self.training:
                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs, use_cache=False)
                    return custom_forward
                hidden_states, _ = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(layer),
                    hidden_states, cos, sin, attention_mask, None,
                    use_reentrant=False,
                )
                present = None
            else:
                hidden_states, present = layer(
                    hidden_states, cos, sin,
                    attention_mask=attention_mask,
                    past_key_value=past_kv,
                    use_cache=use_cache,
                )

            if use_cache:
                present_key_values += (present,)

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, present_key_values, all_hidden_states] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=present_key_values,
            hidden_states=all_hidden_states,
        )


class CALIForCausalLM(CALIPreTrainedModel, GenerationMixin):


    def __init__(self, config: CALIConfig):
        super().__init__(config)
        self.model = CALIModel(config)
        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        if config.tie_embeddings:
            self.lm_head.weight = self.model.embed.weight
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed

    def set_input_embeddings(self, value):
        self.model.embed = value

    def get_tied_weights(self):
        return {"lm_head.weight": "model.embed.weight"} if self.config.tie_embeddings else {}

    def get_output_embeddings(self):
        return self.lm_head

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

    def get_decoder(self):
        return self.model

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        past_key_values=None,
        inputs_embeds=None,
        labels=None,
        use_cache=None,
        output_hidden_states=None,
        return_dict=None,
        **kwargs,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        logits = self.lm_head(outputs[0])

        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, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
        )

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

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        return tuple(
            tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past)
            for layer_past in past_key_values
        )