rnn_shakespeare_tinygrad.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3facbc61-cb3c-4d7d-bf52-eac3bb4635ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "%pylab is deprecated, use %matplotlib inline and import the required libraries.\n",
      "Populating the interactive namespace from numpy and matplotlib\n",
      "50256 Under the canopy.\n",
      "\n",
      "<|endoftext|>Third Servingman:\n"
     ]
    }
   ],
   "source": [
    "# https://arxiv.org/pdf/2410.01201\n",
    "# Were RNNs All We Needed?\n",
    "%pylab inline\n",
    "from tinygrad import Tensor, fetch, nn, TinyJit\n",
    "from typing import Tuple\n",
    "from tqdm import trange\n",
    "base = \"https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/\"\n",
    "X_train = Tensor(fetch(base+\"tiny_shakespeare_train.bin\"))[0x400:].bitcast('uint16').to(None)\n",
    "X_test = Tensor(fetch(base+\"tiny_shakespeare_val.bin\"))[0x400:].bitcast('uint16').to(None)\n",
    "import tiktoken\n",
    "enc = tiktoken.get_encoding(\"gpt2\")\n",
    "print(X_train.max().item(), enc.decode(X_train[0:10].numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a7a816fe-e51c-4ea6-9498-133a619dda4b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def g(x:Tensor) -> Tensor:\n",
    "  return (x>=0).detach().where(x+0.5, x.sigmoid())\n",
    "def log_g(x:Tensor) -> Tensor:\n",
    "  return (x>=0).detach().where((x.relu()+0.5).log(), (-x).softplus())\n",
    "\n",
    "def logcumsumexp(x, dim): return x.log().cumsum(dim).exp()\n",
    "def parallel_scan_log(log_coeffs, log_values):\n",
    "  # log_coeffs: (batch_size, seq_len, input_size)\n",
    "  # log_values: (batch_size, seq_len+1, input_size)\n",
    "  a_star = log_coeffs.cumsum(-1).pad2d((0,0,1,0))\n",
    "  # NOTE: typo in paper here\n",
    "  log_h0_plus_b_star = logcumsumexp(log_values - a_star, dim=1)\n",
    "  log_h = a_star + log_h0_plus_b_star\n",
    "  return log_h.exp()\n",
    "\n",
    "# TODO: we need shape types (einsum like) for Tensors\n",
    "class MinGRU:\n",
    "  def __init__(self, input_size, hidden_size):\n",
    "    self.linear_z = nn.Linear(input_size, hidden_size)\n",
    "    self.linear_h = nn.Linear(input_size, hidden_size)\n",
    "\n",
    "  def single(self, x_t:Tensor, h_prev:Tensor) -> Tensor:\n",
    "    # x_t: (batch_size, input_size)\n",
    "    # h_prev: (batch_size, hidden_size)\n",
    "    z_t = self.linear_z(x_t).sigmoid()\n",
    "    h_tilde = g(self.linear_h(x_t))\n",
    "    h_t = (1-z_t)*h_prev + z_t*h_tilde\n",
    "    return h_t\n",
    "\n",
    "  def __call__(self, x:Tensor, h_0:Tensor):\n",
    "    # x: (batch_size, seq_len, input_size)\n",
    "    # h_0: (batch_size, 1, hidden_size)\n",
    "    k = self.linear_z(x)\n",
    "    log_z = -(-k).softplus()\n",
    "    log_coeffs = -k.softplus()\n",
    "    log_h_0 = log_g(h_0)\n",
    "    log_tilde_h = log_g(self.linear_h(x))\n",
    "    h = parallel_scan_log(log_coeffs, Tensor.cat(log_h_0, log_z + log_tilde_h, dim=1))\n",
    "    return h\n",
    "\n",
    "tokens = 50257\n",
    "hidden_size = 384\n",
    "class Model:\n",
    "  def __init__(self):\n",
    "    self.embedding = nn.Embedding(tokens, 384)\n",
    "    self.gru = MinGRU(384, hidden_size)\n",
    "    self.mlp1 = nn.Linear(384, 384)\n",
    "    self.mlp2 = nn.Linear(384, tokens)\n",
    "    self.reset()\n",
    "\n",
    "  def reset(self): self.h_prev = Tensor.zeros(1, hidden_size, requires_grad=False)\n",
    "  def get_logits(self, x:Tensor) -> Tuple[Tensor, Tensor]:\n",
    "    # TODO: fix bug in embedding shape\n",
    "    h_t = self.gru.single(self.embedding(x)[:, 0], self.h_prev)\n",
    "    logits = self.mlp2(self.mlp1(self.h_prev).relu())\n",
    "    return logits, h_t\n",
    "      \n",
    "  def single(self, x:Tensor) -> Tensor:\n",
    "    logits, self.h_prev = self.get_logits(x)\n",
    "    return logits[0].exp().softmax(-1).multinomial()\n",
    "      \n",
    "  def __call__(self, x:Tensor):\n",
    "    h_prev = Tensor.zeros(x.shape[0], 1, hidden_size)\n",
    "    pc = self.gru(self.embedding(x), h_prev)[:, 1:]\n",
    "    return self.mlp2(self.mlp1(pc).relu())\n",
    "\n",
    "# TODO: it seems like there's an issue with the learning rate and a mean somewhere, this is 100x bigger than the paper\n",
    "model = Model()\n",
    "optim = nn.optim.AdamW(nn.state.get_parameters(model), lr=1e-1)\n",
    "losses = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3039ad72-0596-4b0b-80a0-9f4f8b476ae2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss: 6.5902 [13650, 166037, 220854]:  85%|███████████████████████████████████████████████████▊         | 85/100 [02:08<00:22,  1.47s/it]"
     ]
    }
   ],
   "source": [
    "batch_size, seq_length = 32, 200\n",
    "\n",
    "@TinyJit\n",
    "@Tensor.train()\n",
    "def train_step() -> Tensor:\n",
    "  sel = Tensor.randint((batch_size,), low=0, high=X_train.shape[0]-seq_length)\n",
    "  X = X_train[sel.reshape(-1, 1)+Tensor.arange(seq_length+1).reshape(1, -1)]\n",
    "  optim.zero_grad()\n",
    "  loss = model(X[:, :-1]).sparse_categorical_crossentropy(X[:, 1:]).backward()\n",
    "  optim.step()\n",
    "  return loss, sel[:3]\n",
    "\n",
    "for i in (t:=trange(100)):\n",
    "  loss, sel = train_step()\n",
    "  losses.append(loss.item())\n",
    "  t.set_description(f\"loss: {losses[-1]:.4f} {sel.tolist()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23a09e6b-8fe5-4c82-9db2-886fa82f97c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.reset()\n",
    "arr = enc.encode(\"hello\")\n",
    "for i in range(10): arr.append(model.single(Tensor([arr[-1]])).item())\n",
    "print(arr, enc.decode(arr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8728a685-cbff-421e-a03b-c63fbb13d202",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot(losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399d046c-9a25-4881-8fe7-cfbdd64c429a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}