tensor_parallelism_ffn_layers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyPqqfmbc7/Xa4PVS039NV3D",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/puneeshkhanna/Tensor-Parallelism/blob/master/tensor_parallelism_ffn_layers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Link - https://medium.com/@puneesh.khanna83/understanding-tensor-parallelism-to-fit-larger-models-on-multiple-devices-d4da1821d41b\n"
      ],
      "metadata": {
        "id": "2C3JMhET9ZC5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Pl8o6JNBXvL"
      },
      "outputs": [],
      "source": [
        "import numpy\n",
        "import torch\n",
        "import torch.nn as nn"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Input of dimensions (batch size, no of words or seq len, embedding dimension or hidden size of each word)\n",
        "input = torch.randn(size=(1, 5, 10), dtype=torch.float32)"
      ],
      "metadata": {
        "id": "Td1StIsVCldQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "embedding_dim = input.size(dim=2)\n",
        "embedding_dim"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z2ZqB8YLTnWI",
        "outputId": "57b3d799-d716-41aa-f49e-a9eaa5f32f89"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "10"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Back to back linear layers output\n",
        "\n",
        "Tranformer architectures have back to back linear layers where the embedding dim\n",
        "first goes from h to 4h and then back from 4h to h\n",
        "\n",
        "```\n",
        "(mlp): MLP(\n",
        "          (dense_h_to_4h): Linear(in_features=4096, out_features=16384, bias=False)          \n",
        "          (dense_4h_to_h): Linear(in_features=16384, out_features=4096, bias=False)\n",
        "        )\n",
        "```\n",
        "\n",
        "In this notebook, below dimensions are used.\n",
        "```\n",
        "(mlp): MLP(\n",
        "          (dense_h_to_4h): Linear(in_features=10, out_features=40, bias=False)          \n",
        "          (dense_4h_to_h): Linear(in_features=40, out_features=19, bias=False)\n",
        "        )\n",
        "```"
      ],
      "metadata": {
        "id": "zdTszgCgW1zY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "linear_h_to_4h = nn.Linear(in_features=embedding_dim, out_features=embedding_dim*4, bias=False)\n",
        "\n",
        "# Input of shape [1,5,10] * W.T of shape [10, 40]\n",
        "output = linear_h_to_4h(input)\n",
        "\n",
        "# input of dimension (1,5,10) gets transformed to dimension (1,5,40)\n",
        "output.shape"
      ],
      "metadata": {
        "id": "SkVpKZr7BcCm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "280e96f1-e6ef-4a4a-a95f-447fa39a7e65"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 40])"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "linear_4h_to_h = nn.Linear(in_features=embedding_dim*4, out_features=embedding_dim, bias=False)\n",
        "\n",
        "# h_to_4h output of shape [1,5,40] * W.T of shape [40,10]\n",
        "final_output = linear_4h_to_h(output)\n",
        "\n",
        "# h_to_4h output of dimension (1,5,40) gets transformed back to dimension (1,5,10)\n",
        "final_output.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rMDkkiqACDmd",
        "outputId": "68ac5f55-03ba-4d7f-91b7-5b1c3202bbaa"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Back to back linear layer outputs with Tensor parallelism\n",
        "\n",
        "Assuming 2 devices, below code depicts how weights will be divided between the 2 devices.\n",
        "\n",
        "Note that all the below code is executed on single device only with comments that which blocks will be executed on first device or second device."
      ],
      "metadata": {
        "id": "rINCD0tnWsFE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "n_devices = 2\n",
        "weight_parallel = int((embedding_dim*4)/n_devices)\n",
        "weight_parallel"
      ],
      "metadata": {
        "id": "z871j0n6UijU",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "afd6e177-7a55-4da0-e18f-662c159b14c8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "20"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### First device - h_to_4h logic"
      ],
      "metadata": {
        "id": "mnzBoHpRUT5a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# First device weights of h_to_4h linear layer will be [out_features, in_feature] = [20, 10]\n",
        "linear_h_to_4h_parallel = nn.Linear(in_features=embedding_dim, out_features=weight_parallel, bias=False)\n",
        "\n",
        "# Set weights of this layer to the first 20 rows of the original weights\n",
        "linear_h_to_4h_parallel.weight.data = linear_h_to_4h.weight[:weight_parallel,:]\n",
        "\n",
        "# Input of shape [1,5,10] * W.T of shape [10, 20] ; note that in W.T - it is actually the columns which are divided; hence this is also known as column parallel linear\n",
        "output_parallel_1 = linear_h_to_4h_parallel(input)\n",
        "\n",
        "output_parallel_1.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cu5a_7vnCP4c",
        "outputId": "60fe8ca2-abef-4a41-fe12-f45dafa0dd79"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 20])"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Second device - h_to_4h logic"
      ],
      "metadata": {
        "id": "-GGjyS7AUae4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Second device weights of h_to_4h linear layer will be [out_features, in_feature] = [20, 10]\n",
        "linear_h_to_4h_parallel = nn.Linear(in_features=embedding_dim, out_features=weight_parallel, bias=False)\n",
        "\n",
        "# Set weights of this layer to the last 20 rows of the original weights\n",
        "linear_h_to_4h_parallel.weight.data = linear_h_to_4h.weight[weight_parallel:,:]\n",
        "\n",
        "# Input[1,5,10] * W.T[10, 20] ; note that in W.T - it is actually the columns which are divided; hence this is also known as column parallel linear\n",
        "output_parallel_2 = linear_h_to_4h_parallel(input)\n",
        "\n",
        "output_parallel_2.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1K-W6FiGEj8Q",
        "outputId": "68d0ce92-7f4f-4166-b1b5-db7132b639e3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 20])"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# If we perform an all gather of output_parallel_1 and output_parallel_2, we will get the actual output of h_to_4h linear layer\n",
        "# but since we have one more linear layer of 4h_to_h, we can continue with the tensor parallel outputs."
      ],
      "metadata": {
        "id": "X4YjmSBoTEIu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### First device - 4h_to_h logic"
      ],
      "metadata": {
        "id": "ozgY5VOPUej6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# First device weights of 4h_to_h linear layer will be [out_features, in_features] = [10, 20]\n",
        "linear_4h_to_h_parallel = nn.Linear(in_features=weight_parallel, out_features=embedding_dim, bias=False)\n",
        "\n",
        "# Set weights of this layer to the first 20 columns of the original weights\n",
        "linear_4h_to_h_parallel.weight.data = linear_4h_to_h.weight[:,:weight_parallel]\n",
        "\n",
        "# output_parallel_1 of shape [1,5,20] * W.T of shape [20,10] ; note that in W.T - it is actually the rows which are divided; hence this is also known as row parallel linear\n",
        "final_output_parallel_1 = linear_4h_to_h_parallel(output_parallel_1)\n",
        "\n",
        "final_output_parallel_1.shape"
      ],
      "metadata": {
        "id": "WYsWuO1sExXP",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d7aa4c73-4c5a-4113-bce9-eb75d2c0b683"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 17
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Second device - 4h_to_h logic"
      ],
      "metadata": {
        "id": "H2jEWsp3UiO_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Second device weights of 4h_to_h linear layer will be [out_features, in_features] = [10, 20]\n",
        "linear_4h_to_h_parallel = nn.Linear(in_features=weight_parallel, out_features=embedding_dim, bias=False)\n",
        "\n",
        "# Set weights of this layer to the last 20 columns of the original weights\n",
        "linear_4h_to_h_parallel.weight.data = linear_4h_to_h.weight[:,weight_parallel:]\n",
        "\n",
        "# output_parallel_2 of shape [1,5,20] * W.T of shape [20,10] ; note that in W.T - it is actually the rows which are divided; hence this is also known as row parallel linear\n",
        "final_output_parallel_2 = linear_4h_to_h_parallel(output_parallel_2)\n",
        "\n",
        "final_output_parallel_2.shape"
      ],
      "metadata": {
        "id": "hb8qCCW7E0JI",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "21f5d495-9a3d-40bb-8f7d-c1d372ea7777"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Add tensor parallel outputs of 4h_to_h using torch.add; it will be an all reduce operation when using actual 2 devices"
      ],
      "metadata": {
        "id": "3ao8emKqUmgg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# In actual tensor parallelism between 2 devices, this will be all reduce operation\n",
        "final_output_tp = torch.add(final_output_parallel_1, final_output_parallel_2)\n",
        "\n",
        "final_output_tp.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kcT8l7QGKpuO",
        "outputId": "a9f3342b-1230-478c-effd-4b8dfdde93ae"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 19
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "final_output_tp.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-L9TVaaYIoYp",
        "outputId": "5a8f30e8-0cb8-4b34-f83b-b7c62a2e859e"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([1, 5, 10])"
            ]
          },
          "metadata": {},
          "execution_count": 20
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Compare final_output and final_output_tp using allclose"
      ],
      "metadata": {
        "id": "0Be7wx4XWowO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(torch.allclose(final_output, final_output_tp, rtol=1e-05, atol=1e-05))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ef3Ll13rIc6A",
        "outputId": "a0bbf34d-5465-4cb2-b4c0-eadf81c31b5d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "True\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "final_output"
      ],
      "metadata": {
        "id": "51vlAHWkJHN0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ff187ada-2b48-46fb-e6a4-a9df0959a3db"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[ 0.0230, -0.1114, -0.1713,  0.5962, -0.1553, -0.1313, -0.2674,\n",
              "          -0.2797,  0.0750, -0.1309],\n",
              "         [-0.2180, -0.1899, -0.3012,  0.3890, -0.1467, -0.1274, -0.1041,\n",
              "           0.0927, -0.2869, -0.1906],\n",
              "         [ 0.0226,  0.1911,  0.1427, -0.6180, -0.0171,  0.3005,  0.3173,\n",
              "           0.2533,  0.0814,  0.2457],\n",
              "         [ 0.5145,  0.0557,  0.3507,  0.1115,  0.2716, -0.3608,  0.1050,\n",
              "          -0.2071, -0.0866,  0.1415],\n",
              "         [-0.0225,  0.0401, -0.4115,  0.1852, -0.2485,  0.1331, -0.4189,\n",
              "          -0.1517, -0.0531,  0.2920]]], grad_fn=<UnsafeViewBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "final_output_tp"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6SFJUEqKWIaS",
        "outputId": "0ce0432a-3b9e-420e-84a2-01736a9b263c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([[[ 0.0230, -0.1114, -0.1713,  0.5962, -0.1553, -0.1313, -0.2674,\n",
              "          -0.2797,  0.0750, -0.1309],\n",
              "         [-0.2180, -0.1899, -0.3012,  0.3890, -0.1467, -0.1274, -0.1041,\n",
              "           0.0927, -0.2869, -0.1906],\n",
              "         [ 0.0226,  0.1911,  0.1427, -0.6180, -0.0171,  0.3005,  0.3173,\n",
              "           0.2533,  0.0814,  0.2457],\n",
              "         [ 0.5145,  0.0557,  0.3507,  0.1115,  0.2716, -0.3608,  0.1050,\n",
              "          -0.2071, -0.0866,  0.1415],\n",
              "         [-0.0225,  0.0401, -0.4115,  0.1852, -0.2485,  0.1331, -0.4189,\n",
              "          -0.1517, -0.0531,  0.2920]]], grad_fn=<AddBackward0>)"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    }
  ]
}