Add source: tmt/experiments/01_baseline.ipynb

tmt/experiments/01_baseline.ipynb

@@ -0,0 +1,110 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["# Experiment 01 — Vanilla Transformer Baseline\n", "Train a standard transformer on the same data as TMT for fair comparison."]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "from torch.optim import AdamW\n",
+    "from tmt.data.dataset import load_text_dataset\n",
+    "from tmt.training.scheduler import cosine_warmup_scheduler\n",
+    "import math\n",
+    "\n",
+    "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "SEQ_LEN = 256\n",
+    "BATCH_SIZE = 16\n",
+    "\n",
+    "loaders = load_text_dataset('wikitext-2', seq_len=SEQ_LEN, batch_size=BATCH_SIZE)\n",
+    "print('Train batches:', len(loaders['train']))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Standard transformer — same param budget as TMT\n",
+    "baseline = nn.Transformer(\n",
+    "    d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,\n",
+    "    dim_feedforward=2048, dropout=0.1, batch_first=True\n",
+    ").to(DEVICE)\n",
+    "print(f'Baseline params: {sum(p.numel() for p in baseline.parameters())/1e6:.2f}M')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Simple GPT-style decoder-only baseline using nn.TransformerDecoder\n",
+    "vocab_size = 50258  # gpt2 tokenizer size\n",
+    "\n",
+    "class BaselineGPT(nn.Module):\n",
+    "    def __init__(self, vocab=vocab_size, d=512, heads=8, layers=6, seq=256):\n",
+    "        super().__init__()\n",
+    "        self.embed = nn.Embedding(vocab, d)\n",
+    "        self.pos = nn.Embedding(seq, d)\n",
+    "        layer = nn.TransformerEncoderLayer(d, heads, dim_feedforward=2048, batch_first=True)\n",
+    "        self.transformer = nn.TransformerEncoder(layer, num_layers=layers)\n",
+    "        self.proj = nn.Linear(d, vocab)\n",
+    "        self.proj.weight = self.embed.weight\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        B, S = x.shape\n",
+    "        pos = torch.arange(S, device=x.device).unsqueeze(0)\n",
+    "        h = self.embed(x) + self.pos(pos)\n",
+    "        mask = nn.Transformer.generate_square_subsequent_mask(S, device=x.device)\n",
+    "        h = self.transformer(h, mask=mask, is_causal=True)\n",
+    "        return self.proj(h)\n",
+    "\n",
+    "baseline = BaselineGPT().to(DEVICE)\n",
+    "print(f'BaselineGPT params: {sum(p.numel() for p in baseline.parameters())/1e6:.2f}M')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "opt = AdamW(baseline.parameters(), lr=3e-4, weight_decay=0.1)\n",
+    "sched = cosine_warmup_scheduler(opt, warmup_steps=200, total_steps=2000)\n",
+    "baseline.train()\n",
+    "\n",
+    "losses = []\n",
+    "for step, batch in enumerate(loaders['train']):\n",
+    "    if step >= 2000:\n",
+    "        break\n",
+    "    ids = batch['input_ids'].to(DEVICE)\n",
+    "    x, y = ids[:, :-1], ids[:, 1:]\n",
+    "    logits = baseline(x)\n",
+    "    loss = nn.functional.cross_entropy(logits.reshape(-1, vocab_size), y.reshape(-1))\n",
+    "    opt.zero_grad(); loss.backward()\n",
+    "    nn.utils.clip_grad_norm_(baseline.parameters(), 1.0)\n",
+    "    opt.step(); sched.step()\n",
+    "    losses.append(loss.item())\n",
+    "    if step % 100 == 0:\n",
+    "        print(f'step={step:4d} loss={loss.item():.4f}')\n",
+    "\n",
+    "baseline_ppl = math.exp(sum(losses[-200:]) / 200)\n",
+    "print(f'\\nBaseline final perplexity: {baseline_ppl:.2f}')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+  "language_info": {"name": "python", "version": "3.10.0"}
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}