flops.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def embedding(seq_len, vocab_size, d_model):\n",
    "    return 2 * seq_len * vocab_size * d_model\n",
    "\n",
    "\n",
    "def attention(seq_len, d_model, key_size, num_heads):\n",
    "    projections = 2 * 3 * seq_len * d_model * (key_size * num_heads)\n",
    "    logits = 2 * seq_len * seq_len * (key_size * num_heads)\n",
    "    softmax = 3 * num_heads * seq_len * seq_len\n",
    "    softmax_query_reduction = 2 * seq_len * seq_len * (key_size * num_heads)\n",
    "    final_layer = 2 * seq_len * (key_size * num_heads) * d_model\n",
    "    return projections + logits + softmax + softmax_query_reduction + final_layer\n",
    "\n",
    "\n",
    "def dense(seq_len, d_model, ffw_size, swiglu=False):\n",
    "    if not swiglu:\n",
    "        return 2 * seq_len * (2 * d_model * ffw_size)\n",
    "    else:\n",
    "        return 2 * seq_len * (3 * d_model * ffw_size)\n",
    "\n",
    "\n",
    "def moe(dense, n_experts, top_k, seq_len, d_model, ffw_size, swiglu=False):\n",
    "    dense_flops = top_k * dense(seq_len, d_model, ffw_size, swiglu)\n",
    "    gate_flops = 3 * seq_len * n_experts\n",
    "    return dense_flops + gate_flops\n",
    "\n",
    "\n",
    "def final_logits(seq_len, d_model, vocab_size):\n",
    "    return 2 * seq_len * d_model * vocab_size\n",
    "\n",
    "\n",
    "def get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads,\n",
    "    ffw_size,\n",
    "    swiglu=False,\n",
    "    **kwargs,\n",
    "):\n",
    "    return (\n",
    "        embedding(seq_len, vocab_size, d_model)\n",
    "        + n_layers\n",
    "        * (\n",
    "            attention(seq_len, d_model, key_size, num_heads)\n",
    "            + dense(seq_len, d_model, ffw_size, swiglu=swiglu)\n",
    "        )\n",
    "        + final_logits(seq_len, d_model, vocab_size)\n",
    "    )\n",
    "\n",
    "\n",
    "def flops_moe(\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    n_layers,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads,\n",
    "    ffw_size,\n",
    "    n_experts,\n",
    "    top_k,\n",
    "    swiglu=False,\n",
    "):\n",
    "    return (\n",
    "        embedding(seq_len, vocab_size, d_model)\n",
    "        + n_layers\n",
    "        * (\n",
    "            attention(seq_len, d_model, key_size, num_heads)\n",
    "            + moe(dense, n_experts, top_k, seq_len, d_model, ffw_size, swiglu=swiglu)\n",
    "        )\n",
    "        + final_logits(seq_len, d_model, vocab_size)\n",
    "    )\n",
    "\n",
    "\n",
    "def parameter_count(\n",
    "    vocab_size,\n",
    "    n_layers,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads,\n",
    "    num_kv_heads,\n",
    "    ffw_size,\n",
    "    n_experts=1,\n",
    "    swiglu_or_geglu=False,\n",
    "    **kwargs,\n",
    "):\n",
    "    mul_factor_ffn = 3 if swiglu_or_geglu else 2\n",
    "    attn = 2 * d_model * num_heads * key_size + 2 * d_model * num_kv_heads * key_size\n",
    "    return vocab_size * d_model + n_layers * (\n",
    "        attn + mul_factor_ffn * n_experts * d_model * ffw_size\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "multiple_of = 256\n",
    "\n",
    "tiny = {\n",
    "    \"d_model\": 384,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 6,\n",
    "    \"num_kv_heads\": 6,\n",
    "    \"ffw_size\": int(8 / 3 * 384),\n",
    "    \"n_layers\": 8,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "tiny[\"ffw_size\"] = multiple_of * (\n",
    "    (tiny[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "mini = {\n",
    "    \"d_model\": 512,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 8,\n",
    "    \"num_kv_heads\": 8,\n",
    "    \"ffw_size\": int(8 / 3 * 512),\n",
    "    \"n_layers\": 10,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "mini[\"ffw_size\"] = multiple_of * (\n",
    "    (mini[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "small = {\n",
    "    \"d_model\": 768,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 12,\n",
    "    \"num_kv_heads\": 12,\n",
    "    \"ffw_size\": int(8 / 3 * 768),\n",
    "    \"n_layers\": 12,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "small[\"ffw_size\"] = multiple_of * (\n",
    "    (small[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "\n",
    "_210M = {\n",
    "    \"d_model\": 768,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 12,\n",
    "    \"num_kv_heads\": 12,\n",
    "    \"ffw_size\": int(8 / 3 * 768),\n",
    "    \"n_layers\": 24,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "_210M[\"ffw_size\"] = multiple_of * (\n",
    "    (_210M[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_flops = []\n",
    "all_params = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "33.454464\n",
      "171834605568\n",
      "iters [3.0, 7.0, 10.0]\n",
      "tokens [0.3, 0.7, 1.0]\n",
      "ratio [9.2, 21.4, 30.6]\n",
      "flops [0.1031007633408, 0.2405684477952, 0.343669211136]\n",
      "flop savings 0.6\n"
     ]
    }
   ],
   "source": [
    "model = tiny\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "# lr 0.002 for cos, 0.001 for wsd\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [2400 / 0.8, 5600 / 0.8, 8000 / 0.8]#, 9600 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "52.99456\n",
      "254882611200\n",
      "iters [3.8, 7.5, 11.2]\n",
      "tokens [0.4, 0.8, 1.2]\n",
      "ratio [7.2, 14.5, 21.7]\n",
      "flops [0.1911619584, 0.3823239168, 0.5734858752]\n",
      "flop savings 0.6\n"
     ]
    }
   ],
   "source": [
    "tiny2 = {\n",
    "    \"d_model\": 512,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 8,\n",
    "    \"num_kv_heads\": 8,\n",
    "    \"ffw_size\": int(8 / 3 * 512),\n",
    "    \"n_layers\": 8,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "tiny2[\"ffw_size\"] = multiple_of * (\n",
    "    (tiny2[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "model = tiny2\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "# lr 0.002 for cos, 0.001 for wsd\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [3000 / 0.8, 6000 / 0.8, 9000 / 0.8]# 12000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "59.810304\n",
      "279079550976\n",
      "iters [7.5, 12.5, 17.5]\n",
      "tokens [0.8, 1.3, 1.8]\n",
      "ratio [12.8, 21.4, 30.0]\n",
      "flops [0.418619326464, 0.69769887744, 0.976778428416]\n",
      "flop savings 0.5733333333333334\n"
     ]
    }
   ],
   "source": [
    "model = mini\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [6000 / 0.8, 10000 / 0.8, 14000 / 0.8]# 18000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# lr 0.002 for cos, 0.001 for wsd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "93.11296\n",
      "409294602240\n",
      "iters [10.0, 17.5, 25.0]\n",
      "tokens [1.0, 1.8, 2.6]\n",
      "ratio [11.0, 19.2, 27.5]\n",
      "flops [0.81858920448, 1.43253110784, 2.0464730112]\n",
      "flop savings 0.5809523809523809\n"
     ]
    }
   ],
   "source": [
    "mini2 = {\n",
    "    \"d_model\": 640,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 10,\n",
    "    \"num_kv_heads\": 10,\n",
    "    \"ffw_size\": int(8 / 3 * 640),\n",
    "    \"n_layers\": 12,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "mini2[\"ffw_size\"] = multiple_of * (\n",
    "    (mini2[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "\n",
    "model = mini2\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [8000 / 0.8, 14000 / 0.8, 20000 / 0.8]# 26000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# lr 0.002 for cos, 0.001 for wsd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "123.532032\n",
      "527392309248\n",
      "iters [15.0, 25.0, 35.0]\n",
      "tokens [1.5, 2.6, 3.6]\n",
      "ratio [12.4, 20.7, 29.0]\n",
      "flops [1.582176927744, 2.63696154624, 3.691746164736]\n",
      "flop savings 0.5733333333333334\n"
     ]
    }
   ],
   "source": [
    "model = small\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [12000 / 0.8, 20000 / 0.8, 28000 / 0.8]# 36000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# lr 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "151.843584\n",
      "624142319616\n",
      "iters [25.0, 37.5, 50.0]\n",
      "tokens [2.6, 3.8, 5.1]\n",
      "ratio [16.9, 25.3, 33.7]\n",
      "flops [3.12071159808, 4.68106739712, 6.24142319616]\n",
      "flop savings 0.5555555555555556\n"
     ]
    }
   ],
   "source": [
    "_151M = {\n",
    "    \"d_model\": 768,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 12,\n",
    "    \"num_kv_heads\": 12,\n",
    "    \"ffw_size\": int(8 / 3 * 768),\n",
    "    \"n_layers\": 16,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "\n",
    "_151M[\"ffw_size\"] = multiple_of * (\n",
    "    (_151M[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "model = _151M\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [20000 / 0.8, 30000 / 0.8, 40000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# double batch size?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "166.1408\n",
      "682936762368\n",
      "iters [25.0, 37.5, 50.0]\n",
      "tokens [2.6, 3.8, 5.1]\n",
      "ratio [15.4, 23.1, 30.8]\n",
      "flops [3.41468381184, 5.12202571776, 6.82936762368]\n",
      "flop savings 0.5555555555555556\n"
     ]
    }
   ],
   "source": [
    "_166M = {\n",
    "    \"d_model\": 896,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 14,\n",
    "    \"num_kv_heads\": 14,\n",
    "    \"ffw_size\": int(8 / 3 * 896),\n",
    "    \"n_layers\": 12,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "\n",
    "_166M[\"ffw_size\"] = multiple_of * (\n",
    "    (_166M[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "model = _166M\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [20000 / 0.8, 30000 / 0.8, 40000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# double batch size?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "208.466688\n",
      "817642340352\n",
      "iters [37.5, 50.0, 62.5]\n",
      "tokens [3.8, 5.1, 6.4]\n",
      "ratio [18.4, 24.6, 30.7]\n",
      "flops [6.13231755264, 8.17642340352, 10.2205292544]\n",
      "flop savings 0.5333333333333333\n"
     ]
    }
   ],
   "source": [
    "model = _210M\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "# iters = [22222, 44444, 66666]\n",
    "iters = [30000 / 0.8, 40000 / 0.8, 50000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{200 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{200 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "359.744512\n",
      "1341445373952\n",
      "iters [25.0, 37.5, 50.0]\n",
      "tokens [5.1, 7.7, 10.2]\n",
      "ratio [14.2, 21.3, 28.5]\n",
      "flops [6.70722686976, 10.06084030464, 13.41445373952]\n",
      "flop savings 0.5555555555555556\n"
     ]
    }
   ],
   "source": [
    "_350M = {\n",
    "    \"d_model\": 1024,\n",
    "    \"key_size\": 64,\n",
    "    \"num_heads\": 16,\n",
    "    \"num_kv_heads\": 16,\n",
    "    \"ffw_size\": int(8 / 3 * 1024),\n",
    "    \"n_layers\": 24,\n",
    "    \"vocab_size\": 50257,\n",
    "    \"swiglu\": True,\n",
    "    \"seq_len\": 512,\n",
    "}\n",
    "\n",
    "_350M[\"ffw_size\"] = multiple_of * (\n",
    "    (_350M[\"ffw_size\"] + multiple_of - 1) // multiple_of\n",
    ")\n",
    "\n",
    "model = _350M\n",
    "\n",
    "n_layers = model[\"n_layers\"]\n",
    "d_model = model[\"d_model\"]\n",
    "key_size = model[\"key_size\"]\n",
    "num_heads = model[\"num_heads\"]\n",
    "num_kv_heads = model[\"num_kv_heads\"]\n",
    "ffw_size = model[\"ffw_size\"]\n",
    "vocab_size = model[\"vocab_size\"]\n",
    "swiglu = model[\"swiglu\"]\n",
    "n_experts = 8\n",
    "top_k = 2\n",
    "seq_len = model[\"seq_len\"]\n",
    "\n",
    "\n",
    "flops = 3 * get_flops(\n",
    "    n_layers,\n",
    "    seq_len,\n",
    "    vocab_size,\n",
    "    d_model,\n",
    "    key_size,\n",
    "    num_heads=num_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu=swiglu,\n",
    ")\n",
    "params = parameter_count(\n",
    "    vocab_size=vocab_size,\n",
    "    n_layers=n_layers,\n",
    "    d_model=d_model,\n",
    "    key_size=key_size,\n",
    "    num_heads=num_heads,\n",
    "    num_kv_heads=num_kv_heads,\n",
    "    ffw_size=ffw_size,\n",
    "    swiglu_or_geglu=swiglu,\n",
    ")\n",
    "print(params / 1e6)\n",
    "print(flops)\n",
    "iters = [20000 / 0.8, 30000 / 0.8, 40000 / 0.8]\n",
    "print(\"iters\", [float(f\"{i / 1e3:.1f}\") for i in iters])\n",
    "print(\"tokens\", [float(f\"{400 * 512 * i / 1e9:.1f}\") for i in iters])\n",
    "print(\"ratio\", [float(f\"{400 * 512 * i / params:.1f}\") for i in iters])\n",
    "flops_all = [flops * 200 * i / 1e18 for i in iters]\n",
    "print(\"flops\", flops_all)\n",
    "all_flops.append(flops_all)\n",
    "all_params.append(params)\n",
    "print(\"flop savings\", (flops_all[-1] + 0.2 * sum(flops_all[:-1])) / sum(flops_all))\n",
    "# double batch size?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm-baselines",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}