GPT_J_inference_on_TPU_driver0_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMdmoVjaI8Na2gK5iVt+q9N",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/mosmos6/MTJ-on-TPU0.2/blob/main/GPT_J_inference_on_TPU_driver0_2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# GPT-J inference on TPU_driver0.2"
      ],
      "metadata": {
        "id": "aS5gzMfZKC3t"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## This is a colab demo to infer with modified mesh-transformer-jax on TPU_driver0.2. You need high memory TPU runtime."
      ],
      "metadata": {
        "id": "1IuSSZ0BC6ZT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Load your data from your google cloud bucket (if it's your case)"
      ],
      "metadata": {
        "id": "xLjLuNxwKOqy"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a7kEy7LuKBNi"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!echo \"deb http://packages.cloud.google.com/apt gcsfuse-bionic main\" > /etc/apt/sources.list.d/gcsfuse.list\n",
        "!curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -\n",
        "!apt -qq update\n",
        "!apt -qq install gcsfuse"
      ],
      "metadata": {
        "id": "zq9tSoxQKRjp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir folderOnColab\n",
        "!gcsfuse --implicit-dirs - your bucket - folderOnColab"
      ],
      "metadata": {
        "id": "8w88X_eMKUPJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Installing dependencies"
      ],
      "metadata": {
        "id": "KolFQx91Kyy5"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "pip install --upgrade pip"
      ],
      "metadata": {
        "id": "fShFmBkRKk7h"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# For jax 0.3.25, you must install this BEFORE installing jax.\n",
        "# However for jax 0.3.5, you must install this after installing jax for some reason.\n",
        "\n",
        "import jax.tools.colab_tpu\n",
        "jax.tools.colab_tpu.setup_tpu()"
      ],
      "metadata": {
        "id": "i7z7aePWD_g9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Extract \"mesh_transformer\" folder and requirements.txt file from [this repo](https://github.com/mosmos6/MTJ-on-TPU0.2), replace those with the original ones in your repo."
      ],
      "metadata": {
        "id": "-KV64LbdDL-E"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/<your repo>\n",
        "!pip install -r <your repo>/requirements.txt\n",
        "!pip install <your repo>/"
      ],
      "metadata": {
        "id": "Xi5Dft_oK3XA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install 'jax[tpu]==0.3.25' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
      ],
      "metadata": {
        "id": "18b20uSvLcnI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Setup model"
      ],
      "metadata": {
        "id": "DckE81IZLkWB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import requests \n",
        "import jax\n",
        "import jax.config\n",
        "import jax.tools\n",
        "import jax.tools.colab_tpu\n",
        "\n",
        "tpu_driver = 'tpu_driver0.2'\n",
        "tpu_addr = os.environ['COLAB_TPU_ADDR']\n",
        "colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]\n",
        "#url = f'http://{colab_tpu_addr}:8475/requestversion/tpu_driver0.1_dev20210607'\n",
        "url = f'http://{colab_tpu_addr}:8475/requestversion/{tpu_driver}'\n",
        "requests.post(url)\n",
        "\n",
        "# The following is required to use TPU Driver as JAX's backend.\n",
        "jax.tools.colab_tpu.TPU_DRIVER_MODE = 1\n",
        "jax.config.FLAGS.jax_xla_backend = \"tpu_driver\"\n",
        "jax.config.FLAGS.jax_backend_target = f\"grpc://{tpu_addr}\""
      ],
      "metadata": {
        "id": "3YN8s3TPLf1V"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "\n",
        "import jax\n",
        "from jax.experimental import maps\n",
        "import numpy as np\n",
        "import optax\n",
        "import transformers\n",
        "\n",
        "from mesh_transformer.checkpoint import read_ckpt\n",
        "from mesh_transformer.sampling import nucleaus_sample\n",
        "from mesh_transformer.transformer_shard import CausalTransformer\n",
        "from mesh_transformer.transformer_shard import CausalTransformerV2"
      ],
      "metadata": {
        "id": "ABY1u_oCLtrY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "params = {\n",
        "  \"layers\": 28,\n",
        "  \"d_model\": 4096,\n",
        "  \"n_heads\": 16,\n",
        "  \"d_head\": 256,\n",
        "  \"n_vocab\": 50400,\n",
        "  \"norm\": \"layernorm\",\n",
        "  \"pe\": \"rotary\",\n",
        "  \"pe_rotary_dims\": 64,\n",
        "\n",
        "  \"seq\": 2048,\n",
        "  \"cores_per_replica\": 8,\n",
        "  \"per_replica_batch\": 1,\n",
        "}\n",
        "\n",
        "per_replica_batch = params[\"per_replica_batch\"]\n",
        "cores_per_replica = params[\"cores_per_replica\"]\n",
        "seq = params[\"seq\"]\n",
        "\n",
        "params[\"sampler\"] = nucleaus_sample\n",
        "\n",
        "# here we \"remove\" the optimizer parameters from the model (as we don't need them for inference)\n",
        "params[\"optimizer\"] = optax.scale(0)\n",
        "\n",
        "mesh_shape = (jax.device_count() // cores_per_replica, cores_per_replica)\n",
        "devices = np.array(jax.devices()).reshape(mesh_shape)\n",
        "\n",
        "global_mesh = maps.Mesh(devices, ('dp', 'mp'))\n",
        "maps.thread_resources.env = maps.ResourceEnv(physical_mesh=global_mesh, loops=())\n",
        "\n",
        "tokenizer = transformers.GPT2TokenizerFast.from_pretrained('gpt2')"
      ],
      "metadata": {
        "id": "3DHUa0PcLwbX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Create the network and load your parameters"
      ],
      "metadata": {
        "id": "o1Fz3UWoL-gw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Here we use read_ckpt instead of read_ckpt_lowmem because lowmem gets stuck forever in inference for some reason.\n",
        "\n",
        "total_batch = per_replica_batch * jax.device_count() // cores_per_replica\n",
        "\n",
        "network = CausalTransformer(params)\n",
        "\n",
        "network.state = read_ckpt(network.state, \"/content/folderOnColab/<your data>/step_10/\", devices.shape[1])\n",
        "\n",
        "local_shards = max(jax.local_device_count() // mesh_shape[1], 1)\n",
        "del network.state[\"opt_state\"]\n",
        "network.state = network.move_xmap(network.state, np.zeros(local_shards))"
      ],
      "metadata": {
        "id": "_uWQfAqEL4zf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#Run model"
      ],
      "metadata": {
        "id": "gDSWO7SuFjdt"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# allow text wrapping in generated output: https://stackoverflow.com/a/61401455\n",
        "from IPython.display import HTML, display\n",
        "\n",
        "def set_css():\n",
        "  display(HTML('''\n",
        "  <style>\n",
        "    pre {\n",
        "        white-space: pre-wrap;\n",
        "    }\n",
        "  </style>\n",
        "  '''))\n",
        "get_ipython().events.register('pre_run_cell', set_css)"
      ],
      "metadata": {
        "id": "ePuuOtyHFlFQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def infer(context, top_p=0.9, temp=0.9, gen_len=100):\n",
        "    tokens = tokenizer.encode(context)\n",
        "\n",
        "    provided_ctx = len(tokens)\n",
        "    pad_amount = seq - provided_ctx\n",
        "\n",
        "    pad_amount = max(pad_amount, 0) \n",
        "\n",
        "    padded_tokens = np.pad(tokens, ((pad_amount, 0),)).astype(np.uint32)[-seq:] \n",
        "    batched_tokens = np.array([padded_tokens] * total_batch) \n",
        "    length = np.ones(total_batch, dtype=np.uint32) * len(tokens)\n",
        "\n",
        "    start = time.time()\n",
        "    output = network.generate(batched_tokens, length, gen_len, {\"top_p\": np.ones(total_batch) * top_p, \"temp\": np.ones(total_batch) * temp})\n",
        "\n",
        "    samples = []\n",
        "    decoded_tokens = output[1][0]\n",
        "\n",
        "    for o in decoded_tokens[:, :, 0]:\n",
        "      samples.append(f\"\\033[1m{context}\\033[0m{tokenizer.decode(o)}\")\n",
        "\n",
        "    print(f\"completion done in {time.time() - start:06}s\")\n",
        "    return samples\n",
        "\n",
        "print(infer(\"Google colab is\")[0])"
      ],
      "metadata": {
        "id": "JIFcS22AF7U0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  { form-width: \"300px\" }\n",
        "top_p = 1 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "temp = 1 #@param {type:\"slider\", min:0, max:1, step:0.1}\n",
        "\n",
        "context = \"\"\"Your context here\"\"\"\n",
        "\n",
        "\n",
        "print(infer(top_p=top_p, temp=temp, gen_len=100, context=context)[0])"
      ],
      "metadata": {
        "id": "WPTjI1bGF-jE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Enjoy!*"
      ],
      "metadata": {
        "id": "6CXkR-52GIg7"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ij2bWNLAHM1U"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}