| """ |
| Derived from Andrej Karpathy's nanochat project. |
| |
| MIT License |
| |
| Copyright (c) 2025 Andrej Karpathy |
| |
| Permission is hereby granted, free of charge, to any person obtaining a copy |
| of this software and associated documentation files (the "Software"), to deal |
| in the Software without restriction, including without limitation the rights |
| to use, copy, modify, merge, publish, distribute, sublicense, and/or sell |
| copies of the Software, and to permit persons to whom the Software is |
| furnished to do so, subject to the following conditions: |
| |
| The above copyright notice and this permission notice shall be included in all |
| copies or substantial portions of the Software. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass(frozen=True) |
| class GPTConfig: |
| block_size: int |
| vocab_size: int |
| n_layer: int |
| n_head: int |
| n_embd: int |
| dropout: float |
|
|
| def to_dict(self) -> dict[str, int | float]: |
| return asdict(self) |
|
|
|
|
| def rms_norm(x: torch.Tensor) -> torch.Tensor: |
| return F.rms_norm(x, (x.size(-1),)) |
|
|
|
|
| class Linear(nn.Linear): |
| """Bias-free linear layer matching the nanochat architectural style.""" |
|
|
| def __init__(self, in_features: int, out_features: int): |
| super().__init__(in_features, out_features, bias=False) |
|
|
|
|
| def apply_rotary_emb( |
| x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor |
| ) -> torch.Tensor: |
| half = x.shape[-1] // 2 |
| x1, x2 = x[..., :half], x[..., half:] |
| return torch.cat((x1 * cos + x2 * sin, x1 * (-sin) + x2 * cos), dim=-1) |
|
|
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| if config.n_embd % config.n_head != 0: |
| raise ValueError("n_embd must be divisible by n_head") |
| self.n_head = config.n_head |
| self.head_dim = config.n_embd // config.n_head |
| if self.head_dim % 2 != 0: |
| raise ValueError("rotary attention requires an even head dimension") |
| self.dropout_p = config.dropout |
| self.c_q = Linear(config.n_embd, config.n_embd) |
| self.c_k = Linear(config.n_embd, config.n_embd) |
| self.c_v = Linear(config.n_embd, config.n_embd) |
| self.c_proj = Linear(config.n_embd, config.n_embd) |
| self.resid_dropout = nn.Dropout(config.dropout) |
| mask = torch.ones(config.block_size, config.block_size, dtype=torch.bool).tril() |
| self.register_buffer( |
| "causal_mask", |
| mask.view(1, 1, config.block_size, config.block_size), |
| persistent=False, |
| ) |
|
|
| def set_dropout(self, p: float) -> None: |
| self.dropout_p = p |
| self.resid_dropout.p = p |
|
|
| def forward( |
| self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor |
| ) -> torch.Tensor: |
| batch, seq_len, channels = x.shape |
| q = self.c_q(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) |
| k = self.c_k(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) |
| v = self.c_v(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| q = rms_norm(apply_rotary_emb(q, cos, sin)) * 1.2 |
| k = rms_norm(apply_rotary_emb(k, cos, sin)) * 1.2 |
|
|
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| mask = self.causal_mask[:, :, :seq_len, :seq_len] |
| att = att.masked_fill(~mask, float("-inf")) |
| att = F.softmax(att.float(), dim=-1).to(dtype=x.dtype) |
| att = F.dropout(att, p=self.dropout_p, training=self.training) |
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(batch, seq_len, channels) |
| return self.resid_dropout(self.c_proj(y)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.c_fc = Linear(config.n_embd, 4 * config.n_embd) |
| self.c_proj = Linear(4 * config.n_embd, config.n_embd) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def set_dropout(self, p: float) -> None: |
| self.dropout.p = p |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.c_fc(x) |
| x = F.relu(x).square() |
| x = self.c_proj(x) |
| return self.dropout(x) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.attn = CausalSelfAttention(config) |
| self.mlp = MLP(config) |
|
|
| def set_dropout(self, p: float) -> None: |
| self.attn.set_dropout(p) |
| self.mlp.set_dropout(p) |
|
|
| def forward( |
| self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor |
| ) -> torch.Tensor: |
| x = x + self.attn(rms_norm(x), cos, sin) |
| x = x + self.mlp(rms_norm(x)) |
| return x |
|
|
|
|
| class DropoutGPT(nn.Module): |
| """Minimal nanochat-style causal Transformer with dynamic dropout control.""" |
|
|
| def __init__(self, config: GPTConfig): |
| super().__init__() |
| self.config = config |
| self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd) |
| self.embedding_dropout = nn.Dropout(config.dropout) |
| self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layer)) |
| self.lm_head = Linear(config.n_embd, config.vocab_size) |
| self.register_buffer( |
| "cos", |
| torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), |
| persistent=False, |
| ) |
| self.register_buffer( |
| "sin", |
| torch.empty(1, 1, config.block_size, config.n_embd // config.n_head // 2), |
| persistent=False, |
| ) |
| self._init_weights() |
| self._init_rotary() |
|
|
| @torch.no_grad() |
| def _init_weights(self) -> None: |
| nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.8) |
| nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001) |
| scale = (3.0 ** 0.5) * (self.config.n_embd ** -0.5) |
| for block in self.blocks: |
| nn.init.uniform_(block.attn.c_q.weight, -scale, scale) |
| nn.init.uniform_(block.attn.c_k.weight, -scale, scale) |
| nn.init.uniform_(block.attn.c_v.weight, -scale, scale) |
| nn.init.zeros_(block.attn.c_proj.weight) |
| nn.init.uniform_(block.mlp.c_fc.weight, -0.4 * scale, 0.4 * scale) |
| nn.init.zeros_(block.mlp.c_proj.weight) |
|
|
| @torch.no_grad() |
| def _init_rotary(self) -> None: |
| head_dim = self.config.n_embd // self.config.n_head |
| channel_range = torch.arange( |
| 0, |
| head_dim, |
| 2, |
| dtype=torch.float32, |
| device=self.cos.device, |
| ) |
| inv_freq = 1.0 / (100000 ** (channel_range / head_dim)) |
| positions = torch.arange( |
| self.config.block_size, |
| dtype=torch.float32, |
| device=self.cos.device, |
| ) |
| freqs = torch.outer(positions, inv_freq) |
| self.cos.copy_(freqs.cos().view(1, 1, self.config.block_size, head_dim // 2)) |
| self.sin.copy_(freqs.sin().view(1, 1, self.config.block_size, head_dim // 2)) |
|
|
| def set_dropout(self, p: float) -> None: |
| if not 0.0 <= p < 1.0: |
| raise ValueError("dropout must satisfy 0 <= p < 1") |
| self.embedding_dropout.p = p |
| for block in self.blocks: |
| block.set_dropout(p) |
|
|
| def num_parameters(self) -> int: |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def forward( |
| self, idx: torch.Tensor, targets: torch.Tensor | None = None |
| ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: |
| _, seq_len = idx.shape |
| if seq_len > self.config.block_size: |
| raise ValueError("sequence length exceeds block_size") |
| cos = self.cos[:, :, :seq_len, :] |
| sin = self.sin[:, :, :seq_len, :] |
| x = self.embedding_dropout(rms_norm(self.token_embedding(idx))) |
| for block in self.blocks: |
| x = block(x, cos, sin) |
| logits = self.lm_head(rms_norm(x)).float() |
| if targets is None: |
| return logits |
| loss = F.cross_entropy( |
| logits.reshape(-1, logits.size(-1)), |
| targets.reshape(-1), |
| ) |
| return logits, loss |
|
|
