| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| 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: |
| variance = x.pow(2).mean(-1, keepdim=True) |
| return x * torch.rsqrt(variance + self.eps) * self.weight |
|
|
| class RotaryEmbedding(nn.Module): |
| def __init__(self, dim: int, max_seq_len: int = 512): |
| super().__init__() |
| inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
| t = torch.arange(max_seq_len, dtype=torch.float32) |
| freqs = torch.outer(t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| self.register_buffer("cos_cached", emb.cos(), persistent=False) |
| self.register_buffer("sin_cached", emb.sin(), persistent=False) |
|
|
| def forward(self, q: torch.Tensor, k: torch.Tensor, seq_len: int): |
| cos = self.cos_cached[:seq_len, :].unsqueeze(0).unsqueeze(0) |
| sin = self.sin_cached[:seq_len, :].unsqueeze(0).unsqueeze(0) |
| q_rot = (q * cos) + (self._rotate_half(q) * sin) |
| k_rot = (k * cos) + (self._rotate_half(k) * sin) |
| return q_rot, k_rot |
|
|
| def _rotate_half(self, x: torch.Tensor) -> torch.Tensor: |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2 :] |
| return torch.cat((-x2, x1), dim=-1) |
|
|
| class Attention(nn.Module): |
| def __init__(self, dim: int, n_heads: int): |
| super().__init__() |
| self.n_heads = n_heads |
| self.head_dim = dim // n_heads |
| self.q_proj = nn.Linear(dim, dim, bias=False) |
| self.k_proj = nn.Linear(dim, dim, bias=False) |
| self.v_proj = nn.Linear(dim, dim, bias=False) |
| self.out_proj = nn.Linear(dim, dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor, rope: RotaryEmbedding) -> torch.Tensor: |
| B, S, C = x.shape |
| q = self.q_proj(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2) |
|
|
| q, k = rope(q, k, S) |
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| out = out.transpose(1, 2).contiguous().view(B, S, C) |
| return self.out_proj(out) |
|
|
| class MLP(nn.Module): |
| def __init__(self, dim: int, hidden_dim: int): |
| super().__init__() |
| self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) |
| self.up_proj = nn.Linear(dim, hidden_dim, bias=False) |
| self.down_proj = nn.Linear(hidden_dim, dim, bias=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, dim: int, n_heads: int, hidden_dim: int): |
| super().__init__() |
| self.attn_norm = RMSNorm(dim) |
| self.attn = Attention(dim, n_heads) |
| self.mlp_norm = RMSNorm(dim) |
| self.mlp = MLP(dim, hidden_dim) |
|
|
| def forward(self, x: torch.Tensor, rope: RotaryEmbedding) -> torch.Tensor: |
| x = x + self.attn(self.attn_norm(x), rope) |
| x = x + self.mlp(self.mlp_norm(x)) |
| return x |
|
|
| class MystCodeModel(nn.Module): |
| def __init__(self, vocab_size: int = 16384, dim: int = 512, n_layers: int = 8, n_heads: int = 8, max_seq_len: int = 512): |
| super().__init__() |
| self.vocab_size = vocab_size |
| self.max_seq_len = max_seq_len |
| self.tok_emb = nn.Embedding(vocab_size, dim) |
| self.rope = RotaryEmbedding(dim // n_heads, max_seq_len) |
| hidden_dim = int(2 * (4 * dim) / 3) |
| self.layers = nn.ModuleList([TransformerBlock(dim, n_heads, hidden_dim) for _ in range(n_layers)]) |
| self.norm = RMSNorm(dim) |
| self.lm_head = nn.Linear(dim, vocab_size, bias=False) |
| self.tok_emb.weight = self.lm_head.weight |
|
|
| def forward(self, idx: torch.Tensor, targets: torch.Tensor = None): |
| B, S = idx.shape |
| x = self.tok_emb(idx) |
| for layer in self.layers: |
| x = layer(x, self.rope) |
| x = self.norm(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1)) |
| return logits, loss
|
|
|
