modeling_cheddar.py

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 CausalLMOutput

from .configuration_cheddar import CheddarConfig


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

    def forward(self, x):
        return F.rms_norm(x, (x.size(-1),), self.weight, self.eps)


def compute_rope(
    head_dim,
    seq_len,
    base,
    device,
):
    if head_dim % 2:
        raise ValueError(
            "RoPE head dimension must be even"
        )

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

    positions = torch.arange(
        seq_len,
        device=device,
        dtype=torch.float32,
    )

    freqs = torch.outer(
        positions,
        inv_freq,
    )

    return freqs.cos(), freqs.sin()


def apply_rope(x, cos, sin):
    dtype = x.dtype
    x = x.float()

    x1 = x[..., ::2]
    x2 = x[..., 1::2]

    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]

    out = torch.stack(
        (
            x1 * cos - x2 * sin,
            x1 * sin + x2 * cos,
        ),
        dim=-1,
    ).flatten(-2)

    return out.to(dtype)


class Attention(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_head = config.n_head
        self.n_kv_head = config.n_kv_head
        self.head_dim = config.head_dim

        q_dim = self.n_head * self.head_dim
        kv_dim = self.n_kv_head * self.head_dim
        self.q_dim, self.kv_dim = q_dim, kv_dim

        self.qkv_proj = nn.Linear(
            config.n_embd,
            q_dim + 2 * kv_dim,
            bias=False,
        )

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

        self.q_norm = RMSNorm(self.head_dim, config.rms_eps) if config.qk_norm else nn.Identity()
        self.k_norm = RMSNorm(self.head_dim, config.rms_eps) if config.qk_norm else nn.Identity()

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

        q, k, v = self.qkv_proj(x).split(
            (self.q_dim, self.kv_dim, self.kv_dim),
            dim=-1,
        )

        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_kv_head, self.head_dim).transpose(1, 2)

        q = self.q_norm(q)
        k = self.k_norm(k)

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

        has_padding = (
            attention_mask is not None
            and not attention_mask.bool().all().item()
        )

        if has_padding:
            if attention_mask.ndim != 2:
                raise ValueError(
                    "attention_mask must have shape "
                    "(batch, sequence)"
                )

            if attention_mask.shape != (B, T):
                raise ValueError(
                    f"attention_mask has shape "
                    f"{tuple(attention_mask.shape)}, "
                    f"expected {(B, T)}"
                )

            causal_mask = torch.ones(
                T,
                T,
                device=x.device,
                dtype=torch.bool,
            ).tril()

            key_mask = attention_mask[
                :, None, None, :
            ].bool()

            attn_mask = (
                causal_mask[None, None, :, :]
                & key_mask
            )

            y = F.scaled_dot_product_attention(
                q,
                k,
                v,
                attn_mask=attn_mask,
                dropout_p=0.0,
                is_causal=False,
                enable_gqa=(
                    self.n_head
                    != self.n_kv_head
                ),
            )

        else:
            y = F.scaled_dot_product_attention(
                q,
                k,
                v,
                dropout_p=0.0,
                is_causal=True,
                enable_gqa=(
                    self.n_head
                    != self.n_kv_head
                ),
            )

        y = (
            y.transpose(1, 2)
            .contiguous()
            .view(B, T, self.q_dim)
        )

        return self.o_proj(y)


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

        self.gate_up_proj = nn.Linear(
            config.n_embd,
            2 * config.mlp_hidden,
            bias=False,
        )

        self.down_proj = nn.Linear(
            config.mlp_hidden,
            config.n_embd,
            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 Block(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.attn_norm = RMSNorm(
            config.n_embd,
            config.rms_eps,
        )
        self.attn = Attention(config)

        self.mlp_norm = RMSNorm(
            config.n_embd,
            config.rms_eps,
        )
        self.mlp = SwiGLU(config)

    def forward(self, x, cos, sin, attention_mask=None):
        x = x + self.attn(
            self.attn_norm(x),
            cos,
            sin,
            attention_mask=attention_mask,
        )
        x = x + self.mlp(
            self.mlp_norm(x)
        )
        return x


class CheddarPreTrainedModel(
    PreTrainedModel
):
    config_class = CheddarConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False

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

            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=0.02,
            )


class CheddarForCausalLM(
    CheddarPreTrainedModel,
    GenerationMixin,
):
    _tied_weights_keys = {
        "lm_head.weight": "tok_emb.weight"
    }

    def __init__(self, config):
        super().__init__(config)

        self.tok_emb = nn.Embedding(
            config.vocab_size,
            config.n_embd,
        )

        self.blocks = nn.ModuleList(
            [
                Block(config)
                for _ in range(
                    config.n_layer
                )
            ]
        )

        self.norm = RMSNorm(
            config.n_embd,
            config.rms_eps,
        )

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

        self.post_init()

    def get_input_embeddings(self):
        return self.tok_emb

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

    def get_output_embeddings(self):
        return self.lm_head

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

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

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError(
                "input_ids must be provided"
            )

        B, T = input_ids.shape

        if T > self.config.max_seq_len:
            raise ValueError(
                f"Sequence length {T} exceeds "
                f"max_seq_len="
                f"{self.config.max_seq_len}"
            )

        if attention_mask is not None:
            if (
                attention_mask.shape
                != input_ids.shape
            ):
                raise ValueError(
                    "attention_mask must have "
                    "the same shape as input_ids"
                )

            attention_mask = (
                attention_mask.to(
                    device=input_ids.device
                )
            )

        x = self.tok_emb(input_ids)

        # Compute RoPE deterministically on
        # the actual device for every forward.
        #
        # Do not store cos/sin as
        # non-persistent buffers: those buffers
        # are not restored from the checkpoint
        # by from_pretrained().
        cos, sin = compute_rope(
            self.config.head_dim,
            T,
            self.config.rope_base,
            x.device,
        )

        for block in self.blocks:
            x = block(
                x,
                cos,
                sin,
                attention_mask=attention_mask,
            )

        x = self.norm(x)

        logits = self.lm_head(x)

        loss = None

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

            shift_labels = (
                labels[..., 1:]
                .contiguous()
            )

            if attention_mask is not None:
                shift_labels = (
                    shift_labels.masked_fill(
                        ~attention_mask[
                            ..., 1:
                        ].bool(),
                        -100,
                    )
                )

            flat_logits = (
                shift_logits.view(
                    -1,
                    shift_logits.size(-1),
                )
            )

            flat_labels = (
                shift_labels.view(-1)
            )

            loss = F.cross_entropy(
                flat_logits,
                flat_labels,
                ignore_index=-100,
            )

            if self.config.z_loss > 0:
                valid = (
                    flat_labels != -100
                )

                if valid.any():
                    log_z = torch.logsumexp(
                        flat_logits[
                            valid
                        ].float(),
                        dim=-1,
                    )

                    loss = (
                        loss
                        + self.config.z_loss
                        * log_z.square().mean()
                    )

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