| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
| from typing import Optional |
| from config import NexusConfig |
|
|
| 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(config: NexusConfig) -> torch.Tensor: |
| dim = config.dim // config.num_heads |
| freqs = 1.0 / (config.rope_theta ** (torch.arange(0, dim, 2).float() / dim)) |
| t = torch.arange(config.max_seq_len) |
| freqs = torch.outer(t, freqs) |
| return torch.polar(torch.ones_like(freqs), freqs) |
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, config: NexusConfig): |
| super().__init__() |
| self.freqs_cis = precompute_freqs_cis(config) |
|
|
| def forward(self, x: torch.Tensor, start_pos: int = 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, config: NexusConfig): |
| super().__init__() |
| self.num_heads = config.num_heads |
| self.num_kv_heads = config.num_kv_heads |
| if self.num_kv_heads is None: |
| self.num_kv_heads = config.num_heads |
| self.head_dim = config.dim // config.num_heads |
| self.num_kv_groups = config.num_heads // self.num_kv_heads |
|
|
| self.wq = nn.Linear(config.dim, config.dim, bias=False) |
| self.wk = nn.Linear(config.dim, self.head_dim * self.num_kv_heads, bias=False) |
| self.wv = nn.Linear(config.dim, self.head_dim * self.num_kv_heads, bias=False) |
| self.wo = nn.Linear(config.dim, config.dim, bias=False) |
| self.rotary = RotaryEmbedding(config) |
|
|
| def forward(self, x: torch.Tensor, start_pos: int = 0, mask: Optional[torch.Tensor] = 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, config: NexusConfig): |
| super().__init__() |
| hidden_dim = int(2 * config.ff_dim / 3) |
| hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of) |
|
|
| self.w1 = nn.Linear(config.dim, hidden_dim, bias=False) |
| self.w2 = nn.Linear(hidden_dim, config.dim, bias=False) |
| self.w3 = nn.Linear(config.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, config: NexusConfig): |
| super().__init__() |
| self.attention = Attention(config) |
| self.feed_forward = FeedForward(config) |
| self.attention_norm = RMSNorm(config.dim, config.norm_eps) |
| self.ff_norm = RMSNorm(config.dim, config.norm_eps) |
|
|
| def forward(self, x: torch.Tensor, start_pos: int = 0, mask: Optional[torch.Tensor] = 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: NexusConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.token_embeddings = nn.Embedding(config.vocab_size, config.dim) |
| self.layers = nn.ModuleList([TransformerBlock(config) 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 |
|
|
| self.apply(self._init_weights) |
|
|
| 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: torch.Tensor, start_pos: int = 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) |
| logits = self.output(x) |
|
|
| return logits |
|
|
| def generate(self, input_ids: torch.Tensor, max_new_tokens: int, |
| temperature: float = 0.7, top_k: int = 50, top_p: float = 0.9): |
| self.eval() |
| generated = [] |
|
|
| for _ in range(max_new_tokens): |
| seq_len = input_ids.shape[1] |
| if seq_len > self.config.max_seq_len: |
| input_ids = input_ids[:, -self.config.max_seq_len:] |
|
|
| with torch.no_grad(): |
| logits = self(input_ids, 0) |
| logits = logits[:, -1, :] / temperature |
|
|
| if top_k > 0: |
| top_k_values, _ = torch.topk(logits, top_k) |
| min_top_k = top_k_values[:, -1].unsqueeze(-1) |
| logits = torch.where(logits < min_top_k, |
| torch.full_like(logits, float('-inf')), logits) |
|
|
| if top_p > 0: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) |
| sorted_indices_to_remove = cumulative_probs > top_p |
| sorted_indices_to_remove[:, 0] = False |
| indices_to_remove = torch.zeros_like(logits, dtype=torch.bool) |
| indices_to_remove = indices_to_remove.scatter(1, sorted_indices, |
| sorted_indices_to_remove) |
| logits = torch.where(indices_to_remove, |
| torch.full_like(logits, float('-inf')), logits) |
|
|
| probs = F.softmax(logits, dim=-1) |
| next_token = torch.multinomial(probs, num_samples=1) |
|
|
| input_ids = torch.cat([input_ids, next_token], dim=-1) |
| generated.append(next_token.item()) |
|
|
| return generated, input_ids |
|
|
| def create_nexus_model(): |
| from config import nexus_config |
| config = nexus_config |
|
|
| model = Nexus(config) |
| total_params = sum(p.numel() for p in model.parameters()) |
| trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) |
|
|
| print(f"[Nexus SmAll] Model created with {total_params/1e6:.1f}M parameters " |
| f"({trainable_params/1e6:.1f}M trainable)") |
|
|
| return model, config |
