notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "copyright"
      },
      "outputs": [],
      "source": [
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "title"
      },
      "source": [
        "# PyTorch distributed training with Vertex AI Reduction Server\n",
        "\n",
        "<table align=\"left\">\n",
        "\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>                                                                                               \n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "overview:custom"
      },
      "source": [
        "## Overview\n",
        "\n",
        "When you run a distributed training job across multiple nodes using GPUs, communicating gradients between nodes can contribute significant latency. Reduction Server is an all-reduce algorithm that can increase throughput and reduce latency for distributed training. This notebook demonstrates how to run a PyTorch distributed training job with Reduction Server on Vertex AI. The training job is created to fine-tune pretrained model `bert-large-cased` from the Hugging Face Transformers library on the `imdb` dataset for sentiment classification.\n",
        "\n",
        "Learn more about [Vertex AI Training](https://cloud.google.com/vertex-ai/docs/training/custom-training) and [Vertex AI Reduction Server](https://cloud.google.com/blog/topics/developers-practitioners/optimize-training-performance-reduction-server-vertex-ai)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "objective:custom,training,online_prediction"
      },
      "source": [
        "### Objective\n",
        "\n",
        "In this notebook, you learn how to create a PyTorch distributed training job that uses PyTorch distributed training framework and tools, and run the training job on the Vertex AI Training service with Reduction Server.\n",
        "\n",
        "This tutorial uses the following Google Cloud ML services and resources:\n",
        "\n",
        "* Vertex AI Training\n",
        "* Cloud Storage\n",
        "\n",
        "\n",
        "The steps performed include:\n",
        "\n",
        "* Create a PyTorch distributed training application\n",
        "* Package the training application with pre-built containers\n",
        "* Create a custom job on Vertex AI with Reduction Server\n",
        "* Submit and monitor the job \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dataset:custom,cifar10,icn"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "In this tutorial, we use [`imdb`](https://huggingface.co/datasets/imdb) dataset from Hugging Face. `imdb` is a large movie review dataset for binary sentiment classification containing a set of 25,000 highly polar movie reviews for training, and 25,000 for testing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "costs"
      },
      "source": [
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI\n",
        "pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage\n",
        "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
        "Calculator](https://cloud.google.com/products/calculator/)\n",
        "to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "install_aip"
      },
      "source": [
        "## Installation\n",
        "\n",
        "Install the following packages required to execute this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1fd00fa70a2a"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "restart"
      },
      "source": [
        "### Colab only: Uncomment the following cell to restart the kernel.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bzPxhxS5lugp"
      },
      "outputs": [],
      "source": [
        "# Automatically restart kernel after installs so that your environment can access the new packages\n",
        "# import IPython\n",
        "\n",
        "# app = IPython.Application.instance()\n",
        "# app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2qpIurSjmpT"
      },
      "source": [
        "## Before you begin\n",
        "\n",
        "### Set up your Google Cloud project\n",
        "\n",
        "**The following steps are required, regardless of your notebook environment.**\n",
        "\n",
        "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
        "\n",
        "2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "3. [Enable the following APIs: Vertex AI API, Cloud Resource Manager API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,cloudresourcemanager.googleapis.com).\n",
        "\n",
        "4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "project_id"
      },
      "source": [
        "#### Set your project ID\n",
        "\n",
        "**If you don't know your project ID**, try the following:\n",
        "* Run `gcloud config list`.\n",
        "* Run `gcloud projects list`.\n",
        "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsePm9c4jmpT"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "\n",
        "# Set the project id\n",
        "! gcloud config set project {PROJECT_ID}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a54f9d7c1876"
      },
      "source": [
        "#### Region\n",
        "\n",
        "You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3aaadaaf9b30"
      },
      "outputs": [],
      "source": [
        "REGION = \"us-central1\"  # @param {type: \"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R7eelnCv6EWn"
      },
      "source": [
        "#### UUID\n",
        "\n",
        "If you are in a live tutorial session, you might be using a shared test account or project. To avoid name collisions between users on resources created, you create a uuid for each instance session, and append it onto the name of resources you create in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dHsjsyb76HaN"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import string\n",
        "\n",
        "\n",
        "# Generate a uuid of a specified length(default=8)\n",
        "def generate_uuid(length: int = 8) -> str:\n",
        "    return \"\".join(random.choices(string.ascii_lowercase + string.digits, k=length))\n",
        "\n",
        "\n",
        "UUID = generate_uuid()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5c0404984792"
      },
      "source": [
        "### Authenticate your Google Cloud account\n",
        "\n",
        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x2n5SeAAjmpU"
      },
      "source": [
        "**1. Vertex AI Workbench**\n",
        "* Do nothing as you are already authenticated."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6FDh38swjmpU"
      },
      "source": [
        "**2. Local JupyterLab instance, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nt8cEM2GjmpU"
      },
      "outputs": [],
      "source": [
        "# ! gcloud auth login"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XUSL_JcpjmpU"
      },
      "source": [
        "**3. Colab, uncomment and run:**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_2zemfGvjmpU"
      },
      "outputs": [],
      "source": [
        "# from google.colab import auth\n",
        "# auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TCPJ38n7jmpU"
      },
      "source": [
        "**4. Service account or other**\n",
        "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bucket:custom"
      },
      "source": [
        "### Create a Cloud Storage bucket\n",
        "\n",
        "Create a storage bucket to store intermediate artifacts such as datasets.\n",
        "\n",
        "When you submit a training job using the Cloud SDK, you upload a Python package\n",
        "containing your training code to a Cloud Storage bucket. Vertex AI runs\n",
        "the code from this package. In this tutorial, Vertex AI also saves the\n",
        "trained model that results from your job in the same bucket. Using this model artifact, you can then\n",
        "create Vertex AI Model resource and use for prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bucket"
      },
      "outputs": [],
      "source": [
        "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "create_bucket"
      },
      "source": [
        "**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Oz8J0vmSlugt"
      },
      "outputs": [],
      "source": [
        "! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v3sbyPBU75CR"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x8qQuI1377Jr"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "import_aip"
      },
      "source": [
        "### Initialize Vertex AI SDK for Python\n",
        "\n",
        "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cNEiwLd0lugu"
      },
      "outputs": [],
      "source": [
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tutorial_start:custom"
      },
      "source": [
        "# Tutorial\n",
        "\n",
        "Now you are ready to start creating the PyTorch distributed training job."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "train_custom_model"
      },
      "source": [
        "### Packaging the training application\n",
        "\n",
        "Before running the training job on Vertex AI, the training application code and any dependencies must be packaged and uploaded to Cloud Storage bucket or Container Registry or Artifact Registry that your Google Cloud project can access. This section shows how to package and stage your application in the cloud.\n",
        "\n",
        "There are two ways to package your application and dependencies and train on Vertex AI:\n",
        "\n",
        "1. [Create a Python source distribution](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container) with the training code and dependencies to use with a [pre-built container](https://cloud.google.com/vertex-ai/docs/training/pre-built-containers) on Vertex AI\n",
        "2. Use [custom containers](https://cloud.google.com/ai-platform/training/docs/custom-containers-training) to package dependencies using Docker containers\n",
        "\n",
        "**This notebook shows the Python source distribution option to run a custom training job on Vertex AI.**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBP_kLymhnYP"
      },
      "source": [
        "#### Recommended training application structure\n",
        "\n",
        "You can structure your training application in any way you like. However, the [following structure](https://cloud.google.com/vertex-ai/docs/training/create-python-pre-built-container#structure) is commonly used in Vertex AI samples, and having your project's organization be similar to the samples can make it easier for you to follow the samples.\n",
        "\n",
        "```\n",
        ".\n",
        "├── python_package\n",
        "│   ├── README.md\n",
        "│   ├── setup.py\n",
        "│   └── trainer\n",
        "│       ├── __init__.py\n",
        "│       └── task.py\n",
        "└── pytorch-distributed-training-reduction-server.ipynb    --> This notebook\n",
        "```\n",
        "\n",
        "1. Main project directory contains your `setup.py` file with the dependencies. \n",
        "2. Use a subdirectory named `trainer` to store your main application module and `scripts` to submit training jobs locally or cloud\n",
        "3. Inside `trainer` directory:\n",
        "    - `task.py` - Main application module 1) initializes PyTorch distributed training environment, and 2) Runs the model training and evaluation experiment, and exports the final model.\n",
        "    - `__init__.py` is required to make Python treat directories containing the file as packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4387867f83f2"
      },
      "source": [
        "#### Define variables for the training application\n",
        "\n",
        "Initialize the variables to define pre-built container image, location of training application and training module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mKPS5aGjaKRy"
      },
      "outputs": [],
      "source": [
        "APP_NAME = \"pytorch-bert\"\n",
        "\n",
        "PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai/training/pytorch-gpu.1-9:latest\"\n",
        ")\n",
        "\n",
        "PYTHON_PACKAGE_APPLICATION_DIR = \"python_package\"\n",
        "\n",
        "source_package_file_name = f\"{PYTHON_PACKAGE_APPLICATION_DIR}/dist/trainer-0.1.tar.gz\"\n",
        "python_package_gcs_uri = (\n",
        "    f\"{BUCKET_URI}/pytorch-on-gcp/{APP_NAME}/train/python_package/trainer-0.1.tar.gz\"\n",
        ")\n",
        "\n",
        "python_module_name = \"trainer.task\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4NqTnsxaAdRp"
      },
      "source": [
        "#### Create file structure of the training application"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O9zLCELGbOjS"
      },
      "outputs": [],
      "source": [
        "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}\n",
        "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/README.md\n",
        "\n",
        "! mkdir {PYTHON_PACKAGE_APPLICATION_DIR}/trainer\n",
        "! touch {PYTHON_PACKAGE_APPLICATION_DIR}/trainer/__init__.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3rpehVa0iFjY"
      },
      "source": [
        "#### Create the `setup.py` file for the training application\n",
        "\n",
        "Following is the `setup.py` file for the training application. The `find_packages()` function inside `setup.py` includes the `trainer` directory in the package as it contains `__init__.py` which tells [Python Setuptools](https://setuptools.readthedocs.io/en/latest/) to include all subdirectories of the parent directory as dependencies. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7QNjlhx7cBnj"
      },
      "outputs": [],
      "source": [
        "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/setup.py\n",
        "\n",
        "import os\n",
        "from setuptools import find_packages\n",
        "from setuptools import setup\n",
        "import setuptools\n",
        "\n",
        "from distutils.command.build import build as _build\n",
        "import subprocess\n",
        "\n",
        "\n",
        "REQUIRED_PACKAGES = [\n",
        "    'transformers==4.28.0',\n",
        "    'datasets',\n",
        "    'evaluate',\n",
        "]\n",
        "\n",
        "setup(\n",
        "    name='trainer',\n",
        "    version='0.1',\n",
        "    install_requires=REQUIRED_PACKAGES,\n",
        "    packages=find_packages(),\n",
        "    include_package_data=True,\n",
        "    description='Vertex AI | Training | PyTorch | Text Classification | Python Package'\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jXlifmTtiTzy"
      },
      "source": [
        "#### Create training application code\n",
        "\n",
        "`task.py` is the main application module. It initializes the PyTorch distributed training environment and runs the model training and evaluation experiment, and exports the final model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wruDyrEPsLbH"
      },
      "outputs": [],
      "source": [
        "%%writefile ./{PYTHON_PACKAGE_APPLICATION_DIR}/trainer/task.py\n",
        "# Copyright 2022 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \\\"License\\\");\n",
        "# you may not use this file except in compliance with the License.\\n\",\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \\\"AS IS\\\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "from datetime import datetime\n",
        "import argparse\n",
        "\n",
        "import torch\n",
        "import torch.distributed as dist\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "import datasets\n",
        "from datasets import ClassLabel, Sequence, load_dataset\n",
        "\n",
        "import transformers\n",
        "from transformers import (\n",
        "    AutoModelForSequenceClassification, \n",
        "    AutoTokenizer,\n",
        "    EvalPrediction, \n",
        "    Trainer, \n",
        "    TrainingArguments,\n",
        "    default_data_collator)\n",
        "\n",
        "from google.cloud import storage\n",
        "\n",
        "\n",
        "def main():\n",
        "\n",
        "  parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
        "  parser.add_argument(\"--epochs\", type=int, help=\"Number of training epochs.\", default=2)\n",
        "  parser.add_argument(\"--batch_size\", type=int, help=\"Training batch size for one process.\", default=32)\n",
        "  parser.add_argument(\"--model_dir\", type=str, help=\"Directory for saving models.\", default=os.environ['AIP_MODEL_DIR'] if 'AIP_MODEL_DIR' in os.environ else \"\")\n",
        "  argv = parser.parse_args()\n",
        "\n",
        "  model_name_or_path = \"bert-large-uncased\"\n",
        "  padding = \"max_length\"\n",
        "  max_seq_length = 128\n",
        "\n",
        "  datasets = load_dataset(\"imdb\")\n",
        "  label_list = datasets[\"train\"].unique(\"label\")\n",
        "  label_to_id = {1: 1, 0: 0, -1: 0}\n",
        "\n",
        "  tokenizer = AutoTokenizer.from_pretrained(\n",
        "      model_name_or_path,\n",
        "      use_fast=True,\n",
        "  )\n",
        "\n",
        "  def preprocess_function(examples):\n",
        "      \"\"\"\n",
        "      Tokenize the input example texts\n",
        "      \"\"\"\n",
        "      args = (examples[\"text\"],)\n",
        "      result = tokenizer(\n",
        "          *args, padding=padding, max_length=max_seq_length, truncation=True\n",
        "      )\n",
        "\n",
        "      # Map labels to IDs (not necessary for GLUE tasks)\n",
        "      if label_to_id is not None and \"label\" in examples:\n",
        "          result[\"label\"] = [label_to_id[example] for example in examples[\"label\"]]\n",
        "\n",
        "      return result\n",
        "\n",
        "  # apply preprocessing function to input examples\n",
        "  datasets = datasets.map(preprocess_function, batched=True, load_from_cache_file=True)\n",
        "\n",
        "  model = AutoModelForSequenceClassification.from_pretrained(\n",
        "      model_name_or_path, \n",
        "      num_labels=len(label_list)\n",
        "  )\n",
        "\n",
        "  ngpus_per_node = torch.cuda.device_count()\n",
        "  world_size = int(os.environ[\"WORLD_SIZE\"])\n",
        "  \n",
        "  # Since we have ngpus_per_node processes per node, the total world_size\n",
        "  # needs to be adjusted accordingly\n",
        "  world_size =  world_size * ngpus_per_node\n",
        "\n",
        "  start = datetime.now().strftime(\"%Y_%m_%d_%H_%M_%S\")\n",
        "  print(f'Starting distributed training: {start}') \n",
        "  \n",
        "  # Use torch.multiprocessing.spawn to launch distributed processes\n",
        "  torch.multiprocessing.spawn(main_worker,\n",
        "    args = (ngpus_per_node, world_size, datasets, model, tokenizer, argv),\n",
        "    nprocs = ngpus_per_node,\n",
        "    join = True)\n",
        "  \n",
        "  end = datetime.now().strftime(\"%Y_%m_%d_%H_%M_%S\")\n",
        "  print(f'Distributed training complete: {end}')\n",
        "\n",
        "def main_worker(local_rank, ngpus_per_node, world_size, datasets, model, tokenizer, argv):\n",
        "\n",
        "  # This is the (global) rank of the current process\n",
        "  rank = int(os.environ[\"RANK\"])\n",
        "  \n",
        "  # For multiprocessing distributed training, rank needs to be the\n",
        "  # global rank among all the processes\n",
        "  rank = rank * ngpus_per_node + local_rank\n",
        "  print (f\"Distributed and Multi-processing. Setting rank for each worker. rank={rank}\")\n",
        "\n",
        "  dist.init_process_group(\n",
        "      backend=\"nccl\", \n",
        "      init_method=\"env://\",\n",
        "      world_size=world_size, \n",
        "      rank=rank)\n",
        "  \n",
        "  per_device_batch_size = int(argv.batch_size / ngpus_per_node)\n",
        "\n",
        "  training_args = TrainingArguments(\n",
        "      output_dir=\"/tmp/output/\",\n",
        "      num_train_epochs=argv.epochs, \n",
        "      per_device_train_batch_size=per_device_batch_size,\n",
        "      per_device_eval_batch_size=per_device_batch_size,\n",
        "      local_rank=local_rank,\n",
        "  )\n",
        "\n",
        "  def compute_metrics(p: EvalPrediction):\n",
        "    preds = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions\n",
        "    preds = np.argmax(preds, axis=1)\n",
        "    return {\"accuracy\": (preds == p.label_ids).astype(np.float32).mean().item()}\n",
        "  \n",
        "  trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=datasets[\"train\"],\n",
        "    eval_dataset=datasets[\"test\"],\n",
        "    compute_metrics=compute_metrics,\n",
        "    tokenizer=tokenizer,\n",
        "    data_collator=default_data_collator,\n",
        "  )\n",
        "\n",
        "  trainer.train()\n",
        "\n",
        "  # Save the trained model locally\n",
        "  model_filename = \"pytorch-birt-model\"\n",
        "  local_path = os.path.join(\"/tmp\", model_filename)\n",
        "  trainer.save_model(local_path)\n",
        "\n",
        "  if (os.path.exists(local_path)):\n",
        "    # Upload the trained model to Cloud storage\n",
        "    model_directory = argv.model_dir\n",
        "    storage_path = os.path.join(model_directory, model_filename)\n",
        "    blob = storage.blob.Blob.from_string(storage_path, client=storage.Client())\n",
        "\n",
        "    files = [f for f in os.listdir(local_path) if os.path.isfile(os.path.join(local_path, f))]\n",
        "    for file in files:\n",
        "      local_file = os.path.join(local_path, file)\n",
        "      blob.upload_from_filename(local_file)\n",
        "  \n",
        "    print(f\"Saved model files in {model_directory}/{model_filename}\")\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uzmJxNOnix9O"
      },
      "source": [
        "#### Create a source distribution\n",
        "\n",
        "Create a source distribution `dist/trainer-0.1.tar.gz`, and upload the source distribution with training application to Cloud Storage bucket, and then validate the source distribution exists on Cloud Storage bucket."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dGgRZ4N-bbW0"
      },
      "outputs": [],
      "source": [
        "! cd {PYTHON_PACKAGE_APPLICATION_DIR} && python3 setup.py sdist --formats=gztar\n",
        "\n",
        "! gsutil cp {source_package_file_name} {python_package_gcs_uri}\n",
        "\n",
        "! gsutil ls -l {python_package_gcs_uri}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gv0PB6WhjKsA"
      },
      "source": [
        "### Run custom training job with Reduction Server on Vertex AI\n",
        "\n",
        "Configure a custom job with the pre-built container image for PyTorch and training code packaged as Python source distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OPWYVIEwcQWX"
      },
      "outputs": [],
      "source": [
        "print(f\"APP_NAME={APP_NAME}\")\n",
        "print(\n",
        "    f\"PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI={PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI}\"\n",
        ")\n",
        "print(f\"python_package_gcs_uri={python_package_gcs_uri}\")\n",
        "print(f\"python_module_name={python_module_name}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iy_OwOByjW1W"
      },
      "source": [
        "#### Create a training job\n",
        "\n",
        "You use [Vertex AI SDK for Python](https://cloud.google.com/vertex-ai/docs/start/client-libraries#client_libraries) to create a custom training job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ups1PMXCcjDB"
      },
      "outputs": [],
      "source": [
        "JOB_NAME = f\"pytorch-birt-reduction-server-{UUID}\"\n",
        "print(f\"JOB_NAME={JOB_NAME}\")\n",
        "\n",
        "job = aiplatform.CustomPythonPackageTrainingJob(\n",
        "    display_name=f\"{JOB_NAME}\",\n",
        "    python_package_gcs_uri=python_package_gcs_uri,\n",
        "    python_module_name=python_module_name,\n",
        "    container_uri=PRE_BUILT_TRAINING_CONTAINER_IMAGE_URI,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YimBZzZmAUwP"
      },
      "source": [
        "#### Define the training cluster worker pool and experiment configuration parameters\n",
        "\n",
        "Reduction Server can be used with any distributed training framework that uses the NVIDIA NCCL library for the all-reduce collective operation. You do not need to change or recompile your training application.\n",
        "\n",
        "The Google Cloud guide [Distributed training](https://cloud.google.com/vertex-ai/docs/training/distributed-training) provides detailed instructions on how to run distributed training jobs on Vertex AI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yIjTAfqvcnJu"
      },
      "outputs": [],
      "source": [
        "# Training cluster worker pool configuration\n",
        "REPLICA_COUNT = 3\n",
        "MACHINE_TYPE = \"n1-standard-16\"\n",
        "ACCELERATOR_TYPE = \"NVIDIA_TESLA_V100\"\n",
        "ACCELERATOR_COUNT = 2\n",
        "\n",
        "# Reduction Server configuration\n",
        "REDUCTION_SERVER_COUNT = 4\n",
        "REDUCTION_SERVER_MACHINE_TYPE = \"n1-highcpu-16\"\n",
        "REDUCTION_SERVER_IMAGE_URI = (\n",
        "    \"us-docker.pkg.dev/vertex-ai-restricted/training/reductionserver:latest\"\n",
        ")\n",
        "ENVIRONMENT_VARIABLES = {\"NCCL_DEBUG\": \"INFO\"}\n",
        "\n",
        "# Training experiment parameters\n",
        "EPOCHS = 2\n",
        "BATCH_SIZE = 32\n",
        "MODEL_DIR = f\"{BUCKET_URI}/{JOB_NAME}\"\n",
        "\n",
        "training_args = [\n",
        "    \"--epochs\",\n",
        "    str(EPOCHS),\n",
        "    \"--batch_size\",\n",
        "    str(BATCH_SIZE),\n",
        "    \"--model_dir\",\n",
        "    MODEL_DIR,\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xErX0T7LHYjX"
      },
      "source": [
        "#### Submit the training job\n",
        "\n",
        "After the training cluster configuration parameters have been defined, use the Vertex AI SDK for Python to submit and monitor a training job.\n",
        "\n",
        "*NOTE: When using Vertex AI SDK for Python for submitting a training job, it creates a Training Pipeline which launches the custom job on Vertex AI Training service.*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XtBA-NTIcr7T"
      },
      "outputs": [],
      "source": [
        "model = job.run(\n",
        "    replica_count=REPLICA_COUNT,\n",
        "    machine_type=MACHINE_TYPE,\n",
        "    accelerator_type=ACCELERATOR_TYPE,\n",
        "    accelerator_count=ACCELERATOR_COUNT,\n",
        "    reduction_server_replica_count=REDUCTION_SERVER_COUNT,\n",
        "    reduction_server_machine_type=REDUCTION_SERVER_MACHINE_TYPE,\n",
        "    reduction_server_container_uri=REDUCTION_SERVER_IMAGE_URI,\n",
        "    environment_variables=ENVIRONMENT_VARIABLES,\n",
        "    args=training_args,\n",
        "    sync=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CA38ygH_jeJe"
      },
      "source": [
        "#### Monitor the training job (optional)\n",
        "\n",
        "You can monitor the custom job launched from Cloud Console following the link [here](https://console.cloud.google.com/vertex-ai/training/training-pipelines/) or use gcloud CLI command [`gcloud beta ai custom-jobs stream-logs`](https://cloud.google.com/sdk/gcloud/reference/beta/ai/custom-jobs/stream-logs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s6LpsQI1WpPz"
      },
      "source": [
        "#### Validate the model artifacts (optional)\n",
        "\n",
        "You can validate the model artifacts written to Cloud Storage bucket by the training application after the job completes successfully."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_138iEP3HD0L"
      },
      "outputs": [],
      "source": [
        "print(f\"Model artifacts are available at {MODEL_DIR}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cleanup:custom"
      },
      "source": [
        "# Cleaning up\n",
        "\n",
        "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
        "\n",
        "Otherwise, you can delete the individual resources you created in this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NNmebHf7lug0"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "delete_custom_job = True\n",
        "delete_bucket = False\n",
        "\n",
        "if delete_custom_job:\n",
        "    try:\n",
        "        job.delete()\n",
        "    except Exception as e:\n",
        "        print(e)\n",
        "\n",
        "if delete_bucket or os.getenv(\"IS_TESTING\"):\n",
        "    ! gsutil rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "pytorch_distributed_training_reduction_server.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}