moe.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "b772b407",
      "metadata": {
        "id": "b772b407"
      },
      "source": [
        "# Mixture of Experts from scratch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "89776a6a",
      "metadata": {
        "id": "89776a6a"
      },
      "source": [
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/antonio-f/mixture-of-experts-from-scratch/blob/main/moe.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e704e55d",
      "metadata": {
        "id": "e704e55d"
      },
      "source": [
        "This is a simple implementation of **Mixture of Experts** (**MoE**) technique applied to language modeling tasks.\n",
        "\n",
        "Evaluation and training of deep models can be computationally expensive and time-consuming. The Conditional Computation approach has been proposed to tackle this problem.\n",
        "**Conditional Computation** refers to a class of algorithms in which each input sample uses a different part of the model such that (on average) the compute, latency or power (depending on our objective) is reduced. It operates by selectively activating only parts of the network at a time."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dea0516a",
      "metadata": {
        "id": "dea0516a"
      },
      "source": [
        "### Loading data"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6def97d6",
      "metadata": {
        "id": "6def97d6"
      },
      "source": [
        "We will use the TinyStories dataset ([info](https://huggingface.co/datasets/roneneldan/TinyStories)), it is is suitable and not overly large."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "tV2MvYTdVwe2",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "tV2MvYTdVwe2",
        "outputId": "1841330a-0f23-4c9d-a5b3-2d83b6951cf0"
      },
      "outputs": [],
      "source": [
        "!wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStories_all_data.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4015770f",
      "metadata": {
        "id": "4015770f"
      },
      "source": [
        "We import some modules providing operating system dependent functionality like operations on files, paths etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "ZKuW2rbWW7G-",
      "metadata": {
        "id": "ZKuW2rbWW7G-"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import glob\n",
        "import json"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8a6b0fb8",
      "metadata": {
        "id": "8a6b0fb8"
      },
      "source": [
        "Now we create `TinyStories` folder and extract data inside it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "7dfd7fc0",
      "metadata": {
        "id": "7dfd7fc0"
      },
      "outputs": [],
      "source": [
        "if not os.path.exists(\"./TinyStories\"):\n",
        "    os.makedirs(\"./TinyStories\")\n",
        "\n",
        "!tar -xzf TinyStories_all_data.tar.gz -C TinyStories"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d505f287",
      "metadata": {
        "id": "d505f287"
      },
      "source": [
        "The following command returns a list of paths like\n",
        "\n",
        "`'TinyStories/data00.json'`\n",
        "\n",
        "`'TinyStories/data01.json'`\n",
        "\n",
        "`'TinyStories/data02.json'`\n",
        "\n",
        ". . .\n",
        "\n",
        "and so on."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "-N5m4Rk4Wp72",
      "metadata": {
        "id": "-N5m4Rk4Wp72"
      },
      "outputs": [],
      "source": [
        "shard_filenames = sorted(glob.glob(os.path.join('TinyStories', \"*.json\")))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6d89f17c",
      "metadata": {
        "id": "6d89f17c"
      },
      "source": [
        "We load each json file into `data`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "1m-He99IXCV4",
      "metadata": {
        "id": "1m-He99IXCV4"
      },
      "outputs": [],
      "source": [
        "with open(shard_filenames[0], \"r\") as f:\n",
        "        data = json.load(f)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0bda234c",
      "metadata": {
        "id": "0bda234c"
      },
      "source": [
        "Let us check the first element of `data`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "5000d128",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5000d128",
        "outputId": "1b2c74ef-930e-4a26-8602-af0ada9a6f0a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'story': '\\n\\nLily and Ben are friends. They like to play in the park. One day, they see a big tree with a swing. Lily wants to try the swing. She runs to the tree and climbs on the swing.\\n\"Push me, Ben!\" she says. Ben pushes her gently. Lily feels happy. She swings higher and higher. She laughs and shouts.\\nBen watches Lily. He thinks she is cute. He wants to swing too. He waits for Lily to stop. But Lily does not stop. She swings faster and faster. She is having too much fun.\\n\"Can I swing too, Lily?\" Ben asks. Lily does not hear him. She is too busy swinging. Ben feels sad. He walks away.\\nLily swings so high that she loses her grip. She falls off the swing. She lands on the ground. She hurts her foot. She cries.\\n\"Ow, ow, ow!\" she says. She looks for Ben. She wants him to help her. But Ben is not there. He is gone.\\nLily feels sorry. She wishes she had shared the swing with Ben. She wishes he was there to hug her. She limps to the tree. She sees something hanging from a branch. It is Ben\\'s hat. He left it for her.\\nLily smiles. She thinks Ben is nice. She puts on his hat. She hopes he will come back. She wants to say sorry. She wants to be friends again.',\n",
              " 'instruction': {'prompt:': 'Write a short story (3-5 paragraphs) which only uses very simple words that a 3 year old child would understand. The story should use the verb \"hang\", the noun \"foot\" and the adjective \"cute\". The story has the following features: the story should contain at least one dialogue. Remember to only use simple words!\\n\\nPossible story:',\n",
              "  'words': ['hang', 'foot', 'cute'],\n",
              "  'features': ['Dialogue']},\n",
              " 'summary': 'Lily and Ben play in the park and Lily gets too caught up in swinging, causing Ben to leave. Lily falls off the swing and hurts herself, but Ben leaves his hat for her as a kind gesture.',\n",
              " 'source': 'GPT-4'}"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "data[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5dac7b43",
      "metadata": {
        "id": "5dac7b43"
      },
      "source": [
        "We collect all stories in the `stories` list."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "wmsSEsysXFXF",
      "metadata": {
        "id": "wmsSEsysXFXF"
      },
      "outputs": [],
      "source": [
        "stories = [x['story'] for x in data]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f7f9b21c",
      "metadata": {
        "id": "f7f9b21c"
      },
      "source": [
        "A sample from `stories` is the following."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "yigoqX6fXYMl",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 103
        },
        "id": "yigoqX6fXYMl",
        "outputId": "5f440e13-2d34-40c0-9f7f-4d4a973b9288"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "\"Once upon a time, there was a little girl named Lily. Lily loved to play in the park with her friends. One day, Lily and her friends were playing hide and seek. Lily found a good hiding spot behind a big tree. As she was hiding, she started to yawn because she was very tired.\\nSuddenly, Lily saw an enormous shadow coming towards her. She got scared and started to cry. It turned out that the shadow was just her friend, Timmy. Timmy had found her hiding spot and was trying to surprise her. \\nLily learned that sometimes things that seem scary are not really scary at all. She also learned that it's important to get enough sleep so you don't yawn during the day. From that day on, Lily made sure to get plenty of rest before playing with her friends.\""
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "stories[42]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7000983c",
      "metadata": {
        "id": "7000983c"
      },
      "source": [
        "All the stories are joined together into the string called `text`. At the end of each story there is a new line `\\n` escape sequence."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "1f30df8b",
      "metadata": {
        "id": "1f30df8b"
      },
      "outputs": [],
      "source": [
        "text = \"\\n\".join(stories)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1340c982",
      "metadata": {
        "id": "1340c982"
      },
      "source": [
        "`text` is a very long string."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "a9e34655",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a9e34655",
        "outputId": "0d79d626-8599-4878-bed5-52cf0acfbdf8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "77586884"
            ]
          },
          "execution_count": 10,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "c596e794",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "c596e794",
        "outputId": "c7179d56-dfbb-4c3b-dcb9-a22fd2103f99"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Lily and Ben are friends. They like to play in the park. One day, they see a big tree with a swing\n"
          ]
        }
      ],
      "source": [
        "print(text[:100])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "28829cf1",
      "metadata": {
        "id": "28829cf1"
      },
      "source": [
        "### Character encoding"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f3885408",
      "metadata": {
        "id": "f3885408"
      },
      "source": [
        "We are going to use PyTorch tensors to store data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "14e943d7",
      "metadata": {
        "id": "14e943d7"
      },
      "outputs": [],
      "source": [
        "import torch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e14f124e",
      "metadata": {
        "id": "e14f124e"
      },
      "source": [
        "`chars` contains all the characters found in the text (joined stories). Its size is 97."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "f8InhsLGbFCj",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f8InhsLGbFCj",
        "outputId": "2214e646-9a72-4dd6-b38a-b97c19dc7c4c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\t\n",
            " !\"$%&'()*+,-./0123456789:;?ABCDEFGHIJKLMNOPQRSTUVWXYZ[]`abcdefghijklmnopqrstuvwxyz|~ éñ–—‘’“”…\n",
            "97\n"
          ]
        }
      ],
      "source": [
        "chars = sorted(list(set(text)))\n",
        "vocab_size = len(chars)\n",
        "print(''.join(chars))\n",
        "print(vocab_size)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "54de4004",
      "metadata": {
        "id": "54de4004"
      },
      "source": [
        "Below, two dictionaries: the first binds characters to integers and the second does the reverse."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "UHOqlwJtbKu2",
      "metadata": {
        "id": "UHOqlwJtbKu2"
      },
      "outputs": [],
      "source": [
        "ctoi = {ch:i for i, ch in enumerate(chars)}\n",
        "itoc = {i:ch for i,ch in enumerate(chars)}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "cec90814",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cec90814",
        "outputId": "f0c5f30c-f7df-40ec-99f5-3e2e11c4d9ae"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'\\t': 0,\n",
              " '\\n': 1,\n",
              " ' ': 2,\n",
              " '!': 3,\n",
              " '\"': 4,\n",
              " '$': 5,\n",
              " '%': 6,\n",
              " '&': 7,\n",
              " \"'\": 8,\n",
              " '(': 9,\n",
              " ')': 10,\n",
              " '*': 11,\n",
              " '+': 12,\n",
              " ',': 13,\n",
              " '-': 14,\n",
              " '.': 15,\n",
              " '/': 16,\n",
              " '0': 17,\n",
              " '1': 18,\n",
              " '2': 19,\n",
              " '3': 20,\n",
              " '4': 21,\n",
              " '5': 22,\n",
              " '6': 23,\n",
              " '7': 24,\n",
              " '8': 25,\n",
              " '9': 26,\n",
              " ':': 27,\n",
              " ';': 28,\n",
              " '?': 29,\n",
              " 'A': 30,\n",
              " 'B': 31,\n",
              " 'C': 32,\n",
              " 'D': 33,\n",
              " 'E': 34,\n",
              " 'F': 35,\n",
              " 'G': 36,\n",
              " 'H': 37,\n",
              " 'I': 38,\n",
              " 'J': 39,\n",
              " 'K': 40,\n",
              " 'L': 41,\n",
              " 'M': 42,\n",
              " 'N': 43,\n",
              " 'O': 44,\n",
              " 'P': 45,\n",
              " 'Q': 46,\n",
              " 'R': 47,\n",
              " 'S': 48,\n",
              " 'T': 49,\n",
              " 'U': 50,\n",
              " 'V': 51,\n",
              " 'W': 52,\n",
              " 'X': 53,\n",
              " 'Y': 54,\n",
              " 'Z': 55,\n",
              " '[': 56,\n",
              " ']': 57,\n",
              " '`': 58,\n",
              " 'a': 59,\n",
              " 'b': 60,\n",
              " 'c': 61,\n",
              " 'd': 62,\n",
              " 'e': 63,\n",
              " 'f': 64,\n",
              " 'g': 65,\n",
              " 'h': 66,\n",
              " 'i': 67,\n",
              " 'j': 68,\n",
              " 'k': 69,\n",
              " 'l': 70,\n",
              " 'm': 71,\n",
              " 'n': 72,\n",
              " 'o': 73,\n",
              " 'p': 74,\n",
              " 'q': 75,\n",
              " 'r': 76,\n",
              " 's': 77,\n",
              " 't': 78,\n",
              " 'u': 79,\n",
              " 'v': 80,\n",
              " 'w': 81,\n",
              " 'x': 82,\n",
              " 'y': 83,\n",
              " 'z': 84,\n",
              " '|': 85,\n",
              " '~': 86,\n",
              " '\\xa0': 87,\n",
              " 'é': 88,\n",
              " 'ñ': 89,\n",
              " '–': 90,\n",
              " '—': 91,\n",
              " '‘': 92,\n",
              " '’': 93,\n",
              " '“': 94,\n",
              " '”': 95,\n",
              " '…': 96}"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "ctoi"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "92bbcd2f",
      "metadata": {
        "id": "92bbcd2f"
      },
      "source": [
        "The encoding function transforms a text `s` into a list of integer (one for each character). Decode works exactly in the reverse order: it takes a list of integers and returns the text composed of the characters obtained decoding these integers. For example\n",
        "\n",
        "`encode(\"Hello, world!\")`\n",
        "\n",
        "returns the list\n",
        "\n",
        "`[37, 63, 70, 70, 73, 13, 2, 81, 73, 76, 70, 62, 3]`.\n",
        "\n",
        "Likewise,\n",
        "\n",
        "`decode([37, 63, 70, 70, 73, 13, 2, 81, 73, 76, 70, 62, 3])`\n",
        "\n",
        "returns the string\n",
        "\n",
        "`'Hello, world!'`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "id": "muglCKTAbNQi",
      "metadata": {
        "id": "muglCKTAbNQi"
      },
      "outputs": [],
      "source": [
        "encode = lambda s: [ctoi[c] for c in s]\n",
        "decode = lambda l: \"\".join([itoc[x] for x in l])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8837c47e",
      "metadata": {
        "id": "8837c47e"
      },
      "source": [
        "We store the encoded text into a tensor named `data` (that is not the variable encountered before)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "id": "9dc05388",
      "metadata": {
        "id": "9dc05388"
      },
      "outputs": [],
      "source": [
        "data = torch.tensor(encode(text), dtype = torch.long)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "id": "a4f27fc3",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "a4f27fc3",
        "outputId": "f99f3f01-4737-4e0f-8f90-6675d5b8cf8f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(torch.Size([77586884]), torch.Tensor)"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "data.shape, type(data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "id": "f67b1902",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f67b1902",
        "outputId": "d9a12499-b006-4d79-f1b9-7114443debe3"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([ 1,  1, 41, 67, 70, 83,  2, 59, 72, 62,  2, 31, 63, 72,  2, 59, 76, 63,\n",
              "         2, 64, 76, 67, 63, 72, 62, 77, 15,  2, 49, 66, 63, 83,  2, 70, 67, 69,\n",
              "        63,  2, 78, 73,  2, 74, 70, 59, 83,  2, 67, 72,  2, 78, 66, 63,  2, 74,\n",
              "        59, 76, 69, 15,  2, 44, 72, 63,  2, 62, 59, 83, 13,  2, 78, 66, 63, 83,\n",
              "         2, 77, 63, 63,  2, 59,  2, 60, 67, 65,  2, 78, 76, 63, 63,  2, 81, 67,\n",
              "        78, 66,  2, 59,  2, 77, 81, 67, 72, 65])"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "data[:100]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aa47c3ca",
      "metadata": {
        "id": "aa47c3ca"
      },
      "source": [
        "### Data splitting"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e7a86910",
      "metadata": {
        "id": "e7a86910"
      },
      "source": [
        "Now it's time to create training and validation datasets.\n",
        "Training data amounts to 90% of all data, the rest is validation data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "id": "66922688",
      "metadata": {
        "id": "66922688"
      },
      "outputs": [],
      "source": [
        "n = int(0.9*len(data))\n",
        "train_data = data[:n]\n",
        "val_data = data[n:]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "7dc4b801",
      "metadata": {
        "id": "7dc4b801"
      },
      "source": [
        "Let's define a temporary block size, setting it equal to 8 for testing purposes only. Subsequently this parameter will be set to 256 because it represents the length of the context - it is the set of data that will be provided to the MoE model from time to time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "id": "62038547",
      "metadata": {
        "id": "62038547"
      },
      "outputs": [],
      "source": [
        "block_size = 8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "id": "78a76ae5",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "78a76ae5",
        "outputId": "ead80c08-d88e-47f7-a180-5d4896e2045f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([ 1,  1, 41, 67, 70, 83,  2, 59, 72])"
            ]
          },
          "execution_count": 22,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# Training data block example\n",
        "train_data[:block_size+1]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "57669c35",
      "metadata": {
        "id": "57669c35"
      },
      "source": [
        "Basically, these language models are trained to guess, given n elements of text - words, parts of words, or like in this character-level case, just characters - the next text element. We are going to train a character-level model so, for example, if the first 8 characters (the context) are `your nam`, the next (the 9th) should be `e` (the target). So we need integers `x` for the training data and integers `y` representing all the targets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "id": "bfd62ded",
      "metadata": {
        "id": "bfd62ded"
      },
      "outputs": [],
      "source": [
        "x = train_data[:block_size]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "id": "ed9341b9",
      "metadata": {
        "id": "ed9341b9"
      },
      "outputs": [],
      "source": [
        "y = train_data[1:block_size+1]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "id": "cc4bde3a",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cc4bde3a",
        "outputId": "1a4ac09a-267a-477a-d18e-7ee15e04ea03"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(tensor([ 1,  1, 41, 67, 70, 83,  2, 59]),\n",
              " tensor([ 1, 41, 67, 70, 83,  2, 59, 72]))"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "x,y"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "91359d43",
      "metadata": {
        "id": "91359d43"
      },
      "source": [
        "Here are some examples of contexts-targets, as `t` varies, based on the two tensors `x` and `y` above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "id": "e2f98640",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e2f98640",
        "outputId": "dad178ad-95cc-4a13-bcb0-e12d3df81a31"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "context tensor([1]) target tensor(1)\n",
            "context tensor([1, 1]) target tensor(41)\n",
            "context tensor([ 1,  1, 41]) target tensor(67)\n",
            "context tensor([ 1,  1, 41, 67]) target tensor(70)\n",
            "context tensor([ 1,  1, 41, 67, 70]) target tensor(83)\n",
            "context tensor([ 1,  1, 41, 67, 70, 83]) target tensor(2)\n",
            "context tensor([ 1,  1, 41, 67, 70, 83,  2]) target tensor(59)\n",
            "context tensor([ 1,  1, 41, 67, 70, 83,  2, 59]) target tensor(72)\n"
          ]
        }
      ],
      "source": [
        "for t in range(block_size):\n",
        "    context = x[:t+1]\n",
        "    target = y[t]\n",
        "    print(\"context\", context, \"target\", target)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bbc5d33e",
      "metadata": {
        "id": "bbc5d33e"
      },
      "source": [
        "For reproducibility, we set a seed for PyTorch. Reproducibility is about limiting the number of sources of nondeterministic behavior for a specific platform, device, and PyTorch release. Often, it is possible to control sources of randomness that can cause multiple executions of your application to behave differently."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "id": "7d8ab5b0",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7d8ab5b0",
        "outputId": "90edfdf2-224b-4480-c320-de4c5cd140eb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "<torch._C.Generator at 0x7c4cc25b0ed0>"
            ]
          },
          "execution_count": 27,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "torch.manual_seed(0)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2bf0bb31",
      "metadata": {
        "id": "2bf0bb31"
      },
      "source": [
        "### Creating batches"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6e9f8fa1",
      "metadata": {
        "id": "6e9f8fa1"
      },
      "source": [
        "We set the batch size to 4 for testing (will be changed later). Batch size is how many independent sequences are going to be processed in parallel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "id": "4f400257",
      "metadata": {
        "id": "4f400257"
      },
      "outputs": [],
      "source": [
        "batch_size = 4"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "349b2197",
      "metadata": {
        "id": "349b2197"
      },
      "source": [
        "The following function splits the data into batches."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "id": "6ef9ed1f",
      "metadata": {
        "id": "6ef9ed1f"
      },
      "outputs": [],
      "source": [
        "def get_batch(split):\n",
        "    # generate a small bunch 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"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "id": "1ed4a1c6",
      "metadata": {
        "id": "1ed4a1c6"
      },
      "outputs": [],
      "source": [
        "xb, yb = get_batch('train')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "id": "abafad30",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "abafad30",
        "outputId": "7e28e5aa-3a9f-4907-9a84-c707f34e3515"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "tensor([[71, 71, 83,  2, 64, 73, 76,  2],\n",
              "        [67, 77,  2, 64, 59, 80, 73, 76],\n",
              "        [59, 72, 65,  2, 59, 72, 62,  2],\n",
              "        [ 2, 81, 73, 79, 70, 62,  2, 78]])"
            ]
          },
          "execution_count": 31,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "yb"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61c5f0ac",
      "metadata": {
        "id": "61c5f0ac"
      },
      "source": [
        "Below, examples of context-target sequences on 4 batches."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "id": "f083be50",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f083be50",
        "outputId": "a3d6f6c5-af1e-4f59-8753-e4a9e10f11c7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "tensor([73])       tensor(71)\n",
            "tensor([73, 71])       tensor(71)\n",
            "tensor([73, 71, 71])       tensor(83)\n",
            "tensor([73, 71, 71, 83])       tensor(2)\n",
            "tensor([73, 71, 71, 83,  2])       tensor(64)\n",
            "tensor([73, 71, 71, 83,  2, 64])       tensor(73)\n",
            "tensor([73, 71, 71, 83,  2, 64, 73])       tensor(76)\n",
            "tensor([73, 71, 71, 83,  2, 64, 73, 76])       tensor(2)\n",
            "\n",
            "tensor([66])       tensor(67)\n",
            "tensor([66, 67])       tensor(77)\n",
            "tensor([66, 67, 77])       tensor(2)\n",
            "tensor([66, 67, 77,  2])       tensor(64)\n",
            "tensor([66, 67, 77,  2, 64])       tensor(59)\n",
            "tensor([66, 67, 77,  2, 64, 59])       tensor(80)\n",
            "tensor([66, 67, 77,  2, 64, 59, 80])       tensor(73)\n",
            "tensor([66, 67, 77,  2, 64, 59, 80, 73])       tensor(76)\n",
            "\n",
            "tensor([77])       tensor(59)\n",
            "tensor([77, 59])       tensor(72)\n",
            "tensor([77, 59, 72])       tensor(65)\n",
            "tensor([77, 59, 72, 65])       tensor(2)\n",
            "tensor([77, 59, 72, 65,  2])       tensor(59)\n",
            "tensor([77, 59, 72, 65,  2, 59])       tensor(72)\n",
            "tensor([77, 59, 72, 65,  2, 59, 72])       tensor(62)\n",
            "tensor([77, 59, 72, 65,  2, 59, 72, 62])       tensor(2)\n",
            "\n",
            "tensor([63])       tensor(2)\n",
            "tensor([63,  2])       tensor(81)\n",
            "tensor([63,  2, 81])       tensor(73)\n",
            "tensor([63,  2, 81, 73])       tensor(79)\n",
            "tensor([63,  2, 81, 73, 79])       tensor(70)\n",
            "tensor([63,  2, 81, 73, 79, 70])       tensor(62)\n",
            "tensor([63,  2, 81, 73, 79, 70, 62])       tensor(2)\n",
            "tensor([63,  2, 81, 73, 79, 70, 62,  2])       tensor(78)\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for b in range(batch_size):\n",
        "    for t in range(block_size):\n",
        "        context = xb[b][:t+1]\n",
        "        target = yb[b][t]\n",
        "        print(context, \"     \", target)\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "759eccde",
      "metadata": {
        "id": "759eccde"
      },
      "source": [
        "### Models"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "df33320b",
      "metadata": {
        "id": "df33320b"
      },
      "source": [
        "Let's import some PyTorch neural networks modules."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "id": "90997e9e",
      "metadata": {
        "id": "90997e9e"
      },
      "outputs": [],
      "source": [
        "import torch.nn as nn\n",
        "from torch.nn import functional as F"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9cd1f110",
      "metadata": {
        "id": "9cd1f110"
      },
      "source": [
        "The core of MoE technique is provided by the following code. The MoE layer is a type of neural network layer that combines the predictions of multiple expert networks based a gating mechanism. The gating mechanism is learned.\n",
        "\n",
        "The `__init__` method initializes the `MoeLayer` class with a list of expert modules (`experts`), a gate module (`gate`), and a parameter `k` (default value 1). The experts are the individual neural networks that form the \"experts\" in the mixture, they are feed-forward neural networks. The gate is another neural network (a linear layer) responsible for producing gate logits, which are used to weight the contributions of the experts. The parameter `k` determines how many experts to select based on the gate logits (gate logits are the values that emerge from the application of gate module operations).\n",
        "\n",
        "Let's move on to discussing the mechanics of the `forward` method. At the beginning, the input tensor `inputs` is flattened (squashed) and passed through the gate module to obtain gate logits. The top-k experts with the highest gate logits are selected using `torch.topk`.\n",
        "\n",
        "The gate logits are then normalized using the softmax function along the second dimension. This results in a probability distribution over the selected experts.\n",
        "\n",
        "The selected experts and their corresponding weights are used to compute the weighted sum of the expert outputs. The final result is a tensor representing the output of the mixture of experts layer.\n",
        "\n",
        "The output tensor is reshaped to match the shape of the input tensor and returned."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "id": "cc99dca5",
      "metadata": {
        "id": "cc99dca5"
      },
      "outputs": [],
      "source": [
        "class MoeLayer(nn.Module):\n",
        "    def __init__(self, experts, gate, k=1):\n",
        "        super().__init__()\n",
        "        assert len(experts) > 0\n",
        "        self.experts = nn.ModuleList(experts)\n",
        "        self.gate = gate\n",
        "        self.k = k\n",
        "\n",
        "    def forward(self, inputs: torch.Tensor):\n",
        "        inputs_squashed = inputs.view(-1, inputs.shape[-1])\n",
        "        gate_logits = self.gate(inputs_squashed)\n",
        "        weights, selected_experts = torch.topk(\n",
        "            gate_logits, self.k\n",
        "        )\n",
        "        weights = nn.functional.softmax(\n",
        "            weights,\n",
        "            dim=1,\n",
        "            dtype=torch.float,\n",
        "        ).type_as(inputs)\n",
        "        results = torch.zeros_like(inputs_squashed)\n",
        "        for i, expert in enumerate(self.experts):\n",
        "            batch_idx, nth_expert = torch.where(selected_experts == i)\n",
        "            results[batch_idx] += weights[batch_idx, nth_expert, None] * expert(\n",
        "                inputs_squashed[batch_idx]\n",
        "            )\n",
        "        return results.view_as(inputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dbe9bb49",
      "metadata": {
        "id": "dbe9bb49"
      },
      "source": [
        "The picture below shows the plain Transformer encoder architecture (left) and its MoE modified version (right). Block module is implemented by the `Block` class, which we will see shortly (actually there are n Block modules, n is coded as `n_layer`)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "cb065da2",
      "metadata": {
        "id": "cb065da2"
      },
      "source": [
        "![Transformer](https://github.com/antonio-f/mixture-of-experts-from-scratch/blob/main/Transf.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4f671d79",
      "metadata": {
        "id": "4f671d79"
      },
      "source": [
        "Below, a more detailed picture highlighting MoE layer (taken from https://arxiv.org/pdf/2101.03961.pdf). \"Router\" represents the gating module, experts are Feed Forward Networks (FFN 1, 2, 3 and 4)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "285a9fd6",
      "metadata": {
        "id": "285a9fd6"
      },
      "source": [
        "![MoE](https://github.com/antonio-f/mixture-of-experts-from-scratch/blob/main/MoE.png?raw=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "39f4ccfa",
      "metadata": {
        "id": "39f4ccfa"
      },
      "source": [
        "Below, the code for the Transformer model (modified to include MoE layer). The Transformer consists of several blocks. So, to implement `Transformer` class, we need to implement the `Block` class first. In turn, to implement the `Block` class, we need `MulitHeadAttention` and `FeedForward` classes (other than `MoeLayer`, already defined). To define `MulitHeadAttention` we need the class `Head`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "id": "b2a585da",
      "metadata": {
        "id": "b2a585da"
      },
      "outputs": [],
      "source": [
        "class Head(nn.Module):\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",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        B, T, C = x.shape\n",
        "        k = self.key(x)\n",
        "        q = self.query(x)\n",
        "        wei = q @ k.transpose(-2, -1) * C**-0.5\n",
        "        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))\n",
        "        wei = F.softmax(wei, dim=-1)\n",
        "        wei = self.dropout(wei)\n",
        "        v = self.value(x)\n",
        "        out = wei @ v\n",
        "        return out\n",
        "\n",
        "class MulitHeadAttention(nn.Module):\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",
        "        x =  torch.cat([head(x) for head in self.heads], dim=-1)\n",
        "        out = self.dropout(self.proj(x))\n",
        "        return out\n",
        "\n",
        "\n",
        "class FeedForward(nn.Module):\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",
        "    def forward(self, x):\n",
        "        return self.net(x)\n",
        "\n",
        "class Block(nn.Module):\n",
        "    def __init__(self, n_embed, n_head, num_experts=4):\n",
        "        super().__init__()\n",
        "        self.sa_head= MulitHeadAttention(n_head, n_embed//n_head)\n",
        "        self.ffw = MoeLayer(\n",
        "            experts=[FeedForward(n_embed) for _ in range(num_experts)],\n",
        "            gate=nn.Linear(n_embed, num_experts, bias=False),\n",
        "        )\n",
        "\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_head(self.ln1(x))\n",
        "        x = x + self.ffw(self.ln2(x))\n",
        "        return x\n",
        "\n",
        "\n",
        "class Transformer(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embed, device=device)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embed, device=device)\n",
        "        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
        "        self.lm_head = nn.Linear(n_embed, vocab_size)\n",
        "\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "\n",
        "        token_emb = self.token_embedding_table(idx)\n",
        "        pos_emb = self.position_embedding_table(torch.arange(T).to(device))\n",
        "        x = token_emb + pos_emb\n",
        "        x = self.blocks(x)\n",
        "        logits = self.lm_head(x)\n",
        "        if targets == 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",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokes):\n",
        "        for _ in range(max_new_tokes):\n",
        "            idx_cond = idx[:, -block_size:]\n",
        "            logits, loss = self(idx_cond)\n",
        "            logits = logits[:, -1, :]\n",
        "            probs = F.softmax(logits, dim = -1)\n",
        "            idx_next = torch.multinomial(probs, num_samples = 1)\n",
        "            idx = torch.cat((idx, idx_next), dim = 1)\n",
        "        return idx"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e08f7997",
      "metadata": {
        "id": "e08f7997"
      },
      "source": [
        "Here are all the necessary hyperparameters. `max_iters` is set to 3000 for testing (it will take some time to train). Probably things start to become significant for values larger than 5000..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "id": "35bc917b",
      "metadata": {
        "id": "35bc917b"
      },
      "outputs": [],
      "source": [
        "# hyperparameters\n",
        "batch_size = 64 # independent sequences processed in parallel\n",
        "block_size = 256 # max context length\n",
        "max_iters = 3000 \n",
        "eval_interval = 100\n",
        "learning_rate = 1e-3\n",
        "eval_iters = 200\n",
        "n_embd = 384\n",
        "n_embed = 384\n",
        "n_head = 6\n",
        "n_layer = 6\n",
        "dropout = 0.0\n",
        "\n",
        "# set device\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1063a544",
      "metadata": {
        "id": "1063a544"
      },
      "source": [
        "### Model training"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "740507d5",
      "metadata": {},
      "source": [
        "Our model is the previously defined Transformer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "id": "24122aa8",
      "metadata": {
        "id": "24122aa8"
      },
      "outputs": [],
      "source": [
        "model = Transformer()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e59af998",
      "metadata": {
        "id": "e59af998"
      },
      "source": [
        "The function below evaluates loss for training and validation data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "id": "18a31597",
      "metadata": {
        "id": "18a31597"
      },
      "outputs": [],
      "source": [
        "@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",
        "            X = X.to(device)\n",
        "            Y = Y.to(device)\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": 39,
      "id": "QtLU4WDSPoIF",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "QtLU4WDSPoIF",
        "outputId": "804a9d4d-2f88-462b-e1e7-62d9f498a664"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'cuda'"
            ]
          },
          "execution_count": 39,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "device"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2bc8643a",
      "metadata": {
        "id": "2bc8643a"
      },
      "source": [
        "Move the model to the device and adopt [AdamW](https://www.fast.ai/posts/2018-07-02-adam-weight-decay.html) optimizer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "id": "ba7349df",
      "metadata": {
        "id": "ba7349df"
      },
      "outputs": [],
      "source": [
        "model = model.to(device)\n",
        "optimizer = torch.optim.AdamW(model.parameters(),lr=1e-4)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f7aca4de",
      "metadata": {
        "id": "f7aca4de"
      },
      "source": [
        "The training loop. If `max_iters` is large, it may take some time to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 41,
      "id": "e43b298a",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e43b298a",
        "outputId": "68f6185f-13a3-437b-c9a5-0dc1456a5dfb"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "step 0: train loss 4.9073, val loss 4.9073\n",
            "step 100: train loss 2.3431, val loss 2.3454\n",
            "step 200: train loss 2.3039, val loss 2.3042\n",
            "step 300: train loss 2.2779, val loss 2.2779\n",
            "step 400: train loss 2.2433, val loss 2.2438\n",
            "step 500: train loss 2.1811, val loss 2.1828\n",
            "step 600: train loss 2.0586, val loss 2.0600\n",
            "step 700: train loss 1.8800, val loss 1.8853\n",
            "step 800: train loss 1.7369, val loss 1.7424\n",
            "step 900: train loss 1.6339, val loss 1.6397\n",
            "step 1000: train loss 1.5603, val loss 1.5576\n",
            "step 1100: train loss 1.4920, val loss 1.4932\n",
            "step 1200: train loss 1.4438, val loss 1.4467\n",
            "step 1300: train loss 1.3997, val loss 1.4049\n",
            "step 1400: train loss 1.3656, val loss 1.3669\n",
            "step 1500: train loss 1.3264, val loss 1.3289\n",
            "step 1600: train loss 1.3024, val loss 1.2976\n",
            "step 1700: train loss 1.2736, val loss 1.2743\n",
            "step 1800: train loss 1.2499, val loss 1.2537\n",
            "step 1900: train loss 1.2261, val loss 1.2253\n",
            "step 2000: train loss 1.2046, val loss 1.2061\n",
            "step 2100: train loss 1.1865, val loss 1.1890\n",
            "step 2200: train loss 1.1698, val loss 1.1704\n",
            "step 2300: train loss 1.1549, val loss 1.1545\n",
            "step 2400: train loss 1.1383, val loss 1.1397\n",
            "step 2500: train loss 1.1250, val loss 1.1214\n",
            "step 2600: train loss 1.1100, val loss 1.1127\n",
            "step 2700: train loss 1.0963, val loss 1.0971\n",
            "step 2800: train loss 1.0880, val loss 1.0880\n",
            "step 2900: train loss 1.0735, val loss 1.0768\n",
            "step 2999: train loss 1.0622, val loss 1.0644\n"
          ]
        }
      ],
      "source": [
        "\n",
        "for iter in range(max_iters):\n",
        "\n",
        "    # print the loss on train and val datasets\n",
        "    if iter % 100 == 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",
        "    xb = xb.to(device)\n",
        "    yb = yb.to(device)\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",
      "id": "8d8447c3",
      "metadata": {
        "id": "8d8447c3"
      },
      "source": [
        "### Model evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b1da22a0",
      "metadata": {
        "id": "b1da22a0"
      },
      "source": [
        "We test our model first encoding some small sequence `d` to get started."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 42,
      "id": "SGZjghiVcR7q",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SGZjghiVcR7q",
        "outputId": "e7837b8f-6e2b-4ac4-b5ee-5498312d2b0f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "a long time ago, there was a she what orn it was drawaying.\n",
            "Lily said on the tress and went fast, what let so deep. So, he said, \"From you, Max! I have full new get special?\" But so atcher amaze her paint and hellped swing that mudre that he every day.\n",
            "One Bunny day, a ball abloove make turn very thought animals alun. Lily asked the field mortor the ground, another of get theree were so aftul, scareful deond again.\n",
            "One day, a mexe, something more sak yurng afr he could the make slove locks? \n",
            "Lily asked her for to man stook \n"
          ]
        }
      ],
      "source": [
        "d = 'a long time ago, there was a '\n",
        "x = torch.tensor(encode(d), dtype = torch.long,device=device).unsqueeze(0)\n",
        "print(decode(model.generate(x, max_new_tokes=500)[0].tolist()))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "92898cb0",
      "metadata": {},
      "source": [
        "### Useful links"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1c637c0a",
      "metadata": {},
      "source": [
        "Mixture of Experts from scratch ([Substack](https://monads.substack.com/p/mixture-of-experts-from-scratch)) ([Wordpress](https://m0nads.wordpress.com/2024/02/08/mixture-of-experts-from-scratch/))\n",
        "\n",
        "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity<br>\n",
        "W. Fedus, B. Zoph, N. Shazeer<br>\n",
        "[arXiv:2101.03961v3](https://arxiv.org/abs/2101.03961) [cs.LG](2021, rev. 2022)\n",
        "\n",
        "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding<br>\n",
        "D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, Z. Chen<br>\n",
        "[arXiv:2006.16668v1](https://arxiv.org/abs/2006.16668) [cs.CL] (2020)\n",
        "\n",
        "TinyStories dataset ([link](https://huggingface.co/datasets/roneneldan/TinyStories))\n",
        "\n",
        "Mixture of Experts Explained ([link](https://huggingface.co/blog/moe))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}