notebooks/community/model_garden/model_garden_pytorch_openllama_peft.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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - OpenLLaMA (PEFT)\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_openllama_peft.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_openllama_peft.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_openllama_peft.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> (A Python-3 GPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates running local inference with prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm), finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), quantizing and deploying OpenLLaMA with AWQ or GPTQ, and evaluating PEFT-finetuned OpenLLaMA in Vertex AI.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Run local inference with prebuilt OpenLLaMA\n",
        "- Deploy prebuilt OpenLLaMA\n",
        "- Deploy prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
        "- Finetune and deploy OpenLLaMA with PEFT\n",
        "- Quantize and deploy OpenLLaMA models with AWQ or GPTQ\n",
        "- Evaluate finetuned OpenLLaMA with PEFT\n",
        "\n",
        "| Models | LoRA |\n",
        "| :- | :- |\n",
        "| [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) | Y |\n",
        "| [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) | Y |\n",
        "| [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) | Y |\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), 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\n",
        "\n",
        "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
        "\n",
        "Running local inference with OpenLLaMA requires a GPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ioensNKM8ned"
      },
      "source": [
        "### Colab only\n",
        "Run the following commands for Colab and skip this section if you are using Workbench."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2707b02ef5df"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    ! pip3 install --upgrade google-cloud-aiplatform\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()\n",
        "\n",
        "    # Restart the notebook kernel after installs.\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)\n",
        "! pip3 install transformers==4.31.0\n",
        "! pip3 install sentencepiece==0.1.99\n",
        "! pip3 install accelerate==0.21.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bb7adab99e41"
      },
      "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.\n",
        "\n",
        "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6c460088b873"
      },
      "source": [
        "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). 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\")."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "855d6b96f291"
      },
      "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.\n",
        "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "import os\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
        "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
        "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
        "\n",
        "# The service account looks like:\n",
        "# '@.iam.gserviceaccount.com'\n",
        "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
        "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
        "# The service account for deploying fine tuned model.\n",
        "SERVICE_ACCOUNT = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e828eb320337"
      },
      "source": [
        "### Initialize Vertex AI API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "12cd25839741"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2cc825514deb"
      },
      "source": [
        "### Define constants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b42bd4fa2b2d"
      },
      "outputs": [],
      "source": [
        "# The pre-built training and serving docker images.\n",
        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00\"\n",
        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:20231130_0948_RC00\"\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20231127_0916_RC00\"\n",
        "VLLM_GPTQ_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:gptq\"\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": "0c250872074f"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "354da31189dc"
      },
      "outputs": [],
      "source": [
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "\n",
        "def create_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Creates a 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\")\n",
        "\n",
        "\n",
        "def deploy_model(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    finetuned_lora_model_path: str,\n",
        "    service_account: str,\n",
        "    task: str,\n",
        "    precision_loading_mode: str = \"float16\",\n",
        "    machine_type: str = \"n1-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count: int = 1,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": model_id,\n",
        "        \"PRECISION_LOADING_MODE\": precision_loading_mode,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "    if finetuned_lora_model_path:\n",
        "        serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
        "        serving_container_ports=[7080],\n",
        "        serving_container_predict_route=\"/predictions/peft_serving\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\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 deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"n1-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
        "    accelerator_count: int = 1,\n",
        "    quantization_method: str = \"\",\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\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=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        \"--gpu-memory-utilization=0.9\",\n",
        "        \"--max-num-batched-tokens=4096\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "    if quantization_method:\n",
        "        vllm_args.append(f\"--quantization={quantization_method}\")\n",
        "    if quantization_method == \"gptq\":\n",
        "        vllm_docker_uri = VLLM_GPTQ_DOCKER_URI\n",
        "    else:\n",
        "        vllm_docker_uri = VLLM_DOCKER_URI\n",
        "\n",
        "    serving_env = {\n",
        "        \"MODEL_ID\": \"openlm-research/open_llama\",\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=[7080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=serving_env,\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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65eaa62632d1"
      },
      "source": [
        "## Run inferences locally with prebuilt OpenLLaMA"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "339601a9500b"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from transformers import LlamaForCausalLM, LlamaTokenizer\n",
        "\n",
        "model_path = \"openlm-research/open_llama_3b\"\n",
        "\n",
        "tokenizer = LlamaTokenizer.from_pretrained(model_path)\n",
        "\n",
        "model = LlamaForCausalLM.from_pretrained(\n",
        "    model_path,\n",
        "    torch_dtype=torch.float16,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "\n",
        "prompt = \"Q: What is the largest animal?\\nA:\"\n",
        "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
        "input_ids = input_ids.to(\"cuda\")\n",
        "generation_output = model.generate(input_ids=input_ids, max_new_tokens=32)\n",
        "print(tokenizer.decode(generation_output[0]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V7VOhhHGpUrj"
      },
      "source": [
        "## Deploy Prebuilt OpenLLaMA with vLLM\n",
        "\n",
        "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
        "\n",
        "vLLM is a highly optimized LLM serving framework that can significantly increase serving throughput. The higher QPS you have, the more performance benefits you get from using vLLM."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4GTNnnuYqrW_"
      },
      "source": [
        "Set the prebuilt model id."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kLsRoc4Kqrkx"
      },
      "outputs": [],
      "source": [
        "prebuilt_model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YI0vaDi6p2fi"
      },
      "outputs": [],
      "source": [
        "# Finds Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "# Sets V100 to deploy open_llama_3b and open_llama_7b.\n",
        "# V100 serving has better throughput and latency performance than L4 serving.\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "accelerator_count = 1\n",
        "\n",
        "# Sets L4 to deploy open_llama_3b and open_llama_7b.\n",
        "# L4 serving is more cost efficient than V100 serving.\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "# Sets 2 V100 to deploy open_llama_13b.\n",
        "# V100 serving has better throughput and latency performance than L4 serving.\n",
        "# machine_type = \"n1-standard-16\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "# Sets 2 L4 to deploy open_llama_13b.\n",
        "# L4 serving is more cost efficient than V100 serving.\n",
        "# machine_type = \"g2-standard-24\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
        "    # vLLM currently does not support OpenLLaMA 3B.\n",
        "    precision_loading_mode = \"float16\"\n",
        "    model_without_peft, endpoint_without_peft = deploy_model(\n",
        "        model_name=get_job_name_with_datetime(prefix=\"openllama-serve\"),\n",
        "        model_id=model_id,\n",
        "        finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        task=\"causal-language-modeling-lora\",\n",
        "        precision_loading_mode=precision_loading_mode,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )\n",
        "else:\n",
        "    model_without_peft, endpoint_without_peft = deploy_model_vllm(\n",
        "        model_name=create_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
        "        model_id=prebuilt_model_id,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWYmYWoqqBuZ"
      },
      "source": [
        "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts. If you are interested in additional serving parameters, please refer to the vLLM GitHub [examples/api_client.py](https://github.com/vllm-project/vllm/blob/main/examples/api_client.py) for more details."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fjO4z3qAp3pK"
      },
      "outputs": [],
      "source": [
        "instance = {\n",
        "    \"prompt\": \"Hi, Google. How are you doing?\",\n",
        "    \"n\": 1,\n",
        "    \"max_tokens\": 32,\n",
        "    \"temperature\": 1.0,\n",
        "    \"top_p\": 1.0,\n",
        "    \"top_k\": 10,\n",
        "}\n",
        "response = endpoint_without_peft.predict(instances=[instance])\n",
        "print(response.predictions[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e70e3519ff8b"
      },
      "source": [
        "## Finetune and deploy OpenLLaMA with PEFT\n",
        "\n",
        "This section demonstrates how to finetune the OpenLLaMA-7b model, merge the finetuned LoRA adapter with the base model, and serve using vLLM."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5qCrm_kJH5cz"
      },
      "source": [
        "Set the base model id."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N3UBLiYrM3sU"
      },
      "outputs": [],
      "source": [
        "model_id = \"openlm-research/open_llama_7b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWGwJHqI7LMs"
      },
      "source": [
        "### Finetune"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KKEYoRfiHDVv"
      },
      "source": [
        "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
        "\n",
        "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes). You can either use a [dataset from huggingface](https://huggingface.co/datasets) or a custom JSONL dataset in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) stored in Cloud Storage. The `template` parameter is optional.\n",
        "\n",
        "In order to make the finetuning efficiently, we enabled quantization for loading pretrained models for finetuning LoRA models. Precision options include `\"4bit\"`, `\"8bit\"`, `\"float16\"` (default) and `\"float32\"`, and the precision can be set via `\"--precision_mode\"`. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. `open_llama_3b` and `open_llama_7b` can be finetuned on **1 V100 (16G)** and **1 L4 (24G)**, and `open_llama_13b` can be finetuned on **1 L4 (24G)**.\n",
        "\n",
        "In this section, the finetuned LoRA adapter will be saved to a GCS bucket specified by the variable `lora_adapter_dir` below; and we merge the LoRa adapter with the base model, and save it to a separate GCS bucket specified by `merged_model_output_dir` below.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0810ef72dd9f"
      },
      "source": [
        "#### [Optional] Finetune with a custom dataset\n",
        "\n",
        "To use a custom dataset, you should supply a `gs://` URI to a JSONL file in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) in the `dataset_name` below.\n",
        "\n",
        "For example, here is one data point from the sample dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`:\n",
        "\n",
        "```json\n",
        "{\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
        "```\n",
        "\n",
        "To use this sample dataset that contains `input_text` and `output_text` fields, set `dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl` and `template` to `vertex_sample`. For advanced usage with custom datatset fields, see [the template example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json) and supply your own JSON template as `gs://` URIs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "65467b361315"
      },
      "outputs": [],
      "source": [
        "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
        "dataset_name = \"Abirate/english_quotes\"  # @param {type:\"string\"}\n",
        "# Optional. Template name or gs:// URI to a custom template.\n",
        "template = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "# Finetunes open_llama_3b and open_llama_7b with 1 V100 (16G).\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "accelerator_count = 1\n",
        "\n",
        "# Finetunes open_llama_3b and open_llama_7b with 1 L4 (24G).\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "# Finetunes open_llama_13b with 1 L4 (24G).\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "# Finetunes open_llama_13b with 1 A100 (40G).\n",
        "# machine_type = \"a2-highgpu-1g\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "replica_count = 1\n",
        "\n",
        "\n",
        "# Setup training job.\n",
        "job_name = create_name_with_datetime(\"openllama-lora-train\")\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "\n",
        "# Create a GCS folder to store the LORA adapter.\n",
        "lora_adapter_dir = create_name_with_datetime(\"openllama-lora-adapter\")\n",
        "lora_output_dir = os.path.join(MODEL_BUCKET, lora_adapter_dir)\n",
        "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Create a GCS folder to store the merged model with the base model and the\n",
        "# finetuned LORA adapter.\n",
        "merged_model_dir = create_name_with_datetime(\"openllama-merged-model\")\n",
        "merged_model_output_dir = os.path.join(MODEL_BUCKET, merged_model_dir)\n",
        "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job.run(\n",
        "    args=[\n",
        "        \"--task=causal-language-modeling-lora\",\n",
        "        f\"--pretrained_model_id={model_id}\",\n",
        "        f\"--dataset_name={dataset_name}\",\n",
        "        f\"--output_dir={lora_output_dir_gcsfuse}\",\n",
        "        f\"--merge_base_and_lora_output_dir={merged_model_output_dir_gcsfuse}\",\n",
        "        \"--lora_rank=16\",\n",
        "        \"--lora_alpha=32\",\n",
        "        \"--lora_dropout=0.05\",\n",
        "        \"--warmup_steps=10\",\n",
        "        \"--max_steps=10\",\n",
        "        \"--learning_rate=2e-4\",\n",
        "        f\"--template={template}\",\n",
        "    ],\n",
        "    replica_count=replica_count,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    boot_disk_size_gb=500,\n",
        ")\n",
        "\n",
        "print(\"The finetuned Lora adapter can be found at: \", lora_output_dir)\n",
        "print(\n",
        "    \"The finetuned Lora adapter merged with the base model can be found at: \",\n",
        "    merged_model_output_dir,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53b8a1ad6def"
      },
      "source": [
        "### [Optional] Hyperparameter tuning\n",
        "\n",
        "You can use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to obtain a better performance by experimenting with different hyperparameters such as learning rates.\n",
        "\n",
        "Define the following specifications:\n",
        "\n",
        "- `worker_pool_specs`: Dictionary specifying the machine type and Docker image.\n",
        "\n",
        "- `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
        "\n",
        "- `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the hyperparameter_metric_tag that you set in your training application code, and the value is the optimization goal.\n",
        "\n",
        "The following 4bit QLoRA experiment results show the effectiveness of hyperparameter tuning evaluated on the ARC Challenge dataset (for reference only):\n",
        "\n",
        "| Model         | Training time | Trials | Parallel Trials | GPU  | ∆arc challenge | ∆hellaswag | ∆truthfulqa_mc | cost     |\n",
        "|---------------|---------------|--------|-----------------|------|----------------|------------|----------------|----------|\n",
        "| Openllama-3b  | 2d 10hrs      | 8      | 1               | L4x1 | +1.62          | +7.32      | +3.34          | \\$29.0232 |\n",
        "| Openllama-7b  | 1d 4hrs       | 8      | 2               | L4x1 | +2.82          | +3.55      | +6.68          | \\$47.8016 |\n",
        "| Openllama-13b | 6d 10hrs      | 8      | 2               | L4x1 | +1.01          | +3.67      | +6.19          | \\$87.9208 |\n",
        "\n",
        "The following example runs 8 trials on `timdettmers/openassistant-guanaco` with different learning rates, and evaluates the model on `arc_challenge` dataset. You can customize the search space by extending the range of learning rates, adding other parameters such as LoRA rank, etc. Please refer to the [hyperparameter tuning documentation](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) for more information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a81402f70641"
      },
      "outputs": [],
      "source": [
        "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
        "dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
        "# Optional. Template name or gs:// URI to a custom template.\n",
        "template = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "hpt_precision_mode = \"4bit\"\n",
        "\n",
        "# Worker pool spec for 4bit finetuning.\n",
        "\n",
        "# Finetunes Openllama 3B / 7B / 13B with 1 L4 (24G).\n",
        "machine_type = \"g2-standard-8\"\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "55ed2a2a9d54"
      },
      "source": [
        "### [Optional] Custom evaluation dataset\n",
        "\n",
        "To obtain a model with better performance on some specific tasks, you might want to run hyperparameter tuning with a custom evaluation dataset. The hyperparameter tuning service will pick the model according to the evaluation dataset and the metrics you selected. You can use any of the following tasks as the `eval_task` in the code cell below:\n",
        "\n",
        "1. The name of a [lm-evaluation-harness task](https://github.com/EleutherAI/lm-evaluation-harness/tree/big-refactor/lm_eval/tasks).\n",
        "\n",
        "2. `custom_likelihood`. Then, add a flag `--eval_dataset_path=<Cloud Storage URI to your JSONL dataset>`. The JSONL file must be in the format in Vertex AI language model's [prepare evaluation dataset](https://cloud.google.com/vertex-ai/docs/generative-ai/models/evaluate-models#classification) page.\n",
        "\n",
        "3. `builtin_eval`. The built-in evaluation loop of the trainer will be used to evaluate the model instead of the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) library. You can supply any eval dataset in the same format as the training dataset by specifying `--eval_dataset_path`, `--eval_split`, `--eval_template`, and `--eval_column`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7864cff27197"
      },
      "outputs": [],
      "source": [
        "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
        "\n",
        "eval_task = \"arc_challenge\"  # @param {type:\"string\"}\n",
        "eval_metric_name = \"acc_norm\"  # @param {type:\"string\"}\n",
        "\n",
        "# Runs 10 training steps as a minimal example. Use 1000 to reproduce the experiment results.\n",
        "max_steps = 10  # @param {type:\"integer\"}\n",
        "# Evaluates the model on 10 examples. Use 10000 to reproduce the experiment results.\n",
        "eval_limit = 10  # @param {type:\"integer\"}\n",
        "\n",
        "flags = {\n",
        "    \"learning_rate\": 1e-5,\n",
        "    \"precision_mode\": hpt_precision_mode,\n",
        "    \"task\": \"instruct-lora\",\n",
        "    \"pretrained_model_id\": model_id,\n",
        "    \"output_dir\": lora_output_dir_gcsfuse,\n",
        "    \"warmup_steps\": 10,\n",
        "    \"max_steps\": max_steps,\n",
        "    \"lora_rank\": 32,\n",
        "    \"lora_alpha\": 64,\n",
        "    \"lora_dropout\": 0.05,\n",
        "    \"dataset_name\": dataset_name,\n",
        "    \"eval_steps\": max_steps + 1,  # Only evaluates in the end.\n",
        "    \"eval_tasks\": eval_task,\n",
        "    \"eval_limit\": eval_limit,\n",
        "    \"eval_metric_name\": eval_metric_name,\n",
        "    \"merge_base_and_lora_output_dir\": merged_model_output_dir_gcsfuse,\n",
        "}\n",
        "\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",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_DOCKER_URI,\n",
        "            \"args\": [\"--{}={}\".format(k, v) for k, v in flags.items()],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "metric_spec = {\"model_performance\": \"maximize\"}\n",
        "parameter_spec = {\n",
        "    \"learning_rate\": hpt.DoubleParameterSpec(min=1e-5, max=1e-4, scale=\"linear\"),\n",
        "}\n",
        "train_job = aiplatform.CustomJob(\n",
        "    display_name=job_name,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
        "    display_name=f\"{job_name}_hpt\",\n",
        "    custom_job=train_job,\n",
        "    metric_spec=metric_spec,\n",
        "    parameter_spec=parameter_spec,\n",
        "    max_trial_count=8,\n",
        "    parallel_trial_count=2,\n",
        ")\n",
        "\n",
        "train_hpt_job.run()\n",
        "\n",
        "print(\"Trained models were saved in: \", lora_output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afad2431f37f"
      },
      "source": [
        "Then, find the best trial from the hyperparameter tuning job."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "138b09ee6313"
      },
      "outputs": [],
      "source": [
        "best_trial_id = max(\n",
        "    train_hpt_job.trials, key=lambda trial: trial.final_measurement.metrics[0].value\n",
        ").id\n",
        "lora_output_dir = os.path.join(lora_output_dir, f\"trial_{best_trial_id}\")\n",
        "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "print(f\"Best trial {best_trial_id} saved model in:\", lora_output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqmCtkGnhDmp"
      },
      "source": [
        "### Deploy with vLLM\n",
        "This section uploads the model to Model Registry and deploys it on the Endpoint. vLLM currently does not support serving finetuned [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b).\n",
        "\n",
        "The model deployment step will take ~15 minutes to complete.\n",
        "\n",
        "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~5.3G, ~8.7G and ~15.2G respectively with the default settings.\n",
        "\n",
        "NOTE: vLLM requires a merged model with the base model and the finetuned LoRA adapter. Based on your business need, if you need the base model and the finetuned LoRA weight to be served separately, please consider using the regular Vertex serving instead.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "# Finds Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "# Sets V100 to deploy open_llama_3b and open_llama_7b.\n",
        "# V100 serving has better throughput and latency performance than L4 serving.\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "accelerator_count = 1\n",
        "\n",
        "# Sets L4 to deploy open_llama_3b and open_llama_7b.\n",
        "# L4 serving is more cost efficient than V100 serving.\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "# Sets 2 V100 to deploy open_llama_13b.\n",
        "# V100 serving has better throughput and latency performance than L4 serving.\n",
        "# machine_type = \"n1-standard-16\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "# Sets 2 L4 to deploy open_llama_13b.\n",
        "# L4 serving is more cost efficient than V100 serving.\n",
        "# machine_type = \"g2-standard-24\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
        "    # vLLM currently does not support OpenLLaMA 3B.\n",
        "    precision_loading_mode = \"float16\"\n",
        "    model_with_peft, endpoint_with_peft = deploy_model(\n",
        "        model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
        "        model_id=model_id,\n",
        "        finetuned_lora_model_path=lora_output_dir,  # This will avoid override finetuning models.\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        task=\"causal-language-modeling-lora\",\n",
        "        precision_loading_mode=precision_loading_mode,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )\n",
        "else:\n",
        "    model_with_peft, endpoint_with_peft = deploy_model_vllm(\n",
        "        model_name=create_name_with_datetime(prefix=\"openllama-peft-serve-vllm\"),\n",
        "        model_id=merged_model_output_dir,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )\n",
        "\n",
        "print(\"endpoint_name:\", endpoint_with_peft.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: After the deployment succeeds, the base model weights will be downloaded on the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts. Parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab04da3ec9a"
      },
      "outputs": [],
      "source": [
        "instance = {\n",
        "    \"prompt\": \"Hi, Google. How are you doing?\",\n",
        "    \"n\": 1,\n",
        "    \"max_tokens\": 32,\n",
        "    \"temperature\": 1.0,\n",
        "    \"top_p\": 1.0,\n",
        "    \"top_k\": 10,\n",
        "}\n",
        "response = endpoint_with_peft.predict(instances=[instance])\n",
        "print(response.predictions[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vhDf9dNNn4bP"
      },
      "source": [
        "### [Optional] Merge a previously trained LoRA adapter with the base model\n",
        "\n",
        "This section demonstrates how to merge a previously trained LoRA adapter with a base model, and save the merged model to a GCS bucket. Please be aware that the LoRA adapter should be trained on the same base model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YHdru1aRqRFF"
      },
      "outputs": [],
      "source": [
        "merge_job_name = create_name_with_datetime(prefix=\"openllama-peft-merge\")\n",
        "\n",
        "# The base model to be merged upon. It can be a huggingface model id, or a GCS\n",
        "# path where the base model was stored.\n",
        "base_model_dir = \"gs://\"  # @param {type:\"string\"}\n",
        "# The previously trained LoRA adapter. It needs to be stored in a GCS path.\n",
        "finetuned_lora_adapter_dir = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# The GCS path to save the merged model\n",
        "merged_model_output_dir = os.path.join(MODEL_BUCKET, merge_job_name)\n",
        "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Worker pool spec.\n",
        "# Merges open_llama_3b and open_llama_7b with 1 V100 (16G).\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "\n",
        "# Merges open_llama_3b and open_llama_7b with 1 L4 (24G).\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "\n",
        "# Merges open_llama_13b with 1 L4 (24G).\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "\n",
        "# Merges open_llama_13b with 1 A100 (40G).\n",
        "# machine_type = \"a2-highgpu-1g\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "\n",
        "worker_pool_specs = [\n",
        "    {\n",
        "        \"machine_spec\": {\n",
        "            \"machine_type\": machine_type,\n",
        "            \"accelerator_type\": accelerator_type,\n",
        "            \"accelerator_count\": 1,\n",
        "        },\n",
        "        \"replica_count\": 1,\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_DOCKER_URI,\n",
        "            \"command\": [],\n",
        "            \"args\": [\n",
        "                \"--task=merge-causal-language-model-lora\",\n",
        "                \"--merge_model_precision_mode=float16\",\n",
        "                \"--pretrained_model_id=%s\" % base_model_dir,\n",
        "                \"--finetuned_lora_model_dir=%s\" % finetuned_lora_adapter_dir,\n",
        "                \"--merge_base_and_lora_output_dir=%s\" % merged_model_output_dir_gcsfuse,\n",
        "            ],\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "merge_custom_job = aiplatform.CustomJob(\n",
        "    display_name=merge_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "\n",
        "merge_custom_job.run()\n",
        "\n",
        "print(\"The merged model is stored at: \", merged_model_output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wi6UZnldXpdi"
      },
      "source": [
        "## Quantize and deploy OpenLLaMA 2 models\n",
        "\n",
        "This section demonstrates post-training quantization of OpenLLaMA models with Vertex Custom Job. Quantization reduces the memory required by a model while attempting to retain the same performance. Two such algorithms to do so are AWQ and GPTQ. Read more about AWQ in the following publication: [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://arxiv.org/abs/2306.00978). Read more about GPTQ in the following publication: [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers\n",
        "](https://arxiv.org/abs/2210.17323)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HTUo8dWSo8t9"
      },
      "source": [
        "### Quantize OpenLLaMA models\n",
        "\n",
        "Quantization reduces the amount of GPU required to serve a model by reducing the bit precision of the weights while minimizing drop in performance. Serving quantized models on VLLM requires models to be quantized to 4 bits. It is recommended to first search if a model has already been quantized and made publicly available: [AWQ](https://huggingface.co/TheBloke?search_models=-awq) and [GPTQ](https://huggingface.co/TheBloke?search_models=-gptq).\n",
        "\n",
        "Quantizing models with AWQ with 1 NVIDIA_L4 GPU will take around\n",
        "20 minutes for OpenLLaMA 3B, 30 minutes for OpenLLaMA 7B, and 1 hour for OpenLLaMA 13B.\n",
        "\n",
        "Quantizing models with GPTQ with 1 NVIDIA_L4 GPU will take around 30 minutes for OpenLLaMA 3B, 45 minutes for OpenLLaMA 7B, and 1.5 hour for OpenLLaMA 13B. Finetuned models can also be quantized, so long as the LoRA weights are merged with the base model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_YB31uAEpusN"
      },
      "outputs": [],
      "source": [
        "# Setup quantization job.\n",
        "\n",
        "# Set `finetuned_model_path` to `merged_model_output_dir` from the previous\n",
        "# section above to quantize the finetuned model, if not set the base model will\n",
        "# be quantized.\n",
        "finetuned_model_path = \"\"  # @param {type:\"string\"}\n",
        "if finetuned_model_path:\n",
        "    prequantized_model_path = finetuned_model_path\n",
        "else:\n",
        "    prequantized_model_path = model_id\n",
        "\n",
        "quantization_method = \"awq\"  # @param [\"awq\", \"gptq\"]\n",
        "quantization_job_name = get_job_name_with_datetime(\n",
        "    f\"openllama-{quantization_method}-quantize\"\n",
        ")\n",
        "\n",
        "quantization_output_dir = os.path.join(MODEL_BUCKET, quantization_job_name)\n",
        "quantization_output_dir_gcsfuse = quantization_output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Worker pool spec.\n",
        "\n",
        "# Sets 1 L4 (24G) to quantize OpenLLaMA model.\n",
        "machine_type = \"g2-standard-16\"\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
        "\n",
        "\n",
        "# Quantization parameters.\n",
        "quantization_precision_mode = \"4bit\"\n",
        "if quantization_method == \"awq\":\n",
        "    awq_dataset_name = \"pileval\"\n",
        "    group_size = 64\n",
        "    quantization_args = [\n",
        "        \"--task=quantize-model\",\n",
        "        f\"--quantization_method={quantization_method}\",\n",
        "        f\"--pretrained_model_id={model_id}\",\n",
        "        f\"--quantization_precision_mode={quantization_precision_mode}\",\n",
        "        f\"--quantization_output_dir={quantization_output_dir_gcsfuse}\",\n",
        "        f\"--quantization_dataset_name={awq_dataset_name}\",\n",
        "        f\"--group_size={group_size}\",\n",
        "    ]\n",
        "else:\n",
        "    # The original datasets used in GPTQ paper [\"wikitext2\",\"c4\",\"c4-new\",\"ptb\",\"ptb-new\"].\n",
        "    gptq_dataset_name = \"c4\"  # @param {type:\"string\"}\n",
        "    gptq_precision_mode = \"4bit\"\n",
        "    group_size = -1\n",
        "    damp_percent = 0.1\n",
        "    desc_act = True\n",
        "    quantization_args = [\n",
        "        \"--task=quantize-model\",\n",
        "        f\"--quantization_method={quantization_method}\",\n",
        "        f\"--pretrained_model_id={model_id}\",\n",
        "        f\"--quantization_precision_mode={quantization_precision_mode}\",\n",
        "        f\"--quantization_output_dir={quantization_output_dir_gcsfuse}\",\n",
        "        f\"--quantization_dataset_name={gptq_dataset_name}\",\n",
        "        f\"--group_size={group_size}\",\n",
        "        f\"--damp_percent={damp_percent}\",\n",
        "        f\"--desc_act={desc_act}\",\n",
        "    ]\n",
        "\n",
        "# Pass quantization arguments and launch 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\": 1,\n",
        "        \"disk_spec\": {\n",
        "            \"boot_disk_type\": \"pd-ssd\",\n",
        "            \"boot_disk_size_gb\": 500,\n",
        "        },\n",
        "        \"container_spec\": {\n",
        "            \"image_uri\": TRAIN_DOCKER_URI,\n",
        "            \"env\": [\n",
        "                {\n",
        "                    \"name\": \"PYTORCH_CUDA_ALLOC_CONF\",\n",
        "                    \"value\": \"max_split_size_mb:32\",\n",
        "                },\n",
        "            ],\n",
        "            \"command\": [],\n",
        "            \"args\": quantization_args,\n",
        "        },\n",
        "    }\n",
        "]\n",
        "\n",
        "print(f\"Quantizing {prequantized_model_path}.\")\n",
        "quantize_job = aiplatform.CustomJob(\n",
        "    display_name=quantization_job_name,\n",
        "    project=PROJECT_ID,\n",
        "    worker_pool_specs=worker_pool_specs,\n",
        "    staging_bucket=STAGING_BUCKET,\n",
        ")\n",
        "quantize_job.run()\n",
        "\n",
        "print(\"Quantized models were saved in: \", quantization_output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ocm3m2Hzr4ln"
      },
      "source": [
        "### Deploy quantized models with Google Cloud Text Moderation\n",
        "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
        "\n",
        "The model deployment step will take 15 minutes to 1 hour to complete, depending on the model sizes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DdINEB5Ur9FX"
      },
      "outputs": [],
      "source": [
        "# Finds Vertex AI prediction supported accelerators and regions in\n",
        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "# Sets 1 L4 (24G) to deploy OpenLLaMA models.\n",
        "machine_type = \"g2-standard-8\"\n",
        "accelerator_type = \"NVIDIA_L4\"\n",
        "accelerator_count = 1\n",
        "\n",
        "\n",
        "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
        "    # vLLM currently does not support OpenLLaMA 3B.\n",
        "    precision_loading_mode = \"float16\"\n",
        "    model_quantized_vllm, endpoint_quantized_vllm = deploy_model(\n",
        "        model_name=get_job_name_with_datetime(prefix=\"openllama-quantized-serve\"),\n",
        "        model_id=quantization_output_dir,\n",
        "        finetuned_lora_model_path=\"\",  # This will avoid override finetuning models.\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        task=\"causal-language-modeling-lora\",\n",
        "        precision_loading_mode=precision_loading_mode,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )\n",
        "else:\n",
        "    model_quantized_vllm, endpoint_quantized_vllm = deploy_model_vllm(\n",
        "        model_name=create_name_with_datetime(prefix=\"openllama-quantized-serve-vllm\"),\n",
        "        model_id=quantization_output_dir,\n",
        "        service_account=SERVICE_ACCOUNT,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y3Iy3q7CsEGe"
      },
      "source": [
        "NOTE: After the deployment succeeds, the model weights will be downloaded on the fly. Thus additional 10 ~ 40 minutes (depending on the model sizes) 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.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
        "\n",
        "Example:\n",
        "\n",
        "```\n",
        "Human: What is a car?\n",
        "Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eGP9sCq9sF0V"
      },
      "outputs": [],
      "source": [
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint_quantized_vllm.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint_quantized_vllm` 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 = endpoint_quantized_vllm.name\n",
        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint_quantized_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "\n",
        "# Overides max_length and top_k parameters during inferences.\n",
        "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
        "# you can reduce the max length, such as set max_length as 20.\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": \"What is a car?\",\n",
        "        \"max_tokens\": 50,\n",
        "        \"temperature\": 1.0,\n",
        "        \"top_p\": 1.0,\n",
        "        \"top_k\": 10,\n",
        "    },\n",
        "]\n",
        "response = endpoint_quantized_vllm.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JmuUk3l1DoEo"
      },
      "source": [
        "## Evaluate PEFT-finetuned OpenLLaMA\n",
        "\n",
        "This section demonstrates how to evaluate the OpenLLaMA model fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob.\n",
        "\n",
        "This example uses the dataset [TruthfulQA](https://arxiv.org/abs/2109.07958). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gM4SXaquDoEo"
      },
      "outputs": [],
      "source": [
        "eval_dataset = \"truthfulqa_mc\"  # @param {type:\"string\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "# Sets L4 to evaluate open_llama_3b and open_llama_7b.\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "accelerator_count = 1\n",
        "\n",
        "# Sets L4 to evaluate open_llama_3b and open_llama_7b.\n",
        "# machine_type = \"g2-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 1\n",
        "\n",
        "# Sets 2 V100 to evaluate open_llama_13b.\n",
        "# machine_type = \"n1-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "# Sets 2 L4 to evaluate open_llama_13b.\n",
        "# machine_type = \"g2-standard-24\"\n",
        "# accelerator_type = \"NVIDIA_L4\"\n",
        "# accelerator_count = 2\n",
        "\n",
        "replica_count = 1\n",
        "\n",
        "# Setup evaluation job.\n",
        "job_name = create_name_with_datetime(prefix=\"openllama-peft-eval\")\n",
        "eval_output_dir = os.path.join(MODEL_BUCKET, 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 `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={merged_model_output_dir_gcsfuse},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": "markdown",
      "metadata": {
        "id": "ItWB0WS__CX-"
      },
      "source": [
        "### Submit evaluation CustomJob"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbyIk99bDoEo"
      },
      "outputs": [],
      "source": [
        "# Pass evaluation arguments and launch 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": "markdown",
      "metadata": {
        "id": "kN0lE2iu_NXN"
      },
      "source": [
        "### Fetch and print evaluation results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "927oRxoADoEp"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "from google.cloud import storage\n",
        "\n",
        "# Fetch evaluation results.\n",
        "storage_client = storage.Client()\n",
        "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\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": "af21a3cff1e0"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Delete custom train and evaluation jobs.\n",
        "train_job.delete()\n",
        "eval_job.delete()\n",
        "quantize_job.delete()\n",
        "\n",
        "# Undeploy models and delete endpoints.\n",
        "endpoint_without_peft.delete(force=True)\n",
        "endpoint_with_peft.delete(force=True)\n",
        "endpoint_quantized_vllm.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model_without_peft.delete()\n",
        "model_with_peft.delete()\n",
        "model_quantized_vllm.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_openllama_peft.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}