modules.py

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

from transformers import PretrainedConfig, PreTrainedModel, AutoConfig, AutoModelForCausalLM
from transformers.modeling_outputs import CausalLMOutputWithPast



def build_rope_cache(head_dim: int, context_length: int, base: float = 10000.0, device=None):
    assert head_dim % 2 == 0, "RoPE rotates 2D planes, so head_dim must be even"

    plane_indices = torch.arange(head_dim // 2, dtype=torch.float32, device=device)
    inverse_frequencies = base ** (-2.0 * plane_indices / head_dim)

    positions = torch.arange(context_length, dtype=torch.float32, device=device)
    angles = positions[:, None] * inverse_frequencies[None, :]

    angles = torch.cat([angles, angles], dim=-1)

    return angles.cos(), angles.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    first_half, second_half = x.chunk(2, dim=-1)
    return torch.cat([-second_half, first_half], dim=-1)


def apply_rope(x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor) -> torch.Tensor:
    seq_len = x.shape[-2]
    cos = rope_cos[:seq_len].to(x.dtype)
    sin = rope_sin[:seq_len].to(x.dtype)
    return x * cos + rotate_half(x) * sin


class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.gate = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.up = nn.Linear(cfg["emb_dim"], cfg["hidden_dim"], bias=False)
        self.down = nn.Linear(cfg["hidden_dim"], cfg["emb_dim"], bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class MoE(nn.Module):

    def __init__(self, cfg: dict[str, int | bool]):
        super().__init__()
        self.n_experts = cfg["n_experts"]
        self.top_k = cfg["top_k"]
        self.experts = nn.ModuleList(
            [FeedForward(cfg) for _ in range(self.n_experts)]
        )
        self.router = nn.Linear(cfg["emb_dim"], self.n_experts, bias=False)

    def forward(self, x: torch.Tensor):
        batch_size, seq_len, emb_dim = x.shape
        tokens = x.reshape(batch_size * seq_len, emb_dim)

        router_logits = self.router(tokens)
        router_probs = torch.softmax(router_logits, dim=-1)

        top_weights, top_experts = torch.topk(router_probs, self.top_k, dim=-1)
        top_weights = top_weights / top_weights.sum(dim=-1, keepdim=True)

        expert_mask = F.one_hot(top_experts, self.n_experts)

        tokens_per_expert = torch.sum(expert_mask, dim=1).float().mean(dim=0) / self.top_k
        prob_per_expert = router_probs.mean(dim=0)

        aux_loss = self.n_experts * torch.sum(tokens_per_expert * prob_per_expert, dim=0)

        self.aux_loss = aux_loss

        output = torch.zeros_like(tokens)

        for expert_idx in range(self.n_experts):
            token_pos, slot_pos = torch.where(top_experts == expert_idx)

            selected_tokens = tokens[token_pos]
            expert_output = self.experts[expert_idx](selected_tokens)

            token_weights = top_weights[token_pos, slot_pos].unsqueeze(1)

            output.index_add_(dim=0, index=token_pos, source=expert_output * token_weights)

        return output.reshape(batch_size, seq_len, emb_dim)


class MultiQueryAttention(nn.Module):
    def __init__(self, cfg: dict[str, int | bool]):
        super().__init__()

        assert cfg["emb_dim"] % cfg["n_heads"] == 0

        self.emb_dim = cfg["emb_dim"]
        self.num_heads = cfg["n_heads"]
        self.head_dim = self.emb_dim // self.num_heads
        self.qkv_bias = cfg["qkv_bias"]
        self.drop_rate = cfg["drop_rate"]

        assert self.head_dim % 2 == 0, "RoPE requires an even head_dim"

        self.query_proj = nn.Linear(self.emb_dim, self.emb_dim, self.qkv_bias)
        self.key_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias)
        self.value_proj = nn.Linear(self.emb_dim, self.head_dim, self.qkv_bias)
        self.out_proj = nn.Linear(self.emb_dim, self.emb_dim, bias=False)

    def forward(self, x: torch.Tensor, rope_cos: torch.Tensor, rope_sin: torch.Tensor):
        batch_size, seq_len, _ = x.shape

        queries = self.query_proj(x).reshape(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = self.key_proj(x).unsqueeze(1)
        values = self.value_proj(x).unsqueeze(1)

        queries = apply_rope(queries, rope_cos, rope_sin)
        keys = apply_rope(keys, rope_cos, rope_sin)

        keys = keys.expand(batch_size, self.num_heads, seq_len, self.head_dim)
        values = values.expand(batch_size, self.num_heads, seq_len, self.head_dim)

        context_vecs = F.scaled_dot_product_attention(queries, keys, values, dropout_p=self.drop_rate if self.training else 0.0, is_causal=True)
        context_vecs = context_vecs.transpose(1, 2).reshape(batch_size, seq_len, self.emb_dim)

        return self.out_proj(context_vecs)


class Transformer(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.attention = MultiQueryAttention(cfg)
        self.ff = MoE(cfg)
        self.norm1 = nn.RMSNorm(cfg["emb_dim"])
        self.norm2 = nn.RMSNorm(cfg["emb_dim"])

    def forward(self, x, rope_cos: torch.Tensor, rope_sin: torch.Tensor):

        shortcut = x
        x = self.norm1(x)
        x = self.attention(x, rope_cos, rope_sin)
        x = x + shortcut

        shortcut = x
        x = self.norm2(x)
        x = self.ff(x)
        x = x + shortcut

        return x


class LLMConfig(PretrainedConfig):

    model_type = "custom_llm"

    def __init__(self,
                 vocab_size: int = 32768,
                 eos_token_id=0,
                 bos_token_id=0,
                 pad_token_id=0,
                 context_length: int = 1024,
                 emb_dim: int = 512,
                 hidden_dim: int = 1024,
                 n_heads: int = 8,
                 n_layers: int = 14,
                 qkv_bias: bool = False,
                 drop_rate: float = 0.0,
                 n_experts: int = 8,
                 top_k: int = 2,
                 rope_base: float = 10000.0,
                 **kwargs):

        self.vocab_size = vocab_size
        self.context_length = context_length
        self.emb_dim = emb_dim
        self.n_heads = n_heads
        self.n_layers = n_layers
        self.qkv_bias = qkv_bias
        self.drop_rate = drop_rate
        self.hidden_dim = hidden_dim
        self.n_experts = n_experts
        self.top_k = top_k
        self.rope_base = rope_base

        kwargs.setdefault("tie_word_embeddings", True)
        super().__init__(
            eos_token_id=eos_token_id,
            bos_token_id=bos_token_id,
            pad_token_id=pad_token_id,
            **kwargs)

    def __getitem__(self, key):
        return getattr(self, key)


class LLM(PreTrainedModel):

    config_class = LLMConfig
    _tied_weights_keys = {"out.weight": "tok_emb.weight"}

    def __init__(self, cfg: LLMConfig):
        super().__init__(cfg)

        self.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["emb_dim"])
        self.trans_blocks = nn.ModuleList(
            [Transformer(cfg) for _ in range(cfg["n_layers"])]
        )
        rope_cos, rope_sin = build_rope_cache(
            head_dim=cfg["emb_dim"] // cfg["n_heads"],
            context_length=cfg["context_length"],
            base=cfg["rope_base"],
        )
        self.register_buffer("rope_cos", rope_cos, persistent=True)
        self.register_buffer("rope_sin", rope_sin, persistent=True)

        self.norm = nn.RMSNorm(cfg["emb_dim"])
        self.out = nn.Linear(cfg["emb_dim"], cfg["vocab_size"], bias=False)
        self.out.weight = self.tok_emb.weight

        self.post_init()

    def forward(self, input_ids: torch.Tensor, **kwargs):
        _, seq_len = input_ids.shape

        assert seq_len <= self.rope_cos.shape[0], (
            f"seq_len {seq_len} exceeds cached context_length {self.rope_cos.shape[0]}"
        )

        x = self.tok_emb(input_ids)
        for block in self.trans_blocks:
            x = block(x, self.rope_cos, self.rope_sin)

        logits = self.out(self.norm(x))
        return CausalLMOutputWithPast(logits=logits)

AutoConfig.register("custom_llm", LLMConfig)
AutoModelForCausalLM.register(LLMConfig, LLM)
LLMConfig.register_for_auto_class()
LLM.register_for_auto_class("AutoModelForCausalLM")