notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "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 (Deployment)\n",
        "\n",
        "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n",
        "  \u003ctd style=\"text-align: center\"\u003e\n",
        "    \u003ca 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_gemma_deployment_on_vertex.ipynb\"\u003e\n",
        "      \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
        "    \u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd style=\"text-align: center\"\u003e\n",
        "    \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb\"\u003e\n",
        "      \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n",
        "    \u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates deploying Gemma 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 [vLLM](https://github.com/vllm-project/vllm), the state-of-the-art open source LLM serving solution on GPU.\n",
        "\n",
        "This notebook also showcases how to use the [Text moderation API](https://cloud.google.com/natural-language/docs/moderating-text) to analyze model predictions against a predefined list of safety attributes.\n",
        "\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy Gemma with Hex-LLM on TPU\n",
        "- Deploy Gemma with [vLLM](https://github.com/vllm-project/vllm) 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": "markdown",
      "metadata": {
        "id": "C_wC61dhpWXj"
      },
      "source": [
        "### Request for TPU quota\n",
        "\n",
        "By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 0. If you would like to use Hex-LLM TPU deployment, please request TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota) if you haven't done so already. Please note that vLLM GPU deployment does not need this step."
      ]
    },
    {
      "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 **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\n",
        "\n",
        "import os\n",
        "import json\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\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",
        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\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, please change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    # Create a unique GCS bucket for this notebook if not specified\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\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, \"gemma\")\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",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "# Enable Vertex AI and Cloud Compute APIs.\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# @markdown ## Access Gemma Models\n",
        "# @markdown For GPU based serving, choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
        "# @markdown or Vertex AI as described below.\n",
        "\n",
        "# @markdown If you already obtained access to Gemma models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
        "# @markdown Alternatively, you can also load the original Gemma models for serving from Vertex AI after accepting the agreement.\n",
        "\n",
        "# @markdown For TPU based serving with Hex-LLM, choose the Kaggle option.\n",
        "\n",
        "# @markdown **Please only select and fill one of the three following sections.**\n",
        "LOAD_MODEL_FROM = \"Hugging Face\"  # @param [\"Hugging Face\", \"Google Cloud\", \"Kaggle\"] {isTemplate:true}\n",
        "\n",
        "# @markdown ---\n",
        "\n",
        "# @markdown ### Access Gemma models on Hugging Face for GPU based serving\n",
        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 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",
        "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
        "    assert (\n",
        "        HF_TOKEN\n",
        "    ), \"Please provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
        "\n",
        "# @markdown *--- Or ---*\n",
        "# @markdown ### Access Gemma models on Vertex AI for GPU based serving\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [Gemma model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/335) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
        "# @markdown 1. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
        "# @markdown 1. After accepting the agreement of Gemma, a `https://` link containing Gemma pretrained and finetuned models will be shared.\n",
        "# @markdown 1. Paste the link in the `VERTEX_AI_MODEL_GARDEN_GEMMA` field below.\n",
        "# @markdown **Note:** This will unzip and copy the Gemma model artifacts to your Cloud Storage bucket, which will take around 1 hour.\n",
        "\n",
        "VERTEX_AI_MODEL_GARDEN_GEMMA = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "\n",
        "\n",
        "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
        "    assert (\n",
        "        VERTEX_AI_MODEL_GARDEN_GEMMA\n",
        "    ), \"Please accept the agreement of Gemma in Vertex AI Model Garden and get the URL to Gemma model artifacts, or select a different model source.\"\n",
        "\n",
        "    # Only use the last part in case a full command is pasted.\n",
        "    signed_url = VERTEX_AI_MODEL_GARDEN_GEMMA.split(\" \")[-1].strip('\"')\n",
        "\n",
        "    ! mkdir -p ./gemma\n",
        "    ! curl -X GET \"{signed_url}\" | tar -xzvf - -C ./gemma/\n",
        "    ! gsutil -m cp -R ./gemma/* {MODEL_BUCKET}\n",
        "\n",
        "    model_path_prefix = MODEL_BUCKET\n",
        "    HF_TOKEN = \"\"\n",
        "else:\n",
        "    model_path_prefix = \"google/\"\n",
        "\n",
        "\n",
        "# @markdown *--- Or ---*\n",
        "# @markdown ### Access Gemma models from Kaggle for TPU based serving\n",
        "# @markdown Kaggle credentials are required for Hex-LLM deployment with TPUs.\n",
        "# @markdown Generate the Kaggle username and key by following [these instructions](https://github.com/Kaggle/kaggle-api?tab=readme-ov-file#api-credentials).\n",
        "# @markdown You will need to review and accept the model license.\n",
        "KAGGLE_USERNAME = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "KAGGLE_KEY = \"\"  # @param {type:\"string\", isTemplate:true}\n",
        "if LOAD_MODEL_FROM == \"Kaggle\":\n",
        "    assert (\n",
        "        KAGGLE_USERNAME and KAGGLE_KEY\n",
        "    ), \"Please provide Kaggle credentials to load models from Kaggle, or select a different model source.\"\n",
        "# @markdown ---\n",
        "\n",
        "# The pre-built serving docker images.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20240426_0936_RC01\"\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -\u003e str:\n",
        "    \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def deploy_model_hexllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
        "    max_num_batched_tokens: int = 11264,\n",
        "    tokens_pad_multiple: int = 1024,\n",
        "    seqs_pad_multiple: int = 32,\n",
        ") -\u003e Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    num_tpu_chips = int(machine_type[-2])\n",
        "    hexllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        \"--log_level=INFO\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor_parallel_size={num_tpu_chips}\",\n",
        "        \"--num_nodes=1\",\n",
        "        \"--use_ray\",\n",
        "        \"--batch_mode=continuous\",\n",
        "        f\"--max_num_batched_tokens={max_num_batched_tokens}\",\n",
        "        f\"--tokens_pad_multiple={tokens_pad_multiple}\",\n",
        "        f\"--seqs_pad_multiple={seqs_pad_multiple}\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"PJRT_DEVICE\": \"TPU\",\n",
        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    if KAGGLE_USERNAME and KAGGLE_KEY:\n",
        "        env_vars[\"KAGGLE_USERNAME\"] = KAGGLE_USERNAME\n",
        "        env_vars[\"KAGGLE_KEY\"] = KAGGLE_KEY\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.entrypoints.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",
        "    )\n",
        "\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-12\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_model_len: int = 8192,\n",
        "    dtype: str = \"bfloat16\",\n",
        ") -\u003e Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with vLLM on GPU in 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=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        \"--gpu-memory-utilization=0.9\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    if HF_TOKEN:\n",
        "        env_vars[\"HF_TOKEN\"] = HF_TOKEN\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=[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",
        "    )\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",
        "def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
        "  \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
        "  service_endpoint = \"aiplatform.googleapis.com\"\n",
        "  quota_list_output = !gcloud alpha services quota list --service=$service_endpoint  --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
        "  # Use '.s' on the command output because it is an SList type.\n",
        "  quota_data = json.loads(quota_list_output.s)\n",
        "  if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
        "    return -1\n",
        "  if len(quota_data[0][\"consumerQuotaLimits\"]) == 0 or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]:\n",
        "    return -1\n",
        "  all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
        "  for region_data in all_regions_data:\n",
        "    if region_data.get('dimensions') and region_data['dimensions']['region'] == region:\n",
        "      if 'effectiveLimit' in region_data:\n",
        "        return int(region_data['effectiveLimit'])\n",
        "      else:\n",
        "        return 0\n",
        "  return -1\n",
        "\n",
        "\n",
        "def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
        "  \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
        "  Args:\n",
        "    accelerator_type: The accelerator type.\n",
        "    is_for_training: Whether the resource is used for training. Set false\n",
        "    for serving use case.\n",
        "  Returns:\n",
        "    The resource id.\n",
        "  \"\"\"\n",
        "  training_accelerator_map = {\n",
        "      \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
        "      \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
        "      \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
        "      \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
        "      \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
        "      \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
        "  }\n",
        "  serving_accelerator_map = {\n",
        "      \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
        "      \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
        "      \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
        "      \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
        "      \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
        "  }\n",
        "  if is_for_training:\n",
        "    if accelerator_type in training_accelerator_map:\n",
        "      return training_accelerator_map[accelerator_type]\n",
        "    else:\n",
        "      raise ValueError(\n",
        "          f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
        "      )\n",
        "  else:\n",
        "    if accelerator_type in serving_accelerator_map:\n",
        "      return serving_accelerator_map[accelerator_type]\n",
        "    else:\n",
        "      raise ValueError(\n",
        "          f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
        "      )\n",
        "\n",
        "\n",
        "def check_quota(project_id:str, region: str, accelerator_type: str,\n",
        "                accelerator_count: int, is_for_training: bool):\n",
        "  \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
        "  resource_id = get_resource_id(accelerator_type, is_for_training)\n",
        "  quota = get_quota(project_id, region, resource_id)\n",
        "  quota_request_instruction = (\"Either use \"\n",
        "            \"a different region or request additional quota. Follow \"\n",
        "            \"instructions here \"\n",
        "            \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
        "            \" to check quota in a region or request additional quota for \"\n",
        "            \"your project.\")\n",
        "  if quota == -1:\n",
        "    raise ValueError(\n",
        "            f\"\"\"Quota not found for: {resource_id} in {region}.\n",
        "            {quota_request_instruction}\"\"\"\n",
        "        )\n",
        "  if quota < accelerator_count:\n",
        "    raise ValueError(\n",
        "            f\"\"\"Quota not enough for {resource_id} in {region}:\n",
        "            {quota} < {accelerator_count}.\n",
        "            {quota_request_instruction}\"\"\"\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8neJc8CnDDpu"
      },
      "source": [
        "## Deploy Gemma 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",
        "To request TPU quota, please follow 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": "E8OiHHNNE_wj"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "# @markdown Set the model ID. Model weights can be loaded from Kaggle [google/gemma](https://www.kaggle.com/models/google/gemma/frameworks/pyTorch) or from a GCS bucket.\n",
        "\n",
        "# @markdown Alternatively, you can specify a GCS folder that contains the original or\n",
        "# @markdown customized PyTorch model checkpoint and tokenizer. In this case, the GCS\n",
        "# @markdown folder is expected to contain \"gemma-2b\", \"gemma-7b\", \"gemma-1.1-2b\" or\n",
        "# @markdown  \"gemma-1.1-7b\" in either the GCS folder path or the name of the model checkpoint\n",
        "# @markdown file is expected to have the suffix \".ckpt\" and the tokenizer file is\n",
        "# @markdown file. The model checkpointis expected to have the name \"tokenizer.model\" (same as\n",
        "# @markdown the Kaggle files). An example structure for the GCS folder\n",
        "# @markdown gs://my-deployment-bucket/pytorch-files is:\n",
        "# @markdown - gs://my-deployment-bucket/pytorch-files/gemma-2b.ckpt\n",
        "# @markdown - gs://my-deployment-bucket/pytorch-files/tokenizer.model\n",
        "\n",
        "# @markdown Select one of the four model variations or enter a GCS folder containing the original or customized PyTorch model checkpoint and tokenizer.\n",
        "MODEL_ID = \"google/gemma-1.1-2b-it\"  # @param [\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {allow-input: true, isTemplate: true}\n",
        "\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",
        "\n",
        "import re\n",
        "\n",
        "if MODEL_ID.startswith(\"gs://\"):\n",
        "    MODEL_BUCKET_URI = re.search(\"gs://(.*?)/\", MODEL_ID).group()\n",
        "    ! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $MODEL_BUCKET_URI\n",
        "    MODEL_ID = MODEL_BUCKET_URI\n",
        "\n",
        "if \"2b\" in MODEL_ID:\n",
        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 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",
        "else:\n",
        "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 7B 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",
        "\n",
        "check_quota(project_id=PROJECT_ID,\n",
        "            region=REGION,\n",
        "            accelerator_type=accelerator_type,\n",
        "            accelerator_count=accelerator_count,\n",
        "            is_for_training=False)\n",
        "\n",
        "# Note that a larger max_num_batched_tokens will require more TPU memory.\n",
        "max_num_batched_tokens = 11264\n",
        "# Multiple of tokens for padding alignment. A higher value can reduce\n",
        "# re-compilation but can also increase the waste in computation.\n",
        "tokens_pad_multiple = 1024\n",
        "# Multiple of sequences for padding alignment. A higher value can reduce\n",
        "# re-compilation but can also increase the waste in computation.\n",
        "seqs_pad_multiple = 32\n",
        "\n",
        "model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
        "    model_id=MODEL_ID,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    max_num_batched_tokens=max_num_batched_tokens,\n",
        "    tokens_pad_multiple=tokens_pad_multiple,\n",
        "    seqs_pad_multiple=seqs_pad_multiple,\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. Note that the first few prompts will take longer to execute.\n",
        "\n",
        "# @markdown Example:\n",
        "\n",
        "# @markdown ```\n",
        "# @markdown \u003e What is a car?\n",
        "# @markdown \u003e 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_hexllm.name` allows us to get the endpoint\n",
        "#   name of the endpoint `endpoint_hexllm` 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_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\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 = 10  # @param {type: \"integer\"}\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 = endpoint_hexllm.predict(instances=instances)\n",
        "\n",
        "prediction = response.predictions[0]\n",
        "print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f615c03d6638"
      },
      "source": [
        "### Build chat applications with Gemma\n",
        "\n",
        "You can build chat applications with the instruction finetuned Gemma models.\n",
        "\n",
        "The instruction tuned Gemma models were trained with a specific formatter that annotates instruction tuning examples with extra information, both during training and inference. The annotations (1) indicate roles in a conversation, and (2) delineate tunes in a conversation. Below we show a sample code snippet for formatting the model prompt using the user and model chat templates for a multi-turn conversation. The relevant tokens are:\n",
        "- `user`: user turn\n",
        "- `model`: model turn\n",
        "- `\u003cstart_of_turn\u003e`: beginning of dialogue turn\n",
        "- `\u003cend_of_turn\u003e`: end of dialogue turn\n",
        "\n",
        "An example set of dialogues is:\n",
        "```\n",
        "\u003cstart_of_turn\u003euser\n",
        "knock knock\u003cend_of_turn\u003e\n",
        "\u003cstart_of_turn\u003emodel\n",
        "who is there\u003cend_of_turn\u003e\n",
        "\u003cstart_of_turn\u003euser\n",
        "LaMDA\u003cend_of_turn\u003e\n",
        "\u003cstart_of_turn\u003emodel\n",
        "LaMDA who?\u003cend_of_turn\u003e\n",
        "```\n",
        "where `\u003cend_of_turn\u003e\\n` is the turn separator and `\u003cstart_of_turn\u003emodel\\n` is the prompt prefix. This means if we would like to prompt the model with a question like, `What is Cramer's Rule?`, we should use:\n",
        "```\n",
        "\u003cstart_of_turn\u003euser\n",
        "What is Cramer's Rule?\u003cend_of_turn\u003e\n",
        "\u003cstart_of_turn\u003emodel\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e59377392346"
      },
      "outputs": [],
      "source": [
        "# Chat templates.\n",
        "USER_CHAT_TEMPLATE = \"\u003cstart_of_turn\u003euser\\n{prompt}\u003cend_of_turn\u003e\\n\"\n",
        "MODEL_CHAT_TEMPLATE = \"\u003cstart_of_turn\u003emodel\\n{prompt}\u003cend_of_turn\u003e\\n\"\n",
        "\n",
        "# Sample formatted prompt.\n",
        "prompt = (\n",
        "    USER_CHAT_TEMPLATE.format(prompt=\"What is a good place for travel in the US?\")\n",
        "    + MODEL_CHAT_TEMPLATE.format(prompt=\"California.\")\n",
        "    + USER_CHAT_TEMPLATE.format(prompt=\"What can I do in California?\")\n",
        "    + \"\u003cstart_of_turn\u003emodel\\n\"\n",
        ")\n",
        "print(\"Chat prompt:\\n\", prompt)\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": 50,\n",
        "        \"temperature\": 1.0,\n",
        "        \"top_p\": 1.0,\n",
        "        \"top_k\": 10,\n",
        "    },\n",
        "]\n",
        "response = endpoint_hexllm.predict(instances=instances)\n",
        "\n",
        "prediction = response.predictions[0]\n",
        "print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YKZ4CBJ2kYaW"
      },
      "source": [
        "## Deploy Gemma models with vLLM on GPU\n",
        "\n",
        "[vLLM](https://github.com/vllm-project/vllm) is a high-throughput GPU Large Language Model (LLM) serving library which implements a number of optimizations including paged attention and continuous batching.\n",
        "\n",
        "Note that V100 GPUs generally offer better throughput and latency performance than L4 GPUs, while L4 GPUs are generally more cost efficient than V100 GPUs. The serving efficiency of L4, V100 and T4 GPUs is inferior to that of A100 GPUs, but L4, V100 and T4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
        "\n",
        "Gemma model weights are stored in bfloat16 precision. L4 and A100 GPUs are needed for vLLM serving at bfloat16 precision. V100 and T4 GPUs can support vLLM serving at float32 and float16 precision, and they are also meaningful deployment configurations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "03d504bcd60b"
      },
      "outputs": [],
      "source": [
        "# @title Deploy\n",
        "MODEL_ID = \"gemma-1.1-2b-it\"  # @param [\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {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\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_T4\", \"NVIDIA_TESLA_A100\"] {isTemplate: true}\n",
        "\n",
        "if \"2b\" in MODEL_ID:\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        # Sets 1 L4 (24G) to deploy Gemma 2B models.\n",
        "        machine_type = \"g2-standard-8\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"bfloat16\"\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "        # Sets 1 V100 (16G) to deploy Gemma 2B models.\n",
        "        machine_type = \"n1-standard-8\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"float32\"\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_T4\":\n",
        "        # Sets 1 T4 (16G) to deploy Gemma 2B models.\n",
        "        machine_type = \"n1-standard-8\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"float32\"\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        # Sets 1 A100 (40G) to deploy Gemma 2B models.\n",
        "        machine_type = \"a2-highgpu-1g\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
        "            % accelerator_type\n",
        "        )\n",
        "elif \"7b\" in MODEL_ID:\n",
        "    if accelerator_type == \"NVIDIA_L4\":\n",
        "        # Sets 1 L4 (24G) to deploy Gemma 7B models.\n",
        "        machine_type = \"g2-standard-12\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"bfloat16\"\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        # Sets 1 A100 (40G) to deploy Gemma 7B models.\n",
        "        machine_type = \"a2-highgpu-1g\"\n",
        "        accelerator_count = 1\n",
        "        vllm_dtype = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            \"Recommended machine settings not found for accelerator type: %s\"\n",
        "            % accelerator_type\n",
        "        )\n",
        "else:\n",
        "    raise ValueError(\n",
        "        \"Recommended machine settings not found for accelerator type: %s\"\n",
        "        % accelerator_type\n",
        "    )\n",
        "\n",
        "check_quota(project_id=PROJECT_ID,\n",
        "            region=REGION,\n",
        "            accelerator_type=accelerator_type,\n",
        "            accelerator_count=accelerator_count,\n",
        "            is_for_training=False)\n",
        "\n",
        "# Note that a larger max_model_len will require more GPU memory.\n",
        "max_model_len = 2048\n",
        "\n",
        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\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_model_len=max_model_len,\n",
        "    dtype=vllm_dtype,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RRR11SWykYaX"
      },
      "source": [
        "Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64). Setting `raw_response` to `True` allows you to obtain raw outputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3f5a1e1de60d"
      },
      "outputs": [],
      "source": [
        "# @title Predict\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Note that the first few prompts will take longer to execute.\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint_vllm.name` allows us to get the endpoint\n",
        "#   name of the endpoint `endpoint_vllm` 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_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"What is a car?\"  # @param {type: \"string\"}\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 = 10  # @param {type: \"integer\"}\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 = endpoint_vllm.predict(instances=instances)\n",
        "\n",
        "prediction = response.predictions[0]\n",
        "print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "104fe2c03812"
      },
      "source": [
        "### Apply chat templates\n",
        "\n",
        "Chat templates can be applied to model predictions generated by the vLLM endpoint as well. You may use the same code snippets as for the Hex-LLM endpoint. They are not repeated here for brevity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @title  Delete the models and endpoints\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
        "# Undeploy models and delete endpoints.\n",
        "endpoint_hexllm.delete(force=True)\n",
        "endpoint_vllm.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model_hexllm.delete()\n",
        "model_vllm.delete()\n",
        "\n",
        "# Delete Cloud Storage objects.\n",
        "delete_bucket = False  # @param {type:\"boolean\", isTemplate: true}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_gemma_deployment_on_vertex.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}