modeling_koliber.py

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

from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .configuration_koliber import KoliberConfig


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

    def forward(self, x):
        if hasattr(F, "rms_norm"):
            return F.rms_norm(x, (self.dim,), self.weight, self.eps)

        dtype = x.dtype
        y = x.float()
        y = y * torch.rsqrt(
            y.square().mean(dim=-1, keepdim=True) + self.eps
        )
        return y.to(dtype) * self.weight


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_positions, theta):
        super().__init__()

        inv_freq = 1.0 / (
            theta
            ** (
                torch.arange(0, head_dim, 2, dtype=torch.float32)
                / head_dim
            )
        )

        positions = torch.arange(
            max_positions,
            dtype=torch.float32,
        )

        freqs = torch.outer(
            positions,
            inv_freq,
        )

        self.register_buffer(
            "cos_cached",
            freqs.cos(),
            persistent=True,
        )

        self.register_buffer(
            "sin_cached",
            freqs.sin(),
            persistent=True,
        )

    def forward(self, seq_len, dtype):
        return (
            self.cos_cached[:seq_len].to(dtype=dtype),
            self.sin_cached[:seq_len].to(dtype=dtype),
        )


def apply_rope(x, cos, sin):
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]

    even = x[..., 0::2]
    odd = x[..., 1::2]

    return torch.stack(
        (
            even * cos - odd * sin,
            even * sin + odd * cos,
        ),
        dim=-1,
    ).flatten(-2)


class GQAAttention(nn.Module):
    def __init__(self, config, rope):
        super().__init__()

        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim

        self.kv_repeat = self.num_heads // self.num_kv_heads

        object.__setattr__(self, "rope", rope)

        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim

        self.qkv_proj = nn.Linear(
            config.hidden_size,
            self.q_size + 2 * self.kv_size,
            bias=False,
        )

        self.o_proj = nn.Linear(
            config.hidden_size,
            config.hidden_size,
            bias=False,
        )

    def forward(self, x):
        b, s, _ = x.shape

        qkv = self.qkv_proj(x)

        q, k, v = qkv.split(
            (
                self.q_size,
                self.kv_size,
                self.kv_size,
            ),
            dim=-1,
        )

        q = (
            q.view(
                b,
                s,
                self.num_heads,
                self.head_dim,
            )
            .transpose(1, 2)
        )

        k = (
            k.view(
                b,
                s,
                self.num_kv_heads,
                self.head_dim,
            )
            .transpose(1, 2)
        )

        v = (
            v.view(
                b,
                s,
                self.num_kv_heads,
                self.head_dim,
            )
            .transpose(1, 2)
        )

        cos, sin = self.rope(
            s,
            q.dtype,
        )

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

        try:
            y = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=True,
                enable_gqa=True,
            )

        except TypeError:
            k = k.repeat_interleave(
                self.kv_repeat,
                dim=1,
            )

            v = v.repeat_interleave(
                self.kv_repeat,
                dim=1,
            )

            y = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=None,
                dropout_p=0.0,
                is_causal=True,
            )

        y = (
            y.transpose(1, 2)
            .contiguous()
            .view(
                b,
                s,
                self.hidden_size,
            )
        )

        return self.o_proj(y)


class SwiGLU(nn.Module):
    def __init__(self, hidden, intermediate):
        super().__init__()

        self.gate_up_proj = nn.Linear(
            hidden,
            2 * intermediate,
            bias=False,
        )

        self.down_proj = nn.Linear(
            intermediate,
            hidden,
            bias=False,
        )

    def forward(self, x):
        gate, up = self.gate_up_proj(x).chunk(2, dim=-1)
        return self.down_proj(
            F.silu(gate) * up
        )


class KoliberDecoderLayer(nn.Module):
    def __init__(self, config, rope):
        super().__init__()

        self.input_layernorm = RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )

        self.self_attn = GQAAttention(
            config,
            rope,
        )

        self.post_attention_layernorm = RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )

        self.mlp = SwiGLU(
            config.hidden_size,
            config.intermediate_size,
        )

    def forward(self, hidden_states):
        hidden_states = (
            hidden_states
            + self.self_attn(
                self.input_layernorm(
                    hidden_states
                )
            )
        )

        hidden_states = (
            hidden_states
            + self.mlp(
                self.post_attention_layernorm(
                    hidden_states
                )
            )
        )

        return hidden_states


class KoliberPreTrainedModel(PreTrainedModel):
    config_class = KoliberConfig
    base_model_prefix = ""
    main_input_name = "input_ids"
    supports_gradient_checkpointing = False

    _tied_weights_keys = {
        "lm_head.weight": "embed_tokens.weight"
    }

    def _init_weights(self, module):
        if isinstance(module, nn.Embedding):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.02,
            )

        elif isinstance(module, nn.Linear):
            nn.init.normal_(
                module.weight,
                mean=0.0,
                std=0.02,
            )


class KoliberForCausalLM(
    KoliberPreTrainedModel,
    GenerationMixin,
):
    def __init__(self, config):
        super().__init__(config)

        self.embed_tokens = nn.Embedding(
            config.vocab_size,
            config.hidden_size,
        )

        self.rope = RotaryEmbedding(
            config.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
        )

        self.layers = nn.ModuleList(
            [
                KoliberDecoderLayer(
                    config,
                    self.rope,
                )
                for _ in range(
                    config.num_hidden_layers
                )
            ]
        )

        self.norm = RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )

        self.lm_head = nn.Linear(
            config.hidden_size,
            config.vocab_size,
            bias=False,
        )

        self.post_init()
        self.tie_weights()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value
        self.tie_weights()

    def get_output_embeddings(self):
        return self.lm_head

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

    def forward_hidden(self, input_ids):
        x = self.embed_tokens(input_ids)

        for layer in self.layers:
            x = layer(x)

        return self.norm(x)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        use_cache=False,
        return_dict=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids is required")

        hidden = self.forward_hidden(input_ids)

        logits = F.linear(
            hidden,
            self.embed_tokens.weight,
        )

        loss = None

        if labels is not None:
            shift_logits = logits[:, :-1, :].float()
            shift_labels = labels[:, 1:]

            loss = F.cross_entropy(
                shift_logits.reshape(
                    -1,
                    shift_logits.shape[-1],
                ),
                shift_labels.reshape(-1),
                ignore_index=-100,
            )

        if return_dict is False:
            if loss is None:
                return (logits,)
            return (loss, logits)

        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids,
        attention_mask=None,
        **kwargs,
    ):
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }