nanogpt-experiments/scaling_laws.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Reproducing some results from [Chinchilla](https://arxiv.org/pdf/2203.15556.pdf). I can't get the numbers to match exactly, please open an Issue if you understand why the results are off by a bit. Current running hypothesis is that this is because I am using the FLOPs = 6\\*N\\*D formula, instead of taking all the Gopher sizes and calculating their FLOPs individually, and using that to interpolate?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpt2 = dict(\n",
    "    seq_len = 1024,\n",
    "    vocab_size = 50257,\n",
    "    d_model = 768,\n",
    "    num_heads = 12,\n",
    "    num_layers = 12,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "123.653376"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def params(seq_len, vocab_size, d_model, num_heads, num_layers, ffw_size=4):\n",
    "    \"\"\" Given GPT config calculate total number of parameters\"\"\" \n",
    "    # token and position embeddings\n",
    "    embeddings = d_model * vocab_size + d_model * seq_len\n",
    "    # transformer blocks\n",
    "    attention = 3*d_model**2 + 3*d_model # weights and biases\n",
    "    attproj = d_model**2 + d_model\n",
    "    ffw = d_model*(ffw_size*d_model) + ffw_size*d_model\n",
    "    ffwproj = ffw_size*d_model*d_model + d_model\n",
    "    layernorms = 2*2*d_model\n",
    "    # dense\n",
    "    ln_f = 2*d_model\n",
    "    dense = d_model*vocab_size # note: no bias here\n",
    "    # note: embeddings are not included in the param count!\n",
    "    total_params = num_layers*(attention + attproj + ffw + ffwproj + layernorms) + ln_f + dense\n",
    "    return total_params\n",
    "\n",
    "params(**gpt2)/1e6 # OpenAI reports gpt2 (small) as having 124M params, good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "766.006788096"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def flops(seq_len, vocab_size, d_model, num_heads, num_layers, ffw_size=4):\n",
    "    \"\"\" \n",
    "    Given GPT config calculate total number of FLOPs, see Chinchilla \n",
    "    paper Appendix F as reference: https://arxiv.org/pdf/2203.15556.pdf\n",
    "    \"\"\" \n",
    "    key_size = d_model // num_heads\n",
    "\n",
    "    # embeddings\n",
    "    embeddings = 2 * seq_len * vocab_size * d_model\n",
    "\n",
    "    # attention\n",
    "    # key, query, value projections\n",
    "    attention = 2 * 3 * seq_len * d_model * (key_size * num_heads)\n",
    "    # key @ query logits\n",
    "    attlogits = 2 * seq_len * seq_len * (key_size * num_heads)\n",
    "    # softmax\n",
    "    attsoftmax = 3 * num_heads * seq_len * seq_len # TODO why?\n",
    "    # softmax @ value reductions\n",
    "    attvalue = 2 * seq_len * seq_len * (key_size * num_heads)\n",
    "    # final linear\n",
    "    attlinear = 2 * seq_len * (key_size * num_heads) * d_model\n",
    "    att = attention + attlogits + attsoftmax + attvalue + attlinear\n",
    "    # feed forward\n",
    "    dense = 2 * seq_len * (d_model * ffw_size + d_model * ffw_size)\n",
    "\n",
    "    # logits\n",
    "    logits = 2 * seq_len * d_model * vocab_size\n",
    "\n",
    "    forward_flops = embeddings + num_layers * (att + dense) + logits\n",
    "    backward_flops = 2 * forward_flops # as in Kaplan et al. 2020\n",
    "    total_flops = forward_flops + backward_flops\n",
    "\n",
    "    return total_flops\n",
    "\n",
    "flops(**gpt2)/1e9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "params: 69.72M\n",
      "approx flops: 428.34B\n",
      "chinchilla flops: 434.11B\n",
      "ratio (chinchilla / approx): 1.01\n"
     ]
    }
   ],
   "source": [
    "# Reproduce Table A4 from Chinchilla paper Appendix, \n",
    "# comparing accurate flops above to approximate flops F = 6*N*D\n",
    "# note Chinchilla uses vocab_size = 32K\n",
    "\n",
    "chin_73M = dict(seq_len = 1024, vocab_size = 32000, d_model = 640, num_heads = 10, num_layers = 10)\n",
    "args = chin_73M\n",
    "\n",
    "D = 1024 # dataset size, cancels anyway\n",
    "N = params(**args) \n",
    "F = flops(**args)\n",
    "\n",
    "approx_flops = 6*D*N\n",
    "chinch_flops = F * (float(D) / args['seq_len'])\n",
    "\n",
    "print(f\"params: {N/1e6:.2f}M\")\n",
    "print(f\"approx flops: {approx_flops/1e9:.2f}B\")\n",
    "print(f\"chinchilla flops: {chinch_flops/1e9:.2f}B\")\n",
    "print(f\"ratio (chinchilla / approx): {chinch_flops / approx_flops:.2f}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well this is awkward because Chinchilla paper claims that number of params for these args are 73M (we see only 70M), and the ratio is supposed to be 1.03 but we get 1.01. TODO stare at more..."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## scaling laws"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.colorbar.Colorbar at 0x7fd5bc48bb80>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABDQAAAHWCAYAAACIWtlUAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8o6BhiAAAACXBIWXMAAA9hAAAPYQGoP6dpAACWeElEQVR4nO3deXxU1f3/8fckkAUkURAIS9gVEAHBNWAFK4JIEap14WsL7tqCirRa6a8qiBRtXXDFnbhRVwQrFgQUEFmUrQUXFGXTJuAGAYQAM/f3B2XqkPkMueFOZubO6/l43Ec7d86ce+5MyHxyvPe8A47jOAIAAAAAAEghGYkeAAAAAAAAgFtMaAAAAAAAgJTDhAYAAAAAAEg5TGgAAAAAAICUw4QGAAAAAABIOUxoAAAAAACAlMOEBgAAAAAASDlMaAAAAAAAgJTDhAYAAAAAAEg5TGgAMRQXFysQCGjdunWJHgoAAIAv+bHe+tvf/qZWrVopMzNTxx13nCSpRYsWuuSSSxI6LsBvmNAAAAAA4HsTJkzQ+eefr2bNmikQCMScXNiyZYuuuuoq1a9fX7Vr19bpp5+uZcuWVeo4b7/9tm666SZ1795dEydO1F/+8hePzgDAgWokegAAAAAAEG933XWXtm3bppNOOkklJSVmu1AopH79+ulf//qXbrzxRh155JF65JFH1LNnTy1dulRHHXVUzOO88847ysjI0FNPPaWsrCyvTwPATzChAQAAAMD35s6dG74647DDDjPbvfrqq1qwYIFeeeUV/epXv5IkXXDBBTr66KN12223adKkSTGPs3nzZuXm5jKZAVQDbjkBXHrkkUfUoUMHZWdnq3Hjxho6dKi2bNkS0ebzzz/Xeeedp4KCAuXk5Khp06a66KKLtHXr1nCbmTNn6tRTT9Xhhx+uww47TG3bttWf/vSnaj4bAACA5FSZmkuSHn74YbVq1Uq5ubk66aST9N5776lnz57q2bNnRLvmzZsrEAgc9LivvvqqGjZsqHPPPTe8r379+rrgggs0depUlZeXm68NBAKaOHGiduzYoUAgoEAgoOLiYrP9l19+qfPPP19169ZVrVq1dMopp2jatGkRbebMmaNAIKCXXnpJf/rTn1RQUKDatWvrnHPO0caNGyPaVqYGBfyEKzQAF0aNGqXRo0erV69e+u1vf6vVq1drwoQJ+vDDD/X++++rZs2a2r17t/r06aPy8nJde+21Kigo0Ndff60333xTW7ZsUX5+vj766CP94he/UKdOnXT77bcrOztba9as0fvvv5/oUwQAAEi4ytRc0r51MYYNG6af/exnuuGGG7Ru3ToNHDhQRxxxhJo2bVqlYy9fvlxdu3ZVRkbkf/s96aST9Pjjj+uzzz5Tx44do772ueee0+OPP64PPvhATz75pCSpW7duUdtu2rRJ3bp1048//qjrrrtO9erV0zPPPKNzzjlHr776qn75y19GtB87dqwCgYD++Mc/avPmzRo/frx69eqlFStWKDc3t1I1KOA7DgDTxIkTHUnO2rVrnc2bNztZWVlO7969nWAwGG7z0EMPOZKcp59+2nEcx1m+fLkjyXnllVfMfu+77z5HkvPNN9/E/RwAAACS2U/rLcdxKl1zlZeXO/Xq1XNOPPFEZ8+ePeF2xcXFjiSnR48e5jFr167tDBkyxHzusssuq7B/2rRpjiRn+vTpMc9nyJAhTu3atSvsb968ecQxhw8f7khy3nvvvfC+bdu2OS1btnRatGgRPvd3333XkeQ0adLEKSsrC7d9+eWXHUnO/fff7zhO5WpQwG+45QSopFmzZmn37t0aPnx4xIz9lVdeqby8vPDlgftnv2fMmKEff/wxal+HH364JGnq1KkKhULxHTgAAEAKqWzNtWTJEn333Xe68sorVaPG/y48v/jii3XEEUdU+fg7d+5UdnZ2hf05OTnh573w1ltv6aSTTtKpp54a3nfYYYfpqquu0rp16/Txxx9HtB88eLDq1KkTfvyrX/1KjRo10ltvvSWpcjUo4DdMaACVtH79eklS27ZtI/ZnZWWpVatW4edbtmypESNG6Mknn9SRRx6pPn366OGHH464d/HCCy9U9+7ddcUVV6hhw4a66KKL9PLLLzO5AQAA0l5la679/9umTZuIdjVq1FCLFi2qfPzc3Nyo62Ts2rUr/LwX1q9fX+EcJal9+/bh53/qwHSVQCCgNm3aaN26dZIqV4MCfsOEBhAH99xzj/7973/rT3/6k3bu3KnrrrtOHTp00FdffSVp3xfhvHnzNGvWLP3mN7/Rv//9b1144YU688wzFQwGEzx6AACA9NWoUaOosa779zVu3Li6h1RpB6tBAb9hQgOopObNm0uSVq9eHbF/9+7dWrt2bfj5/Tp27Kg///nPmjdvnt577z19/fXXevTRR8PPZ2Rk6IwzztC9996rjz/+WGPHjtU777yjd999N/4nAwAAkKQqW3Pt/981a9ZEtNu7d2/4qoWqOO6447Rs2bIKV84uXrxYtWrV0tFHH13lvn+qefPmFc5Rkj799NPw8z/1+eefRzx2HEdr1qypcDXKwWpQwE+Y0AAqqVevXsrKytIDDzwgx3HC+5966ilt3bpV/fr1kySVlZVp7969Ea/t2LGjMjIywpcvfv/99xX6P+644yQpZhQYAACA31W25jrhhBNUr149PfHEExG11wsvvKAffvihysf/1a9+pU2bNmny5Mnhfd9++61eeeUV9e/fP+r6GlVx9tln64MPPtDChQvD+3bs2KHHH39cLVq00DHHHBPR/tlnn9W2bdvCj1999VWVlJSob9++kipXgwJ+Q2wrUEn169fXyJEjNXr0aJ111lk655xztHr1aj3yyCM68cQT9etf/1qS9M4772jYsGE6//zzdfTRR2vv3r167rnnlJmZqfPOO0+SdPvtt2vevHnq16+fmjdvrs2bN+uRRx5R06ZNIxaGAgAASDeVrbmysrI0atQoXXvttfr5z3+uCy64QOvWrVNxcbFat26tQCAQ0e8//vEP/etf/5Ik7dmzR//+9791xx13SJLOOeccderUSdK+CY1TTjlFl156qT7++GMdeeSReuSRRxQMBjV69GjPzvPmm2/W3//+d/Xt21fXXXed6tatq2eeeUZr167Va6+9ViE2tm7dujr11FN16aWXatOmTRo/frzatGmjK6+8UlLlalDAb5jQAFwYNWqU6tevr4ceekg33HCD6tatq6uuukp/+ctfwnnonTt3Vp8+ffSPf/xDX3/9tWrVqqXOnTvrn//8p0455RRJ+740161bp6efflrffvutjjzySPXo0UOjR48mIxwAAKS9ytRckjRs2DA5jqN77rlHf/jDH9S5c2e98cYbuu6668KpJPu99tpreuaZZ8KPly9fruXLl0uSmjZtGp7QyMzM1FtvvaUbb7xRDzzwgHbu3KkTTzxRxcXFURfxrKqGDRtqwYIF+uMf/6gHH3xQu3btUqdOnfSPf/wjfBXKT/3pT3/Sv//9b40bN07btm3TGWecoUceeUS1atWSVLkaFPCbgPPT67gAAAAAIIWFQiHVr19f5557rp544olED+eQzZkzR6effrpeeeUV/epXv0r0cICkwhoaAAAAAFLSrl27dOB/n3322Wf1/fffq2fPnokZFIBqwy0nAAAAAFLSokWLdMMNN+j8889XvXr1tGzZMj311FM69thjdf755yd6eADijAkNAAAAACmpRYsWKiws1AMPPKDvv/9edevW1eDBg3XnnXcqKysr0cMDEGcJveVk3rx56t+/vxo3bqxAIKApU6ZEPD958mT17t1b9erVUyAQ0IoVKyrV7yuvvKJ27dopJydHHTt21FtvveX94AEAAGCizkN1aNGihd544w2VlpZq9+7dKi0t1dNPP60GDRokemie6dmzpxzHYf0MIIqETmjs2LFDnTt31sMPP2w+f+qpp+quu+6qdJ8LFizQoEGDdPnll2v58uUaOHCgBg4cqFWrVnk1bAAAABwEdR4AIN6SJuUkEAjo9ddf18CBAys8t27dOrVs2VLLly/XcccdF7OfCy+8UDt27NC
      "text/plain": [
       "<Figure size 1500x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def L(N, D):\n",
    "    \"\"\" \n",
    "    Approximates loss given N parameters and D dataset size (in tokens),\n",
    "    per Chinchilla paper.\n",
    "    \"\"\"\n",
    "    E = 1.69 # entropy of natural language, limit of infinite model on infinite data\n",
    "    A = 406.4\n",
    "    B = 410.7\n",
    "    alpha = 0.34\n",
    "    beta = 0.28\n",
    "    return A / (N ** alpha) + B / (D ** beta) + E\n",
    "\n",
    "ns = 10 ** np.arange(7, 11, step=2**-4) # model sizes from 10M to 100B\n",
    "ds = 10 ** np.arange(9, 12, step=2**-4) # dataset sizes from 1B to 1T\n",
    "plt.figure(figsize=(15, 5))\n",
    "# plot a heatmap of loss as a function of model size and dataset size\n",
    "plt.subplot(121)\n",
    "plt.imshow(np.array([[L(n, d) for d in ds] for n in ns]), extent=[9, 12, 7, 11], origin='lower')\n",
    "plt.xlabel('log10(dataset size)')\n",
    "plt.ylabel('log10(model size)')\n",
    "plt.title('loss')\n",
    "plt.colorbar()\n",
    "# plot the compute for each point, which is a deterministic function: flops = 6*N*D\n",
    "plt.subplot(122)\n",
    "plt.imshow(np.log10(np.array([[6*n*d for d in ds] for n in ns])), extent=[9, 12, 7, 11], origin='lower')\n",
    "plt.xlabel('log10(dataset size)')\n",
    "plt.ylabel('log10(model size)')\n",
    "plt.title('log10 flops')\n",
    "plt.colorbar()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok so given any N,D we can estimate both: 1) the loss, and 2) the total flops. Now we want to solve the following problem: Given a specific budget of flops C, find: N_opt, D_opt = argmin_{FLOPs(N,D) = C} L(N, D). i.e. how big of a model should we train and for how many tokens?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "best model size: 316.23M\n",
      "best dataset size: 10.12B\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'loss')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS8AAAEqCAYAAABEE9ZrAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8o6BhiAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAuW0lEQVR4nO3deViU9f7/8ecMy7Ajm4AspqIoKoqaCpor5m7W6ZfHLKy0Tp0sy7IT1Tl1Wg6ezJZz8tt6ytZjqannWC7pSS3FDTTRVMIFEFkUkWWQbeb+/QFOoiAIw9wzzPtxXXNdcnPPzHuY5tXn/sxn0SiKoiCEEDZGq3YBQgjREhJeQgibJOElhLBJEl5CCJsk4SWEsEkSXkIImyThJYSwSY5qF2BpRqORM2fO4OnpiUajUbscIcQVFEWhtLSUTp06odU23r6yu/A6c+YMYWFhapchhGhCdnY2oaGhjf7e7sLL09MTqP3DeHl5qVxNO6bXQ6dOtf8+cwbc3dWtR9iMkpISwsLCTJ/VxthdeF26VPTy8pLwaksODr/928tLwktct6a6daTDXghhkyS8hBA2ScJLCGGTJLyEEDZJwksIYZMkvIQQNknCSwjR5j7ecZK1B3Ioqag222NKeAkh2lS1wciSTenMX36ArMJysz2uhJcQok0dyL5AWWUNvu7ORAWbb2C4hJcQok39mH4WgLhufmi15lsMQcJLCNGmfsw4B8CI7gFmfVwJLyFEmykur+bn7AsADO/ub9bHlvASQrSZ5BPnMCrQLcCdTh1czfrYEl5CiDaz/dfaS8abzHzJCBJeQog29JMpvMx7yQgSXkKINpJZqCfrfDmOWg1DuvqZ/fElvIQQbeLHulbXgM4+eOjMv+6phJcQok2YLhkjzH/JCBJeQog2UGMwsuN4XXj1MH9nPUh4CSHawMGcYkoravB2daJviHebPIeElxDC7LYdq50SNCzCDwczTgm6nISXEMLsttbNZxzVo2ObPYeElxDCrM7rqzh4+gIAI9qovwskvIQQZvbjr2dRFOgZ5EmQt0ubPY+ElxDCrC71d42MbLtWF0h4CSHMyGhU2GaB/i6Q8BJCmNHhMyUU6qtwd3ZgYGefNn0uCS8hhNlsPVYAwLAIf5wd2zZeJLyEEGZjumSMbNtLRpDwEkKYSXF5NalZRUDbd9aDhJcQwkx+yqhdNbV7Rw9CzLxqakMkvIQQZnGpv2tkGw5MvZyElxCi1YxGhR+OWa6/CyS8hBBmkJZTzLmySjx0jgzu4muR55TwEkK02pajtZeMI3q0/RCJSyS8hBCt9r+j+QCMttAlI0h4CSFaKb+kgkM5JWg0luvvAgkvIUQr/VB3ydgvtAMBnjqLPa+ElxCiVS71d43tablWF0h4CSFaoaLaYNolaEwvCS8hhI3YdaKQi9UGgrxciAr2suhzS3gJIVrsUn/X6J4d0WjaZqONxkh4CSFaRFEU1fq7QMJLCNFCx/JLOV10EZ2jlrgIP4s/v9WE16JFi9BoNDz22GONnrNs2TI0Gk29m4tL2y3wL4Ro3KbDtQNTb+ruj5uzo8Wf3/LP2IC9e/fy3nvvER0d3eS5Xl5eHDt2zPSzpa+zhRC1vv+lNrzGRQWq8vyqt7zKysqYNWsWH3zwAT4+Ta95rdFoCAoKMt0CA9X5wwlhz3KLL5KWU4xGA2N62ml4Pfzww0yePJn4+PhmnV9WVkbnzp0JCwvjlltu4fDhw9c8v7KykpKSkno3IUTrbK5rdQ0I97HoqPrLqRpey5cvJzU1laSkpGadHxkZyUcffcTatWv5/PPPMRqNxMXFcfr06Ubvk5SUhLe3t+kWFhZmrvKFsFub6sLrZpUuGUHF8MrOzmb+/Pl88cUXze50j42NJSEhgf79+zNy5Ei++eYbAgICeO+99xq9T2JiIsXFxaZbdna2uV6CEHappKKaXScKAfX6u0DFDvuUlBQKCgoYMGCA6ZjBYGD79u28/fbbVFZW4uDgcM3HcHJyIiYmhoyMjEbP0el06HTqNGuFaI+2HjtLtUGhW4A7XQM8VKtDtfAaO3YsaWlp9Y7de++99OzZkz/96U9NBhfUhl1aWhqTJk1qqzKFEFf47VvGIFXrUC28PD096dOnT71j7u7u+Pn5mY4nJCQQEhJi6hN78cUXGTp0KBEREVy4cIHFixeTmZnJ3LlzLV6/EPaoqsbI1rpR9WpeMoKVjPNqTFZWFlrtb91yRUVF3H///eTl5eHj48PAgQPZuXMnUVFRKlYphP3YfbKQ0soa/D10xIR1ULUWjaIoiqoVWFhJSQne3t4UFxfj5WXZWfB2Ra8Hj7r+kLIycHdXtx5hFs+uTuOL3VnMHBxG0m1NDypvieZ+RlUf5yWEsA0Go8LGuilBE/oEq1yNhJcQoplSMos4V1aJl4sjsV0tPxH7ShJeQohm2XAoD4D4qECLbW92LepXIISweoqisPFwbXhN6K3uEIlLJLyEEE06eLqYnAsXcXN2YESPALXLASS8hBDNsL7uknF0z464ODU9gNwSJLyEENekKAobDuUCMLGPdVwygoSXEKIJx/JLOVVYjrOj1qI7YjdFwksIcU3r02ovGUd0D8BDZz2TciS8hBDXtN4KLxlBwksIcQ3p+aWk55fh7KAlXuWJ2FeS8BJCNGrdwdpW14ge/ni7OqlcTX0SXkKIBimKwrqDZwCYEt1J5WquJuElhGjQkdxSTpzV4+yoZWwv6/mW8RIJLyFEg75Nq211jY4MwNPFui4ZQcJLCNGA2kvG2v4ua7xkBAkvIUQDDp8pIbOwHBcnLWN6Wt8lI0h4CSEa8N+6jvqxPQNxt6KBqZeT8BJC1KMoCt/WXTJOjlZ/xdTGSHgJIerZn32B00W1y9+MtqK5jFeS8BJC1LN2fw4A43sH4epsHcvfNETCSwhhUmMwmr5lnNbfOr9lvETCSwhh8lPGOQr1Vfi5OzM8wl/tcq5JwksIYfKfA7XfMk6ODsbJwbrjwbqrE0JYzMUqg2mTjVv6h6hcTdMkvIQQAGw+ko++ykCYrysDwjuoXU6TJLyEEACsrbtkvKVfCBqNRuVqmibhJYTgQnkV29ILALjFyr9lvETCSwjBt2m5VBsUooK96B7oqXY5zSLhJYRgVcppAKbH2EarCyS8hLB7J8/pSc26gFYD023gW8ZLJLyEsHPfpNa2ukb0CKCjl4vK1TSfhJcQdsxoVPgmtXYu4+8GhKpczfWR8BLCju06WUjOhYt4ujgyzsq2NmuKhJcQdmxVSm2ra0p0MC5O1ruCREMkvISwU+VVNabdsG3tkhEkvISwWxsO5VFeZeAGPzcGdvZRu5zrJuElhJ1aVfct420DQm1iOtCVJLyEsEPZ58vZkVGIRgO3xtjO2K7LSXgJYYdW7MsGYHiEP2G+bipX0zISXkLYGYNRYUXddKA7BoWpXE3LWU14LVq0CI1Gw2OPPXbN81asWEHPnj1xcXGhb9++fPfdd5YpUIh2YvuvZ8ktrqCDmxM397atsV2Xs4rw2rt3L++99x7R0dHXPG/nzp3MnDmTOXPmsH//fqZPn8706dM5dOiQhSoVwvZ9vbf2knF6/xB0jrY1tutyqodXWVkZs2bN4oMPPsDH59pf17711ltMmDCBhQsX0qtXL1566SUGDBjA22+/baFqhbBthWWVbD6SD8CMG233khGsILwefvhhJk+eTHx8fJPnJicnX3Xe+PHjSU5ObvQ+lZWVlJSU1LsJYa9W78+h2qAQHepNr2AvtctpFUc1n3z58uWkpqayd+/eZp2fl5dHYGD9a/TAwEDy8vIavU9SUhJ//etfW1WnEO2Boigsr7tktPVWF6jY8srOzmb+/Pl88cUXuLi03TIciYmJFBcXm27Z2dlt9lxCWLPUrCIyCspwcdIytZ/tLDrYGNVaXikpKRQUFDBgwADTMYPBwPbt23n77beprKzEwaF+Z2JQUBD5+fn1juXn5xMUFNTo8+h0OnQ6nXmLF8IGfbErC4Cp0Z3wcnFSuZrWU63lNXbsWNLS0jhw4IDpNmjQIGbNmsWBAweuCi6A2NhYtmzZUu/Y999/T2xsrKXKFsI
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "c = 1.92e19 # target compute budget (usually know this because we know how many GPU for how long go brrr)\n",
    "# sweep model sizes from 10M to 100B\n",
    "ns = 10 ** np.arange(7, 11, step=2**-4)\n",
    "# using C = 6*N*D, solve for D that maintains the compute budget c\n",
    "ds = c / (6 * ns)\n",
    "# evaluate the loss in each case\n",
    "losses = L(ns, ds)\n",
    "# find the argmin\n",
    "best = np.argmin(losses)\n",
    "print(f\"best model size: {ns[best]/1e6:.2f}M\")\n",
    "print(f\"best dataset size: {ds[best]/1e9:.2f}B\")\n",
    "# plot the loss\n",
    "plt.figure(figsize=(3,3))\n",
    "plt.plot(ns, losses)\n",
    "plt.xscale('log')\n",
    "# plot a vertical bar at the best model size\n",
    "plt.axvline(ns[best], color='red')\n",
    "plt.xlabel('model size')\n",
    "plt.ylabel('loss')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the plot above, basically the models on the left of best are too small and trained for too long. The models on the right of best are way too large and trained for too little. The model at the red line is just right.\n",
    "\n",
    "Now, the Chinchilla paper says that best model size is 400M params and 8B tokens, so this disagrees and there is some calculations problem. TODO figure out and fix..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "compute budget 1.000000e+17: best model size: 29.43M, best dataset size: 0.57B\n",
      "compute budget 1.778279e+17: best model size: 36.52M, best dataset size: 0.81B\n",
      "compute budget 3.162278e+17: best model size: 48.70M, best dataset size: 1.08B\n",
      "compute budget 5.623413e+17: best model size: 60.43M, best dataset size: 1.55B\n",
      "compute budget 1.000000e+18: best model size: 80.58M, best dataset size: 2.07B\n",
      "compute budget 1.778279e+18: best model size: 107.46M, best dataset size: 2.76B\n",
      "compute budget 3.162278e+18: best model size: 133.35M, best dataset size: 3.95B\n",
      "compute budget 5.623413e+18: best model size: 177.83M, best dataset size: 5.27B\n",
      "compute budget 1.000000e+19: best model size: 220.67M, best dataset size: 7.55B\n",
      "compute budget 1.778279e+19: best model size: 294.27M, best dataset size: 10.07B\n",
      "compute budget 3.162278e+19: best model size: 392.42M, best dataset size: 13.43B\n",
      "compute budget 5.623413e+19: best model size: 486.97M, best dataset size: 19.25B\n",
      "compute budget 1.000000e+20: best model size: 649.38M, best dataset size: 25.67B\n",
      "compute budget 1.778279e+20: best model size: 865.96M, best dataset size: 34.23B\n",
      "compute budget 3.162278e+20: best model size: 1074.61M, best dataset size: 49.05B\n",
      "compute budget 5.623413e+20: best model size: 1433.01M, best dataset size: 65.40B\n",
      "compute budget 1.000000e+21: best model size: 1778.28M, best dataset size: 93.72B\n",
      "compute budget 1.778279e+21: best model size: 2371.37M, best dataset size: 124.98B\n",
      "compute budget 3.162278e+21: best model size: 3162.28M, best dataset size: 166.67B\n",
      "compute budget 5.623413e+21: best model size: 3924.19M, best dataset size: 238.84B\n",
      "compute budget 1.000000e+22: best model size: 5232.99M, best dataset size: 318.49B\n",
      "compute budget 1.778279e+22: best model size: 6493.82M, best dataset size: 456.40B\n",
      "compute budget 3.162278e+22: best model size: 8659.64M, best dataset size: 608.62B\n",
      "compute budget 5.623413e+22: best model size: 11547.82M, best dataset size: 811.61B\n",
      "compute budget 1.000000e+23: best model size: 14330.13M, best dataset size: 1163.05B\n",
      "compute budget 1.778279e+23: best model size: 19109.53M, best dataset size: 1550.95B\n",
      "compute budget 3.162278e+23: best model size: 23713.74M, best dataset size: 2222.54B\n",
      "compute budget 5.623413e+23: best model size: 31622.78M, best dataset size: 2963.80B\n",
      "compute budget 1.000000e+24: best model size: 42169.65M, best dataset size: 3952.29B\n",
      "compute budget 1.778279e+24: best model size: 52329.91M, best dataset size: 5663.68B\n",
      "compute budget 3.162278e+24: best model size: 69783.06M, best dataset size: 7552.64B\n",
      "compute budget 5.623413e+24: best model size: 93057.20M, best dataset size: 10071.61B\n",
      "compute budget 1.000000e+25: best model size: 115478.20M, best dataset size: 14432.74B\n",
      "compute budget 1.778279e+25: best model size: 153992.65M, best dataset size: 19246.37B\n",
      "compute budget 3.162278e+25: best model size: 191095.30M, best dataset size: 27580.28B\n",
      "compute budget 5.623413e+25: best model size: 254829.67M, best dataset size: 36778.90B\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.legend.Legend at 0x7fd5bc240070>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAA1gAAAEqCAYAAAAI6vFLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8o6BhiAAAACXBIWXMAAA9hAAAPYQGoP6dpAACGMklEQVR4nO3deXxU1fn48c+dyWSdzCQh+waBsEsSCKK4AUJVbLVg9autC2JLa8XlJ2ILtS60KlaqpVX6xdoq2trWrdpWW79WSsBWrCWQgMiWECD7npnMZJ2Z+/tjkoFhAmS/M8nzfr3mleTOmTtnMg+XPHPOeY6iqqqKEEIIIYQQQogB02ndASGEEEIIIYQYKSTBEkIIIYQQQohBIgmWEEIIIYQQQgwSSbCEEEIIIYQQYpBIgiWEEEIIIYQQg0QSLCGEEEIIIYQYJJJgCSGEEEIIIcQgCdK6A/7K5XJRUVFBZGQkiqJo3R0hhBBCCCGERlRVpbm5meTkZHS6s49RSYJ1BhUVFaSlpWndDSGEEEIIIYSfKC0tJTU19axtJME6zaZNm9i0aRMOhwOA999/n4iICI17JQbC6XRSXFzMhAkT0Ov1WndHjFISh0JrEoNCaxKDQmsDiUG73c6Xv/xlIiMjz9lWUVVV7W8nRzKr1YrZbKa2thaTyaR1d8QAOJ1OCgoKyMnJkQu60IzEodCaxKDQmsSg0NpAYtBqtRIXF4fFYjlnbiAjWKfpHsFyOp0A7N+/H6PRqHGvxEB0f1qh1+vlgi40I3EotCYxKLQmMSi0NpAYtNlsvW4rCdZpVq5cycqVKz0jWFlZWTKCFeC6k2X5xExoSeJQaE1iUGhNYlBobSAxaLVae91WyrQLIYQQQgghxCCREazTnD5FcO/evTJFMMA5nU6OHTsGIJ+YCc04nU4OHz4MSBwKbUgMCq1JDAqtDSQG+zJFUIpcnEH3FMGGhoYepwh2dHRw/PhxXC6XBr0TfdXR0UFwcPCQPofZbCYhIUH2TRM9ksXdQmsSg0JrEoNCawMtchETEyNFLoaKqqpUV1cTFBREUlLSOTcbE9pSVZW2tjZCQ0OHJPlRVZWWlhZqa2sBSExMHPTnEEIIIYQQgUESrNP0doqgXq8nKSkJQEaxAoCqqp7bUAgLC8NsNlNdXU1VVdWQPIcIbDI1RmhNYlBoTWJQ9JbibCfccoSIxgPYo6dijzlvUM47XFMEJcE6TW+qCLa1tXH8+HEiIyMJCwvTqKeit7qTqrCwsCGdvqcoCk1NTYwdO5bQ0NAhex4RmKR6ltCaxKDQmsSg6JHqgvoilPJ8KM93f63Zj+JyAOCa8x3UnFsG5amGq4qgJFjn0FOdfL1ej6Io6HQ6WW8TQBRFGdL3qzseZH8PcSbdsSHxIbQiMSi0JjEosNVAeT6U7YLyXVC+B9otvu0i4iF1NrqUWTCI8dLfGOxLe0mwhBBCCCGEEIOvowUqC90JVfkuKMsHywnfdkFhkJwDKbmQOhtSZoM5FQJ0IEMSrHNwOp2e4cRTjw31mh4xeLrfo6F+r7rjoaeYEaI7LiQ2hFYkBoXWJAZHONUFdYfdU/wqdru/Vu9HUb3fbxUFYiehpuRCymz31/ipoDstLRmCGgcDicG+PEYSrNP0tshFUFAQra2tUuDiDK666iqysrJ4+umne9X+d7/7Hd///vcpLy/v1/MZjUb+8Ic/cM011/R4f1tbW7/O2xft7e10dnZy8ODBIX8uEXhkcbfQmsSg0JrE4MhjaK3BVJuPqS6fyNo9GDoafdp0hsRgj5rqLlYRPQW7eTIuQ8TJBlVOqPp8WPorRS400pciF2FhYVLM4Az0ej1BQUGEh4f3qn33HlW9bX+6iooKoqOjCQkJ8blvuIpc6HQ6DAYDmZmZEhfChyzuFlqTGBRakxgcAdqscPxfKEfzUEq2o9Qf8bpb7Zrqp6bkoibnQkouOlMKkYpCpEZdPpUUufATZytyMdRFEwJdX34/3e36+/vsLpk/WP3pj+7zy+JdcSayuFtoTWJQaE1iMMA4OqDsv3A0z30rz4dTp/wpOkieBRMWwPj5KKlzICgYf/7reDiKXMgOuYNAVVVaOhya3Pqyrmj+/Pncc889/L//9/+Ijo4mISGBF198EbvdzvLly4mMjCQzM5O///3vXo/bvn07c+bMISQkhKSkJNasWYPD4fDcb7fbue222zAajSQlJfHMM8/4PHd7ezurV68mJSWFiIgILrjgAvLy8nrd946ODu6++26SkpIIDQ1l7NixrF+/3nO/oii8++67ADz22GNeCbBOp8NoNLJlyxbAvW/Z+vXrycjIICwsjOzsbN56661e90UIIYQQYkRSVaj+Anb+El67AX4yDrZcDTuehrLP3MnVmEw4/1tw42vwvRJYsRUu/yGMuwSCgrV+BX5BRrAGQWunk2mP/J8mz/3Fj64kPLj3b+Mrr7zC9773PT777DNef/11vvvd7/LOO++wdOlSfvCDH/Czn/2MW2+9lRMnThAeHk55eTlXX301t99+O6+++ioHDx5kxYoVhIaG8thjjwHw4IMPsn37dv785z8THx/PD37wA3bv3k1OTo7nee+++26++OIL/vjHP5KcnMw777zDVVddxb59+5g4ceI5+/2LX/yCv/zlL7zxxhukp6dTWlpKaWlpj21Xr17NnXfe6fn5d7/7HY8++iizZ88GYP369fzud79j8+bNTJw4kR07dnDLLbcQFxfHvHnzev27FEIIIYQIeNYKOLodjm5zj1LZqr3vD4+F8fNgvHuUiqg0LXoZUCTBOofeVBHUspJgX58/Ozubhx56CIA1a9bw1FNPERsby7e+9S0AHn74Yf73f/+XwsJCLrzwQjZt2kRaWhrPPfcciqIwefJkysvLWbNmDQ8//DAtLS385je/4be//S2XX345AFu2bCEtLc3TtxMnTvDyyy9z/PhxkpOTAXjggQf44IMPeOmll3jyySfPWenv+PHjTJw4kYsvvhhFUUhPT/dp3/18ERERRES4F09++umnPPzww7zwwgtMnz6dtrY2nnzySf7xj38wd+5cADIyMvj444954YUXuOyyy3r9uzydVBEUZyPVs4TWJAaF1iQG/YDqgtpDKBW73Zv6ntiJUnfIu0lQGIydi5oxDzVjPiRMd08F7BbA759UEdRIf6oIqqrKf79/6XB3FQC1s50WR0ev2jqdTqZOnUpLS4vnWExMDJMnT/Yci4x0L0EsLS0lKyuLzz//nPPPP5/W1lbPY3Jzc7HZbBw5coSmpiY6OjrIysrynCM0NJSJEyficDhoaWlh165dOJ1OJk+e7NWf9vZ2zGYzLS0tdHS4X8OpfTvVTTfdxLXXXsvkyZNZtGgRixcvZuHChT7nO/XxpaWlLFmyhHvvvZcvf/nLtLa28sUXX9DS0sIVV1zh9diOjg6ys7PP+Py9IVUExdlI9SyhNYlBoTWJweEX1FZPROMBIpoOENF0kIimQ+gd3n/rqCi0RE3CGpuLNS4Xe/R0VH3XVL8qJ1Tt1aDnQ0OqCGqkv1UEI3o6mZ/R6/WEhYV5VerT6XSEh4f7VO8zGAyEh4cTFBTkUw2w+zWHhYV5yp/3dN7ux3V2dqLX69m1a5dPMBuNRsLDw89ZRfCiiy7i6NGj/P3vf+ejjz7itttuY9GiRbz55pueNiEhIZ7H2+12brrpJi666CKefPJJ2traCAsL8yTO7733HikpKV7Pcerj+0OqCIqzkepZQmsSg0JrEoNDrMMOlYUo5bu69qLKR7FW+DRTDeGQ1FXpL2U2jLuE0LBoQoH44e/1sJIqgn5ipFUR7KnPZzs2depU3n77bc8xgE8++YTIyEjS0tIYM2YMBoOBzz77jLFjxwLQ2NjI4cOHmTdvHoqiMGvWLJxOJ7W1tVx6ac8jfb2pImg2m7npppu46aabuOGGG7jqqqtobGwkJibGq8+qqnLrrbficrn47W9/i06n89w/ffp0QkJCKC0tZf78+X387Z2dVBEU5yLVs4TWJAaF1iQGB4nLCbUH3VX9yna5v9Z84Z4CeCpFB3FTIWUWpM6GlNkocVNAH+TXlf6G0nB
      "text/plain": [
       "<Figure size 1000x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# sweep over compute budgets from 1e17 to 1e26\n",
    "cs = 10 ** np.arange(17, 26, step=2**-2)\n",
    "best_ns = []\n",
    "best_ds = []\n",
    "for c in cs:\n",
    "    ns = 10 ** np.arange(7, 14, step=2**-5)\n",
    "    ds = c / (6 * ns)\n",
    "    losses = L(ns, ds)\n",
    "    best = np.argmin(losses)\n",
    "    best_ns.append(ns[best])\n",
    "    best_ds.append(ds[best])\n",
    "    print(f\"compute budget {c:e}: best model size: {ns[best]/1e6:.2f}M, best dataset size: {ds[best]/1e9:.2f}B\")\n",
    "\n",
    "# plot both the model size and dataset size as a function of compute budget, on one plot\n",
    "plt.figure(figsize=(10,3))\n",
    "plt.plot(cs, best_ns, label='model size')\n",
    "plt.plot(cs, best_ds, label='dataset size')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel('compute budget')\n",
    "plt.ylabel('model size / dataset size')\n",
    "plt.grid(True, which=\"both\", ls=\"-\", color='k', alpha=0.2)\n",
    "plt.legend()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch2",
   "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.10.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "7f5833218766b48e6e35e4452ee875aac0e2188d05bbe5298f2c62b79f08b222"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}