| """Compact decoder-only Transformer for the Part 2 chat experiment.""" |
| from dataclasses import asdict, dataclass |
| import math |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| @dataclass |
| class ModelConfig: |
| vocab_size: int = 8192 |
| d_model: int = 192 |
| n_layers: int = 6 |
| n_heads: int = 6 |
| n_kv_heads: int = 2 |
| d_ff: int = 512 |
| max_seq_len: int = 512 |
| dropout: float = 0.0 |
|
|
|
|
| def parameter_count(config: ModelConfig) -> int: |
| |
| d, h, kv, f, v = config.d_model, config.n_heads, config.n_kv_heads, config.d_ff, config.vocab_size |
| attn = d * d + 2 * d * (d * kv // h) + d * d |
| |
| layer = attn + 3 * d * f + 2 * d |
| return v * d + config.n_layers * layer + d |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, c: ModelConfig): |
| super().__init__() |
| assert c.d_model % c.n_heads == 0 and c.n_heads % c.n_kv_heads == 0 |
| self.head_dim = c.d_model // c.n_heads |
| self.n_heads, self.n_kv_heads = c.n_heads, c.n_kv_heads |
| self.q = nn.Linear(c.d_model, c.d_model, bias=False) |
| kv_dim = c.n_kv_heads * self.head_dim |
| self.k = nn.Linear(c.d_model, kv_dim, bias=False) |
| self.v = nn.Linear(c.d_model, kv_dim, bias=False) |
| self.out = nn.Linear(c.d_model, c.d_model, bias=False) |
| self.dropout = c.dropout |
|
|
| def forward(self, x): |
| b, t, _ = x.shape |
| q = self.q(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2) |
| k = self.k(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) |
| v = self.v(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2) |
| if self.n_kv_heads != self.n_heads: |
| repeat = self.n_heads // self.n_kv_heads |
| k, v = k.repeat_interleave(repeat, dim=1), v.repeat_interleave(repeat, dim=1) |
| |
| half = self.head_dim // 2 |
| freq = torch.arange(half, device=x.device, dtype=torch.float32) |
| inv_freq = 1.0 / (10000 ** (freq / half)) |
| pos = torch.arange(t, device=x.device, dtype=torch.float32) |
| angle = torch.outer(pos, inv_freq) |
| cos, sin = angle.cos()[None, None], angle.sin()[None, None] |
|
|
| def rotate(z): |
| z1, z2 = z[..., :half], z[..., half:2 * half] |
| return torch.cat((z1 * cos - z2 * sin, z1 * sin + z2 * cos, z[..., 2 * half:]), dim=-1).to(z.dtype) |
|
|
| q, k = rotate(q), rotate(k) |
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True, |
| dropout_p=self.dropout if self.training else 0.0) |
| return self.out(y.transpose(1, 2).contiguous().view(b, t, -1)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.attn_norm, self.ffn_norm = RMSNorm(c.d_model), RMSNorm(c.d_model) |
| self.attn = Attention(c) |
| self.gate = nn.Linear(c.d_model, c.d_ff, bias=False) |
| self.up = nn.Linear(c.d_model, c.d_ff, bias=False) |
| self.down = nn.Linear(c.d_ff, c.d_model, bias=False) |
| self.dropout = nn.Dropout(c.dropout) |
|
|
| def forward(self, x): |
| x = x + self.dropout(self.attn(self.attn_norm(x))) |
| x = x + self.dropout(self.down(F.silu(self.gate(self.ffn_norm(x))) * self.up(self.ffn_norm(x)))) |
| return x |
|
|
|
|
| class TinyChatLM(nn.Module): |
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.config = config |
| self.token_embedding = nn.Embedding(config.vocab_size, config.d_model) |
| self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)]) |
| self.norm = RMSNorm(config.d_model) |
| self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False) |
| self.lm_head.weight = self.token_embedding.weight |
| self.apply(self._init_weights) |
| for name, p in self.named_parameters(): |
| if name.endswith("out.weight") or name.endswith("down.weight"): |
| nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layers)) |
|
|
| @staticmethod |
| def _init_weights(module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def forward(self, ids, targets=None, loss_mask=None): |
| x = self.token_embedding(ids) |
| for block in self.blocks: |
| x = block(x) |
| logits = self.lm_head(self.norm(x)) |
| loss = None |
| if targets is not None: |
| per_token = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="none") |
| if loss_mask is not None: |
| weights = loss_mask.reshape(-1).to(per_token.dtype) |
| loss = (per_token * weights).sum() / weights.sum().clamp_min(1) |
| else: |
| loss = per_token.mean() |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, ids, max_new_tokens=100, temperature=0.8, top_k=40, eos_id=None): |
| self.eval() |
| for _ in range(max_new_tokens): |
| ctx = ids[:, -self.config.max_seq_len:] |
| logits, _ = self(ctx) |
| logits = logits[:, -1, :] / max(temperature, 1e-5) |
| if top_k: |
| cutoff = torch.topk(logits, min(top_k, logits.size(-1))).values[:, -1:] |
| logits = logits.masked_fill(logits < cutoff, -float("inf")) |
| nxt = torch.multinomial(F.softmax(logits, dim=-1), 1) |
| ids = torch.cat((ids, nxt), dim=1) |
| if eos_id is not None and (nxt == eos_id).all(): |
| break |
| return ids |
|
|
|
|
| def config_dict(config): |
| return asdict(config) |
|
|
