notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "language": "python",
      "metadata": {
        "cellView": "form",
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - Gemma 2 (Deployment)\n",
        "\n",
        "<table><tbody><tr>\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%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma2_deployment_on_vertex.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying Gemma 2 models\n",
        " * on TPU using **Hex-LLM**, a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel serving solution built with **XLA** that is being developed by Google Cloud, and\n",
        " * on GPU using **TGI** ([text-generation-inference](https://github.com/huggingface/text-generation-inference)), the state-of-the-art open source LLM serving solution on GPU.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy Gemma 2 with Hex-LLM on TPU\n",
        "- Deploy Gemma with [TGI](https://github.com/huggingface/text-generation-inference) on GPU\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), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and 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": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "C_wC61dhpWXj"
      },
      "outputs": [],
      "source": [
        "# @title Request for TPU quota\n",
        "\n",
        "# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Import the necessary packages\n",
        "\n",
        "# Upgrade Vertex AI SDK.\n",
        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import os\n",
        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "models, endpoints = {}, {}\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma2\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
        "\n",
        "# @markdown ## Access Gemma 2 Models\n",
        "\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 2 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
        "\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "assert (\n",
        "    HF_TOKEN\n",
        "), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
        "\n",
        "model_path_prefix = \"google/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8neJc8CnDDpu"
      },
      "source": [
        "## Deploy Gemma 2 models with Hex-LLM on TPU\n",
        "\n",
        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
        "\n",
        "Refer to the \"Request for TPU quota\" section for TPU quota."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E8OiHHNNE_wj"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
        "\n",
        "# The pre-built serving docker images.\n",
        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:gemma2\"\n",
        "\n",
        "# @markdown Select one of the four model variations.\n",
        "MODEL_ID = \"gemma-2-2b-it\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
        "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
        "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
        "if \"2b\" in model_id:\n",
        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2 2B models.\n",
        "    machine_type = \"ct5lp-hightpu-1t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 1\n",
        "elif \"9b\" in model_id:\n",
        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 2 9B models.\n",
        "    machine_type = \"ct5lp-hightpu-4t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 4\n",
        "else:\n",
        "    # Sets ct5lp-hightpu-8t (8 TPU chips) to deploy Gemma 2 27B models.\n",
        "    machine_type = \"ct5lp-hightpu-8t\"\n",
        "    accelerator_type = \"TPU_V5e\"\n",
        "    # Note: 1 TPU V5 chip has only one core.\n",
        "    accelerator_count = 8\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=TPU_DEPLOYMENT_REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "# Server parameters.\n",
        "tensor_parallel_size = accelerator_count\n",
        "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
        "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
        "\n",
        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "# Endpoint configurations.\n",
        "min_replica_count = 1\n",
        "max_replica_count = 1\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    tensor_parallel_size: int = 1,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
        "    tpu_topology: str = \"1x1\",\n",
        "    hbm_utilization_factor: float = 0.6,\n",
        "    max_running_seqs: int = 256,\n",
        "    max_model_len: int = 4096,\n",
        "    endpoint_id: str = \"\",\n",
        "    min_replica_count: int = 1,\n",
        "    max_replica_count: int = 1,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
        "    if endpoint_id:\n",
        "        aip_endpoint_name = (\n",
        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
        "        )\n",
        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "    else:\n",
        "        endpoint = aiplatform.Endpoint.create(\n",
        "            display_name=f\"{model_name}-endpoint\",\n",
        "            location=TPU_DEPLOYMENT_REGION,\n",
        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "        )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    if not tensor_parallel_size:\n",
        "        tensor_parallel_size = int(machine_type[-2])\n",
        "\n",
        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
        "\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
        "        f\"--num_hosts={num_hosts}\",\n",
        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
        "        f\"--max_running_seqs={max_running_seqs}\",\n",
        "        f\"--max_model_len={max_model_len}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
        "    except:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
        "        serving_container_args=hexllm_args,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "        location=TPU_DEPLOYMENT_REGION,\n",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        min_replica_count=min_replica_count,\n",
        "        max_replica_count=max_replica_count,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    tensor_parallel_size=tensor_parallel_size,\n",
        "    hbm_utilization_factor=hbm_utilization_factor,\n",
        "    max_running_seqs=max_running_seqs,\n",
        "    min_replica_count=min_replica_count,\n",
        "    max_replica_count=max_replica_count,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nkUaMxIus6Pv"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown > What is a car?\n",
        "# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
        "#   name of the endpoint `endpoint` created in the cell above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint:\n",
        "# endpoint_name = endpoint_without_peft.name\n",
        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
        "max_tokens = 50  # @param {type: \"integer\"}\n",
        "temperature = 1.0  # @param {type: \"number\"}\n",
        "top_p = 1.0  # @param {type: \"number\"}\n",
        "top_k = 1  # @param {type: \"integer\"}\n",
        "\n",
        "# Overrides parameters for inferences.\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "    },\n",
        "]\n",
        "response = endpoints[\"hexllm_tpu\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFfqpQm8BNwZ"
      },
      "source": [
        "## Deploy Gemma models with TGI on GPU\n",
        "\n",
        "[TGI](https://github.com/huggingface/text-generation-inference) stands for Text Generation Inference. It's a powerful library designed specifically for running large language models on GPUs efficiently. TGI utilizes techniques like \"paged attention\" and \"continuous batching\" to improve the speed and throughput of LLMs.\n",
        "\n",
        "Currently, only L4 GPUs are demonstrated in this notebook. Functionality on other GPU types will be confirmed and added in the future.\n",
        "\n",
        "Gemma2 2B, 9B and 27B models require at least 1, 2, and 4 L4 GPUs respectively for deployment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TBNJYZMlBNwZ"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "\n",
        "# The pre-built serving docker image.\n",
        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
        "\n",
        "MODEL_ID = \"gemma-2-2b\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
        "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
        "\n",
        "# @markdown Finds Vertex AI prediction supported accelerators and regions in\n",
        "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
        "\n",
        "if \"2b\" in MODEL_ID:\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        # Sets 1 L4 (24G) to deploy Gemma 2 2B models.\n",
        "        machine_type = \"g2-standard-12\"\n",
        "        accelerator_count = 1\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
        "            % accelerator_type\n",
        "        )\n",
        "elif \"9b\" in MODEL_ID:\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        # Sets 2 L4 (24G) to deploy Gemma 2 9B models.\n",
        "        machine_type = \"g2-standard-24\"\n",
        "        accelerator_count = 2\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
        "            % accelerator_type\n",
        "        )\n",
        "elif \"27b\" in MODEL_ID:\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        # Sets 4 L4 (24G) to deploy Gemma 2 27B models.\n",
        "        machine_type = \"g2-standard-48\"\n",
        "        accelerator_count = 4\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
        "            % accelerator_type\n",
        "        )\n",
        "else:\n",
        "    raise ValueError(\"Recommended machine settings not found for model: %s\" % MODEL_ID)\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "# Note that larger token counts will require more GPU memory. For example, if you'd\n",
        "# like to increase the `max_total_tokens` and `max_batch_prefill_tokens` to 8192,\n",
        "# you may need 1 L4 for 2b model, 4 L4s for the 9b model, and 8 L4s for the 27b model.\n",
        "max_input_length = 1562\n",
        "max_total_tokens = 2048\n",
        "max_batch_prefill_tokens = 2048\n",
        "\n",
        "\n",
        "def deploy_model_tgi(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_input_length: int = 2047,\n",
        "    max_total_tokens: int = 2048,\n",
        "    max_batch_prefill_tokens: int = 2048,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=TGI_DOCKER_URI,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "    )\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",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
        "    model_id=model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_input_length=max_input_length,\n",
        "    max_total_tokens=max_total_tokens,\n",
        "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZCwACkjuBNwZ"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
        "\n",
        "# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the finetuning outcome:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown ### Human: How would the Future of AI in 10 Years look?### Assistant: Predicting the future is always a challenging task, but here are some possible ways that AI could evolve over the next 10 years: Continued advancements in deep learning: Deep learning has been one of the main drivers of recent AI breakthroughs, and we can expect continued advancements in this area. This may include improvements to existing algorithms, as well as the development of new architectures that are better suited to specific types of data and tasks. Increased use of AI in healthcare: AI has the potential to revolutionize healthcare, by improving the accuracy of diagnoses, developing new treatments, and personalizing patient care. We can expect to see continued investment in this area, with more healthcare providers and researchers using AI to improve patient outcomes. Greater automation in the workplace: Automation is already transforming many industries, and AI is likely to play an increasingly important role in this process. We can expect to see more jobs being automated, as well as the development of new types of jobs that require a combination of human and machine skills. More natural and intuitive interactions with technology: As AI becomes more advanced, we can expect to see more natural and intuitive ways of interacting with technology. This may include voice and gesture recognition, as well as more sophisticated chatbots and virtual assistants. Increased focus on ethical considerations: As AI becomes more powerful, there will be a growing need to consider its ethical implications. This may include issues such as bias in AI algorithms, the impact of automation on employment, and the use of AI in surveillance and policing. Overall, the future of AI in 10 years is likely to be shaped by a combination of technological advancements, societal changes, and ethical considerations. While there are many exciting possibilities for AI in the future, it will be important to carefully consider its potential impact on society and to work towards ensuring that its benefits are shared fairly and equitably.\n",
        "# @markdown ```\n",
        "\n",
        "# @markdown Click \"Show Code\" to see more details.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"How would the Future of AI in 10 Years look?\"  # @param {type: \"string\"}\n",
        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
        "max_tokens = 128  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 0.9  # @param {type:\"number\"}\n",
        "top_k = 1  # @param {type:\"integer\"}\n",
        "\n",
        "# Overrides max_tokens and top_k parameters during inferences.\n",
        "instances = [\n",
        "    {\n",
        "        \"inputs\": f\"### Human: {prompt}### Assistant: \",\n",
        "        \"parameters\": {\n",
        "            \"max_new_tokens\": max_tokens,\n",
        "            \"temperature\": temperature,\n",
        "            \"top_p\": top_p,\n",
        "            \"top_k\": top_k,\n",
        "        },\n",
        "    },\n",
        "]\n",
        "response = endpoints[\"tgi\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma2_deployment_on_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}