notebooks/community/model_garden/model_garden_pytorch_dolly_v2.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2023 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": "2bd716bf3e39"
      },
      "source": [
        "# Vertex AI Model Garden - Dolly V2\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td>\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.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/community/model_garden/model_garden_pytorch_dolly_v2.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/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/model_garden/model_garden_pytorch_dolly_v2.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",
        "    (a Python-3 CPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8cd12648da4"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates downloading and deploying the pre-trained [Dolly-V2-3b](https://huggingface.co/databricks/dolly-v2-3b), [Dolly-V2-7b](https://huggingface.co/databricks/dolly-v2-7b) & [Dolly-V2-12b](https://huggingface.co/databricks/dolly-v2-12b) model on Vertex AI for online prediction. It also demonstrates the evaluation of popular benchmark datasets through Vertex CustomJobs using [EleutherAI’s evaluation harness](https://github.com/EleutherAI/lm-evaluation-harness).\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
        "- Run predictions by giving instructions or by QnA.\n",
        "- Run the evaluation on any of the benchmark datasets (hellaswag, openbookqa, etc...) as per your requirement/choice and get the evaluation results.\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing).\n",
        "\n",
        "Use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "264c07757582"
      },
      "source": [
        "## Setup environment\n",
        "\n",
        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FepKNEq7kjtm"
      },
      "source": [
        "### Import the required packages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4zrBU0yzkFrD"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import storage"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d73ffa0c0b83"
      },
      "source": [
        "### Colab only"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "!pip3 install --upgrade google-cloud-aiplatform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b60a4d7100bf"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth as google_auth\n",
        "\n",
        "google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0f826ff482a2"
      },
      "source": [
        "### Setup Google Cloud project\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",
        "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
        "\n",
        "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8958ebc71868"
      },
      "source": [
        "Fill following variables for experiments environment:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9db30f827a65"
      },
      "outputs": [],
      "source": [
        "# Cloud project id.\n",
        "PROJECT_ID = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The region you want to launch jobs in.\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The Cloud Storage bucket for storing experiments output. Fill it without the 'gs://' prefix.\n",
        "GCS_BUCKET = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "BUCKET_URI = \"gs://\" + GCS_BUCKET"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92f16e22c20b"
      },
      "source": [
        "Initialize Vertex AI API:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1680c257acfb"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=GCS_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ca48b699d17"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "de9882ea89ea"
      },
      "outputs": [],
      "source": [
        "# The pre-built serving and evaluation docker image.\n",
        "# The model artifacts are embedded within the container, except for model weights which will be downloaded during deployment.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240424_0916_RC00\"\n",
        "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20231011_0934_RC00\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10188266a5cd"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cac4478ae098"
      },
      "outputs": [],
      "source": [
        "def deploy_model(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_model_len: int = 1024,\n",
        "):\n",
        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=8080\",\n",
        "        \"--swap-space=4\",\n",
        "        \"--gpu-memory-utilization=0.9\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        \"--max-num-batched-tokens=2048\",\n",
        "        \"--dtype=float16\",\n",
        "        \"--disable-log-stats\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        # service_account=SERVICE_ACCOUNT,\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
        "    print(\"To load this existing endpoint from a different session:\")\n",
        "    print(\"from google.cloud import aiplatform\")\n",
        "    print(\n",
        "        f'endpoint = aiplatform.Endpoint(\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}\")'\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2d72ecdb8c9"
      },
      "source": [
        "## Upload and deploy models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9448c5f545fa"
      },
      "source": [
        "This section uploads the pre-trained model to Model Registry and deploys it on the Endpoint with 1 A100 GPU.\n",
        "\n",
        "The model deployment step will take ~15 minutes to complete."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b4b46c28d8b1"
      },
      "outputs": [],
      "source": [
        "MODEL_ID = \"databricks/dolly-v2-3b\"  # @param [\"databricks/dolly-v2-3b\", \"databricks/dolly-v2-7b\", \"databricks/dolly-v2-12b\"]\n",
        "\n",
        "machine_type = \"g2-standard-8\"\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "\n",
        "if \"12b\" in MODEL_ID:\n",
        "    accelerator_count = 2\n",
        "    machine_type = \"g2-standard-24\"\n",
        "else:\n",
        "    accelerator_count = 1\n",
        "\n",
        "model, endpoint = deploy_model(\n",
        "    model_name=get_job_name_with_datetime(\"dolly\"),\n",
        "    model_id=MODEL_ID,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 10 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ohPQGAawXtAf"
      },
      "source": [
        "## Send prediction request"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6be655247cb1"
      },
      "outputs": [],
      "source": [
        "input_text = \"Explain to me the difference between nuclear fission and fusion.\"\n",
        "\n",
        "instances = [\n",
        "    {\"prompt\": input_text},\n",
        "]\n",
        "preds = endpoint.predict(instances=instances).predictions\n",
        "print(preds)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-P3IVO8Ykrx6"
      },
      "source": [
        "## Evaluation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hrYUiP7LQFw-"
      },
      "source": [
        "This section demonstrates how to evaluate the Dolly V2 models using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex Custom Job. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
        "\n",
        "This example uses the dataset [HellaSwag](https://allenai.org/data/hellaswag). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "honqZVqIpiTX"
      },
      "source": [
        "### Define the variables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "micAdZTFk3_X"
      },
      "outputs": [],
      "source": [
        "# Define the machine_type, accelerator_type, accelerator_count and benchmark dataset.\n",
        "\n",
        "eval_dataset = \"hellaswag\"  # @param [\"arc_challenge\", \"arc_easy\", \"boolq\", \"hellaswag\", \"openbookqa\", \"piqa\", \"winogrande\"]\n",
        "\n",
        "# Worker pool spec.\n",
        "# Find Vertex AI supported accelerators and regions in:\n",
        "# https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
        "\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "# Sets V100 to deploy Dolly V2 as an example.\n",
        "if \"12b\" in MODEL_ID:\n",
        "    accelerator_count = 2  # for dolly-v2-12b\n",
        "else:\n",
        "    accelerator_count = 1  # for dolly-v2-3b & dolly-v2-7b\n",
        "\n",
        "replica_count = 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZOAIZNMyl2uB"
      },
      "outputs": [],
      "source": [
        "# Setup evaluation job.\n",
        "job_name = get_job_name_with_datetime(prefix=\"dolly-v2\")\n",
        "eval_output_dir = os.path.join(BUCKET_URI, job_name)\n",
        "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g0t0RBixIw0P"
      },
      "outputs": [],
      "source": [
        "# Prepare evaluation command that runs the evaluation harness.\n",
        "# Set `trust_remote_code = True` because evaluating the model requires\n",
        "# executing code from the model repository.\n",
        "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
        "eval_command = [\n",
        "    \"python\",\n",
        "    \"main.py\",\n",
        "    \"--model\",\n",
        "    \"hf-causal-experimental\",\n",
        "    \"--model_args\",\n",
        "    f\"pretrained={MODEL_ID},trust_remote_code=True,use_accelerate=True,device_map_option=auto\",\n",
        "    \"--tasks\",\n",
        "    f\"{eval_dataset}\",\n",
        "    \"--output_path\",\n",
        "    f\"{eval_output_dir_gcsfuse}\",\n",
        "]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_Gjzbb8pI5_V"
      },
      "outputs": [],
      "source": [
        "# Run the evaluation job.\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": accelerator_type,\n",
        "            \"accelerator_count\": accelerator_count,\n",
        "        },\n",
        "        \"replica_count\": replica_count,\n",
        "        \"disk_spec\": {\n",
        "            \"boot_disk_size_gb\": 500,\n",
        "        },\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": EVAL_DOCKER_URI,\n",
        "            \"command\": eval_command,\n",
        "            \"args\": [],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "eval_job = aiplatform.CustomJob(\n",
        "    display_name=job_name,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    base_output_dir=eval_output_dir,\n",
        ")\n",
        "\n",
        "eval_job.run()\n",
        "\n",
        "print(\"Evaluation results were saved in:\", eval_output_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gxi5nZOUpZMR"
      },
      "outputs": [],
      "source": [
        "# Fetch evaluation results.\n",
        "storage_client = storage.Client()\n",
        "BUCKET_NAME = GCS_BUCKET\n",
        "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
        "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
        "blob = bucket.blob(RESULT_FILE_PATH)\n",
        "raw_result = blob.download_as_string()\n",
        "\n",
        "# Print evaluation results.\n",
        "result = json.loads(raw_result)\n",
        "result_formatted = json.dumps(result, indent=2)\n",
        "print(f\"Evaluation result:\\n{result_formatted}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uXLWt43CXMgj"
      },
      "source": [
        "## Clean up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2ccf3714dbe9"
      },
      "outputs": [],
      "source": [
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()\n",
        "\n",
        "# Delete evaluation jobs.\n",
        "eval_job.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_dolly_v2.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}