import math from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel, PretrainedConfig, AutoConfig, AutoModel, AutoModelForCausalLM class NexusSmAllConfig(PretrainedConfig): model_type = "nexus_small" def __init__( self, vocab_size=50304, max_seq_len=512, dim=768, num_layers=10, num_heads=12, num_kv_heads=4, multiple_of=256, ff_dim=2048, norm_eps=1e-6, rope_theta=500000.0, **kwargs, ): super().__init__(**kwargs) self.vocab_size = vocab_size self.max_seq_len = max_seq_len self.dim = dim self.num_layers = num_layers self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.multiple_of = multiple_of self.ff_dim = ff_dim self.norm_eps = norm_eps self.rope_theta = rope_theta AutoConfig.register("nexus_small", NexusSmAllConfig) 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): rms = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) return (x.float() * rms * self.weight.float()).type_as(x) def precompute_freqs_cis(max_seq_len, dim, num_heads, rope_theta): head_dim = dim // num_heads freqs = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) t = torch.arange(max_seq_len) freqs = torch.outer(t, freqs) return torch.polar(torch.ones_like(freqs), freqs) class RotaryEmbedding(nn.Module): def __init__(self, max_seq_len, dim, num_heads, rope_theta): super().__init__() self.freqs_cis = precompute_freqs_cis(max_seq_len, dim, num_heads, rope_theta) def forward(self, x, start_pos=0): _, seq_len, _, head_dim = x.shape freqs_cis = self.freqs_cis[start_pos : start_pos + seq_len, : head_dim // 2].to(x.device) freqs_cis = freqs_cis.view(1, seq_len, 1, head_dim // 2) x_shaped = x.float().reshape(*x.shape[:-1], -1, 2) x_complex = torch.complex(x_shaped[..., 0], x_shaped[..., 1]) x_rotated = x_complex * freqs_cis x_out = torch.stack([x_rotated.real, x_rotated.imag], dim=-1).reshape_as(x_shaped) return x_out.reshape_as(x).type_as(x) class Attention(nn.Module): def __init__(self, dim, num_heads, num_kv_heads, max_seq_len, rope_theta): super().__init__() self.num_heads = num_heads self.num_kv_heads = num_kv_heads or num_heads self.head_dim = dim // num_heads self.num_kv_groups = num_heads // self.num_kv_heads self.wq = nn.Linear(dim, dim, bias=False) self.wk = nn.Linear(dim, self.head_dim * self.num_kv_heads, bias=False) self.wv = nn.Linear(dim, self.head_dim * self.num_kv_heads, bias=False) self.wo = nn.Linear(dim, dim, bias=False) self.rotary = RotaryEmbedding(max_seq_len, dim, num_heads, rope_theta) def forward(self, x, start_pos=0, mask=None): bsz, seqlen, _ = x.shape q = self.wq(x).view(bsz, seqlen, self.num_heads, self.head_dim).transpose(1, 2) k = self.wk(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.wv(x).view(bsz, seqlen, self.num_kv_heads, self.head_dim).transpose(1, 2) q = self.rotary(q, start_pos) k = self.rotary(k, start_pos) if self.num_kv_groups > 1: k = k[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim) k = k.reshape(bsz, self.num_heads, seqlen, self.head_dim) v = v[:, :, None, :, :].expand(bsz, self.num_kv_heads, self.num_kv_groups, seqlen, self.head_dim) v = v.reshape(bsz, self.num_heads, seqlen, self.head_dim) scale = 1.0 / math.sqrt(self.head_dim) attn_weights = torch.matmul(q, k.transpose(-2, -1)) * scale if mask is not None: attn_weights = attn_weights + mask attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(q) attn_output = torch.matmul(attn_weights, v) attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, seqlen, -1) return self.wo(attn_output) class FeedForward(nn.Module): def __init__(self, dim, ff_dim, multiple_of): super().__init__() hidden_dim = int(2 * ff_dim / 3) hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) self.w1 = nn.Linear(dim, hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, dim, bias=False) self.w3 = nn.Linear(dim, hidden_dim, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class TransformerBlock(nn.Module): def __init__(self, dim, num_heads, num_kv_heads, ff_dim, multiple_of, norm_eps, max_seq_len, rope_theta): super().__init__() self.attention = Attention(dim, num_heads, num_kv_heads, max_seq_len, rope_theta) self.feed_forward = FeedForward(dim, ff_dim, multiple_of) self.attention_norm = RMSNorm(dim, norm_eps) self.ff_norm = RMSNorm(dim, norm_eps) def forward(self, x, start_pos=0, mask=None): h = x + self.attention(self.attention_norm(x), start_pos, mask) out = h + self.feed_forward(self.ff_norm(h)) return out class Nexus(nn.Module): def __init__(self, config): super().__init__() self.config = config self.token_embeddings = nn.Embedding(config.vocab_size, config.dim) self.layers = nn.ModuleList([ TransformerBlock( config.dim, config.num_heads, config.num_kv_heads, config.ff_dim, config.multiple_of, config.norm_eps, config.max_seq_len, config.rope_theta, ) for _ in range(config.num_layers) ]) self.norm = RMSNorm(config.dim, config.norm_eps) self.output = nn.Linear(config.dim, config.vocab_size, bias=False) self.token_embeddings.weight = self.output.weight def forward(self, input_ids, start_pos=0): _, seqlen = input_ids.shape mask = torch.full( (1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=torch.float32, device=input_ids.device, ) mask = torch.triu(mask, diagonal=start_pos + 1).type_as(input_ids) x = self.token_embeddings(input_ids) for layer in self.layers: x = layer(x, start_pos, mask) x = self.norm(x) return self.output(x) class NexusForCausalLM(PreTrainedModel): config_class = NexusSmAllConfig base_model_prefix = "nexus_small" supports_gradient_checkpointing = False _no_split_modules = ["TransformerBlock"] def __init__(self, config): super().__init__(config) self.model = Nexus(config) self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False) with torch.no_grad(): self.lm_head.weight = self.model.output.weight def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, input_ids, attention_mask=None, **kwargs): return self.model(input_ids) def prepare_inputs_for_generation(self, input_ids, **kwargs): return {"input_ids": input_ids} AutoModel.register(NexusSmAllConfig, NexusForCausalLM) AutoModelForCausalLM.register(NexusSmAllConfig, NexusForCausalLM)