makeMoE_from_Scratch_with_Expert_Capacity.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fe612761-a75c-4c61-83ad-a68e4e991645",
     "showTitle": false,
     "title": ""
    },
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/AviSoori1x/makeMoE/blob/main/makeMoE_from_Scratch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "8e9b80fc-12cf-41a9-a0de-354f678b412b",
     "showTitle": false,
     "title": ""
    },
    "id": "90vgVgmDkRJQ"
   },
   "source": [
    "#### Sparse mixture of experts language model from scratch inspired by (and largely based on) Andrej Karpathy's makemore (https://github.com/karpathy/makemore) :)\n",
    "\n",
    "This is a from scratch implementation of a sparse mixture of experts language model. This is inspired by and largely based on Andrej Karpathy's project 'makemore' and borrows most of the re-usable components from that implementation. Just like makemore, makeMoE is also an autoregressive character-level language model but uses the aforementioned sparse mixture of experts architecture.\n",
    "\n",
    "Just like makemore, pytorch is the only requirement (so I hope the from scratch claim is justified).\n",
    "\n",
    "Significant Changes from the makemore architecture\n",
    "\n",
    "- Sparse mixture of experts instead of the solitary feed forward neural net.\n",
    "- Top-k gating and noisy top-k gating implementations.\n",
    "- initialization - Kaiming He initialization is used here but the point of this notebook is to be hackable so you can swap in Xavier Glorot etc. and take it for a spin.\n",
    "\n",
    "Unchanged from makemore\n",
    "- The dataset, preprocessing (tokenization), and the language modeling task Andrej chose originally - generate Shakespeare-like text\n",
    "- Casusal self attention implementation\n",
    "- Training loop\n",
    "- Inference logic\n",
    "\n",
    "Publications heavily referenced for this implementation:\n",
    "- Mixtral of experts: https://arxiv.org/pdf/2401.04088.pdf\n",
    "- Outrageosly Large Neural Networks: The Sparsely-Gated Mixture-Of-Experts layer: https://arxiv.org/pdf/1701.06538.pdf\n",
    "\n",
    "\n",
    "This notebook walks through the intuition for the entire model architecture and how everything comes together\n",
    "\n",
    "The code was entirely developed on Databricks using a single A100 for compute. If you're running this on Databricks, you can scale this on an arbitrarily large GPU cluster with no issues in the cloud provider of your choice\n",
    "\n",
    "I chose to use mlflow (which comes pre-installed in Databricks. You can pip install easily elsewhere) as I find it helpful to track and log all the metrics necessary. This is entirely optional.\n",
    "\n",
    "Please note that the implementation emphasizes readability and hackability vs performance, so there are many ways in which you could improve this. Please try and let me know\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "2f4a58a8-bd4c-40de-a4a9-95457842db0b",
     "showTitle": false,
     "title": ""
    },
    "id": "hywLNfb0kRJT"
   },
   "source": [
    "![mixture of experts overview](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/moe.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "35b3daa3-3b3b-47af-b3e7-be95878f9e06",
     "showTitle": false,
     "title": ""
    },
    "id": "RQAnP6_RkRJU"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[43mNote: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages.\u001B[0m\nCollecting mlflow\n  Downloading mlflow-2.11.1-py3-none-any.whl (19.7 MB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/19.7 MB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[91m╸\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.3/19.7 MB\u001B[0m \u001B[31m8.2 MB/s\u001B[0m eta \u001B[36m0:00:03\u001B[0m\n\u001B[2K     \u001B[91m━\u001B[0m\u001B[91m╸\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.8/19.7 MB\u001B[0m \u001B[31m12.5 MB/s\u001B[0m eta \u001B[36m0:00:02\u001B[0m\n\u001B[2K     \u001B[91m━━━\u001B[0m\u001B[90m╺\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m1.6/19.7 MB\u001B[0m \u001B[31m16.1 MB/s\u001B[0m eta \u001B[36m0:00:02\u001B[0m\n\u001B[2K     \u001B[91m━━━━━\u001B[0m\u001B[91m╸\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m2.7/19.7 MB\u001B[0m \u001B[31m19.8 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━\u001B[0m\u001B[90m╺\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m4.1/19.7 MB\u001B[0m \u001B[31m23.7 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━\u001B[0m\u001B[90m╺\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m6.0/19.7 MB\u001B[0m \u001B[31m29.2 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m\u001B[90m━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m8.6/19.7 MB\u001B[0m \u001B[31m35.8 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m\u001B[90m━━━━━━━━━━━━━━━\u001B[0m \u001B[32m12.3/19.7 MB\u001B[0m \u001B[31m65.1 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m\u001B[90m━━━━━\u001B[0m \u001B[32m17.1/19.7 MB\u001B[0m \u001B[31m112.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m\u001B[91m╸\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m136.0 MB/s\u001B[0m eta \u001B[36m0:00:01\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m19.7/19.7 MB\u001B[0m \u001B[31m21.7 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hRequirement already satisfied: scipy<2 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.11.1)\nRequirement already satisfied: sqlalchemy<3,>=1.4.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.4.39)\nRequirement already satisfied: packaging<24 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (23.2)\nRequirement already satisfied: markdown<4,>=3.3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.4.1)\nRequirement already satisfied: matplotlib<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.7.2)\nRequirement already satisfied: gitpython<4,>=3.1.9 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.1.27)\nRequirement already satisfied: sqlparse<1,>=0.4.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (0.4.2)\nCollecting graphene<4\n  Downloading graphene-3.3-py2.py3-none-any.whl (128 kB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/128.2 kB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m128.2/128.2 kB\u001B[0m \u001B[31m18.0 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hCollecting docker<8,>=4.0.0\n  Downloading docker-7.0.0-py3-none-any.whl (147 kB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/147.6 kB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m147.6/147.6 kB\u001B[0m \u001B[31m18.1 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hRequirement already satisfied: pyyaml<7,>=5.1 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (6.0)\nRequirement already satisfied: pytz<2025 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2022.7)\nRequirement already satisfied: pyarrow<16,>=4.0.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (14.0.1)\nRequirement already satisfied: protobuf<5,>=3.12.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (4.24.1)\nRequirement already satisfied: scikit-learn<2 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.3.0)\nRequirement already satisfied: numpy<2 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.23.5)\nRequirement already satisfied: gunicorn<22 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (20.1.0)\nRequirement already satisfied: pandas<3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.0.3)\nRequirement already satisfied: importlib-metadata!=4.7.0,<8,>=3.7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (6.0.0)\nRequirement already satisfied: click<9,>=7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (8.0.4)\nRequirement already satisfied: entrypoints<1 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (0.4)\nRequirement already satisfied: Flask<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.2.5)\nRequirement already satisfied: cloudpickle<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.2.1)\nRequirement already satisfied: requests<3,>=2.17.3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.31.0)\nCollecting querystring-parser<2\n  Downloading querystring_parser-1.2.4-py2.py3-none-any.whl (7.9 kB)\nRequirement already satisfied: Jinja2<4,>=2.11 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.1.2)\nCollecting alembic!=1.10.0,<2\n  Downloading alembic-1.13.1-py3-none-any.whl (233 kB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/233.4 kB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m233.4/233.4 kB\u001B[0m \u001B[31m35.8 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hRequirement already satisfied: Mako in /databricks/python3/lib/python3.11/site-packages (from alembic!=1.10.0,<2->mlflow) (1.2.0)\nRequirement already satisfied: typing-extensions>=4 in /databricks/python3/lib/python3.11/site-packages (from alembic!=1.10.0,<2->mlflow) (4.7.1)\nRequirement already satisfied: urllib3>=1.26.0 in /databricks/python3/lib/python3.11/site-packages (from docker<8,>=4.0.0->mlflow) (1.26.16)\nRequirement already satisfied: Werkzeug>=2.2.2 in /databricks/python3/lib/python3.11/site-packages (from Flask<4->mlflow) (2.2.3)\nRequirement already satisfied: itsdangerous>=2.0 in /databricks/python3/lib/python3.11/site-packages (from Flask<4->mlflow) (2.0.1)\nRequirement already satisfied: gitdb<5,>=4.0.1 in /databricks/python3/lib/python3.11/site-packages (from gitpython<4,>=3.1.9->mlflow) (4.0.11)\nCollecting graphql-relay<3.3,>=3.1\n  Downloading graphql_relay-3.2.0-py3-none-any.whl (16 kB)\nCollecting graphql-core<3.3,>=3.1\n  Downloading graphql_core-3.2.3-py3-none-any.whl (202 kB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/202.9 kB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m202.9/202.9 kB\u001B[0m \u001B[31m32.5 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hCollecting aniso8601<10,>=8\n  Downloading aniso8601-9.0.1-py2.py3-none-any.whl (52 kB)\n\u001B[?25l     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m0.0/52.8 kB\u001B[0m \u001B[31m?\u001B[0m eta \u001B[36m-:--:--\u001B[0m\n\u001B[2K     \u001B[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001B[0m \u001B[32m52.8/52.8 kB\u001B[0m \u001B[31m9.5 MB/s\u001B[0m eta \u001B[36m0:00:00\u001B[0m\n\u001B[?25hRequirement already satisfied: setuptools>=3.0 in /databricks/python3/lib/python3.11/site-packages (from gunicorn<22->mlflow) (68.0.0)\nRequirement already satisfied: zipp>=0.5 in /databricks/python3/lib/python3.11/site-packages (from importlib-metadata!=4.7.0,<8,>=3.7.0->mlflow) (3.11.0)\nRequirement already satisfied: MarkupSafe>=2.0 in /databricks/python3/lib/python3.11/site-packages (from Jinja2<4,>=2.11->mlflow) (2.1.1)\nRequirement already satisfied: cycler>=0.10 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (0.11.0)\nRequirement already satisfied: pyparsing<3.1,>=2.3.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (3.0.9)\nRequirement already satisfied: kiwisolver>=1.0.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (1.4.4)\nRequirement already satisfied: contourpy>=1.0.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (1.0.5)\nRequirement already satisfied: python-dateutil>=2.7 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (2.8.2)\nRequirement already satisfied: pillow>=6.2.0 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (9.4.0)\nRequirement already satisfied: fonttools>=4.22.0 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (4.25.0)\nRequirement already satisfied: tzdata>=2022.1 in /databricks/python3/lib/python3.11/site-packages (from pandas<3->mlflow) (2022.1)\nRequirement already satisfied: six in /usr/lib/python3/dist-packages (from querystring-parser<2->mlflow) (1.16.0)\nRequirement already satisfied: idna<4,>=2.5 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (3.4)\nRequirement already satisfied: certifi>=2017.4.17 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (2023.7.22)\nRequirement already satisfied: charset-normalizer<4,>=2 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (2.0.4)\nRequirement already satisfied: joblib>=1.1.1 in /databricks/python3/lib/python3.11/site-packages (from scikit-learn<2->mlflow) (1.2.0)\nRequirement already satisfied: threadpoolctl>=2.0.0 in /databricks/python3/lib/python3.11/site-packages (from scikit-learn<2->mlflow) (2.2.0)\nRequirement already satisfied: greenlet!=0.4.17 in /databricks/python3/lib/python3.11/site-packages (from sqlalchemy<3,>=1.4.0->mlflow) (2.0.1)\nRequirement already satisfied: smmap<6,>=3.0.1 in /databricks/python3/lib/python3.11/site-packages (from gitdb<5,>=4.0.1->gitpython<4,>=3.1.9->mlflow) (5.0.0)\nInstalling collected packages: aniso8601, querystring-parser, graphql-core, graphql-relay, docker, alembic, graphene, mlflow\nSuccessfully installed alembic-1.13.1 aniso8601-9.0.1 docker-7.0.0 graphene-3.3 graphql-core-3.2.3 graphql-relay-3.2.0 mlflow-2.11.1 querystring-parser-1.2.4\n\u001B[43mNote: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages.\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "#Using mlflow is entirely optional. I personally like to use MLFlow to track and log everything. If you're using Databricks, it comes pre-installed.\n",
    "%pip install mlflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5e1a3e38-8717-42ec-9bbc-71d3712c1c68",
     "showTitle": false,
     "title": ""
    },
    "id": "V521QQ_qkRJV"
   },
   "outputs": [],
   "source": [
    "#Import the necessary packages and set seed for reproducibility. For this notebook, pytorch is all you need\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "torch.manual_seed(42)\n",
    "#Optional\n",
    "import mlflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "faf99ef2-39bb-46fc-b772-05d6d0482bbc",
     "showTitle": false,
     "title": ""
    },
    "id": "-4r_QNRRkRJV"
   },
   "source": [
    "Next few sections, downloading the data, preprocessing it and self attention are directly from makemore. I have elaborated a little on self attention and added visual aids to understand the process a bit better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "45143d84-28c7-463d-9fb5-e21122842600",
     "showTitle": false,
     "title": ""
    },
    "id": "2GhDw0yWkRJV",
    "outputId": "22b79649-9819-4c65-a46a-d0796d6c9085"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-17 20:52:21--  https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/input.txt\r\nResolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\r\nConnecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\r\nHTTP request sent, awaiting response... 200 OK\r\nLength: 1115394 (1.1M) [text/plain]\r\nSaving to: ‘input.txt.5’\r\n\r\n\rinput.txt.5           0%[                    ]       0  --.-KB/s               \rinput.txt.5         100%[===================>]   1.06M  --.-KB/s    in 0.03s   \r\n\r\n2024-03-17 20:52:22 (32.4 MB/s) - ‘input.txt.5’ saved [1115394/1115394]\r\n\r\n"
     ]
    }
   ],
   "source": [
    "# Downloading the tiny shakespeare dataset\n",
    "!wget https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/input.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "192e830a-762d-4573-9484-70a58deb1fec",
     "showTitle": false,
     "title": ""
    },
    "id": "3sPAL1AKkRJW"
   },
   "outputs": [],
   "source": [
    "# read it in to inspect it\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2d7181b7-f5e5-4ab5-bdd8-74c507c798ad",
     "showTitle": false,
     "title": ""
    },
    "id": "wNkF3RYLkRJX",
    "outputId": "671be677-be3d-49e0-d391-569559a99ae2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of dataset in characters:  1115394\n"
     ]
    }
   ],
   "source": [
    "print(\"length of dataset in characters: \", len(text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "68032e07-8625-4750-a340-bc8f4eed2458",
     "showTitle": false,
     "title": ""
    },
    "id": "AHIwr-yxkRJX",
    "outputId": "fa10ee0f-c8cb-4ba1-9fb1-79f9fbf336fe"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\nBefore we proceed any further, hear me speak.\n\nAll:\nSpeak, speak.\n\nFirst Citizen:\nYou are all resolved rather to die than to famish?\n\nAll:\nResolved. resolved.\n\nFirst Citizen:\nFirst, you know Caius Marcius is chief enemy to the people.\n\nAll:\nWe know't, we know't.\n\nFirst Citizen:\nLet us kill him, and we'll have corn at our own price.\nIs't a verdict?\n\nAll:\nNo more talking on't; let it be done: away, away!\n\nSecond Citizen:\nOne word, good citizens.\n\nFirst Citizen:\nWe are accounted poor citizens, the patricians good.\nWhat authority surfeits on would relieve us: if they\nwould yield us but the superfluity, while it were\nwholesome, we might guess they relieved us humanely;\nbut they think we are too dear: the leanness that\nafflicts us, the object of our misery, is as an\ninventory to particularise their abundance; our\nsufferance is a gain to them Let us revenge this with\nour pikes, ere we become rakes: for the gods know I\nspeak this in hunger for bread, not in thirst for revenge.\n\n\n"
     ]
    }
   ],
   "source": [
    "# let's look at the first 1000 characters\n",
    "print(text[:1000])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b6995ad6-c9ac-4a21-9da0-ebbd3273c991",
     "showTitle": false,
     "title": ""
    },
    "id": "DHGayz7mkRJY",
    "outputId": "54b52e55-b6b9-4bdc-974a-5c48a43e1497"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n65\n"
     ]
    }
   ],
   "source": [
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "print(''.join(chars))\n",
    "print(vocab_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "43002fa3-ffd3-416c-9aaf-0a03b19c7bc1",
     "showTitle": false,
     "title": ""
    },
    "id": "pzn11WJckRJY",
    "outputId": "cd7b3ff1-680b-4036-b3c8-b597a361aeb3"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 47, 47, 1, 58, 46, 43, 56, 43]\nhii there\n"
     ]
    }
   ],
   "source": [
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "print(encode(\"hii there\"))\n",
    "print(decode(encode(\"hii there\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b4609fc4-09c7-4a39-8367-e9ee39d440ed",
     "showTitle": false,
     "title": ""
    },
    "id": "YbBGz0O2kRJY",
    "outputId": "49e94dde-6e1f-4a9e-a18b-9b8f1c51de86"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115394]) torch.int64\ntensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n        58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50, 50,\n         1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1, 58,\n        53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51, 47,\n        57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43, 42,\n         8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57, 58,\n         1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1, 63,\n        53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56, 41,\n        47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51, 63,\n         1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0, 13,\n        50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,  1,\n        49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47, 58,\n        47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57,  1, 49, 47, 50, 50,  1,\n        46, 47, 51,  6,  1, 39, 52, 42,  1, 61, 43,  5, 50, 50,  1, 46, 39, 60,\n        43,  1, 41, 53, 56, 52,  1, 39, 58,  1, 53, 59, 56,  1, 53, 61, 52,  1,\n        54, 56, 47, 41, 43,  8,  0, 21, 57,  5, 58,  1, 39,  1, 60, 43, 56, 42,\n        47, 41, 58, 12,  0,  0, 13, 50, 50, 10,  0, 26, 53,  1, 51, 53, 56, 43,\n         1, 58, 39, 50, 49, 47, 52, 45,  1, 53, 52,  5, 58, 11,  1, 50, 43, 58,\n         1, 47, 58,  1, 40, 43,  1, 42, 53, 52, 43, 10,  1, 39, 61, 39, 63,  6,\n         1, 39, 61, 39, 63,  2,  0,  0, 31, 43, 41, 53, 52, 42,  1, 15, 47, 58,\n        47, 64, 43, 52, 10,  0, 27, 52, 43,  1, 61, 53, 56, 42,  6,  1, 45, 53,\n        53, 42,  1, 41, 47, 58, 47, 64, 43, 52, 57,  8,  0,  0, 18, 47, 56, 57,\n        58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 35, 43,  1, 39, 56, 43,  1,\n        39, 41, 41, 53, 59, 52, 58, 43, 42,  1, 54, 53, 53, 56,  1, 41, 47, 58,\n        47, 64, 43, 52, 57,  6,  1, 58, 46, 43,  1, 54, 39, 58, 56, 47, 41, 47,\n        39, 52, 57,  1, 45, 53, 53, 42,  8,  0, 35, 46, 39, 58,  1, 39, 59, 58,\n        46, 53, 56, 47, 58, 63,  1, 57, 59, 56, 44, 43, 47, 58, 57,  1, 53, 52,\n         1, 61, 53, 59, 50, 42,  1, 56, 43, 50, 47, 43, 60, 43,  1, 59, 57, 10,\n         1, 47, 44,  1, 58, 46, 43, 63,  0, 61, 53, 59, 50, 42,  1, 63, 47, 43,\n        50, 42,  1, 59, 57,  1, 40, 59, 58,  1, 58, 46, 43,  1, 57, 59, 54, 43,\n        56, 44, 50, 59, 47, 58, 63,  6,  1, 61, 46, 47, 50, 43,  1, 47, 58,  1,\n        61, 43, 56, 43,  0, 61, 46, 53, 50, 43, 57, 53, 51, 43,  6,  1, 61, 43,\n         1, 51, 47, 45, 46, 58,  1, 45, 59, 43, 57, 57,  1, 58, 46, 43, 63,  1,\n        56, 43, 50, 47, 43, 60, 43, 42,  1, 59, 57,  1, 46, 59, 51, 39, 52, 43,\n        50, 63, 11,  0, 40, 59, 58,  1, 58, 46, 43, 63,  1, 58, 46, 47, 52, 49,\n         1, 61, 43,  1, 39, 56, 43,  1, 58, 53, 53,  1, 42, 43, 39, 56, 10,  1,\n        58, 46, 43,  1, 50, 43, 39, 52, 52, 43, 57, 57,  1, 58, 46, 39, 58,  0,\n        39, 44, 44, 50, 47, 41, 58, 57,  1, 59, 57,  6,  1, 58, 46, 43,  1, 53,\n        40, 48, 43, 41, 58,  1, 53, 44,  1, 53, 59, 56,  1, 51, 47, 57, 43, 56,\n        63,  6,  1, 47, 57,  1, 39, 57,  1, 39, 52,  0, 47, 52, 60, 43, 52, 58,\n        53, 56, 63,  1, 58, 53,  1, 54, 39, 56, 58, 47, 41, 59, 50, 39, 56, 47,\n        57, 43,  1, 58, 46, 43, 47, 56,  1, 39, 40, 59, 52, 42, 39, 52, 41, 43,\n        11,  1, 53, 59, 56,  0, 57, 59, 44, 44, 43, 56, 39, 52, 41, 43,  1, 47,\n        57,  1, 39,  1, 45, 39, 47, 52,  1, 58, 53,  1, 58, 46, 43, 51,  1, 24,\n        43, 58,  1, 59, 57,  1, 56, 43, 60, 43, 52, 45, 43,  1, 58, 46, 47, 57,\n         1, 61, 47, 58, 46,  0, 53, 59, 56,  1, 54, 47, 49, 43, 57,  6,  1, 43,\n        56, 43,  1, 61, 43,  1, 40, 43, 41, 53, 51, 43,  1, 56, 39, 49, 43, 57,\n        10,  1, 44, 53, 56,  1, 58, 46, 43,  1, 45, 53, 42, 57,  1, 49, 52, 53,\n        61,  1, 21,  0, 57, 54, 43, 39, 49,  1, 58, 46, 47, 57,  1, 47, 52,  1,\n        46, 59, 52, 45, 43, 56,  1, 44, 53, 56,  1, 40, 56, 43, 39, 42,  6,  1,\n        52, 53, 58,  1, 47, 52,  1, 58, 46, 47, 56, 57, 58,  1, 44, 53, 56,  1,\n        56, 43, 60, 43, 52, 45, 43,  8,  0,  0])\n"
     ]
    }
   ],
   "source": [
    "# let's now encode the entire text dataset and store it into a torch.Tensor\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "print(data.shape, data.dtype)\n",
    "print(data[:1000]) # the 1000 characters we looked at earier will to the GPT look like this"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "88f7cb0d-02ff-42b0-92a5-a505dc3f8f25",
     "showTitle": false,
     "title": ""
    },
    "id": "hoLIeA7YkRJZ"
   },
   "outputs": [],
   "source": [
    "# Let's now split up the data into train and validation sets\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6b554ddf-50f4-441b-8acf-10b81a508b7e",
     "showTitle": false,
     "title": ""
    },
    "id": "VY55nr6EkRJZ",
    "outputId": "1a5685e9-66bb-40ed-f19a-f0cb09e8e51c"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "block_size = 8\n",
    "train_data[:block_size+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "22ba4512-309d-4895-a908-2ef3efa317bc",
     "showTitle": false,
     "title": ""
    },
    "id": "5YbgrB9HkRJZ",
    "outputId": "05577a5f-10d5-495b-82e8-c3278b4fdea0"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when input is tensor([18]) the target: 47\nwhen input is tensor([18, 47]) the target: 56\nwhen input is tensor([18, 47, 56]) the target: 57\nwhen input is tensor([18, 47, 56, 57]) the target: 58\nwhen input is tensor([18, 47, 56, 57, 58]) the target: 1\nwhen input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\nwhen input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\nwhen input is tensor([18, 47, 56, 57, 58,  1, 15, 47]) the target: 58\n"
     ]
    }
   ],
   "source": [
    "x = train_data[:block_size]\n",
    "y = train_data[1:block_size+1]\n",
    "for t in range(block_size):\n",
    "    context = x[:t+1]\n",
    "    target = y[t]\n",
    "    print(f\"when input is {context} the target: {target}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bf386bff-0f63-4358-82fc-6c7d02c37321",
     "showTitle": false,
     "title": ""
    },
    "id": "Oaxhage8kRJZ"
   },
   "outputs": [],
   "source": [
    "batch_size = 4 # how many independent sequences will we process in parallel?\n",
    "block_size = 8 # what is the maximum context length for predictions?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "99acd85c-233f-4f2d-a062-028dbcde9960",
     "showTitle": false,
     "title": ""
    },
    "id": "HfpkIUNdkRJZ",
    "outputId": "1ad2bc10-2b4a-458b-a6ee-a169728ec90d"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([250930, 237205, 974116, 383898])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "ix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e46dc826-9f39-4aed-b2d2-f9ea401136de",
     "showTitle": false,
     "title": ""
    },
    "id": "faoGVPG3kRJa",
    "outputId": "dcf8b139-d8f8-4565-a962-62d094c5da9f"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([[42,  1, 58, 46, 59, 57,  1, 21],\n",
       "        [54, 56, 47, 43, 57, 58, 11,  0],\n",
       "        [49, 47, 52, 45, 12,  1, 58, 46],\n",
       "        [58, 46, 53, 59, 58,  1, 56, 43]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2886aedf-200e-40bd-9a9d-4658cf6c509b",
     "showTitle": false,
     "title": ""
    },
    "id": "hkllYFCPkRJa",
    "outputId": "9350713c-be0f-492a-c873-bd279226668d"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([[ 1, 58, 46, 59, 57,  1, 21,  1],\n",
       "        [56, 47, 43, 57, 58, 11,  0, 37],\n",
       "        [47, 52, 45, 12,  1, 58, 46, 53],\n",
       "        [46, 53, 59, 58,  1, 56, 43, 42]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a486fc04-ed29-456f-918b-5f8395e455cb",
     "showTitle": false,
     "title": ""
    },
    "id": "tWrajECBkRJa"
   },
   "source": [
    "The following code block clearly shows the autoregressive nature of the prediction and how the context is a rolling windows over a 1 dimentional arrangement of tokens (characters in this case)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "49a86e10-ac37-4b92-8f18-775cd4853fdc",
     "showTitle": false,
     "title": ""
    },
    "id": "xjgtxxztkRJa",
    "outputId": "1be0ebc2-b22d-4136-e77d-ecdafded66d2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\ntorch.Size([4, 8])\ntensor([[ 6,  0, 14, 43, 44, 53, 56, 43],\n        [39,  1, 42, 59, 43,  1, 39, 52],\n        [47, 41, 43,  1, 39, 52, 42,  1],\n        [53, 44,  1, 50, 43, 58,  1, 58]])\ntargets:\ntorch.Size([4, 8])\ntensor([[ 0, 14, 43, 44, 53, 56, 43,  1],\n        [ 1, 42, 59, 43,  1, 39, 52, 42],\n        [41, 43,  1, 39, 52, 42,  1, 42],\n        [44,  1, 50, 43, 58,  1, 58, 46]])\n----\nwhen input is [6] the target: 0\nwhen input is [6, 0] the target: 14\nwhen input is [6, 0, 14] the target: 43\nwhen input is [6, 0, 14, 43] the target: 44\nwhen input is [6, 0, 14, 43, 44] the target: 53\nwhen input is [6, 0, 14, 43, 44, 53] the target: 56\nwhen input is [6, 0, 14, 43, 44, 53, 56] the target: 43\nwhen input is [6, 0, 14, 43, 44, 53, 56, 43] the target: 1\nwhen input is [39] the target: 1\nwhen input is [39, 1] the target: 42\nwhen input is [39, 1, 42] the target: 59\nwhen input is [39, 1, 42, 59] the target: 43\nwhen input is [39, 1, 42, 59, 43] the target: 1\nwhen input is [39, 1, 42, 59, 43, 1] the target: 39\nwhen input is [39, 1, 42, 59, 43, 1, 39] the target: 52\nwhen input is [39, 1, 42, 59, 43, 1, 39, 52] the target: 42\nwhen input is [47] the target: 41\nwhen input is [47, 41] the target: 43\nwhen input is [47, 41, 43] the target: 1\nwhen input is [47, 41, 43, 1] the target: 39\nwhen input is [47, 41, 43, 1, 39] the target: 52\nwhen input is [47, 41, 43, 1, 39, 52] the target: 42\nwhen input is [47, 41, 43, 1, 39, 52, 42] the target: 1\nwhen input is [47, 41, 43, 1, 39, 52, 42, 1] the target: 42\nwhen input is [53] the target: 44\nwhen input is [53, 44] the target: 1\nwhen input is [53, 44, 1] the target: 50\nwhen input is [53, 44, 1, 50] the target: 43\nwhen input is [53, 44, 1, 50, 43] the target: 58\nwhen input is [53, 44, 1, 50, 43, 58] the target: 1\nwhen input is [53, 44, 1, 50, 43, 58, 1] the target: 58\nwhen input is [53, 44, 1, 50, 43, 58, 1, 58] the target: 46\n"
     ]
    }
   ],
   "source": [
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    return x, y\n",
    "\n",
    "xb, yb = get_batch('train')\n",
    "print('inputs:')\n",
    "print(xb.shape)\n",
    "print(xb)\n",
    "print('targets:')\n",
    "print(yb.shape)\n",
    "print(yb)\n",
    "\n",
    "print('----')\n",
    "\n",
    "for b in range(batch_size): # batch dimension\n",
    "    for t in range(block_size): # time dimension\n",
    "        context = xb[b, :t+1]\n",
    "        target = yb[b,t]\n",
    "        print(f\"when input is {context.tolist()} the target: {target}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "dde3273f-0519-4108-ba84-dfd99e020722",
     "showTitle": false,
     "title": ""
    },
    "id": "RlON_gNikRJa"
   },
   "source": [
    "### Understanding the intuition of Causal Scaled Dot Product Self Attention\n",
    "\n",
    "This code is borrowed from Andrej Karpathy's excellent makemore repository linked in the repo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e435d0cf-1383-446a-9026-cd80b4266019",
     "showTitle": false,
     "title": ""
    },
    "id": "uBVWP40SkRJa"
   },
   "source": [
    "![scaled dot product self attention](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/self_attention.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "97660589-1719-48c0-ad6f-4f7c2888348a",
     "showTitle": false,
     "title": ""
    },
    "id": "i-jFgMntkRJa"
   },
   "source": [
    "The provided code demonstrates self-attention's mechanics and fundamental concepts, specifically focusing on the classic scaled dot product self-attention. In this variant, the query, key, and value matrices all originate from the same input sequence. To ensure the integrity of the autoregressive language generation process, particularly in a decoder-only model, the code implements masking. This masking technique is crucial as it obscures any information following the current token's position, thereby directing the model's attention to only the preceding parts of the sequence. Such an attention mechanism is known as causal self-attention. It's important to note that the Sparse Mixture of Experts model isn't restricted to decoder-only Transformer architectures. In fact, much of the significant work in this field, particularly that by Shazeer et al, revolves around the T5 architecture, which encompasses both encoder and decoder components in the Transformer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6f82ca41-a301-4a92-aed9-ba7ac3a2bf88",
     "showTitle": false,
     "title": ""
    },
    "id": "lxMSgZWGkRJb",
    "outputId": "8d968403-a14f-4f3c-f2ce-9099a2b34f1e"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 16])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(1337)\n",
    "B,T,C = 4,8,32 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "\n",
    "# let's see a single Head perform self-attention\n",
    "head_size = 16\n",
    "key = nn.Linear(C, head_size, bias=False)\n",
    "query = nn.Linear(C, head_size, bias=False)\n",
    "value = nn.Linear(C, head_size, bias=False)\n",
    "k = key(x)   # (B, T, 16)\n",
    "q = query(x) # (B, T, 16)\n",
    "wei =  q @ k.transpose(-2, -1) # (B, T, 16) @ (B, 16, T) ---> (B, T, T)\n",
    "\n",
    "tril = torch.tril(torch.ones(T, T))\n",
    "#wei = torch.zeros((T,T))\n",
    "wei = wei.masked_fill(tril == 0, float('-inf'))\n",
    "wei = F.softmax(wei, dim=-1) #B,T,T\n",
    "\n",
    "v = value(x) #B,T,H\n",
    "out = wei @ v # (B,T,T) @ (B,T,H) -> (B,T,H)\n",
    "#The output from this final matrix product is subsequently passsed through a linear layer as shown in the diagram above\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "49c278ec-19db-4c5d-b4a3-3bdc45c5a443",
     "showTitle": false,
     "title": ""
    },
    "id": "cA3iXggEkRJb"
   },
   "source": [
    "Generalizing and Modularizing code for causal self attention and multi-head causal self attention. Multi-head self attention applied multiple attention heads in parallel, each focusing on a separate section of the channel (the embedding dimension)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "608c6c9f-fb93-43ed-9580-5e782fd90d61",
     "showTitle": false,
     "title": ""
    },
    "id": "909nX3PHkRJb"
   },
   "outputs": [],
   "source": [
    "#Causal scaled dot product self-Attention Head\n",
    "\n",
    "n_embd = 64\n",
    "n_head = 4\n",
    "n_layer = 4\n",
    "head_size = 16\n",
    "dropout = 0.1\n",
    "\n",
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)   # (B,T,C)\n",
    "        q = self.query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B,T,C)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "6e8b31af-f45a-4066-8288-fb0d9c8e2aff",
     "showTitle": false,
     "title": ""
    },
    "id": "T3MoVK_WkRJb"
   },
   "outputs": [],
   "source": [
    "#Multi-Headed Self Attention\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embd, n_embd)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "16267e9a-008b-46e3-82ce-2ae41396a1a1",
     "showTitle": false,
     "title": ""
    },
    "id": "T-w53_mSkRJb",
    "outputId": "052e4478-b284-43c7-924c-fcf13deea320"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "torch.Size([4, 8, 64])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Confirming that what's output from multi head attention is the original embedding size\n",
    "B,T,C = 4,8,64 # batch, time, channels\n",
    "x = torch.randn(B,T,C)\n",
    "mha = MultiHeadAttention(4,16)\n",
    "mha(x).shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "5f7ff128-7fe5-4a91-b9f2-208e2132e505",
     "showTitle": false,
     "title": ""
    },
    "id": "wNlJTtfhkRJb"
   },
   "source": [
    "### Creating an Expert module i.e. a simple Multi Layer Perceptron"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "f6e422a5-57c1-4b2f-b7b9-2757e109848a",
     "showTitle": false,
     "title": ""
    },
    "id": "zv3fGRpbkRJb"
   },
   "source": [
    "In the Sparse Mixture of Experts (MoE) architecture, the self-attention mechanism within each transformer block remains unchanged. However, a notable alteration occurs in the structure of each block: the standard feed-forward neural network is replaced with several sparsely activated feed-forward networks, known as experts. \"Sparse activation\" refers to the process where each token in the sequence is routed to only a limited number of these experts – typically one or two – out of the total pool available. This modification allows for specialized processing of different parts of the input data, enabling the model to handle a wider range of complexities efficiently."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "efe9fdcc-82eb-4047-9233-ad3cfe8759b1",
     "showTitle": false,
     "title": ""
    },
    "id": "7Kz0Y_P0kRJc"
   },
   "source": [
    "![experts](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/experts.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a2f0f382-ab4a-45e1-9dce-27ae0d3da641",
     "showTitle": false,
     "title": ""
    },
    "id": "a-9CYWXgkRJc"
   },
   "outputs": [],
   "source": [
    "#Expert module\n",
    "class Expert(nn.Module):\n",
    "    \"\"\" An MLP is a simple linear layer followed by a non-linearity i.e. each Expert \"\"\"\n",
    "\n",
    "    def __init__(self, n_embd):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embd, 4 * n_embd),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embd, n_embd),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "a7764385-26e9-4d75-9aa7-ce011023e24e",
     "showTitle": false,
     "title": ""
    },
    "id": "qderdEuykRJc"
   },
   "source": [
    "### Top-k Gating Intuition through an Example"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "d3fca4df-4c47-4e9a-98cd-08cf8ccf7726",
     "showTitle": false,
     "title": ""
    },
    "id": "VxJv5y44kRJc"
   },
   "source": [
    "![top k gating](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/topk.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "8e494b86-cdb2-4f2a-8824-5fa2ef4b2606",
     "showTitle": false,
     "title": ""
    },
    "id": "n6DuhY0DkRJc"
   },
   "source": [
    "The gating network, also known as the router, determines which expert network receives the output for each token from the multi-head attention. Let's consider a simple example: suppose there are 4 experts, and the token is to be routed to the top 2 experts. Initially, we input the token into the gating network through a linear layer. This layer projects the input tensor from a shape of (2, 4, 32) — representing (Batch size, Tokens, n_embed, where n_embed is the channel dimension of the input) — to a new shape of (2, 4, 4), which corresponds to (Batch size, Tokens, num_experts), where num_experts is the count of expert networks. Following this, we determine the top k=2 highest values and their respective indices along the last dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "621916ff-2290-4e2f-9fd7-5181ed98d540",
     "showTitle": false,
     "title": ""
    },
    "id": "pNAuFDDvkRJc",
    "outputId": "a9d313d5-1dd1-4df8-9c9a-dc6d0b90d4ed"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(tensor([[[ 0.9558,  0.1610],\n",
       "          [ 0.8659, -0.1494],\n",
       "          [ 0.8765,  0.7202],\n",
       "          [ 0.9496, -0.6609]],\n",
       " \n",
       "         [[ 0.4419, -0.2500],\n",
       "          [ 1.2602,  0.8430],\n",
       "          [ 0.8570,  0.7822],\n",
       "          [ 0.7376,  0.2561]]], grad_fn=<TopkBackward0>),\n",
       " tensor([[[2, 0],\n",
       "          [3, 2],\n",
       "          [3, 0],\n",
       "          [1, 2]],\n",
       " \n",
       "         [[1, 3],\n",
       "          [1, 2],\n",
       "          [1, 2],\n",
       "          [0, 1]]]))"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Understanding how gating works\n",
    "num_experts = 4\n",
    "top_k=2\n",
    "n_embed=32\n",
    "\n",
    "\n",
    "#Example multi-head attention output for a simple illustrative example, consider n_embed=32, context_length=4 and batch_size=2\n",
    "mh_output = torch.randn(2, 4, n_embed)\n",
    "\n",
    "topkgate_linear = nn.Linear(n_embed, num_experts) # nn.Linear(32, 4)\n",
    "\n",
    "logits = topkgate_linear(mh_output)\n",
    "top_k_logits, top_k_indices = logits.topk(top_k, dim=-1)  # Get top-k experts\n",
    "top_k_logits, top_k_indices"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0f135ff7-0aa3-4b6d-ab5e-42399c48427b",
     "showTitle": false,
     "title": ""
    },
    "id": "EKwAyJxrkRJd"
   },
   "source": [
    "Obtain the sparse gating output by only keeping the top k values in their respective index along the last dimension. Fill the rest with '-inf' and pass through a softmax activation. This pushes '-inf' values to zero, makes the top two values more accentuated and sum to 1. This summation to 1 helps with the weighting of expert outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "735e160a-ef1e-424d-b6d9-09f63ea99ec1",
     "showTitle": false,
     "title": ""
    },
    "id": "IiVejzOpkRJd",
    "outputId": "2ff61217-56f8-4835-d310-2169240fb4a1"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([[[ 0.1610,    -inf,  0.9558,    -inf],\n",
       "         [   -inf,    -inf, -0.1494,  0.8659],\n",
       "         [ 0.7202,    -inf,    -inf,  0.8765],\n",
       "         [   -inf,  0.9496, -0.6609,    -inf]],\n",
       "\n",
       "        [[   -inf,  0.4419,    -inf, -0.2500],\n",
       "         [   -inf,  1.2602,  0.8430,    -inf],\n",
       "         [   -inf,  0.8570,  0.7822,    -inf],\n",
       "         [ 0.7376,  0.2561,    -inf,    -inf]]], grad_fn=<ScatterBackward0>)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zeros = torch.full_like(logits, float('-inf')) #full_like clones a tensor and fills it with a specified value (like infinity) for masking or calculations.\n",
    "sparse_logits = zeros.scatter(-1, top_k_indices, top_k_logits)\n",
    "sparse_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "9146e6f9-4eee-4a8b-8338-55072719ed59",
     "showTitle": false,
     "title": ""
    },
    "id": "HFgRxDF4kRJh",
    "outputId": "cde71bbf-ba67-4972-eef1-2fba48a106d5"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "tensor([[[0.3111, 0.0000, 0.6889, 0.0000],\n",
       "         [0.0000, 0.0000, 0.2660, 0.7340],\n",
       "         [0.4610, 0.0000, 0.0000, 0.5390],\n",
       "         [0.0000, 0.8335, 0.1665, 0.0000]],\n",
       "\n",
       "        [[0.0000, 0.6664, 0.0000, 0.3336],\n",
       "         [0.0000, 0.6028, 0.3972, 0.0000],\n",
       "         [0.0000, 0.5187, 0.4813, 0.0000],\n",
       "         [0.6181, 0.3819, 0.0000, 0.0000]]], grad_fn=<SoftmaxBackward0>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gating_output= F.softmax(sparse_logits, dim=-1)\n",
    "gating_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fe558b12-e443-4b62-9a85-c59120456352",
     "showTitle": false,
     "title": ""
    },
    "id": "dGJrq2uqkRJh"
   },
   "source": [
    "### Generalizing and Modularizing above code and adding noisy top-k Gating for load balancing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "45516b59-d814-4853-a34e-d36aae9f04eb",
     "showTitle": false,
     "title": ""
    },
    "id": "TKp4DqwYkRJh"
   },
   "outputs": [],
   "source": [
    "# First define the top k router module\n",
    "class TopkRouter(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k):\n",
    "        super(TopkRouter, self).__init__()\n",
    "        self.top_k = top_k\n",
    "        self.linear =nn.Linear(n_embed, num_experts)\n",
    "\n",
    "    def forward(self, mh_ouput):\n",
    "        # mh_ouput is the output tensor from multihead self attention block\n",
    "        logits = self.linear(mh_output)\n",
    "        top_k_logits, indices = logits.topk(self.top_k, dim=-1)\n",
    "        zeros = torch.full_like(logits, float('-inf'))\n",
    "        sparse_logits = zeros.scatter(-1, indices, top_k_logits)\n",
    "        router_output = F.softmax(sparse_logits, dim=-1)\n",
    "        return router_output, indices\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c500844f-0866-4bbf-acef-c0c1d4979721",
     "showTitle": false,
     "title": ""
    },
    "id": "KjkouzwkkRJh",
    "outputId": "4e5cd861-4f75-4e98-c3cf-5697fb5f7451"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(torch.Size([2, 4, 4]),\n",
       " tensor([[[0.0000, 0.4249, 0.5751, 0.0000],\n",
       "          [0.3467, 0.6533, 0.0000, 0.0000],\n",
       "          [0.3970, 0.0000, 0.6030, 0.0000],\n",
       "          [0.0000, 0.0000, 0.7713, 0.2287]],\n",
       " \n",
       "         [[0.4043, 0.5957, 0.0000, 0.0000],\n",
       "          [0.0000, 0.5281, 0.0000, 0.4719],\n",
       "          [0.7053, 0.0000, 0.2947, 0.0000],\n",
       "          [0.0000, 0.4602, 0.5398, 0.0000]]], grad_fn=<SoftmaxBackward0>),\n",
       " tensor([[[2, 1],\n",
       "          [1, 0],\n",
       "          [2, 0],\n",
       "          [2, 3]],\n",
       " \n",
       "         [[1, 0],\n",
       "          [1, 3],\n",
       "          [0, 2],\n",
       "          [2, 1]]]))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Testing this out:\n",
    "num_experts = 4\n",
    "top_k = 2\n",
    "n_embd = 32\n",
    "\n",
    "mh_output = torch.randn(2, 4, n_embd)  # Example input\n",
    "top_k_gate = TopkRouter(n_embd, num_experts, top_k)\n",
    "gating_output, indices = top_k_gate(mh_output)\n",
    "gating_output.shape, gating_output, indices\n",
    "#And it works!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9fa02d0c-3688-4d01-811d-d0b2b851ab33",
     "showTitle": false,
     "title": ""
    },
    "id": "tAouN-GwkRJi"
   },
   "source": [
    "Althought the mixtral paper released recently does not make any mention of it, I believe Noisy top-k Gating is an important tool in training MoE models. Essentially, you don't want all the tokens to be sent to the same set of 'favored' experts. You want a fine balance of exploitation and exploration. For this purpose, to load balance, it is helpful to add standard normal noise to the logits from the gating linear layer. This makes training more efficient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e05b3306-b89f-4ebc-901b-f16398a925c2",
     "showTitle": false,
     "title": ""
    },
    "id": "aWeueE83kRJi"
   },
   "source": [
    "![noisy top-k gating](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/noisytopkgating.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "dda4d805-373c-48f7-9037-da08fbc06e64",
     "showTitle": false,
     "title": ""
    },
    "id": "ZBrN-w3JkRJi"
   },
   "outputs": [],
   "source": [
    "#Changing the above to accomodate noisy top-k gating\n",
    "class NoisyTopkRouter(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k):\n",
    "        super(NoisyTopkRouter, self).__init__()\n",
    "        self.top_k = top_k\n",
    "        #layer for router logits\n",
    "        self.topkroute_linear = nn.Linear(n_embed, num_experts)\n",
    "        self.noise_linear =nn.Linear(n_embed, num_experts)\n",
    "\n",
    "\n",
    "    def forward(self, mh_output):\n",
    "        # mh_ouput is the output tensor from multihead self attention block\n",
    "        logits = self.topkroute_linear(mh_output)\n",
    "\n",
    "        #Noise logits\n",
    "        noise_logits = self.noise_linear(mh_output)\n",
    "\n",
    "        #Adding scaled unit gaussian noise to the logits\n",
    "        noise = torch.randn_like(logits)*F.softplus(noise_logits)\n",
    "        noisy_logits = logits + noise\n",
    "\n",
    "        top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1)\n",
    "        zeros = torch.full_like(noisy_logits, float('-inf'))\n",
    "        sparse_logits = zeros.scatter(-1, indices, top_k_logits)\n",
    "        router_output = F.softmax(sparse_logits, dim=-1)\n",
    "        return router_output, indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a01a9d6b-fedb-427d-b0da-c3b2a75a8643",
     "showTitle": false,
     "title": ""
    },
    "id": "7Q6KcH9AkRJi",
    "outputId": "7aea5328-ba96-4c98-b550-042ecf440700"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "(torch.Size([2, 4, 8]),\n",
       " tensor([[[0.0000, 0.4304, 0.0000, 0.5696, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000, 0.4109, 0.0000, 0.0000, 0.0000, 0.5891],\n",
       "          [0.0000, 0.3708, 0.0000, 0.0000, 0.0000, 0.0000, 0.6292, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000, 0.0000, 0.4542, 0.0000, 0.5458, 0.0000]],\n",
       " \n",
       "         [[0.0000, 0.0000, 0.1747, 0.0000, 0.0000, 0.0000, 0.8253, 0.0000],\n",
       "          [0.0000, 0.0000, 0.5008, 0.4992, 0.0000, 0.0000, 0.0000, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000, 0.0000, 0.6160, 0.0000, 0.3840, 0.0000],\n",
       "          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1952, 0.0000, 0.8048]]],\n",
       "        grad_fn=<SoftmaxBackward0>),\n",
       " tensor([[[3, 1],\n",
       "          [7, 3],\n",
       "          [6, 1],\n",
       "          [6, 4]],\n",
       " \n",
       "         [[6, 2],\n",
       "          [2, 3],\n",
       "          [4, 6],\n",
       "          [7, 5]]]))"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Testing this out, again:\n",
    "num_experts = 8\n",
    "top_k = 2\n",
    "n_embd = 16\n",
    "\n",
    "mh_output = torch.randn(2, 4, n_embd)  # Example input\n",
    "noisy_top_k_gate = NoisyTopkRouter(n_embd, num_experts, top_k)\n",
    "gating_output, indices = noisy_top_k_gate(mh_output)\n",
    "gating_output.shape, gating_output, indices\n",
    "#It works!!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "076fa004-a165-42a7-b729-0bca8ad39418",
     "showTitle": false,
     "title": ""
    },
    "id": "XyKjpR-dkRJi"
   },
   "source": [
    "\n",
    "### Creating a sparse Mixture of Experts module\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6747b8de-0086-4cb0-8fbd-46ee95457eb9",
     "showTitle": false,
     "title": ""
    },
    "id": "UsRCy7i3kRJi"
   },
   "source": [
    "The primary aspect of this process involves the gating network's output. After acquiring these results, the top k values are selectively multiplied with the outputs from the corresponding top-k experts for a given token. This selective multiplication forms a weighted sum, which constitutes the SparseMoe block's output. The critical and challenging part of this process is to avoid unnecessary multiplications. It's essential to conduct forward passes only for the top_k experts and then compute this weighted sum. Performing forward passes for each expert would defeat the purpose of employing a sparse MoE, as it would no longer be sparse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d6809b3f-4be9-4859-b39e-24fcdd6c8d86",
     "showTitle": false,
     "title": ""
    },
    "id": "7dDUHU_IkRJi"
   },
   "outputs": [],
   "source": [
    "class SparseMoE(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k):\n",
    "        super(SparseMoE, self).__init__()\n",
    "        self.router = NoisyTopkRouter(n_embed, num_experts, top_k)\n",
    "        self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])\n",
    "        self.top_k = top_k\n",
    "\n",
    "    def forward(self, x):\n",
    "        gating_output, indices = self.router(x)\n",
    "        final_output = torch.zeros_like(x)\n",
    "\n",
    "        # Reshape inputs for batch processing\n",
    "        flat_x = x.view(-1, x.size(-1))\n",
    "        flat_gating_output = gating_output.view(-1, gating_output.size(-1))\n",
    "\n",
    "        # Process each expert in parallel\n",
    "        for i, expert in enumerate(self.experts):\n",
    "            # Create a mask for the inputs where the current expert is in top-k\n",
    "            expert_mask = (indices == i).any(dim=-1)\n",
    "            flat_mask = expert_mask.view(-1)\n",
    "\n",
    "            if flat_mask.any():\n",
    "                expert_input = flat_x[flat_mask]\n",
    "                expert_output = expert(expert_input)\n",
    "\n",
    "                # Extract and apply gating scores\n",
    "                gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1)\n",
    "                weighted_output = expert_output * gating_scores\n",
    "\n",
    "                # Update final output additively by indexing and adding\n",
    "                final_output[expert_mask] += weighted_output.squeeze(1)\n",
    "\n",
    "        return final_output\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "06239630-0a1c-47c9-976c-7770f3d82e18",
     "showTitle": false,
     "title": ""
    },
    "id": "q8kDLI1ukRJj",
    "outputId": "69e7b330-cd47-4094-99c2-7928f8b60d0d"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the final output: torch.Size([4, 8, 16])\ntensor([[[ 0.1048,  0.1091, -0.0641, -0.0028,  0.0128,  0.1013,  0.0984,\n           0.2361,  0.1531,  0.0541,  0.0117,  0.1020,  0.1968, -0.1290,\n           0.0773, -0.1298],\n         [-0.1560, -0.0066,  0.0165,  0.0999,  0.4629,  0.1595,  0.0000,\n          -0.2844, -0.0649,  0.1817,  0.0575,  0.1993, -0.0075, -0.1548,\n           0.1300, -0.1441],\n         [ 0.0521, -0.4851,  0.1351, -0.4214,  0.2063,  0.0647, -0.1414,\n          -0.0012, -0.0691, -0.0890, -0.1402, -0.0714, -0.5327,  0.1769,\n           0.0707,  0.0029],\n         [ 0.2238,  0.1036,  0.0479,  0.1971,  0.3933, -0.2208,  0.0695,\n           0.1594,  0.2549, -0.3261, -0.2240, -0.1214, -0.1737, -0.3955,\n          -0.0866, -0.3569],\n         [ 0.2065, -0.3067,  0.2048, -0.1736, -0.0092,  0.0581, -0.3415,\n          -0.4282,  0.4316, -0.1257, -0.0764,  0.0938, -0.3932,  0.0290,\n           0.2679, -0.2724],\n         [ 0.1621, -0.2511,  0.2731, -0.1064, -0.0121,  0.1772,  0.0242,\n          -0.4435,  0.1737, -0.0327, -0.1157,  0.0934, -0.1550, -0.1154,\n          -0.0055, -0.1116],\n         [-0.1760, -0.0041,  0.0457, -0.2011,  0.0762,  0.2753,  0.1200,\n          -0.2905, -0.0817,  0.1967, -0.0857,  0.1158, -0.0188, -0.0775,\n           0.2569, -0.0667],\n         [ 0.1873, -0.0629, -0.0265,  0.1152,  0.0093,  0.2998, -0.3798,\n          -0.2109,  0.0456,  0.0654,  0.0126,  0.0643,  0.1700, -0.1483,\n           0.1320, -0.0765]],\n\n        [[ 0.1749, -0.4452,  0.0375,  0.4030, -0.0589, -0.0557,  0.2202,\n          -0.2070,  0.2774, -0.1397, -0.3506,  0.1670,  0.0994, -0.1768,\n           0.0687,  0.0075],\n         [ 0.4911,  0.2073, -0.2944, -0.0929, -0.1730, -0.3299,  0.1641,\n           0.2849,  0.2833,  0.4697, -0.1884,  0.1149, -0.5412, -0.1892,\n           0.2454,  0.2670],\n         [-0.0258,  0.2709,  0.1265,  0.0305,  0.1662, -0.3579, -0.2588,\n          -0.0602,  0.1777,  0.3465,  0.0067, -0.1786,  0.1058,  0.0810,\n          -0.0162, -0.4533],\n         [ 0.2528,  0.1140,  0.4796,  0.0720,  0.1378, -0.2780,  0.1419,\n          -0.0351, -0.0084, -0.2817, -0.3290, -0.7016, -0.0943, -0.0448,\n          -0.4251, -0.3939],\n         [ 0.2818,  0.2597, -0.1666,  0.2278, -0.2895,  0.0326, -0.2536,\n          -0.3955,  0.3245, -0.2268,  0.0388, -0.0151, -0.0036, -0.1326,\n          -0.0476, -0.2762],\n         [-0.1127,  0.0482,  0.0661,  0.0378, -0.0070, -0.1059, -0.0311,\n          -0.0311, -0.1317, -0.0468, -0.0651, -0.0385,  0.0007, -0.1871,\n          -0.2204, -0.1286],\n         [-0.1871,  0.2141,  0.0721,  0.2008,  0.2033,  0.3226, -0.1237,\n          -0.3688,  0.0246,  0.0856, -0.1324,  0.2939, -0.2533,  0.1236,\n           0.2845, -0.0519],\n         [-0.1128, -0.1132,  0.0178,  0.2444, -0.0283,  0.1765, -0.0097,\n          -0.2997,  0.0694, -0.3719,  0.0000,  0.2617, -0.0845, -0.1418,\n           0.2802, -0.1770]],\n\n        [[ 0.3124, -0.0377,  0.0320, -0.0836,  0.0323, -0.0945, -0.1557,\n           0.1118,  0.0728, -0.0322, -0.2244,  0.3234, -0.0721, -0.0891,\n           0.3090, -0.0016],\n         [ 0.3595,  0.3793,  0.1171,  0.1065, -0.3076, -0.1722,  0.0586,\n           0.0023, -0.0287, -0.2390, -0.0453,  0.0524,  0.1820, -0.2429,\n           0.2909, -0.0966],\n         [ 0.1624, -0.2464,  0.1155,  0.2456,  0.1685,  0.2580,  0.1276,\n          -0.2908,  0.0963, -0.0264,  0.0522, -0.2774,  0.0813, -0.1886,\n           0.1026, -0.1154],\n         [ 0.3913, -0.3193,  0.1475, -0.5435, -0.0627, -0.2530, -0.0132,\n          -0.1300, -0.1488, -0.3335, -0.0973,  0.1744, -0.2725,  0.0305,\n          -0.0482, -0.1040],\n         [ 0.3110, -0.0463, -0.0096,  0.1301,  0.0410,  0.3452, -0.1977,\n          -0.3983,  0.0018, -0.2717, -0.0178,  0.1880, -0.1838, -0.0089,\n           0.0236, -0.0066],\n         [ 0.0556, -0.1992,  0.3387,  0.1902,  0.2421,  0.2000,  0.1573,\n          -0.3408,  0.4080,  0.0645, -0.1038,  0.0946,  0.0000, -0.0125,\n           0.1616, -0.2892],\n         [ 0.0638,  0.0880, -0.0223,  0.3853, -0.0637, -0.0483,  0.0840,\n          -0.0579,  0.0000, -0.1013, -0.2518, -0.0564, -0.0226, -0.1857,\n          -0.1649, -0.2703],\n         [ 0.1150,  0.2131,  0.1994,  0.1988, -0.2313, -0.2971, -0.0044,\n          -0.0937,  0.1935,  0.1543, -0.1097,  0.0407,  0.1786, -0.0691,\n          -0.0980, -0.3126]],\n\n        [[ 0.1079,  0.1596,  0.3365,  0.2271, -0.0702,  0.2927, -0.0158,\n          -0.0169,  0.0789,  0.0316,  0.2629, -0.1192, -0.0590, -0.0774,\n           0.1162,  0.0315],\n         [-0.0151, -0.1179,  0.0791,  0.6538, -0.0188, -0.3190,  0.2474,\n          -0.4115,  0.0786, -0.3171, -0.0767,  0.0731,  0.4347, -0.2968,\n          -0.2395,  0.1059],\n         [ 0.2747, -0.0377, -0.1290,  0.2299, -0.0701, -0.2335, -0.1921,\n           0.2715, -0.1677,  0.1207,  0.1637,  0.0312,  0.1185,  0.0470,\n           0.0718, -0.2176],\n         [ 0.4119, -0.2886, -0.1818, -0.5075,  0.2650, -0.3419,  0.2648,\n           0.0146, -0.2189,  0.1048,  0.3014,  0.2350, -0.4225, -0.1668,\n           0.1802, -0.1162],\n         [-0.0671,  0.0747, -0.2080,  0.1087, -0.0698, -0.0266,  0.0999,\n          -0.3547,  0.1129, -0.0540, -0.0605, -0.1544, -0.1951, -0.0874,\n           0.1364, -0.1657],\n         [ 0.3648, -0.2365,  0.0735, -0.2014,  0.0190, -0.0966, -0.0088,\n          -0.1481,  0.0574, -0.3550,  0.3277,  0.0671,  0.3020, -0.2079,\n          -0.1634, -0.2814],\n         [ 0.5786, -0.0495,  0.0461, -0.0771,  0.1385, -0.2441,  0.0118,\n           0.0471, -0.1069, -0.1958, -0.0022, -0.1445,  0.1712,  0.1308,\n          -0.0314, -0.0228],\n         [ 0.2225, -0.2066, -0.0654,  0.0314, -0.0041, -0.0515, -0.0731,\n          -0.1248,  0.1275,  0.2082, -0.0565, -0.0173, -0.2220, -0.0030,\n          -0.0262,  0.1187]]], grad_fn=<IndexPutBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "#Let's test this out\n",
    "num_experts = 8\n",
    "top_k = 2\n",
    "n_embd = 16\n",
    "dropout=0.1\n",
    "\n",
    "mh_output = torch.randn(4, 8, n_embd)  # Example multi-head attention output\n",
    "sparse_moe = SparseMoE(n_embd, num_experts, top_k)\n",
    "final_output = sparse_moe(mh_output)\n",
    "print(\"Shape of the final output:\", final_output.shape)\n",
    "print(final_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "7476ca07-a315-4108-aa77-46173a703ca2",
     "showTitle": false,
     "title": ""
    },
    "id": "l7GoxCi2kRJj"
   },
   "source": [
    "To emphasize, it's important to recognize that the magnitudes of the top_k experts output from the Router/ gating network, as illustrated in the code above, are also significant. These top_k indices identify the experts that are activated, and the magnitude of the values in those top_k dimensions determines their respective weighting. This concept of weighted summation is further highlighted in the diagram below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "d99b5dce-301e-4380-8263-b5cfb4136ab2",
     "showTitle": false,
     "title": ""
    },
    "id": "e9a_oQ2akRJj"
   },
   "source": [
    "![sparse MoE](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/sparseMoEfinal.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "9096728c-3d39-48c5-88bc-e7f17f91a1a0",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Introducing Expert Capacity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "11be1b2f-a099-4e55-9263-fa48f0e2aa4f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "When pretraining mixture of experts language models, or any large language model, the process typically spans several GPUs and often many machines. The way training is parallelized across these hardware resources is critical for balancing the computational load. However, if certain experts or a set of experts become overly favored—reflecting a bias towards exploitation over exploration—it can lead not only to potential performance issues in the model but also to an imbalance in the computational load across the cluster.\n",
    "\n",
    "The Switch Transformer implementation uses expert capacity to circumvent this.The expert capacity determines how many tokens each expert is responsible for during the training or inference process. It is defined based on the number of tokens in the batch and the number of available experts, often adjusted by a capacity factor. This factor allows for flexibility in the allocation, providing a buffer to account for variations in the data distribution and ensuring that no single expert becomes a bottleneck due to being overloaded. Hardware failures are common in training these large models for many weeks if not months, so this is quite important. \n",
    "\n",
    "Here's how expert capacity is generally calculated:\n",
    "\n",
    "Expert Capacity = (Tokens per batch / Number of experts) × Capacity factor\n",
    "Where:\n",
    "\n",
    "Tokens per batch is the total number of tokens present in a batch that needs to be processed.\n",
    "Number of experts is the total count of experts available in the MoE layer to process the data.\n",
    "Capacity factor is a multiplier used to adjust the base capacity (tokens per batch divided by the number of experts). A capacity factor greater than 1 allows each expert to handle a buffer above the evenly distributed share, accommodating imbalances in token assignment. The general range for this value is 1-1.25\n",
    "\n",
    "The following code block makes a slight adjustment to implement a simple version of expert capacity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2352f25c-bc72-4ce8-9fb0-aa560943e4d8",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class SparseMoE(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k, capacity_factor=1.0):\n",
    "        super(SparseMoE, self).__init__()\n",
    "        self.router = NoisyTopkRouter(n_embed, num_experts, top_k)\n",
    "        self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])\n",
    "        self.top_k = top_k\n",
    "        self.capacity_factor = capacity_factor\n",
    "        self.num_experts = num_experts\n",
    "    \n",
    "    def forward(self, x):\n",
    "    # Assuming x has shape [batch_size, seq_len, n_embd]\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        gating_output, indices = self.router(x)\n",
    "        final_output = torch.zeros_like(x)\n",
    "\n",
    "        # Flatten the batch and sequence dimensions to treat each token independently\n",
    "        flat_x = x.view(-1, x.size(-1))  # Now shape [batch_size * seq_len, n_embd]\n",
    "        flat_gating_output = gating_output.view(-1, gating_output.size(-1))\n",
    "\n",
    "        tokens_per_batch = batch_size * seq_len * self.top_k\n",
    "        expert_capacity = int((tokens_per_batch / self.num_experts) * self.capacity_factor)\n",
    "\n",
    "        updates = torch.zeros_like(flat_x)\n",
    "\n",
    "        for i, expert in enumerate(self.experts):\n",
    "            expert_mask = (indices == i).any(dim=-1)\n",
    "            flat_mask = expert_mask.view(-1)\n",
    "            selected_indices = torch.nonzero(flat_mask).squeeze(-1)\n",
    "\n",
    "            limited_indices = selected_indices[:expert_capacity] if selected_indices.numel() > expert_capacity else selected_indices\n",
    "            if limited_indices.numel() > 0:\n",
    "                expert_input = flat_x[limited_indices]\n",
    "                expert_output = expert(expert_input)\n",
    "\n",
    "                gating_scores = flat_gating_output[limited_indices, i].unsqueeze(1)\n",
    "                weighted_output = expert_output * gating_scores\n",
    "\n",
    "                updates.index_add_(0, limited_indices, weighted_output)\n",
    "\n",
    "        # Reshape updates to match the original dimensions of x\n",
    "        final_output += updates.view(batch_size, seq_len, -1)\n",
    "\n",
    "        return final_output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7b30c5c7-3b70-48ef-910c-e99337f3a2cf",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the final output: torch.Size([4, 8, 16])\ntensor([[[ 3.8829e-02, -6.1598e-02,  2.4070e-02, -5.5337e-01,  9.7003e-02,\n          -1.1740e-01,  1.4278e-01, -1.9145e-01,  4.1449e-02, -1.6396e-01,\n          -3.9448e-02, -7.1892e-02, -3.8065e-02,  1.5322e-01,  7.6931e-03,\n           3.0947e-01],\n         [-5.5785e-02,  6.3108e-02,  2.8818e-02,  1.2226e-03,  4.2681e-01,\n          -2.5701e-02, -3.2161e-02,  3.0531e-02, -8.8392e-02, -6.3882e-01,\n           9.0863e-02, -1.1883e-01, -6.2885e-02, -4.7814e-02, -2.5941e-02,\n           1.7633e-01],\n         [ 4.3926e-03, -5.8481e-02,  4.6390e-02, -5.1557e-01, -8.0752e-02,\n          -1.4103e-01, -1.1572e-01,  3.7564e-01, -1.8460e-02, -1.5760e-01,\n          -6.3076e-02,  1.3879e-01,  3.8779e-01, -5.6332e-02, -6.1815e-02,\n           1.9131e-01],\n         [-1.6490e-01, -2.0238e-01,  1.7096e-01,  3.3168e-01,  6.8391e-02,\n           3.4293e-01, -2.7043e-01,  5.1552e-01, -4.1907e-02, -5.7373e-02,\n          -4.2664e-02, -2.5897e-01,  4.9681e-02, -2.2637e-01, -2.5592e-01,\n          -4.3623e-01],\n         [ 1.1192e-01, -2.0017e-01,  1.5898e-01,  2.1665e-02, -1.5533e-02,\n           3.9343e-01,  3.3780e-01,  1.1301e-01, -1.6729e-01, -8.0209e-02,\n           2.3421e-01,  2.0428e-01,  6.2819e-02, -4.3882e-02, -4.9752e-01,\n          -3.3788e-01],\n         [ 3.0292e-01,  1.2632e-01,  7.0440e-02, -2.6895e-01, -2.7104e-01,\n          -1.0157e-01, -4.1019e-02, -1.3302e-01,  2.1430e-01, -1.2813e-01,\n           9.9809e-02,  7.7734e-03,  1.2782e-01, -9.7855e-02,  1.8376e-01,\n          -2.2943e-01],\n         [ 1.5236e-01,  3.1608e-02,  3.3809e-01, -8.0975e-02,  9.6067e-02,\n          -1.0134e-01, -3.4674e-01, -2.8861e-01, -1.4704e-01, -1.4246e-02,\n           1.7942e-01, -4.1562e-02, -1.4488e-02,  7.8012e-03, -2.3404e-01,\n          -3.0936e-01],\n         [ 1.8530e-01,  9.3396e-02, -1.7914e-03, -1.8016e-01, -6.9005e-02,\n          -5.5640e-02,  2.2173e-01,  6.4840e-02, -2.7157e-01, -1.4801e-01,\n           1.7874e-01, -1.8243e-01, -2.3770e-01,  2.5014e-01,  5.1655e-02,\n           2.5591e-01]],\n\n        [[ 7.8846e-02, -7.9402e-02,  4.2252e-04, -2.3439e-01, -4.6693e-01,\n           1.0347e-01, -2.1222e-01,  0.0000e+00, -4.8033e-01, -5.0988e-02,\n           3.3131e-01, -1.3476e-01,  4.0150e-02,  1.0579e-01,  1.6469e-02,\n          -2.7733e-01],\n         [ 2.2507e-01, -5.6083e-02, -2.6389e-01,  5.6235e-02, -1.4721e-01,\n           2.3475e-01,  1.2690e-01,  7.3004e-02, -4.0002e-02, -7.6029e-02,\n           1.2575e-01,  4.3562e-02, -2.3968e-01,  1.5215e-01, -3.3064e-01,\n           8.5029e-02],\n         [ 5.6429e-03,  8.0912e-02,  1.1466e-02,  1.1759e-01, -2.5671e-01,\n           1.6194e-01, -2.1085e-01, -7.2070e-01, -3.0520e-01,  9.9465e-02,\n           1.1205e-01, -8.0660e-02,  2.4345e-01,  1.7447e-02, -1.8797e-01,\n          -2.2256e-01],\n         [-9.1357e-02,  4.2767e-02, -1.3787e-01,  2.5129e-01, -1.4351e-01,\n           1.2715e-01, -2.9797e-02,  3.8052e-01,  8.9910e-02, -7.9867e-02,\n           1.5351e-01, -1.3429e-01,  1.8843e-01,  1.3405e-01, -1.1784e-01,\n          -7.6185e-02],\n         [-2.0273e-01, -2.7053e-01,  1.2292e-01,  2.0311e-01,  1.8657e-04,\n          -4.8982e-02, -2.0292e-01,  1.8525e-01,  2.1668e-01,  1.1099e-01,\n           4.1379e-02,  0.0000e+00, -1.6045e-01, -1.1878e-01, -1.6615e-01,\n          -1.9525e-02],\n         [ 8.1091e-02,  4.1072e-01, -4.0442e-01,  2.4271e-02, -4.0198e-01,\n          -4.8897e-02, -2.4575e-01, -1.4247e-01, -2.5412e-01,  3.1725e-03,\n           3.5476e-03, -2.7825e-01,  2.2469e-01,  1.2127e-01,  2.1675e-01,\n          -4.8169e-02],\n         [ 3.6590e-02,  1.9505e-01,  9.7517e-02, -9.5200e-02,  2.7198e-01,\n           7.9089e-02, -5.0379e-02,  2.3437e-01,  1.3500e-01, -1.4875e-01,\n           4.2517e-01, -6.2765e-02,  1.9841e-02, -1.1475e-01,  6.0201e-03,\n          -1.2854e-01],\n         [ 2.4987e-03, -6.9246e-02,  1.3846e-03, -2.8647e-01, -2.0792e-02,\n           1.3569e-01, -4.7075e-01,  2.2852e-01, -3.2622e-01, -1.5090e-01,\n           6.8632e-02,  2.3949e-02,  9.8145e-02,  2.1269e-01,  2.6956e-01,\n          -5.6896e-02]],\n\n        [[-1.2407e-02,  1.2390e-01,  5.4870e-02, -3.2936e-01, -9.1524e-02,\n           4.3641e-02,  3.4771e-02, -1.5105e-02, -1.7580e-01,  2.3287e-02,\n           3.5996e-01, -2.4644e-01, -2.9213e-01,  3.5063e-02, -4.4453e-02,\n           9.1115e-02],\n         [ 1.6025e-01, -2.1792e-01,  7.6735e-02, -3.1049e-01,  8.0726e-02,\n           1.7411e-01,  5.1778e-01, -1.6498e-02, -3.6364e-01,  7.8156e-02,\n          -6.5910e-03, -2.2307e-01, -1.7757e-01,  2.6452e-01,  7.0265e-02,\n          -1.1863e-01],\n         [-9.7254e-03, -6.0539e-02, -1.4289e-01, -1.8740e-01, -1.6023e-02,\n           2.7862e-01, -2.2776e-02,  1.7756e-01, -3.1297e-01, -2.4593e-02,\n          -5.0128e-03, -1.0907e-01,  6.4314e-02,  4.5431e-01, -2.2536e-02,\n           1.3231e-01],\n         [-6.5999e-02, -1.2334e-01,  2.2360e-01, -1.3675e-03, -4.6372e-03,\n          -2.2655e-01, -8.8361e-02, -2.5459e-01, -1.5959e-01,  8.2173e-02,\n          -3.5941e-02,  4.7639e-02, -2.0930e-01,  4.2207e-02,  3.6008e-01,\n           8.8476e-02],\n         [ 3.2619e-02,  3.4652e-01, -1.4679e-01, -1.3125e-01,  3.1324e-01,\n          -1.3008e-02, -2.7615e-01,  2.0545e-01, -6.3104e-02,  8.7131e-02,\n           2.0929e-01,  2.9709e-02,  9.2231e-02,  4.1322e-01, -5.6908e-01,\n          -3.9222e-03],\n         [ 1.8726e-01, -4.8134e-02,  6.0983e-02, -1.7235e-01, -9.3733e-02,\n           6.1528e-02, -4.0874e-01, -1.5726e-01, -2.0188e-01, -3.8612e-01,\n          -3.0585e-01, -3.0549e-01,  1.6624e-01,  8.5448e-02,  7.2860e-02,\n           1.9309e-01],\n         [ 1.2793e-01, -1.1460e-01, -1.3877e-01,  4.1917e-01,  2.8715e-01,\n          -1.7335e-01, -1.2060e-01,  6.6951e-01, -1.4283e-01,  1.6194e-01,\n           7.6279e-01,  6.6734e-02, -1.0459e-01,  1.1360e-01, -4.1824e-01,\n           0.0000e+00],\n         [-1.6618e-01, -1.3762e-01,  1.0223e-01, -5.6128e-02, -1.8210e-02,\n           2.0556e-01, -1.4219e-01,  1.3737e-01, -6.0496e-02,  1.0694e-02,\n           1.3115e-01, -3.7184e-02, -1.2207e-01, -3.6441e-02, -1.8151e-01,\n           1.4794e-01]],\n\n        [[ 5.5805e-02, -9.6174e-02, -5.0972e-02, -1.2038e-01,  2.4002e-01,\n           2.7006e-01, -1.2706e-01,  3.7488e-01,  1.2984e-01, -3.1188e-01,\n           1.6299e-01, -8.6462e-04,  1.7678e-03,  4.7784e-02, -3.1683e-01,\n           6.0736e-02],\n         [ 3.5152e-01, -2.0748e-02, -3.3030e-01, -1.4243e-01, -3.9356e-01,\n           3.6193e-02, -1.6062e-01,  5.3316e-02,  7.5414e-02, -2.1863e-02,\n           2.7194e-01, -1.7353e-01,  1.9504e-02,  1.5616e-01, -1.1125e-01,\n          -2.1409e-01],\n         [-6.1641e-02, -5.5712e-02, -1.7302e-01,  2.6501e-01,  5.1760e-02,\n          -2.2489e-01, -3.0600e-01,  1.8294e-01,  7.5629e-02,  2.1215e-02,\n           3.0244e-01, -1.4149e-01, -1.2973e-01, -5.8397e-02, -2.9961e-02,\n           1.3260e-01],\n         [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n           0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n           0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,\n           0.0000e+00],\n         [-3.0740e-02,  6.8911e-03, -1.3946e-01,  0.0000e+00, -9.9590e-02,\n           7.2630e-02, -1.5939e-01,  1.4464e-01,  4.0463e-03,  1.2086e-01,\n           9.5772e-02,  0.0000e+00,  1.9205e-02,  1.4282e-02, -7.4317e-02,\n          -4.6120e-02],\n         [-8.0736e-02,  9.0654e-02,  2.5160e-01,  1.3211e-01, -9.7253e-02,\n           3.4490e-02, -9.7416e-02,  4.5538e-02, -1.5034e-01,  8.2626e-02,\n           1.8621e-01, -3.3723e-02, -1.3703e-01,  5.5660e-02, -1.8671e-01,\n          -6.6035e-02],\n         [-2.1099e-01,  0.0000e+00,  1.0384e-01,  2.2306e-01,  4.2068e-01,\n          -2.4380e-01, -3.9424e-02,  4.6948e-02,  1.0136e-01,  0.0000e+00,\n           2.6688e-01,  4.2314e-02, -1.4266e-01,  8.2734e-02,  7.9087e-02,\n           1.6459e-01],\n         [ 0.0000e+00,  1.1030e-01,  1.6229e-02,  1.3760e-01, -5.8265e-02,\n           2.7115e-02, -1.0889e-01,  1.4462e-01,  1.7505e-01,  1.5387e-02,\n           7.3504e-02, -1.1557e-01,  4.1865e-02, -5.6453e-04, -2.2887e-02,\n          -8.7004e-02]]], grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "#Let's test this out\n",
    "num_experts = 8\n",
    "top_k = 2\n",
    "n_embd = 16\n",
    "dropout=0.1\n",
    "\n",
    "mh_output = torch.randn(4, 8, n_embd)  # Example multi-head attention output\n",
    "sparse_moe = SparseMoE(n_embd, num_experts, top_k)\n",
    "final_output = sparse_moe(mh_output)\n",
    "print(\"Shape of the final output:\", final_output.shape)\n",
    "print(final_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "da5f3be4-f155-4d6c-bcbd-2f88a088261f",
     "showTitle": false,
     "title": ""
    },
    "id": "vMZCda7dkRJj"
   },
   "source": [
    "### Putting it all together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0eaf71cd-c77e-40c7-b5be-e364e91685cf",
     "showTitle": false,
     "title": ""
    },
    "id": "f8yczkFHkRJj"
   },
   "outputs": [],
   "source": [
    "#First defining hyperparameters and boiler plate code. Imports and data preparation code is repeated for convenience\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.nn import init\n",
    "\n",
    "# hyperparameters\n",
    "batch_size = 16 # how many independent sequences will we process in parallel?\n",
    "block_size = 32 # what is the maximum context length for predictions?\n",
    "max_iters = 5000\n",
    "eval_interval = 100\n",
    "learning_rate = 1e-3\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "eval_iters = 400\n",
    "head_size = 16\n",
    "n_embed = 128\n",
    "n_head = 8\n",
    "n_layer = 8\n",
    "dropout = 0.1\n",
    "num_experts = 8\n",
    "top_k = 2\n",
    "# ------------\n",
    "\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "with open('input.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# here are all the unique characters that occur in this text\n",
    "chars = sorted(list(set(text)))\n",
    "vocab_size = len(chars)\n",
    "# create a mapping from characters to integers\n",
    "stoi = { ch:i for i,ch in enumerate(chars) }\n",
    "itos = { i:ch for i,ch in enumerate(chars) }\n",
    "encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers\n",
    "decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string\n",
    "\n",
    "# Train and test splits\n",
    "data = torch.tensor(encode(text), dtype=torch.long)\n",
    "n = int(0.9*len(data)) # first 90% will be train, rest val\n",
    "train_data = data[:n]\n",
    "val_data = data[n:]\n",
    "\n",
    "# data loading\n",
    "def get_batch(split):\n",
    "    # generate a small batch of data of inputs x and targets y\n",
    "    data = train_data if split == 'train' else val_data\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([data[i:i+block_size] for i in ix])\n",
    "    y = torch.stack([data[i+1:i+block_size+1] for i in ix])\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    return x, y\n",
    "\n",
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = get_batch(split)\n",
    "            logits, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ee1180f7-5004-4425-87fe-9a81a17b9024",
     "showTitle": false,
     "title": ""
    },
    "id": "QfxJ6B2fkRJj"
   },
   "outputs": [],
   "source": [
    "class Head(nn.Module):\n",
    "    \"\"\" one head of self-attention \"\"\"\n",
    "\n",
    "    def __init__(self, head_size):\n",
    "        super().__init__()\n",
    "        self.key = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.query = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.value = nn.Linear(n_embed, head_size, bias=False)\n",
    "        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B,T,C = x.shape\n",
    "        k = self.key(x)   # (B,T,C)\n",
    "        q = self.query(x) # (B,T,C)\n",
    "        # compute attention scores (\"affinities\")\n",
    "        wei = q @ k.transpose(-2,-1) * C**-0.5 # (B, T, C) @ (B, C, T) -> (B, T, T)\n",
    "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # (B, T, T)\n",
    "        wei = F.softmax(wei, dim=-1) # (B, T, T)\n",
    "        wei = self.dropout(wei)\n",
    "        # perform the weighted aggregation of the values\n",
    "        v = self.value(x) # (B,T,C)\n",
    "        out = wei @ v # (B, T, T) @ (B, T, C) -> (B, T, C)\n",
    "        return out\n",
    "\n",
    "#Multi-Headed Self Attention\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
    "\n",
    "    def __init__(self, num_heads, head_size):\n",
    "        super().__init__()\n",
    "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
    "        self.proj = nn.Linear(n_embed, n_embed)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
    "        out = self.dropout(self.proj(out))\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "03611a92-aaa2-4e0d-9755-cba56f96c794",
     "showTitle": false,
     "title": ""
    },
    "id": "y35jVCZYkRJk"
   },
   "outputs": [],
   "source": [
    "#Expert module\n",
    "class Expert(nn.Module):\n",
    "    \"\"\" An MLP is a simple linear layer followed by a non-linearity i.e. each Expert \"\"\"\n",
    "\n",
    "    def __init__(self, n_embed):\n",
    "        super().__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(n_embed, 4 * n_embed),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(4 * n_embed, n_embed),\n",
    "            nn.Dropout(dropout),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "#noisy top-k gating\n",
    "class NoisyTopkRouter(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k):\n",
    "        super(NoisyTopkRouter, self).__init__()\n",
    "        self.top_k = top_k\n",
    "        #layer for router logits\n",
    "        self.topkroute_linear = nn.Linear(n_embed, num_experts)\n",
    "        self.noise_linear =nn.Linear(n_embed, num_experts)\n",
    "\n",
    "\n",
    "    def forward(self, mh_output):\n",
    "        # mh_ouput is the output tensor from multihead self attention block\n",
    "        logits = self.topkroute_linear(mh_output)\n",
    "\n",
    "        #Noise logits\n",
    "        noise_logits = self.noise_linear(mh_output)\n",
    "\n",
    "        #Adding scaled unit gaussian noise to the logits\n",
    "        noise = torch.randn_like(logits)*F.softplus(noise_logits)\n",
    "        noisy_logits = logits + noise\n",
    "\n",
    "        top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1)\n",
    "        zeros = torch.full_like(noisy_logits, float('-inf'))\n",
    "        sparse_logits = zeros.scatter(-1, indices, top_k_logits)\n",
    "        router_output = F.softmax(sparse_logits, dim=-1)\n",
    "        return router_output, indices\n",
    "\n",
    "#Now create the sparse mixture of experts module\n",
    "\n",
    "\n",
    "class SparseMoE(nn.Module):\n",
    "    def __init__(self, n_embed, num_experts, top_k, capacity_factor=1.0):\n",
    "        super(SparseMoE, self).__init__()\n",
    "        self.router = NoisyTopkRouter(n_embed, num_experts, top_k)\n",
    "        self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])\n",
    "        self.top_k = top_k\n",
    "        self.capacity_factor = capacity_factor\n",
    "        self.num_experts = num_experts\n",
    "    \n",
    "    def forward(self, x):\n",
    "    # Assuming x has shape [batch_size, seq_len, n_embd]\n",
    "        batch_size, seq_len, _ = x.shape\n",
    "        gating_output, indices = self.router(x)\n",
    "        final_output = torch.zeros_like(x)\n",
    "\n",
    "        # Flatten the batch and sequence dimensions to treat each token independently\n",
    "        flat_x = x.view(-1, x.size(-1))  \n",
    "        flat_gating_output = gating_output.view(-1, gating_output.size(-1))\n",
    "\n",
    "        tokens_per_batch = batch_size * seq_len * self.top_k\n",
    "        expert_capacity = int((tokens_per_batch / self.num_experts) * self.capacity_factor)\n",
    "\n",
    "        updates = torch.zeros_like(flat_x)\n",
    "\n",
    "        for i, expert in enumerate(self.experts):\n",
    "            expert_mask = (indices == i).any(dim=-1)\n",
    "            flat_mask = expert_mask.view(-1)\n",
    "            selected_indices = torch.nonzero(flat_mask).squeeze(-1)\n",
    "\n",
    "            limited_indices = selected_indices[:expert_capacity] if selected_indices.numel() > expert_capacity else selected_indices\n",
    "            if limited_indices.numel() > 0:\n",
    "                expert_input = flat_x[limited_indices]\n",
    "                expert_output = expert(expert_input)\n",
    "\n",
    "                gating_scores = flat_gating_output[limited_indices, i].unsqueeze(1)\n",
    "                weighted_output = expert_output * gating_scores\n",
    "\n",
    "                updates.index_add_(0, limited_indices, weighted_output)\n",
    "\n",
    "        # Reshape updates to match the original dimensions of x\n",
    "        final_output += updates.view(batch_size, seq_len, -1)\n",
    "\n",
    "        return final_output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bfdff2bb-092f-41c8-9a33-c84e6f8d6633",
     "showTitle": false,
     "title": ""
    },
    "id": "jGiuRsSgkRJk"
   },
   "outputs": [],
   "source": [
    "#First create a self attention + mixture of experts block, that may be repeated several number of times\n",
    "#Copy pasting key architecture variables for clarity\n",
    "\n",
    "class Block(nn.Module):\n",
    "    \"\"\" Mixture of Experts Transformer block: communication followed by computation (multi-head self attention + SparseMoE) \"\"\"\n",
    "\n",
    "    def __init__(self, n_embed, n_head, num_experts, top_k):\n",
    "        # n_embed: embedding dimension, n_head: the number of heads we'd like\n",
    "        super().__init__()\n",
    "        head_size = n_embed // n_head\n",
    "        self.sa = MultiHeadAttention(n_head, head_size)\n",
    "        self.smoe = SparseMoE(n_embed, num_experts, top_k)\n",
    "        self.ln1 = nn.LayerNorm(n_embed)\n",
    "        self.ln2 = nn.LayerNorm(n_embed)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.sa(self.ln1(x))\n",
    "        x = x + self.smoe(self.ln2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "2d32a276-d0cc-4808-90d7-62441771af44",
     "showTitle": false,
     "title": ""
    },
    "id": "RpyZBA71kRJk"
   },
   "outputs": [],
   "source": [
    "#Finally putting it all together to crease a sparse mixture of experts language model\n",
    "class SparseMoELanguageModel(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        # each token directly reads off the logits for the next token from a lookup table\n",
    "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed)\n",
    "        self.position_embedding_table = nn.Embedding(block_size, n_embed)\n",
    "        self.blocks = nn.Sequential(*[Block(n_embed, n_head=n_head, num_experts=num_experts,top_k=top_k) for _ in range(n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(n_embed) # final layer norm\n",
    "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
    "\n",
    "    def forward(self, idx, targets=None):\n",
    "        B, T = idx.shape\n",
    "\n",
    "        # idx and targets are both (B,T) tensor of integers\n",
    "        tok_emb = self.token_embedding_table(idx) # (B,T,C)\n",
    "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,C)\n",
    "        x = tok_emb + pos_emb # (B,T,C)\n",
    "        x = self.blocks(x) # (B,T,C)\n",
    "        x = self.ln_f(x) # (B,T,C)\n",
    "        logits = self.lm_head(x) # (B,T,vocab_size)\n",
    "\n",
    "        if targets is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C = logits.shape\n",
    "            logits = logits.view(B*T, C)\n",
    "            targets = targets.view(B*T)\n",
    "            loss = F.cross_entropy(logits, targets)\n",
    "\n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, idx, max_new_tokens):\n",
    "        # idx is (B, T) array of indices in the current context\n",
    "        for _ in range(max_new_tokens):\n",
    "            # crop idx to the last block_size tokens\n",
    "            idx_cond = idx[:, -block_size:]\n",
    "            # get the predictions\n",
    "            logits, loss = self(idx_cond)\n",
    "            # focus only on the last time step\n",
    "            logits = logits[:, -1, :] # becomes (B, C)\n",
    "            # apply softmax to get probabilities\n",
    "            probs = F.softmax(logits, dim=-1) # (B, C)\n",
    "            # sample from the distribution\n",
    "            idx_next = torch.multinomial(probs, num_samples=1) # (B, 1)\n",
    "            # append sampled index to the running sequence\n",
    "            idx = torch.cat((idx, idx_next), dim=1) # (B, T+1)\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "622ba7ce-3f20-4820-8982-93f3d3b7be09",
     "showTitle": false,
     "title": ""
    },
    "id": "bBzQXmfykRJk"
   },
   "source": [
    "Kaiming He initialization is used here because of presence of ReLU activations in the experts. Feel free to experiment with Glorot initialization which is more commonly used in transformers. Jeremy Howard's Fastai Part 2 has an excellent lecture that implements these from scratch: https://course.fast.ai/Lessons/lesson17.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a6d3c057-08ee-4c1b-8013-6a88b2eadac5",
     "showTitle": false,
     "title": ""
    },
    "id": "guGaJqHbkRJk"
   },
   "outputs": [],
   "source": [
    "\n",
    "def kaiming_init_weights(m):\n",
    "    if isinstance (m, (nn.Linear)):\n",
    "        init.kaiming_normal_(m.weight)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "5b4d9525-8405-4a51-adda-661aba004e57",
     "showTitle": false,
     "title": ""
    },
    "id": "nJGGkXz4kRJl",
    "outputId": "8518ec23-caa0-4167-88c6-65a7c905743a"
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "SparseMoELanguageModel(\n",
       "  (token_embedding_table): Embedding(65, 128)\n",
       "  (position_embedding_table): Embedding(32, 128)\n",
       "  (blocks): Sequential(\n",
       "    (0): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (1): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (2): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (3): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (4): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (5): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (6): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (7): Block(\n",
       "      (sa): MultiHeadAttention(\n",
       "        (heads): ModuleList(\n",
       "          (0-7): 8 x Head(\n",
       "            (key): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (query): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (value): Linear(in_features=128, out_features=16, bias=False)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (proj): Linear(in_features=128, out_features=128, bias=True)\n",
       "        (dropout): Dropout(p=0.1, inplace=False)\n",
       "      )\n",
       "      (smoe): SparseMoE(\n",
       "        (router): NoisyTopkRouter(\n",
       "          (topkroute_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "          (noise_linear): Linear(in_features=128, out_features=8, bias=True)\n",
       "        )\n",
       "        (experts): ModuleList(\n",
       "          (0-7): 8 x Expert(\n",
       "            (net): Sequential(\n",
       "              (0): Linear(in_features=128, out_features=512, bias=True)\n",
       "              (1): ReLU()\n",
       "              (2): Linear(in_features=512, out_features=128, bias=True)\n",
       "              (3): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "      (ln1): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "      (ln2): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (ln_f): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
       "  (lm_head): Linear(in_features=128, out_features=65, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = SparseMoELanguageModel()\n",
    "model.apply(kaiming_init_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6adf1d04-e668-4d14-b691-161ea4e4dccf",
     "showTitle": false,
     "title": ""
    },
    "id": "_cb1-L8ckRJl"
   },
   "source": [
    "I have used mlflow to track and log the metrics I care about and the training hyperparameters. The training loop in the next cell includes this mlflow code. If you prefer to just train without using mlflow, the subsequent cell has code without the mlflow code. However, I find it very convenient to track parameters and metrics, particularly when experimenting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b8968247-0d7b-4460-b96b-06743b31c55d",
     "showTitle": false,
     "title": ""
    },
    "id": "WTG1Fv4SkRJl",
    "outputId": "38318015-63a9-4959-cfe0-cab305625f49"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.996545 M parameters\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unexpected internal error when monkey patching `PreTrainedModel.from_pretrained`: Failed to import transformers.modeling_utils because of the following error (look up to see its traceback):\nFailed to import transformers.generation.utils because of the following error (look up to see its traceback):\ncannot import name 'allow_in_graph' from partially initialized module 'torch._dynamo' (most likely due to a circular import) (/databricks/python/lib/python3.11/site-packages/torch/_dynamo/__init__.py)\nUnexpected internal error when monkey patching `Trainer.train`: Failed to import transformers.trainer because of the following error (look up to see its traceback):\nFailed to import transformers.integrations.integration_utils because of the following error (look up to see its traceback):\ncannot import name 'allow_in_graph' from partially initialized module 'torch._dynamo' (most likely due to a circular import) (/databricks/python/lib/python3.11/site-packages/torch/_dynamo/__init__.py)\n"
     ]
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 5.3552, val loss 5.3479\nstep 100: train loss 2.7525, val loss 2.7736\nstep 200: train loss 2.5523, val loss 2.5550\nstep 300: train loss 2.4612, val loss 2.4580\nstep 400: train loss 2.3730, val loss 2.3827\nstep 500: train loss 2.3048, val loss 2.3181\nstep 600: train loss 2.2294, val loss 2.2478\nstep 700: train loss 2.1690, val loss 2.1955\nstep 800: train loss 2.1177, val loss 2.1481\nstep 900: train loss 2.0738, val loss 2.1403\nstep 1000: train loss 2.0255, val loss 2.0911\nstep 1100: train loss 1.9962, val loss 2.0733\nstep 1200: train loss 1.9641, val loss 2.0472\nstep 1300: train loss 1.9384, val loss 2.0331\nstep 1400: train loss 1.9026, val loss 2.0157\nstep 1500: train loss 1.8949, val loss 1.9966\nstep 1600: train loss 1.8581, val loss 1.9857\nstep 1700: train loss 1.8488, val loss 1.9719\nstep 1800: train loss 1.8278, val loss 1.9528\nstep 1900: train loss 1.8124, val loss 1.9621\nstep 2000: train loss 1.7862, val loss 1.9346\nstep 2100: train loss 1.7787, val loss 1.9316\nstep 2200: train loss 1.7708, val loss 1.9130\nstep 2300: train loss 1.7569, val loss 1.9206\nstep 2400: train loss 1.7475, val loss 1.8879\nstep 2500: train loss 1.7310, val loss 1.8834\nstep 2600: train loss 1.7142, val loss 1.8605\nstep 2700: train loss 1.7074, val loss 1.8784\nstep 2800: train loss 1.7036, val loss 1.8659\nstep 2900: train loss 1.6860, val loss 1.8584\nstep 3000: train loss 1.6883, val loss 1.8537\nstep 3100: train loss 1.6791, val loss 1.8366\nstep 3200: train loss 1.6685, val loss 1.8455\nstep 3300: train loss 1.6611, val loss 1.8146\nstep 3400: train loss 1.6556, val loss 1.8127\nstep 3500: train loss 1.6467, val loss 1.8172\nstep 3600: train loss 1.6428, val loss 1.8097\nstep 3700: train loss 1.6480, val loss 1.8142\nstep 3800: train loss 1.6298, val loss 1.7940\nstep 3900: train loss 1.6171, val loss 1.7921\nstep 4000: train loss 1.6206, val loss 1.8008\nstep 4100: train loss 1.6155, val loss 1.7896\nstep 4200: train loss 1.6126, val loss 1.7892\nstep 4300: train loss 1.6082, val loss 1.7895\nstep 4400: train loss 1.5957, val loss 1.7725\nstep 4500: train loss 1.5955, val loss 1.7819\nstep 4600: train loss 1.5819, val loss 1.7706\nstep 4700: train loss 1.5816, val loss 1.7621\nstep 4800: train loss 1.5814, val loss 1.7790\nstep 4900: train loss 1.5757, val loss 1.7675\nstep 4999: train loss 1.5713, val loss 1.7645\n"
     ]
    }
   ],
   "source": [
    "#Using MLFlow\n",
    "m = model.to(device)\n",
    "# print the number of parameters in the model\n",
    "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "#mlflow.set_experiment(\"makeMoE\")\n",
    "with mlflow.start_run():\n",
    "    #If you use mlflow.autolog() this will be automatically logged. I chose to explicitly log here for completeness\n",
    "    params = {\"batch_size\": batch_size , \"block_size\" : block_size, \"max_iters\": max_iters, \"eval_interval\": eval_interval,\n",
    "              \"learning_rate\": learning_rate, \"device\": device, \"eval_iters\": eval_iters, \"dropout\" : dropout, \"num_experts\": num_experts, \"top_k\": top_k }\n",
    "    mlflow.log_params(params)\n",
    "    for iter in range(max_iters):\n",
    "\n",
    "        # every once in a while evaluate the loss on train and val sets\n",
    "        if iter % eval_interval == 0 or iter == max_iters - 1:\n",
    "            losses = estimate_loss()\n",
    "            print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "            metrics = {\"train_loss\": float(losses['train']), \"val_loss\": float(losses['val'])}\n",
    "            mlflow.log_metrics(metrics, step=iter)\n",
    "\n",
    "\n",
    "        # sample a batch of data\n",
    "        xb, yb = get_batch('train')\n",
    "\n",
    "        # evaluate the loss\n",
    "        logits, loss = model(xb, yb)\n",
    "        optimizer.zero_grad(set_to_none=True)\n",
    "        loss.backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "1ed96085-c292-4624-a2cc-be8aad38df79",
     "showTitle": false,
     "title": ""
    },
    "id": "jFBVfgfekRJl"
   },
   "source": [
    "Logging train and validation losses gives you a good indication of how the training is going. The plot shows that I probably should have stopped around 4500 steps (when the validation loss jumps up a bit)\n",
    "\n",
    "![mlflow_dash](https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/images/mlflow_dash.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6360e1b7-94c4-4ef1-a850-9bc93f49a083",
     "showTitle": false,
     "title": ""
    },
    "id": "WP_2lcRUkRJl"
   },
   "outputs": [],
   "source": [
    "#Not using MLflow\n",
    "m = model.to(device)\n",
    "# print the number of parameters in the model\n",
    "print(sum(p.numel() for p in m.parameters())/1e6, 'M parameters')\n",
    "\n",
    "# create a PyTorch optimizer\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for iter in range(max_iters):\n",
    "\n",
    "    # every once in a while evaluate the loss on train and val sets\n",
    "    if iter % eval_interval == 0 or iter == max_iters - 1:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step {iter}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "\n",
    "    # sample a batch of data\n",
    "    xb, yb = get_batch('train')\n",
    "\n",
    "    # evaluate the loss\n",
    "    logits, loss = model(xb, yb)\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optimizer.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "8aa6e4c4-c688-4985-a3b8-e2af1f771e54",
     "showTitle": false,
     "title": ""
    },
    "id": "W4yshpXMkRJl",
    "outputId": "e69af1c7-2e2f-475c-eeba-b20dec8532c2"
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\nae\nold, buem neran dose: bo'ld bua woold nor may siver!\n\nLyYORTOLNSIA:\nyou, what, recooatusan that thou mody!\n\nBuNhRUTUO SINE:\n\nFrotsef'O Rich Lenger:\n\nreemendsseou:\n groung at yo thes Eave hers how had? Whht blse\nthe tronighnest with of that but mind son\nWhat a gentle, wo thave rear amore off this the poma thesrixe.\n\nFUY BLOKE VFIY BOLLUCVGge:Dlectst yook; a goad, them.\n\nCOMINIALE:\nvilny an? with! Priviini modeisus!\nThAsesd uproe--\n\nCORuTHy.\nFRray;\nYoo,e sads yoush higher waslif you minhim! he'f over hrins miet.\n\nCAMILLO:\nHase dost coe must nou, and live thoug fel!\n\nJuCeIvg:\nI o? no there, arm you. cand his cannow thmake a stopworn fothin aling:\nSive lauge meot an to the sort uae: ne\nIus ast begent agrody thugh. Some shimRs of dele, so now.\nWhot, Whut not hase one's presier broberes\nThatt him sent my nainher's preself,\nBouviu. Asfriet me mays server yet bew wall if wragaft its:\nBut it lover faminzs, havH, livNa rhav,\nHow shat add nobty quars that gost lamor.\n\nKING RICHARR LIEDWD IRDkRkol'Er,\nMysolvilgoo- yeary, and cond me you pass: noth and never\nFattion's your foARar-give.\nBuegh fortain asiumb, hisnow brou;\nOre bleat at leedo'tung. Trius me so, god'd.\n\nFARUTIROza your Jausy. WOisPry, there\ncome a Divitiove, in'these: ao shed thes\nas thessertiend too! and her, ses'gions to herr!\n\nGLOULEEEOTROMROEE:\nBnoy, mad yot deed, maith, right--we and welcomm-not- manurse:\nI hast spinl give twake yor foe;\nYour and furteher? and velage gembe aman:\nWeat nonge show but but in wh eyer:\nBeoess\npe'tion tha bpooe kew wikno bott\nAzgiven egefcer'd in latter of Llors\nThers I, punibola birt sabeimiing,\nFor men sticce the waatef hill wouth?\nMishn!\n\nKukest Lefeit:\nHave lest tomene it so\nAnd suppupoit; he wear ny anot in.\n\nLAOLY:\nYet, as merry, I will dow, name, he  now;\nWhom hwse didg thee ? uts;\nWhurt, 'I,ow't, but cold oat colk.\n\nNo OLeVTs Marniust:\nDidsioek'ill him!- then, spray;s,\nClest thenou,denned friars andieggo. That Ise, be we thind tone, my courtague,\nAnd fever senp trout of the\n"
     ]
    }
   ],
   "source": [
    "# generate from the model. Not great. Not too bad either\n",
    "context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
    "print(decode(m.generate(context, max_new_tokens=2000)[0].tolist()))"
   ]
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "makeMoE_from_Scratch_with_Expert_Capacity",
   "widgets": {}
  },
  "colab": {
   "include_colab_link": true,
   "provenance": []
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}