makeMoE_with_Expert.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Sparse mixture of experts instead of the solitary feed forward neural net. \n",
    "- Top-k gating and noisy top-k gating implementations."
   ]
  },
  {
   "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": null,
   "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": [
    {
     "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\n",
      "Collecting 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)\n",
      "Requirement already satisfied: sqlalchemy<3,>=1.4.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.4.39)\n",
      "Requirement already satisfied: packaging<24 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (23.2)\n",
      "Requirement already satisfied: markdown<4,>=3.3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.4.1)\n",
      "Requirement already satisfied: matplotlib<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.7.2)\n",
      "Requirement already satisfied: gitpython<4,>=3.1.9 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.1.27)\n",
      "Requirement already satisfied: sqlparse<1,>=0.4.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (0.4.2)\n",
      "Collecting 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)\n",
      "Requirement already satisfied: pytz<2025 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2022.7)\n",
      "Requirement already satisfied: pyarrow<16,>=4.0.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (14.0.1)\n",
      "Requirement already satisfied: protobuf<5,>=3.12.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (4.24.1)\n",
      "Requirement already satisfied: scikit-learn<2 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.3.0)\n",
      "Requirement already satisfied: numpy<2 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (1.23.5)\n",
      "Requirement already satisfied: gunicorn<22 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (20.1.0)\n",
      "Requirement already satisfied: pandas<3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.0.3)\n",
      "Requirement already satisfied: importlib-metadata!=4.7.0,<8,>=3.7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (6.0.0)\n",
      "Requirement already satisfied: click<9,>=7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (8.0.4)\n",
      "Requirement already satisfied: entrypoints<1 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (0.4)\n",
      "Requirement already satisfied: Flask<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.2.5)\n",
      "Requirement already satisfied: cloudpickle<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.2.1)\n",
      "Requirement already satisfied: requests<3,>=2.17.3 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (2.31.0)\n",
      "Collecting querystring-parser<2\n",
      "  Downloading querystring_parser-1.2.4-py2.py3-none-any.whl (7.9 kB)\n",
      "Requirement already satisfied: Jinja2<4,>=2.11 in /databricks/python3/lib/python3.11/site-packages (from mlflow) (3.1.2)\n",
      "Collecting 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)\n",
      "Requirement already satisfied: typing-extensions>=4 in /databricks/python3/lib/python3.11/site-packages (from alembic!=1.10.0,<2->mlflow) (4.7.1)\n",
      "Requirement already satisfied: urllib3>=1.26.0 in /databricks/python3/lib/python3.11/site-packages (from docker<8,>=4.0.0->mlflow) (1.26.16)\n",
      "Requirement already satisfied: Werkzeug>=2.2.2 in /databricks/python3/lib/python3.11/site-packages (from Flask<4->mlflow) (2.2.3)\n",
      "Requirement already satisfied: itsdangerous>=2.0 in /databricks/python3/lib/python3.11/site-packages (from Flask<4->mlflow) (2.0.1)\n",
      "Requirement 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)\n",
      "Collecting graphql-relay<3.3,>=3.1\n",
      "  Downloading graphql_relay-3.2.0-py3-none-any.whl (16 kB)\n",
      "Collecting 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)\n",
      "Requirement 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)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /databricks/python3/lib/python3.11/site-packages (from Jinja2<4,>=2.11->mlflow) (2.1.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (0.11.0)\n",
      "Requirement already satisfied: pyparsing<3.1,>=2.3.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (3.0.9)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (1.4.4)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (1.0.5)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (2.8.2)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (9.4.0)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /databricks/python3/lib/python3.11/site-packages (from matplotlib<4->mlflow) (4.25.0)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /databricks/python3/lib/python3.11/site-packages (from pandas<3->mlflow) (2022.1)\n",
      "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from querystring-parser<2->mlflow) (1.16.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (3.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (2023.7.22)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow) (2.0.4)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /databricks/python3/lib/python3.11/site-packages (from scikit-learn<2->mlflow) (1.2.0)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /databricks/python3/lib/python3.11/site-packages (from scikit-learn<2->mlflow) (2.2.0)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /databricks/python3/lib/python3.11/site-packages (from sqlalchemy<3,>=1.4.0->mlflow) (2.0.1)\n",
      "Requirement 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)\n",
      "Installing collected packages: aniso8601, querystring-parser, graphql-core, graphql-relay, docker, alembic, graphene, mlflow\n",
      "Successfully 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": null,
   "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--2024-03-17 20:52:21--  https://raw.githubusercontent.com/AviSoori1x/makeMoE/main/input.txt\r\n",
      "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\r\n",
      "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\r\n",
      "HTTP request sent, awaiting response... 200 OK\r\n",
      "Length: 1115394 (1.1M) [text/plain]\r\n",
      "Saving 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\n",
      "2024-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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You are all resolved rather to die than to famish?\n",
      "\n",
      "All:\n",
      "Resolved. resolved.\n",
      "\n",
      "First Citizen:\n",
      "First, you know Caius Marcius is chief enemy to the people.\n",
      "\n",
      "All:\n",
      "We know't, we know't.\n",
      "\n",
      "First Citizen:\n",
      "Let us kill him, and we'll have corn at our own price.\n",
      "Is't a verdict?\n",
      "\n",
      "All:\n",
      "No more talking on't; let it be done: away, away!\n",
      "\n",
      "Second Citizen:\n",
      "One word, good citizens.\n",
      "\n",
      "First Citizen:\n",
      "We are accounted poor citizens, the patricians good.\n",
      "What authority surfeits on would relieve us: if they\n",
      "would yield us but the superfluity, while it were\n",
      "wholesome, we might guess they relieved us humanely;\n",
      "but they think we are too dear: the leanness that\n",
      "afflicts us, the object of our misery, is as an\n",
      "inventory to particularise their abundance; our\n",
      "sufferance is a gain to them Let us revenge this with\n",
      "our pikes, ere we become rakes: for the gods know I\n",
      "speak 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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n",
      "65\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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[46, 47, 47, 1, 58, 46, 43, 56, 43]\n",
      "hii 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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1115394]) torch.int64\n",
      "tensor([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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "when input is tensor([18]) the target: 47\n",
      "when input is tensor([18, 47]) the target: 56\n",
      "when input is tensor([18, 47, 56]) the target: 57\n",
      "when input is tensor([18, 47, 56, 57]) the target: 58\n",
      "when input is tensor([18, 47, 56, 57, 58]) the target: 1\n",
      "when input is tensor([18, 47, 56, 57, 58,  1]) the target: 15\n",
      "when input is tensor([18, 47, 56, 57, 58,  1, 15]) the target: 47\n",
      "when 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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs:\n",
      "torch.Size([4, 8])\n",
      "tensor([[ 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]])\n",
      "targets:\n",
      "torch.Size([4, 8])\n",
      "tensor([[ 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",
      "----\n",
      "when input is [6] the target: 0\n",
      "when input is [6, 0] the target: 14\n",
      "when input is [6, 0, 14] the target: 43\n",
      "when input is [6, 0, 14, 43] the target: 44\n",
      "when input is [6, 0, 14, 43, 44] the target: 53\n",
      "when input is [6, 0, 14, 43, 44, 53] the target: 56\n",
      "when input is [6, 0, 14, 43, 44, 53, 56] the target: 43\n",
      "when input is [6, 0, 14, 43, 44, 53, 56, 43] the target: 1\n",
      "when input is [39] the target: 1\n",
      "when input is [39, 1] the target: 42\n",
      "when input is [39, 1, 42] the target: 59\n",
      "when input is [39, 1, 42, 59] the target: 43\n",
      "when input is [39, 1, 42, 59, 43] the target: 1\n",
      "when input is [39, 1, 42, 59, 43, 1] the target: 39\n",
      "when input is [39, 1, 42, 59, 43, 1, 39] the target: 52\n",
      "when input is [39, 1, 42, 59, 43, 1, 39, 52] the target: 42\n",
      "when input is [47] the target: 41\n",
      "when input is [47, 41] the target: 43\n",
      "when input is [47, 41, 43] the target: 1\n",
      "when input is [47, 41, 43, 1] the target: 39\n",
      "when input is [47, 41, 43, 1, 39] the target: 52\n",
      "when input is [47, 41, 43, 1, 39, 52] the target: 42\n",
      "when input is [47, 41, 43, 1, 39, 52, 42] the target: 1\n",
      "when input is [47, 41, 43, 1, 39, 52, 42, 1] the target: 42\n",
      "when input is [53] the target: 44\n",
      "when input is [53, 44] the target: 1\n",
      "when input is [53, 44, 1] the target: 50\n",
      "when input is [53, 44, 1, 50] the target: 43\n",
      "when input is [53, 44, 1, 50, 43] the target: 58\n",
      "when input is [53, 44, 1, 50, 43, 58] the target: 1\n",
      "when input is [53, 44, 1, 50, 43, 58, 1] the target: 58\n",
      "when 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": null,
   "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": [
    {
     "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": null,
   "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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the final output: torch.Size([4, 8, 16])\n",
      "tensor([[[ 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": null,
   "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": null,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7b30c5c7-3b70-48ef-910c-e99337f3a2cf",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of the final output: torch.Size([4, 8, 16])\n",
      "tensor([[[ 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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": null,
   "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": [
    {
     "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8.996545 M parameters\n"
     ]
    },
    {
     "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):\n",
      "Failed to import transformers.generation.utils because of the following error (look up to see its traceback):\n",
      "cannot 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",
      "Unexpected internal error when monkey patching `Trainer.train`: Failed to import transformers.trainer because of the following error (look up to see its traceback):\n",
      "Failed to import transformers.integrations.integration_utils because of the following error (look up to see its traceback):\n",
      "cannot 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: train loss 5.3552, val loss 5.3479\n",
      "step 100: train loss 2.7525, val loss 2.7736\n",
      "step 200: train loss 2.5523, val loss 2.5550\n",
      "step 300: train loss 2.4612, val loss 2.4580\n",
      "step 400: train loss 2.3730, val loss 2.3827\n",
      "step 500: train loss 2.3048, val loss 2.3181\n",
      "step 600: train loss 2.2294, val loss 2.2478\n",
      "step 700: train loss 2.1690, val loss 2.1955\n",
      "step 800: train loss 2.1177, val loss 2.1481\n",
      "step 900: train loss 2.0738, val loss 2.1403\n",
      "step 1000: train loss 2.0255, val loss 2.0911\n",
      "step 1100: train loss 1.9962, val loss 2.0733\n",
      "step 1200: train loss 1.9641, val loss 2.0472\n",
      "step 1300: train loss 1.9384, val loss 2.0331\n",
      "step 1400: train loss 1.9026, val loss 2.0157\n",
      "step 1500: train loss 1.8949, val loss 1.9966\n",
      "step 1600: train loss 1.8581, val loss 1.9857\n",
      "step 1700: train loss 1.8488, val loss 1.9719\n",
      "step 1800: train loss 1.8278, val loss 1.9528\n",
      "step 1900: train loss 1.8124, val loss 1.9621\n",
      "step 2000: train loss 1.7862, val loss 1.9346\n",
      "step 2100: train loss 1.7787, val loss 1.9316\n",
      "step 2200: train loss 1.7708, val loss 1.9130\n",
      "step 2300: train loss 1.7569, val loss 1.9206\n",
      "step 2400: train loss 1.7475, val loss 1.8879\n",
      "step 2500: train loss 1.7310, val loss 1.8834\n",
      "step 2600: train loss 1.7142, val loss 1.8605\n",
      "step 2700: train loss 1.7074, val loss 1.8784\n",
      "step 2800: train loss 1.7036, val loss 1.8659\n",
      "step 2900: train loss 1.6860, val loss 1.8584\n",
      "step 3000: train loss 1.6883, val loss 1.8537\n",
      "step 3100: train loss 1.6791, val loss 1.8366\n",
      "step 3200: train loss 1.6685, val loss 1.8455\n",
      "step 3300: train loss 1.6611, val loss 1.8146\n",
      "step 3400: train loss 1.6556, val loss 1.8127\n",
      "step 3500: train loss 1.6467, val loss 1.8172\n",
      "step 3600: train loss 1.6428, val loss 1.8097\n",
      "step 3700: train loss 1.6480, val loss 1.8142\n",
      "step 3800: train loss 1.6298, val loss 1.7940\n",
      "step 3900: train loss 1.6171, val loss 1.7921\n",
      "step 4000: train loss 1.6206, val loss 1.8008\n",
      "step 4100: train loss 1.6155, val loss 1.7896\n",
      "step 4200: train loss 1.6126, val loss 1.7892\n",
      "step 4300: train loss 1.6082, val loss 1.7895\n",
      "step 4400: train loss 1.5957, val loss 1.7725\n",
      "step 4500: train loss 1.5955, val loss 1.7819\n",
      "step 4600: train loss 1.5819, val loss 1.7706\n",
      "step 4700: train loss 1.5816, val loss 1.7621\n",
      "step 4800: train loss 1.5814, val loss 1.7790\n",
      "step 4900: train loss 1.5757, val loss 1.7675\n",
      "step 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": null,
   "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": null,
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "ae\n",
      "old, buem neran dose: bo'ld bua woold nor may siver!\n",
      "\n",
      "LyYORTOLNSIA:\n",
      "you, what, recooatusan that thou mody!\n",
      "\n",
      "BuNhRUTUO SINE:\n",
      "\n",
      "Frotsef'O Rich Lenger:\n",
      "\n",
      "reemendsseou:\n",
      " groung at yo thes Eave hers how had? Whht blse\n",
      "the tronighnest with of that but mind son\n",
      "What a gentle, wo thave rear amore off this the poma thesrixe.\n",
      "\n",
      "FUY BLOKE VFIY BOLLUCVGge:Dlectst yook; a goad, them.\n",
      "\n",
      "COMINIALE:\n",
      "vilny an? with! Priviini modeisus!\n",
      "ThAsesd uproe--\n",
      "\n",
      "CORuTHy.\n",
      "FRray;\n",
      "Yoo,e sads yoush higher waslif you minhim! he'f over hrins miet.\n",
      "\n",
      "CAMILLO:\n",
      "Hase dost coe must nou, and live thoug fel!\n",
      "\n",
      "JuCeIvg:\n",
      "I o? no there, arm you. cand his cannow thmake a stopworn fothin aling:\n",
      "Sive lauge meot an to the sort uae: ne\n",
      "Ius ast begent agrody thugh. Some shimRs of dele, so now.\n",
      "Whot, Whut not hase one's presier broberes\n",
      "Thatt him sent my nainher's preself,\n",
      "Bouviu. Asfriet me mays server yet bew wall if wragaft its:\n",
      "But it lover faminzs, havH, livNa rhav,\n",
      "How shat add nobty quars that gost lamor.\n",
      "\n",
      "KING RICHARR LIEDWD IRDkRkol'Er,\n",
      "Mysolvilgoo- yeary, and cond me you pass: noth and never\n",
      "Fattion's your foARar-give.\n",
      "Buegh fortain asiumb, hisnow brou;\n",
      "Ore bleat at leedo'tung. Trius me so, god'd.\n",
      "\n",
      "FARUTIROza your Jausy. WOisPry, there\n",
      "come a Divitiove, in'these: ao shed thes\n",
      "as thessertiend too! and her, ses'gions to herr!\n",
      "\n",
      "GLOULEEEOTROMROEE:\n",
      "Bnoy, mad yot deed, maith, right--we and welcomm-not- manurse:\n",
      "I hast spinl give twake yor foe;\n",
      "Your and furteher? and velage gembe aman:\n",
      "Weat nonge show but but in wh eyer:\n",
      "Beoess\n",
      "pe'tion tha bpooe kew wikno bott\n",
      "Azgiven egefcer'd in latter of Llors\n",
      "Thers I, punibola birt sabeimiing,\n",
      "For men sticce the waatef hill wouth?\n",
      "Mishn!\n",
      "\n",
      "Kukest Lefeit:\n",
      "Have lest tomene it so\n",
      "And suppupoit; he wear ny anot in.\n",
      "\n",
      "LAOLY:\n",
      "Yet, as merry, I will dow, name, he  now;\n",
      "Whom hwse didg thee ? uts;\n",
      "Whurt, 'I,ow't, but cold oat colk.\n",
      "\n",
      "No OLeVTs Marniust:\n",
      "Didsioek'ill him!- then, spray;s,\n",
      "Clest thenou,denned friars andieggo. That Ise, be we thind tone, my courtague,\n",
      "And 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
}