ares_core/model.py

from __future__ import annotations

import math
from typing import List, Optional, Tuple

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

from .config import AresConfig

KVCache = Tuple[torch.Tensor, torch.Tensor]


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:
        dtype = x.dtype
        x_float = x.float()
        rms = torch.rsqrt(x_float.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (x_float * rms).to(dtype) * self.weight


def precompute_rope(seq_len: int, head_dim: int, theta: float, device=None) -> Tuple[torch.Tensor, torch.Tensor]:
    if head_dim % 2 != 0:
        raise ValueError("RoPE requires an even head_dim")
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(seq_len, device=device).float()
    freqs = torch.einsum("i,j->ij", t, inv_freq)
    return freqs.cos(), freqs.sin()


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Apply rotary embeddings.

    x: [batch, time, heads, head_dim]
    cos/sin: [time, head_dim // 2]
    """
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    cos = cos[None, :, None, :].to(dtype=x.dtype, device=x.device)
    sin = sin[None, :, None, :].to(dtype=x.dtype, device=x.device)
    y_even = x_even * cos - x_odd * sin
    y_odd = x_even * sin + x_odd * cos
    return torch.stack((y_even, y_odd), dim=-1).flatten(-2)


def repeat_kv(x: torch.Tensor, repeats: int) -> torch.Tensor:
    """Repeat grouped KV heads to full query-head count.

    Input: [batch, time, kv_heads, head_dim]
    Output: [batch, time, kv_heads * repeats, head_dim]
    """
    if repeats == 1:
        return x
    b, t, h, d = x.shape
    return x[:, :, :, None, :].expand(b, t, h, repeats, d).reshape(b, t, h * repeats, d)


class CausalSelfAttention(nn.Module):
    def __init__(self, cfg: AresConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.n_heads = cfg.n_heads
        self.n_kv_heads = cfg.n_kv_heads
        self.head_dim = cfg.head_dim
        self.kv_repeats = cfg.n_heads // cfg.n_kv_heads

        self.wq = nn.Linear(cfg.d_model, cfg.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(cfg.d_model, cfg.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(cfg.n_heads * self.head_dim, cfg.d_model, bias=False)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        kv_cache: Optional[KVCache] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
        b, t, _ = x.shape
        q = self.wq(x).view(b, t, self.n_heads, self.head_dim)
        k = self.wk(x).view(b, t, self.n_kv_heads, self.head_dim)
        v = self.wv(x).view(b, t, self.n_kv_heads, self.head_dim)

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

        if kv_cache is not None:
            past_k, past_v = kv_cache
            k = torch.cat([past_k, k], dim=1)
            v = torch.cat([past_v, v], dim=1)

        present: Optional[KVCache] = (k, v) if use_cache else None

        k_full = repeat_kv(k, self.kv_repeats)
        v_full = repeat_kv(v, self.kv_repeats)

        # [B, H, T, D]
        q = q.transpose(1, 2)
        k_full = k_full.transpose(1, 2)
        v_full = v_full.transpose(1, 2)

        scores = (q @ k_full.transpose(-2, -1)) / math.sqrt(self.head_dim)
        total_k = k_full.size(-2)
        past_len = total_k - t
        q_positions = past_len + torch.arange(t, device=x.device)[:, None]
        k_positions = torch.arange(total_k, device=x.device)[None, :]
        causal = k_positions <= q_positions
        scores = scores.masked_fill(~causal[None, None, :, :], torch.finfo(scores.dtype).min)
        att = F.softmax(scores.float(), dim=-1).to(dtype=x.dtype)
        att = self.dropout(att)
        y = att @ v_full
        y = y.transpose(1, 2).contiguous().view(b, t, self.n_heads * self.head_dim)
        return self.wo(y), present


class SwiGLU(nn.Module):
    def __init__(self, cfg: AresConfig):
        super().__init__()
        self.w1 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w3 = nn.Linear(cfg.d_model, cfg.d_ff, bias=False)
        self.w2 = nn.Linear(cfg.d_ff, cfg.d_model, bias=False)
        self.dropout = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(self.dropout(F.silu(self.w1(x)) * self.w3(x)))


class TransformerBlock(nn.Module):
    def __init__(self, cfg: AresConfig):
        super().__init__()
        self.attn_norm = RMSNorm(cfg.d_model)
        self.attn = CausalSelfAttention(cfg)
        self.ffn_norm = RMSNorm(cfg.d_model)
        self.ffn = SwiGLU(cfg)

    def forward(
        self,
        x: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        kv_cache: Optional[KVCache] = None,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[KVCache]]:
        h, present = self.attn(self.attn_norm(x), cos, sin, kv_cache=kv_cache, use_cache=use_cache)
        x = x + h
        x = x + self.ffn(self.ffn_norm(x))
        return x, present


class AresForCausalLM(nn.Module):
    def __init__(self, cfg: AresConfig):
        super().__init__()
        cfg.validate()
        self.cfg = cfg
        self.tok_embeddings = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.layers = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.d_model)
        self.lm_head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
        if cfg.tie_embeddings:
            self.lm_head.weight = self.tok_embeddings.weight
        self.apply(self._init_weights)

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(
        self,
        input_ids: torch.Tensor,
        targets: Optional[torch.Tensor] = None,
        start_pos: int = 0,
        kv_cache: Optional[List[Optional[KVCache]]] = None,
        use_cache: bool = False,
    ):
        b, t = input_ids.shape
        if t + start_pos > self.cfg.max_seq_len:
            raise ValueError(f"Sequence length {t + start_pos} exceeds max_seq_len={self.cfg.max_seq_len}")
        if kv_cache is None:
            kv_cache = [None] * len(self.layers)

        x = self.tok_embeddings(input_ids)
        cos_all, sin_all = precompute_rope(t + start_pos, self.cfg.head_dim, self.cfg.rope_theta, device=input_ids.device)
        cos = cos_all[start_pos : start_pos + t]
        sin = sin_all[start_pos : start_pos + t]

        new_cache: List[Optional[KVCache]] = []
        for layer, layer_cache in zip(self.layers, kv_cache):
            x, present = layer(x, cos, sin, kv_cache=layer_cache, use_cache=use_cache)
            new_cache.append(present)
        x = self.norm(x)
        logits = self.lm_head(x)

        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return {"logits": logits, "loss": loss, "kv_cache": new_cache if use_cache else None}

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 64,
        temperature: float = 0.8,
        top_k: int = 50,
        eos_id: Optional[int] = None,
    ) -> torch.Tensor:
        self.eval()
        if input_ids.dim() == 1:
            input_ids = input_ids[None, :]
        generated = input_ids
        cache = None
        start_pos = 0
        next_input = input_ids
        for _ in range(max_new_tokens):
            out = self(next_input, start_pos=start_pos, kv_cache=cache, use_cache=True)
            logits = out["logits"][:, -1, :]
            cache = out["kv_cache"]
            # The next forward pass will process the sampled token at the position
            # immediately after the tokens that are already present in the cache.
            next_start_pos = start_pos + next_input.size(1)
            if temperature <= 0:
                next_token = torch.argmax(logits, dim=-1, keepdim=True)
            else:
                logits = logits / temperature
                if top_k and top_k > 0:
                    values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                    cutoff = values[:, -1, None]
                    logits = torch.where(logits < cutoff, torch.full_like(logits, -float("inf")), logits)
                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            generated = torch.cat([generated, next_token], dim=1)
            next_input = next_token
            start_pos = next_start_pos
            if eos_id is not None and int(next_token[0, 0]) == int(eos_id):
                break
            if generated.size(1) >= self.cfg.max_seq_len:
                break
        return generated


def count_parameters(model: nn.Module) -> int:
    return sum(p.numel() for p in model.parameters())