ares/model.py

import math, torch
import torch.nn as nn
import torch.nn.functional as F
from .config import AresConfig

class RMSNorm(nn.Module):
 def __init__(self,d,eps=1e-5): super().__init__();self.weight=nn.Parameter(torch.ones(d));self.eps=eps
 def forward(self,x): return x*torch.rsqrt(x.pow(2).mean(-1,keepdim=True)+self.eps)*self.weight

def rope(x, positions, theta):
 # x B,H,T,D; RoPE on even-dimensional head features
 d=x.size(-1);freq=1.0/(theta**(torch.arange(0,d,2,device=x.device,dtype=x.dtype)/d));ang=positions.to(x.dtype)[None,None,:,None]*freq
 c,s=ang.cos(),ang.sin();a,b=x[...,::2],x[...,1::2];return torch.stack((a*c-b*s,a*s+b*c),-1).flatten(-2)
class Attention(nn.Module):
 def __init__(self,c):
  super().__init__();assert c.dim%c.n_heads==0 and c.n_heads%c.n_kv_heads==0;self.h,self.kh,self.d=c.n_heads,c.n_kv_heads,c.dim//c.n_heads
  self.q=nn.Linear(c.dim,c.dim,bias=False);self.drop=c.dropout;self.k=nn.Linear(c.dim,self.kh*self.d,bias=False);self.v=nn.Linear(c.dim,self.kh*self.d,bias=False);self.o=nn.Linear(c.dim,c.dim,bias=False);self.theta=c.rope_theta
 def forward(self,x,cache=None):
  b,t,_=x.shape; past=0 if cache is None else cache[0].size(2);pos=torch.arange(past,past+t,device=x.device)
  q=rope(self.q(x).view(b,t,self.h,self.d).transpose(1,2),pos,self.theta);k=rope(self.k(x).view(b,t,self.kh,self.d).transpose(1,2),pos,self.theta);v=self.v(x).view(b,t,self.kh,self.d).transpose(1,2)
  if cache is not None:k=torch.cat((cache[0],k),2);v=torch.cat((cache[1],v),2)
  k=k.repeat_interleave(self.h//self.kh,1);v=v.repeat_interleave(self.h//self.kh,1)
  # is_causal correctly handles full prefill; generation uses one query and all keys.
  y=F.scaled_dot_product_attention(q,k,v,dropout_p=self.drop if self.training else 0.0,is_causal=(t>1 and past==0));return self.o(y.transpose(1,2).reshape(b,t,-1)),(k[:,::self.h//self.kh],v[:,::self.h//self.kh])
class MLP(nn.Module):
 def __init__(self,c): super().__init__();h=int(8*c.dim*c.ffn_multiplier/3)//64*64;self.w1=nn.Linear(c.dim,h,bias=False);self.w3=nn.Linear(c.dim,h,bias=False);self.w2=nn.Linear(h,c.dim,bias=False);self.drop=c.dropout
 def forward(self,x):return F.dropout(self.w2(F.silu(self.w1(x))*self.w3(x)),p=self.drop,training=self.training)
class Block(nn.Module):
 def __init__(self,c):super().__init__();self.n1=RMSNorm(c.dim,c.rms_norm_eps);self.a=Attention(c);self.n2=RMSNorm(c.dim,c.rms_norm_eps);self.m=MLP(c)
 def forward(self,x,cache=None):a,k=self.a(self.n1(x),cache);return x+a+self.m(self.n2(x+a)),k
class AresTransformer(nn.Module):
 def __init__(self,c=AresConfig()):
  super().__init__();self.config=c;self.embed=nn.Embedding(c.vocab_size,c.dim);self.blocks=nn.ModuleList([Block(c) for _ in range(c.n_layers)]);self.norm=RMSNorm(c.dim,c.rms_norm_eps);self.lm_head=nn.Linear(c.dim,c.vocab_size,bias=False)
  if c.tie_embeddings:self.lm_head.weight=self.embed.weight
 def forward(self,ids,targets=None,kv_cache=None):
  x=self.embed(ids);new=[]
  for i,b in enumerate(self.blocks):x,k=b(x,None if kv_cache is None else kv_cache[i]);new.append(k)
  logits=self.lm_head(self.norm(x));loss=F.cross_entropy(logits[:,:-1].reshape(-1,logits.size(-1)),targets[:,1:].reshape(-1)) if targets is not None else None
  return logits,loss,new
 @torch.no_grad()
 def generate(self,ids,max_new=128,temperature=.8):
  cache=None
  for _ in range(max_new):
   logits,_,cache=self(ids[:,-self.config.max_seq_len:] if cache is None else ids[:,-1:],kv_cache=cache);p=F.softmax(logits[:,-1]/max(temperature,1e-5),-1);ids=torch.cat((ids,torch.multinomial(p,1)),1)
  return ids