| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| class RotaryPositionalEmbedding(nn.Module): | |
| def __init__(self, dim, max_seq_len=2048, base=10000): | |
| super().__init__() | |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) | |
| self.register_buffer('inv_freq', inv_freq) | |
| self.max_seq_len = max_seq_len | |
| def forward(self, seq_len, device): | |
| t = torch.arange(seq_len, device=device).type_as(self.inv_freq) | |
| freqs = torch.einsum('i,j->ij', t, self.inv_freq) | |
| emb = torch.cat((freqs, freqs), dim=-1) | |
| return emb.cos(), emb.sin() | |
| def apply_rotary_pos_emb(q, k, cos, sin): | |
| def rotate_half(x): | |
| x1, x2 = x.chunk(2, dim=-1) | |
| return torch.cat((-x2, x1), dim=-1) | |
| q_embed = (q * cos) + (rotate_half(q) * sin) | |
| k_embed = (k * cos) + (rotate_half(k) * sin) | |
| return q_embed, k_embed | |
| class MultiHeadSelfAttention(nn.Module): | |
| def __init__(self, d_model, n_heads, dropout=0.1, max_seq_len=2048): | |
| super().__init__() | |
| assert d_model % n_heads == 0 | |
| self.d_model = d_model | |
| self.n_heads = n_heads | |
| self.d_k = d_model // n_heads | |
| self.q_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.k_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.v_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.out_linear = nn.Linear(d_model, d_model, bias=False) | |
| self.dropout = nn.Dropout(dropout) | |
| self.attn_dropout = nn.Dropout(dropout) | |
| self.rope = RotaryPositionalEmbedding(self.d_k, max_seq_len) | |
| self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') | |
| def forward(self, x, mask=None): | |
| batch_size, seq_len, d_model = x.size() | |
| Q = self.q_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) | |
| K = self.k_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) | |
| V = self.v_linear(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) | |
| cos, sin = self.rope(seq_len, x.device) | |
| cos = cos[None, None, :, :] | |
| sin = sin[None, None, :, :] | |
| Q, K = apply_rotary_pos_emb(Q, K, cos, sin) | |
| if self.flash and mask is None: | |
| context = F.scaled_dot_product_attention(Q, K, V, attn_mask=None, dropout_p=0.0, is_causal=True) | |
| else: | |
| scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) | |
| if mask is not None: | |
| scores = scores.masked_fill(mask == 0, float('-inf')) | |
| attn_weights = F.softmax(scores, dim=-1) | |
| attn_weights = self.attn_dropout(attn_weights) | |
| context = torch.matmul(attn_weights, V) | |
| context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, d_model) | |
| output = self.out_linear(context) | |
| return self.dropout(output) | |
| class FeedForward(nn.Module): | |
| def __init__(self, d_model, d_ff, dropout=0.1): | |
| super().__init__() | |
| self.linear1 = nn.Linear(d_model, d_ff) | |
| self.linear2 = nn.Linear(d_ff, d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x): | |
| return self.linear2(self.dropout(F.gelu(self.linear1(x)))) | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim, eps=1e-6): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x): | |
| norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) | |
| return x * norm * self.weight | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, d_model, n_heads, d_ff, dropout=0.1, max_seq_len=2048, use_swiglu=False): | |
| super().__init__() | |
| self.attention = MultiHeadSelfAttention(d_model, n_heads, dropout, max_seq_len) | |
| self.feed_forward = FeedForward(d_model, d_ff, dropout) | |
| self.norm1 = RMSNorm(d_model) | |
| self.norm2 = RMSNorm(d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, mask=None): | |
| x = x + self.attention(self.norm1(x), mask) | |
| x = x + self.feed_forward(self.norm2(x)) | |
| return x | |
| class MTPMiniModel(nn.Module): | |
| def __init__(self, vocab_size, d_model=512, n_layers=8, n_heads=8, | |
| d_ff=2048, max_seq_len=512, dropout=0.2, use_swiglu=False): | |
| super().__init__() | |
| self.vocab_size = vocab_size | |
| self.d_model = d_model | |
| self.max_seq_len = max_seq_len | |
| self.token_embedding = nn.Embedding(vocab_size, d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| self.blocks = nn.ModuleList([ | |
| TransformerBlock(d_model, n_heads, d_ff, dropout, max_seq_len, use_swiglu) | |
| for _ in range(n_layers) | |
| ]) | |
| self.norm_f = RMSNorm(d_model) | |
| self.lm_head = nn.Linear(d_model, vocab_size, bias=False) | |
| self.lm_head.weight = self.token_embedding.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) | |
| if module.bias is not None: | |
| torch.nn.init.zeros_(module.bias) | |
| elif isinstance(module, nn.Embedding): | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) | |
| def forward(self, input_ids, targets=None, use_eos_weight=False): | |
| batch_size, seq_len = input_ids.size() | |
| mask = torch.tril(torch.ones(seq_len, seq_len, device=input_ids.device)).view(1, 1, seq_len, seq_len) | |
| x = self.dropout(self.token_embedding(input_ids)) | |
| for block in self.blocks: | |
| x = block(x, mask) | |
| x = self.norm_f(x) | |
| logits = self.lm_head(x) | |
| loss = None | |
| if targets is not None: | |
| if use_eos_weight: | |
| weights = torch.ones(self.vocab_size, device=logits.device) | |
| weights[3] = 2.0 | |
| loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1), weight=weights, label_smoothing=0.1) | |
| else: | |
| loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1), label_smoothing=0.1) | |
| return logits, loss | |
| def count_parameters(self): | |
| return sum(p.numel() for p in self.parameters() if p.requires_grad) | |