| """ |
| Self-contained 63.8M Parameter GPT-Style Decoder Transformer |
| - Pre-LayerNorm Architecture |
| - Flash Attention via PyTorch F.scaled_dot_product_attention |
| - Tied Token Embedding & LM Head Weights |
| """ |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from config import GPTConfig |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| assert config.d_model % config.n_head == 0, "d_model must be divisible by n_head" |
| self.n_head = config.n_head |
| self.d_model = config.d_model |
| self.head_dim = config.d_model // config.n_head |
| self.dropout_p = config.dropout |
|
|
| self.c_attn = nn.Linear(config.d_model, 3 * config.d_model, bias=config.bias) |
| self.c_proj = nn.Linear(config.d_model, config.d_model, bias=config.bias) |
| self.resid_dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, T, C = x.size() |
| qkv = self.c_attn(x) |
| q, k, v = qkv.split(self.d_model, dim=2) |
|
|
| k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| dropout_p = self.dropout_p if self.training else 0.0 |
| y = F.scaled_dot_product_attention( |
| q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True |
| ) |
|
|
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.resid_dropout(self.c_proj(y)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.c_fc = nn.Linear(config.d_model, config.d_ffn, bias=config.bias) |
| self.gelu = nn.GELU() |
| self.c_proj = nn.Linear(config.d_ffn, config.d_model, bias=config.bias) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.dropout(self.c_proj(self.gelu(self.c_fc(x)))) |
|
|
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.ln_1 = nn.LayerNorm(config.d_model, elementwise_affine=config.bias) |
| self.attn = CausalSelfAttention(config) |
| self.ln_2 = nn.LayerNorm(config.d_model, elementwise_affine=config.bias) |
| self.mlp = MLP(config) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = x + self.attn(self.ln_1(x)) |
| x = x + self.mlp(self.ln_2(x)) |
| return x |
|
|
|
|
| class SmallGPT(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.transformer = nn.ModuleDict(dict( |
| wte = nn.Embedding(config.vocab_size, config.d_model), |
| wpe = nn.Embedding(config.context_length, config.d_model), |
| drop = nn.Dropout(config.dropout), |
| h = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)]), |
| ln_f = nn.LayerNorm(config.d_model, elementwise_affine=config.bias), |
| )) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
|
|
| |
| self.transformer.wte.weight = self.lm_head.weight |
|
|
| |
| self.apply(self._init_weights) |
| for pn, p in self.named_parameters(): |
| if pn.endswith('c_proj.weight'): |
| torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)) |
|
|
| 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, idx: torch.Tensor, targets: torch.Tensor = None): |
| b, t = idx.size() |
| pos = torch.arange(0, t, dtype=torch.long, device=idx.device) |
|
|
| tok_emb = self.transformer.wte(idx) |
| pos_emb = self.transformer.wpe(pos) |
| x = self.transformer.drop(tok_emb + pos_emb) |
|
|
| for block in self.transformer.h: |
| x = block(x) |
|
|
| x = self.transformer.ln_f(x) |
|
|
| if targets is not None: |
| logits = self.lm_head(x) |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) |
| else: |
| logits = self.lm_head(x[:, [-1], :]) |
| loss = None |
|
|
| return logits, loss |
|
|
| def get_num_params(self) -> int: |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def configure_optimizers(self, weight_decay: float, learning_rate: float, betas: tuple, device_type: str): |
| decay_params = [p for n, p in self.named_parameters() if p.requires_grad and p.dim() >= 2] |
| nodecay_params = [p for n, p in self.named_parameters() if p.requires_grad and p.dim() < 2] |
|
|
| optim_groups = [ |
| {'params': decay_params, 'weight_decay': weight_decay}, |
| {'params': nodecay_params, 'weight_decay': 0.0} |
| ] |
|
|
| fused_available = 'fused' in torch.optim.AdamW.__init__.__code__.co_varnames |
| use_fused = fused_available and device_type == 'cuda' |
| extra_args = dict(fused=True) if use_fused else dict() |
| return torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, eps=1e-8, **extra_args) |
|
|
| @torch.no_grad() |
| def generate(self, idx: torch.Tensor, max_new_tokens: int = 80, temperature: float = 0.8, top_k: int = 40) -> torch.Tensor: |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx if idx.size(1) <= self.config.context_length else idx[:, -self.config.context_length:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1, :] / max(temperature, 1e-5) |
|
|
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float('Inf') |
|
|
| probs = F.softmax(logits, dim=-1) |
| idx_next = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat((idx, idx_next), dim=1) |
|
|
| return idx |
|
|
|
|
| |
| GPT = SmallGPT |
|
|
