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 style=\"text-align: center\">\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\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Freduction_server%2Fpytorch_distributed_training_reduction_server.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td> \n",
        "  <td style=\"text-align: center\">\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\"><br>\n",
        "      Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>  \n",
        "  <td style=\"text-align: center\">\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\"><br>\n",
        "      View on GitHub\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": "f0316df526f8"
      },
      "source": [
        "## Get started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2c2cb2109a0"
      },
      "source": [
        "### Install Vertex AI SDK for Python and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6406a27bfea8"
      },
      "outputs": [],
      "source": [
        "! pip3 install --upgrade --quiet google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ff555b32bab8"
      },
      "source": [
        "### Restart runtime (Colab only)\n",
        "\n",
        "To use the newly installed packages, you must restart the runtime on Google Colab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f09b4dff629a"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ee775571c2b5"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92e68cfc3a90"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "Authenticate your environment on Google Colab.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "46604f70e831"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4f872cd812d0"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "294fe4e5a671"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}"
      ]
    },
    {
      "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": [
        "**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 {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
      ]
    },
    {
      "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": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=BUCKET_URI)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R7eelnCv6EWn"
      },
      "source": [
        "#### UUID\n",
        "\n",
        "If you're 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": "tutorial_start:custom"
      },
      "source": [
        "## Tutorial\n",
        "\n",
        "Now you're 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\", verification_mode='no_checks')\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
}