contrib/colab/LazyTensor_Basics.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.10"
    },
    "colab": {
      "name": "LazyTensor-Basics.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "06e73109"
      },
      "source": [
        "# Setup PyTorch/XLA Environment"
      ],
      "id": "06e73109"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WJ5qowFp1yJP"
      },
      "source": [
        "!pip install cloud-tpu-client==0.10 torch==1.9.0 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl"
      ],
      "id": "WJ5qowFp1yJP",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5VDH-vKP1q9u"
      },
      "source": [
        "import os\n",
        "assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'"
      ],
      "id": "5VDH-vKP1q9u",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8aa70226"
      },
      "source": [
        "import torch\n",
        "import torch_xla\n",
        "import torch_xla.core.xla_model as xm\n",
        "import torch_xla.debug.metrics as met\n",
        "\n",
        "import torch.nn as nn"
      ],
      "id": "8aa70226",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e55acb83"
      },
      "source": [
        "## LazyTensor Basics \n",
        "\n",
        "This colab is a companion to the blog post titled \"Understanding Lazy Tensor System Performance\".\n",
        "\n",
        "For illustration of lazy tensor behavior, let's perform some operations with XLA tensor(s), and examine the resulting HLO Graph:"
      ],
      "id": "e55acb83"
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "20d93911"
      },
      "source": [
        "dev = xm.xla_device()\n",
        "\n",
        "x1 = torch.rand((3, 3)).to(dev)\n",
        "x2 = torch.rand((3, 8)).to(dev)\n",
        "\n",
        "y1 = torch.einsum('bs,st->bt', x1, x2)\n",
        "y1 = y1 + x2\n",
        "print(torch_xla._XLAC._get_xla_tensors_text([y1]))"
      ],
      "id": "20d93911",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that XLA Tensors are \"Lazy\", i.e. The operations have been recorded but no computation/execution actually is done until required.\n",
        "\n",
        "The execution is done when a LazyTensor Barrier is inserted.\n",
        "The easiest way to insert a barrier is mark_step() call:"
      ],
      "metadata": {
        "id": "RdNpPCL_eYUp"
      },
      "id": "RdNpPCL_eYUp"
    },
    {
      "cell_type": "code",
      "source": [
        "xm.mark_step()\n",
        "print(torch_xla._XLAC._get_xla_tensors_text([x1]))\n",
        "print(y1.device)"
      ],
      "metadata": {
        "id": "Pc_GPV7zgV2b"
      },
      "id": "Pc_GPV7zgV2b",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dynamic Graph\n",
        "Now let's create a method which executes operations on xla tensors followed by a mark_step call. Optionally we also introduce a dynamic structure with these tensors. We then execute this method with and without the dynamic structure and measure the run time."
      ],
      "metadata": {
        "id": "jBthtD2MzSyD"
      },
      "id": "jBthtD2MzSyD"
    },
    {
      "cell_type": "code",
      "source": [
        "def dummy_step(x, y, loss, acc=False):\n",
        "  z = torch.einsum('bs,st->bt', y, x)\n",
        "  step_loss = z.sum().view(1,)\n",
        "  if acc: \n",
        "    loss = torch.cat((loss, step_loss))\n",
        "  else:\n",
        "    loss = step_loss\n",
        "  xm.mark_step()\n",
        "  return loss"
      ],
      "metadata": {
        "id": "tWcwhuglRYAG"
      },
      "id": "tWcwhuglRYAG",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "def measure_time(acc=False):\n",
        "  exec_times = []\n",
        "  iter_count = 100\n",
        "  x = torch.rand((512, 8)).to(dev)\n",
        "  y = torch.rand((512, 512)).to(dev)\n",
        "  loss = torch.zeros(1).to(dev)\n",
        "  for i in range(iter_count):\n",
        "    tic = time.time()\n",
        "    loss = dummy_step(x, y, loss, acc=acc)\n",
        "    toc = time.time()\n",
        "    exec_times.append(toc - tic)\n",
        "  return exec_times"
      ],
      "metadata": {
        "id": "YGAf-6c0ZxoK"
      },
      "id": "YGAf-6c0ZxoK",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dyn = measure_time(acc=True) # acc= True Results in dynamic graph\n",
        "st = measure_time(acc=False) # Static graph, computation shape, inputs and output shapes don't change"
      ],
      "metadata": {
        "id": "I5UFoYwDndpU"
      },
      "id": "I5UFoYwDndpU",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.plot(st, label = 'static graph')\n",
        "plt.plot(dyn, label = 'dynamic graph')\n",
        "plt.legend()\n",
        "plt.title('Execution time in seconds')"
      ],
      "metadata": {
        "id": "4cnVF11ykWxG"
      },
      "id": "4cnVF11ykWxG",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that dynamic graph execution times are consistently higher for same computation because of the compilation cost incurred in every iteration. Static graph curve benefits from compilation cache and quickly stablizes to a faster execution time."
      ],
      "metadata": {
        "id": "1bkkZ5vN0f8u"
      },
      "id": "1bkkZ5vN0f8u"
    }
  ]
}