| from dataclasses import dataclass
|
|
|
| import torch
|
| import torch.nn as nn
|
|
|
|
|
| @dataclass
|
| class TransformerConfig:
|
| model_dim: int
|
| num_heads: int
|
| N_ctx: int
|
| use_moe: bool
|
| num_experts : int
|
| active_experts : int
|
| shared_experts : int
|
| expert_width : int
|
| vocab_size : int
|
| layers : int
|
|
|
| class MHA(nn.Module):
|
| def __init__(self, num_heads, model_dim):
|
| super().__init__()
|
| self.num_heads = num_heads
|
| self.model_dim = model_dim
|
| self.head_dim = model_dim // num_heads
|
|
|
| self.q = nn.Linear(in_features=model_dim, out_features=model_dim)
|
| self.k = nn.Linear(in_features=model_dim, out_features=model_dim)
|
| self.v = nn.Linear(in_features=model_dim, out_features=model_dim)
|
| self.out_proj = nn.Linear(in_features=model_dim, out_features=model_dim)
|
|
|
| def forward(self, x, mask=None):
|
| q = self.q(x)
|
| k = self.k(x)
|
| v = self.v(x)
|
|
|
| q_heads = q.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
| k_heads = k.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
| v_heads = v.reshape(x.shape[0], x.shape[1], self.num_heads, self.head_dim).permute(0, 2, 1, 3)
|
|
|
| attention_scores = (q_heads @ torch.transpose(k_heads, 3, 2)) / self.head_dim ** (1/2)
|
| if mask is not None:
|
| attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))
|
|
|
| attention_scores = torch.softmax(attention_scores, dim=-1)
|
| self_attention = attention_scores @ v_heads
|
|
|
| mha_output = self_attention.permute(0, 2, 1, 3).reshape(x.shape[0], x.shape[1], self.model_dim)
|
| output = self.out_proj(mha_output)
|
| return output
|
|
|
| class RMSNorm(nn.Module):
|
| def __init__(self, model_dim):
|
| super().__init__()
|
| self.model_dim = model_dim
|
| self.gamma = nn.Parameter(torch.ones(model_dim))
|
|
|
| def forward(self, x):
|
| eps = 1e-5
|
| rms = (((torch.sum(torch.square(x), dim=-1, keepdim=True)) / self.model_dim) + eps) ** (1/2)
|
| norm = self.gamma * (x / rms)
|
| return norm
|
|
|
| class SinusoidalPositionalEmbeddings(nn.Module):
|
| def __init__(self, model_dim, N=10000):
|
| super().__init__()
|
| self.d = model_dim
|
| self.N = N
|
| self.register_buffer('k', torch.arange(0, model_dim, step=2))
|
|
|
| def forward(self, x):
|
| t = torch.arange(x.shape[1], device=x.device).unsqueeze(1)
|
| sin = torch.sin(t/self.N ** (self.k/self.d))
|
| cos = torch.cos(t/self.N ** (self.k/self.d))
|
| total_angles = torch.cat([sin, cos], dim=-1).unsqueeze(0)
|
| return x + total_angles
|
|
|
| class Block(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| self.config = config
|
| self.self_attn = MHA(config.num_heads, config.model_dim)
|
| self.norms = nn.ModuleList([RMSNorm(config.model_dim) for _ in range(2)])
|
| if config.use_moe:
|
| self.ffn = MoE(self.config)
|
| else:
|
| self.ffn = MLP(self.config)
|
|
|
| def forward(self, x, mask=None, pad_mask=None):
|
| x = x + self.self_attn(self.norms[0](x), mask)
|
| h, aux, counts = self.ffn(self.norms[1](x), pad_mask)
|
| return x + h, aux, counts
|
|
|
| class Expert(nn.Module):
|
| def __init__(self, config):
|
| super().__init__()
|
| self.expert = nn.Sequential(
|
| nn.Linear(config.model_dim, config.expert_width * config.model_dim),
|
| nn.GELU(),
|
| nn.Linear(config.expert_width * config.model_dim, config.model_dim)
|
| )
|
|
|
| def forward(self, x):
|
| x = self.expert(x)
|
| return x
|
|
|
| class MoE(nn.Module):
|
| def __init__(self, config: TransformerConfig):
|
| super().__init__()
|
| self.config = config
|
| self.router = nn.Sequential(
|
| nn.Linear(config.model_dim, config.num_experts - config.shared_experts),
|
| nn.Softmax(-1)
|
| )
|
| self.experts = nn.ModuleList([Expert(config) for _ in range(config.num_experts)])
|
|
|
| def forward(self, x, pad_mask=None):
|
| B, N, D = x.shape
|
| x_flat = x.reshape(-1, D)
|
| probs = self.router(x_flat)
|
| top_w, top_idx = torch.topk(probs, self.config.active_experts)
|
| out = torch.zeros_like(x_flat)
|
|
|
|
|
| if self.config.shared_experts:
|
| for i in range(self.config.shared_experts):
|
| out += self.experts[i](x_flat)
|
|
|
| real = pad_mask.reshape(-1).bool()
|
| r_probs = probs[real]
|
| r_idx = top_idx[real]
|
| num_routed = self.config.num_experts - self.config.shared_experts
|
| counts = torch.bincount(r_idx.reshape(-1), minlength=num_routed)
|
|
|
| for e in range(self.config.num_experts - self.config.shared_experts):
|
| e_mask = (top_idx == e).any(dim=-1)
|
| if not e_mask.any():
|
| continue
|
| e_wt = (top_w * (top_idx == e)).sum(dim=-1)
|
| out[e_mask] += e_wt[e_mask, None] * self.experts[self.config.shared_experts + e](x_flat[e_mask])
|
|
|
| f = counts / r_idx.numel()
|
| P = r_probs.mean(dim=0)
|
| aux = num_routed * (f * P).sum()
|
| return out.reshape(B, N, D), aux, counts
|
|
|
|
|
|
|
| class MLP(nn.Module):
|
| def __init__(self, config: TransformerConfig):
|
| super().__init__()
|
| self.mlp = nn.Sequential(
|
| nn.Linear(config.model_dim, 4 * config.model_dim),
|
| nn.GELU(),
|
| nn.Linear(4 * config.model_dim, config.model_dim)
|
| )
|
|
|
| def forward(self, x, pad_mask=None):
|
| x = self.mlp(x)
|
| return x, 0, None
|
|
|
|
|
| class Transformer(nn.Module):
|
| def __init__(self, config: TransformerConfig):
|
| super().__init__()
|
| self.config = config
|
| self.token_embedding = nn.Embedding(config.vocab_size, config.model_dim)
|
| self.pos_embedding = SinusoidalPositionalEmbeddings(config.model_dim)
|
| self.blocks = nn.ModuleList([Block(config) for _ in range(config.layers)])
|
| self.norm = RMSNorm(config.model_dim)
|
| self.lm_head = nn.Linear(config.model_dim, config.vocab_size, bias=False)
|
| self.lm_head.weight = self.token_embedding.weight
|
| nn.init.normal_(self.token_embedding.weight, std=0.02)
|
|
|
| def forward(self, input_ids, attention_mask, device, last_positions=None):
|
| N = input_ids.shape[-1]
|
| causal_mask = torch.tril(torch.ones(N, N, device=device, dtype=torch.bool))
|
| pad_mask = attention_mask
|
| attn_mask = causal_mask & pad_mask[:, None, None, :].bool()
|
|
|
| x = self.token_embedding(input_ids)
|
| x = self.pos_embedding(x)
|
|
|
| aux_total = 0
|
| all_counts = []
|
| for block in self.blocks:
|
| x, aux, counts = block(x, attn_mask, pad_mask)
|
| aux_total = aux_total + aux
|
| all_counts.append(counts)
|
|
|
| if last_positions is not None:
|
| x = x[torch.arange(x.shape[0], device=x.device), last_positions]
|
| logits = self.lm_head(self.norm(x))
|
|
|
| return logits, aux_total, all_counts
|
|
|
|
|
