notebooks/community/model_garden/model_garden_pytorch_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 - 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_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_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_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 CPU notebook is recommended)\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook demonstrates finetuning models with performance efficient finetuning libaries ([PEFT](https://github.com/huggingface/peft)), and running inferences with containers in Vertex AI. There are various models supported in [PEFT](https://github.com/huggingface/peft). This notebook shows examples with some models, such as OpenLLaMA, Falcon-instruct, BERT, RoBERTa-large, and XLM-RoBERTa-large, etc.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Finetune causal language models with PEFT and run inferences with Vertex AI, supporting\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",
        "- Finetune instruct models with PEFT and run inferences with Vertex AI, supporting\n",
        "\n",
        "| Models | LoRA |\n",
        "| :- | :- |\n",
        "| [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct) | Y |\n",
        "| [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) | Y |\n",
        "\n",
        "\n",
        "- Finetune sequence classification models with PEFT and run inferences with Vertex AI, supporting\n",
        "\n",
        "| Models | LoRA |\n",
        "| :- | :- |\n",
        "| [bert-base-uncase](https://huggingface.co/bert-base-uncased) | Y |\n",
        "| [RoBERTa-large](https://huggingface.co/roberta-large) | Y |\n",
        "| [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) | 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."
      ]
    },
    {
      "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",
        "    # Install gdown for downloading example training images.\n",
        "    ! pip3 install gdown\n",
        "\n",
        "    # Restart the notebook kernel after installs.\n",
        "    import IPython\n",
        "\n",
        "    app = IPython.Application.instance()\n",
        "    app.kernel.do_shutdown(True)"
      ]
    },
    {
      "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",
        "BUCKET_URI = \"\"  # @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:20231020_0936_RC00\"\n",
        "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:20231108_1540_RC00\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0c250872074f"
      },
      "source": [
        "### Define common functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "354da31189dc"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from datetime import datetime\n",
        "\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str):\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "def deploy_model(\n",
        "    model_name,\n",
        "    model_id,\n",
        "    finetuned_lora_model_path,\n",
        "    service_account,\n",
        "    task,\n",
        "    machine_type=\"n1-standard-8\",\n",
        "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
        "):\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",
        "        \"FINETUNED_LORA_MODEL_PATH\": finetuned_lora_model_path,\n",
        "        \"TASK\": task,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\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=1,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e70e3519ff8b"
      },
      "source": [
        "## Casual Language Modeling + PEFT\n",
        "\n",
        "This section demonstrates how to finetune OpenLLaMA with PEFT LoRA, including [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)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWGwJHqI7LMs"
      },
      "source": [
        "### Finetune with LoRA"
      ]
    },
    {
      "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_3b\"  # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
      ]
    },
    {
      "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.\n",
        "\n",
        "In order to make the finetunig efficiently, we enabled quantization (8bits) when loading pretrained models for finetuning LoRA models. 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. In theory, open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bd995e7aa529"
      },
      "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. The `template` parameter is optional.\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",
        "# machine_type = \"n1-standard-8\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "machine_type = \"a2-highgpu-1g\"\n",
        "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "replica_count = 1\n",
        "accelerator_count = 1\n",
        "\n",
        "# Setup training job.\n",
        "job_name = get_job_name_with_datetime(\"openllama-lora-train\")\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
        "output_dir_gcsfuse = 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={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(\"Trained models were saved in: \", output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jqmCtkGnhDmp"
      },
      "source": [
        "### Run inferences with serving images\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 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 separately with the default settings. We use V100 in deployments for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bf55e38815dc"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
        "    model_id=model_id,\n",
        "    finetuned_lora_model_path=output_dir,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    task=\"causal-language-modeling-lora\",\n",
        ")\n",
        "print(\"endpoint_name:\", endpoint.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "80b3fd2ace09"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4ab04da3ec9a"
      },
      "outputs": [],
      "source": [
        "# # Loads an existing endpoint as below.\n",
        "# endpoint_name = endpoint.name\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": \"Hi Google.\",\n",
        "        \"max_tokens\": 50,\n",
        "        \"temperature\": 1.0,\n",
        "        \"top_p\": 1.0,\n",
        "        \"top_k\": 10,\n",
        "    },\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "af21a3cff1e0"
      },
      "source": [
        "### Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# Delete custom train jobs.\n",
        "train_job.delete()\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W2USaHtQbE-l"
      },
      "source": [
        "## Instruct + PEFT\n",
        "\n",
        "This section demonstrates how to finetune Falcon-instruct with PEFT LoRA, including [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PwluMqvYbLu9"
      },
      "source": [
        "### Finetune with LoRA"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ap1pQd-ckAAb"
      },
      "source": [
        "Set the base model id."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vVTfG2bpbR90"
      },
      "outputs": [],
      "source": [
        "model_id = \"tiiuae/falcon-7b-instruct\"  # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJe_us8mkZEV"
      },
      "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 [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco). 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",
        "The peak GPU memory usages are ~11G and ~34G for finetuning LoRA models for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) separately with default training parameters and the example dataset. In theory, falcon-7b-instruct can be finetuned on 1 P100/V100, and falcon-40b-instruct can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4c128f63be2a"
      },
      "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": "-i4M7mWPbV8s"
      },
      "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",
        "# Worker pool spec.\n",
        "machine_type = \"a2-highgpu-1g\"\n",
        "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "replica_count = 1\n",
        "accelerator_count = 1\n",
        "\n",
        "# Setup training job.\n",
        "job_name = get_job_name_with_datetime(\"falcon-instruct-lora-train\")\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "max_steps = 10\n",
        "train_job.run(\n",
        "    args=[\n",
        "        \"--task=instruct-lora\",\n",
        "        f\"--pretrained_model_id={model_id}\",\n",
        "        f\"--dataset_name={dataset_name}\",\n",
        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
        "        \"--lora_rank=64\",\n",
        "        \"--lora_alpha=16\",\n",
        "        \"--lora_dropout=0.1\",\n",
        "        \"--warmup_ratio=0.03\",\n",
        "        f\"--max_steps={max_steps}\",\n",
        "        \"--max_seq_length=512\",\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(\"Trained models were saved in: \", output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oCgivsdllZ1x"
      },
      "source": [
        "### Run inferences with serving images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lNIFvielcyV"
      },
      "source": [
        "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 complete.\n",
        "\n",
        "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) with LoRA weights are ~15.5G and ~38.2G separately with the default settings. We use V100 in deployments for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "inrBXzoxk53t"
      },
      "outputs": [],
      "source": [
        "# # If deploy finetuned falcon-40b-instruct models, please set\n",
        "# machine_type = \"a2-highgpu-1g\",\n",
        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "\n",
        "model, endpoint = deploy_model(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"falcon-peft-serve\"),\n",
        "    model_id=model_id,\n",
        "    finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    task=\"instruct-lora\",\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        ")\n",
        "print(\"endpoint_name: \", endpoint.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_dmO3XooliAs"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint.\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": "JrKsrffGl14T"
      },
      "outputs": [],
      "source": [
        "# # Loads an existing endpoint as below.\n",
        "# endpoint_name = endpoint.name\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\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.predict(instances=instances)\n",
        "\n",
        "for prediction in response.predictions:\n",
        "    print(prediction)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uk-t69nFl9rr"
      },
      "source": [
        "### Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CcIX2WAJmAI7"
      },
      "outputs": [],
      "source": [
        "# Delete custom train jobs.\n",
        "train_job.delete()\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "og_s64QVmJDb"
      },
      "source": [
        "## Sequence Classification + PEFT\n",
        "\n",
        "This section demonstrates how to finetune sequence classification models with PEFT LoRA, including [bert-base-uncased](https://huggingface.co/bert-base-uncased), [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QqJcvCb2mn8l"
      },
      "source": [
        "### Finetune with PEFT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c0E7dfqloeBB"
      },
      "source": [
        "Set the base model id."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eru79ot6mtNG"
      },
      "outputs": [],
      "source": [
        "model_id = \"xlm-roberta-large\"  # @param [\"bert-base-uncased\", \"roberta-large\", \"xlm-roberta-large\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y9n-qheaotS9"
      },
      "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 [glue](https://huggingface.co/datasets/glue).\n",
        "\n",
        "The peak GPU memory usages are ~8.8G and ~15G for finetuning LoRA models for [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) separately with default training parameters and the example dataset. In theory, RoBERTa-large can be finetuned on 1 P100/V100, and XLM-RoBERTa-large can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OHeSYEKinfP3"
      },
      "outputs": [],
      "source": [
        "dataset_name = \"glue\"  # @param {type:\"string\"}\n",
        "\n",
        "# Worker pool spec.\n",
        "# # Please switch to A100 or other powerful machines if you run out of memory.\n",
        "# machine_type = \"a2-highgpu-1g\"\n",
        "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
        "machine_type = \"n1-standard-8\"\n",
        "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
        "replica_count = 1\n",
        "accelerator_count = 1\n",
        "\n",
        "# Setup training job.\n",
        "job_name = get_job_name_with_datetime(os.path.basename(model_id) + \"-lora-train\")\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        ")\n",
        "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
        "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job.run(\n",
        "    args=[\n",
        "        \"--task=sequence-classification-lora\",\n",
        "        f\"--pretrained_model_id={model_id}\",\n",
        "        f\"--dataset_name={dataset_name}\",\n",
        "        f\"--output_dir={output_dir_gcsfuse}\",\n",
        "        \"--lora_rank=8\",\n",
        "        \"--lora_alpha=16\",\n",
        "        \"--lora_dropout=0.1\",\n",
        "        \"--num_epochs=20\",\n",
        "        \"--batch_size=32\",\n",
        "        \"--learning_rate=3e-4\",\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(\"Trained models were saved in: \", output_dir)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kJuHaBCqoacg"
      },
      "source": [
        "### Run inferences with serving images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jw1GFXydoft0"
      },
      "source": [
        "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 complete.\n",
        "\n",
        "The peak GPU memory usages for [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) with LoRA weights are ~2G and ~2.4G separately with the default settings. We use V100 in deployments for simplicity."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A5j0a3AZoj3h"
      },
      "outputs": [],
      "source": [
        "model, endpoint = deploy_model(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"sequence-classification-peft-serve\"),\n",
        "    model_id=model_id,\n",
        "    finetuned_lora_model_path=output_dir,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    task=\"sequence-classification-lora\",\n",
        ")\n",
        "print(\"endpoint_name:\", endpoint.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6dR3-Hc7oqu-"
      },
      "source": [
        "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint.\n",
        "\n",
        "Once deployment succeeds, you can send requests to the endpoint with text prompts."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xAmULPegoyXl"
      },
      "outputs": [],
      "source": [
        "# # Loads an existing endpoint as below.\n",
        "# endpoint_name = endpoint.name\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
        "instances = [\n",
        "    {\"prompt\": \"The cat sat on the mat.\"},\n",
        "]\n",
        "response = endpoint.predict(instances=instances)\n",
        "labels = [int(item) for item in response.predictions]\n",
        "print(labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Vg_b46Jo6C1"
      },
      "source": [
        "### Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vNiTDL_VpAlp"
      },
      "outputs": [],
      "source": [
        "# Delete custom train jobs.\n",
        "train_job.delete()\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "model.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_peft.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}