{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Tensor Parallelism’s Communication Boundary\n",
        "\n",
        "A paired matrix partition shows where tensor parallelism needs communication, and why the nonlinearity determines that boundary.\n",
        "\n",
        "Original CPU forward-pass illustration using ReLU; not a distributed Megatron-LM implementation or GPU scaling test.\n",
        "\n",
        "The code uses Python 3 and the standard library. Saved output is from the checked CPU run.\n"
      ],
      "id": "megatron-tensor-parallelism-co-0"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "source": [
        "def row_times_matrix(row, matrix):\n",
        "    return [sum(row[i] * matrix[i][j] for i in range(len(row)))\n",
        "            for j in range(len(matrix[0]))]\n",
        "\n",
        "x = [1, 2]\n",
        "w1 = [[1, -1, 2, 0], [0, 1, -1, 3]]\n",
        "w2 = [[1, 0], [0, 1], [2, 1], [1, -1]]\n",
        "hidden = [max(0, v) for v in row_times_matrix(x, w1)]\n",
        "dense = row_times_matrix(hidden, w2)\n",
        "partials = []\n",
        "for columns in ((0, 1), (2, 3)):\n",
        "    local_w1 = [[row[j] for j in columns] for row in w1]\n",
        "    local_hidden = [max(0, v) for v in row_times_matrix(x, local_w1)]\n",
        "    local_w2 = [w2[j] for j in columns]\n",
        "    partials.append(row_times_matrix(local_hidden, local_w2))\n",
        "split = [sum(part[j] for part in partials) for j in range(2)]\n",
        "assert split == dense == [7, -5]\n",
        "print(f\"dense={dense}; split={split}\")\n",
        "tokens, width, expanded, element_bytes = 128, 1024, 4096, 2\n",
        "hidden_bytes = tokens * expanded * element_bytes\n",
        "output_bytes = tokens * width * element_bytes\n",
        "print(f\"expanded_tensor_bytes={hidden_bytes}\")\n",
        "print(f\"output_reduction_tensor_bytes={output_bytes}\")\n",
        "assert hidden_bytes == 4 * output_bytes\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dense=[7, -5]; split=[7, -5]\n",
            "expanded_tensor_bytes=1048576\n",
            "output_reduction_tensor_bytes=262144\n"
          ]
        }
      ],
      "id": "megatron-tensor-parallelism-co-1"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Sources\n",
        "\n",
        "- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism: https://arxiv.org/html/1909.08053v4\n"
      ],
      "id": "megatron-tensor-parallelism-co-2"
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}