_01_discord_01.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "2d553b43",
      "metadata": {
        "id": "2d553b43"
      },
      "source": [
        "Discussion on discord:\n",
        "\n",
        "_Did anyone else find that indexing into the weight matrix was slower than using the \"redundant\" one-hot encoded vectors? The perf difference was around 8x for me on GPU.\n",
        "I'm guessing it's faster to let the GPU rip through the matmul than have to do many random memory lookups._\n",
        "\n",
        "https://discord.com/channels/1020383067459821711/1029849849765564528/1056756317777305710"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "id": "88a37e75",
      "metadata": {
        "id": "88a37e75"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "69652fe0",
      "metadata": {
        "id": "69652fe0"
      },
      "outputs": [],
      "source": [
        "row_dimensions = 703 # Trigram word model input: (.., .a, [...], .z, aa, [...], az)\n",
        "col_dimensions = 27 # Trigram word model output: (., a, [...], z)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "db288660",
      "metadata": {
        "id": "db288660"
      },
      "outputs": [],
      "source": [
        "# Simulated weight matrix\n",
        "W = torch.randn([row_dimensions, col_dimensions], requires_grad=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "22d4e69d",
      "metadata": {
        "id": "22d4e69d"
      },
      "outputs": [],
      "source": [
        "# simulated X matrix, consisting of 1000 random integers between 0 and 703\n",
        "X = torch.randint(low=0, high = row_dimensions, size=(1000000,))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "8b257035",
      "metadata": {
        "id": "8b257035"
      },
      "outputs": [],
      "source": [
        "# One hot encoded x_enc matrix\n",
        "x_enc = F.one_hot(X, num_classes=row_dimensions).float()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "ca6fcee5",
      "metadata": {
        "id": "ca6fcee5"
      },
      "outputs": [],
      "source": [
        "# Alternative1: using matrix multiplication between one-hot encoded matrix and W\n",
        "a1 = x_enc @ W"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "id": "7ea0180e",
      "metadata": {
        "id": "7ea0180e"
      },
      "outputs": [],
      "source": [
        "# Alternative2: directly indexing into W with original X vector\n",
        "a2 = W[X]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "id": "3c476ba8",
      "metadata": {
        "id": "3c476ba8"
      },
      "outputs": [],
      "source": [
        "# Make sure the results are the same\n",
        "assert torch.allclose(a1, a2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "id": "fbb284ae",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fbb284ae",
        "outputId": "565ad0a7-cf0b-4e40-f716-59e174cef9ec"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "963 ms ± 13.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "x_enc @ W"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "id": "3f001c33",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3f001c33",
        "outputId": "550f8075-f02d-4c13-bcda-db4d23eb3b43"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "34.8 ms ± 4.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "W[X]"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e919a466",
      "metadata": {
        "id": "e919a466"
      },
      "source": [
        "### Directly indexing into W is roughly 10 times faster then using matrix multiplication"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "id": "80aa3033",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "80aa3033",
        "outputId": "bcebad12-7d90-40bd-b63c-fd0fc85bad46"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3.05 s ± 104 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "x_enc = F.one_hot(X, num_classes=row_dimensions).float()\n",
        "x_enc @ W"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9b1f1f3d",
      "metadata": {
        "id": "9b1f1f3d"
      },
      "source": [
        "### If we also consider the duration of creating the one_hot encoded matrix, the difference is even larger, it's 1300 vs 35 $\\mu$s, a factor 30!"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "124072df",
      "metadata": {
        "id": "124072df"
      },
      "source": [
        "### EDIT: Also when including the backward pass, indexing into W is faster then one-hot encoding"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "id": "dacc047b",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dacc047b",
        "outputId": "5e802bf5-c268-474d-b8c7-4e3cd85efc6f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.28 s ± 9.25 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "W.grad = None\n",
        "a1 = x_enc @ W\n",
        "a1.sum().backward()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "id": "78bc6ea4",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "78bc6ea4",
        "outputId": "033cd624-4790-44ac-b6c3-8ef84f2d80ed"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "58 ms ± 536 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "W.grad = None\n",
        "a2 = W[X]\n",
        "a2.sum().backward()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### EDIT: Try on GPU"
      ],
      "metadata": {
        "id": "Z3TP-o_2T4jZ"
      },
      "id": "Z3TP-o_2T4jZ"
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "id": "61dabdeb",
      "metadata": {
        "id": "61dabdeb"
      },
      "outputs": [],
      "source": [
        "x_enc = x_enc.cuda()\n",
        "W = W.cuda()\n",
        "X = X.cuda()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "id": "160b5559",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "160b5559",
        "outputId": "61f87f4d-ddc8-4aa1-d5d1-515be85c1d61"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "48.3 ms ± 1.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ],
      "source": [
        "%%timeit\n",
        "W.grad = None\n",
        "a1 = x_enc @ W\n",
        "a1.sum().backward()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "W.grad = None\n",
        "a2 = W[X]\n",
        "a2.sum().backward()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "m8MYFdrPTKyf",
        "outputId": "1cd45dbc-2140-4cb5-c53a-150a871a3a14"
      },
      "id": "m8MYFdrPTKyf",
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "5.06 ms ± 21.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "KRAEC0m2TMD6"
      },
      "id": "KRAEC0m2TMD6",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.8"
    },
    "colab": {
      "provenance": []
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}