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())