diff --git a/notebooks/community/CODEOWNERS b/notebooks/community/CODEOWNERS
index e21ec41d6..5fc8312c9 100644
--- a/notebooks/community/CODEOWNERS
+++ b/notebooks/community/CODEOWNERS
@@ -107,7 +107,6 @@
/notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb @KCFindstr
/notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb @KCFindstr
/notebooks/community/model_garden/model_garden_pytorch_open_clip.ipynb @lydhr
-/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb @genquan9
/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb @xiangxu-google
/notebooks/community/model_garden/model_garden_pytorch_nllb.ipynb @weigary
/notebooks/community/model_garden/model_garden_pytorch_vicuna.ipynb @dstnluong-google
diff --git a/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb b/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb
deleted file mode 100644
index 911042667..000000000
--- a/notebooks/community/model_garden/model_garden_pytorch_llama2_peft.ipynb
+++ /dev/null
@@ -1,2748 +0,0 @@
-{
- "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 - LLaMA2 (PEFT)\n",
- "\n",
- "
\n",
- " \n",
- " \n",
- " Run in Colab\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " \n",
- " View on GitHub\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " \n",
- "Open in Vertex AI Workbench\n",
- " (A Python-3 CPU notebook is recommended)\n",
- " | \n",
- "
"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "3de7470326a2"
- },
- "source": [
- "## Overview\n",
- "\n",
- "This notebook demonstrates downloading and deploying prebuilt [LLaMA2 models](https://huggingface.co/meta-llama). It demonstrates deployment of pre-built LLaMA2 models with vLLM. It demonstrates finetuning and deploying LLaMA2 models with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)) in Vertex AI, quantizating LLaMA2 models using either AWQ or GPTQ to reduce the GPU memory requirements, and the evaluation of popular benchmark datasets through Vertex CustomJobs using [EleutherAI’s evaluation harness](https://github.com/EleutherAI/lm-evaluation-harness). This notebook uses [Text moderation APIs](https://cloud.google.com/natural-language/docs/moderating-text) to analyze predictions against a list of safety attributes.\n",
- "\n",
- "### Objective\n",
- "\n",
- "- Download and deploy prebuilt LLaMA2 models\n",
- "- Deploy LLaMA2 with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
- "- Finetune and deploy LLaMA2 models with PEFT\n",
- "- Quantize and deploy LLaMA2 models with AWQ or GPTQ\n",
- "- Evaluate the LLaMA2 models on any of the benchmark datasets\n",
- "\n",
- "### Costs\n",
- "\n",
- "This tutorial uses billable components of Google Cloud:\n",
- "\n",
- "* Vertex AI\n",
- "* Cloud Storage\n",
- "* Cloud NL APIs\n",
- "\n",
- "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), [Cloud NL API pricing](https://cloud.google.com/natural-language/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "264c07757582"
- },
- "source": [
- "## Before you begin"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "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",
- " ! pip3 install ipython pandas[output_formatting] google-cloud-language==2.10.0\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)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "038512c9338a"
- },
- "source": [
- "### Workbench only\n",
- "If you are using Workbench, you should find that the neccesary dependencies are already pre-installed. If this is not the case or if you have previously modified the existing libraries, you may install the dependencies using the following commands:\n",
- "```\n",
- "! pip3 install --upgrade google-cloud-aiplatform\n",
- "! pip3 install ipython pandas[output_formatting] google-cloud-language==2.10.0\n",
- "```"
- ]
- },
- {
- "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, Compute Engine API and Cloud Natural Language API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component,language.googleapis.com).\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": "35Dvbzb0hH3-"
- },
- "source": [
- "### Import the necessary packages"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "nLsuvskfhOv4"
- },
- "outputs": [],
- "source": [
- "import json\n",
- "import os\n",
- "import sys\n",
- "from datetime import datetime\n",
- "from typing import Tuple\n",
- "\n",
- "from google.cloud import aiplatform, language, storage"
- ]
- },
- {
- "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",
- "# Region for launching jobs.\n",
- "REGION = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# 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",
- "! gcloud services enable language.googleapis.com\n",
- "\n",
- "\n",
- "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
- "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
- "BASE_MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"base_model\")\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\"}\n",
- "\n",
- "if \"google.colab\" in sys.modules:\n",
- " from google.colab import auth\n",
- "\n",
- " auth.authenticate_user(project_id=PROJECT_ID)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e828eb320337"
- },
- "source": [
- "### Initialize Vertex AI API"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "12cd25839741"
- },
- "outputs": [],
- "source": [
- "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, serving and evaluation docker images.\n",
- "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231211_0936_RC00\"\n",
- "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:20231129_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": [
- "def get_job_name_with_datetime(prefix: str) -> str:\n",
- " \"\"\"Gets the job name with date time when triggering training or deployment\n",
- " jobs in Vertex AI.\n",
- " \"\"\"\n",
- " return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\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",
- " 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",
- " )\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 moderate_text(text: str) -> language.ModerateTextResponse:\n",
- " \"\"\"Calls Vertex AI APIs to analyze text moderations.\"\"\"\n",
- " client = language.LanguageServiceClient()\n",
- " document = language.Document(\n",
- " content=text,\n",
- " type_=language.Document.Type.PLAIN_TEXT,\n",
- " )\n",
- " return client.moderate_text(document=document)\n",
- "\n",
- "\n",
- "def show_text_moderation(text: str, response: language.ModerateTextResponse) -> None:\n",
- " \"\"\"Shows text moderation results.\"\"\"\n",
- " import pandas as pd\n",
- "\n",
- " def confidence(category: language.ClassificationCategory) -> float:\n",
- " return category.confidence\n",
- "\n",
- " columns = [\"category\", \"confidence\"]\n",
- " categories = sorted(response.moderation_categories, key=confidence, reverse=True)\n",
- " data = ((category.name, category.confidence) for category in categories)\n",
- " df = pd.DataFrame(columns=columns, data=data)\n",
- "\n",
- " print(f\"Text analyzed:\\n{text}\")\n",
- " print(df.to_markdown(index=False, tablefmt=\"presto\", floatfmt=\".0%\"))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ivs2RK093c8X"
- },
- "source": [
- "## Access LLaMA2 pretrained and finetuned models\n",
- "The original models from Meta are converted into the Hugging Face format for finetuning and serving in Vertex AI.\n",
- "\n",
- "Accept the model agreement to access the models:\n",
- "1. Navigate to the Vertex AI > Model Garden page in the Google Cloud console\n",
- "2. Find the LLaMA2 model card and click on \"VIEW DETAILS\"\n",
- "3. Review the agreement on the model card page\n",
- "4. After clicking the agreement of LLaMA2, a Cloud Storage bucket containing LLaMA2 pretrained and finetuned models will be shared\n",
- "5. Paste the Cloud Storage bucket link below and assign it to `VERTEX_AI_MODEL_GARDEN_LLAMA2`"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Jwn4PcTf4EMt"
- },
- "outputs": [],
- "source": [
- "VERTEX_AI_MODEL_GARDEN_LLAMA2 = \"\" # This will be shared once click the agreement of LLaMA2 in Vertex AI Model Garden.\n",
- "assert (\n",
- " VERTEX_AI_MODEL_GARDEN_LLAMA2\n",
- "), \"Please click the agreement of LLaMA2 in Vertex AI Model Garden, and get the GCS path of LLaMA2 model artifacts.\"\n",
- "print(\n",
- " \"Copy LLaMA2 model artifacts from\",\n",
- " VERTEX_AI_MODEL_GARDEN_LLAMA2,\n",
- " \"to \",\n",
- " BASE_MODEL_BUCKET,\n",
- ")\n",
- "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA2/* $BASE_MODEL_BUCKET"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "2MjaORIIFDVu"
- },
- "source": [
- "Set the base model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "E8OiHHNNE_wj"
- },
- "outputs": [],
- "source": [
- "base_model_name = \"llama2-7b-chat-hf\" # @param [\"llama2-7b-hf\", \"llama2-7b-chat-hf\", \"llama2-13b-hf\", \"llama2-13b-chat-hf\", \"llama2-70b-hf\", \"llama2-70b-chat-hf\"]\n",
- "model_id = os.path.join(BASE_MODEL_BUCKET, base_model_name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "8neJc8CnDDpu"
- },
- "source": [
- "## Deploy prebuilt LLaMA2 models with Google Cloud Text Moderation\n",
- "\n",
- "This section deploys prebuilt LLaMA2 models on Vertex AI. V100 GPUs are used for demonstration. Note that V100 serving generally offers better throughput and latency performance than L4 serving, while L4 serving is generally more cost efficient than V100 serving. The serving efficiency of V100 and L4 GPUs is inferior to that of A100 GPUs, but V100 and L4 GPUs are nevertheless good serving solutions if you do not have A100 quota."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "dHFW7yvjaVFV"
- },
- "source": [
- "We use the PEFT serving images to deploy prebuilt LLaMA2 models, by setting finetuning LoRA model paths as empty. 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": "Uak1pyEeExYM"
- },
- "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 (16G) to deploy LLaMA2 7B models.\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 1 L4 (24G) to deploy LLaMA2 7B models.\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",
- "# If A100 is not available, you may serve LLaMA2 13B models with multiple V100s\n",
- "# or L4s. Please keep in mind that the efficiency of serving with multiple\n",
- "# V100s or L4s is inferior to serving with 1 A100.\n",
- "# Sets 2 V100 (16G) to deploy LLaMA2 13B models.\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 = 2\n",
- "\n",
- "# Sets 2 L4 (24G) to deploy LLaMA2 13B models.\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",
- "# Sets A100 (40G) to deploy LLaMA2 13B models.\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 8 L4 (24G) to deploy LLaMA2 70B models.\n",
- "# If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA 2 70B\n",
- "# models with 8 L4 (24G) GPUs.\n",
- "# Note that with the default timeout threshold of Vertex endpoints, you should\n",
- "# set a `max_length` configuration of around 1,000 tokens or fewer. If you need\n",
- "# longer generated sequences, please file a request with Vertex to allowlist\n",
- "# your project for a longer timeout threshold with Vertex endpoints.\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "# Sets 4 A100 (40G) to deploy LLaMA2 70B models.\n",
- "# machine_type = \"a2-highgpu-4g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 4\n",
- "\n",
- "# The supported precision loading types are \"4bit\", \"8bit\", \"float16\" and \"float32\".\n",
- "# Quantization on prebuilt models directly might get poor results.\n",
- "# Suggests to use float16 for prebuild LLaMA2 model deployments.\n",
- "precision_loading_mode = \"float16\"\n",
- "\n",
- "model_without_peft, endpoint_without_peft = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"llama2-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",
- "print(\"endpoint_name:\", endpoint_without_peft.name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "sGKIjgmDFRW2"
- },
- "source": [
- "NOTE: The prebuilt model weights will be downloaded on the fly from $BASE_MODEL_BUCKET after the deployment succeeds. 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": "rDHsCOqvFYBi"
- },
- "outputs": [],
- "source": [
- "# Loads an existing endpoint instance using the endpoint name:\n",
- "# - Using `endpoint_name = endpoint_without_peft.name` allows us to get the\n",
- "# endpoint name of the endpoint `endpoint_without_peft` 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_without_peft.name\n",
- "# # endpoint_name = \"\" # @param {type:\"string\"}\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint_without_peft = 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",
- "# If you are using L4 GPUs to serve LLaMA2 70B models, you should set\n",
- "# max_length to around 1,000 tokens or fewer. If you need longer generated\n",
- "# sequences, please file a request with Vertex to allowlist your project for a\n",
- "# longer timeout threshold with Vertex endpoints.\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_without_peft.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "aUMsewPDj_pS"
- },
- "source": [
- "### Moderate model predictions\n",
- "\n",
- "Text moderation analyzes a document against a list of safety attributes, which include \"harmful categories\" and topics that may be considered sensitive."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "EE_GCSVVkBWj"
- },
- "outputs": [],
- "source": [
- "for generated_text in response.predictions:\n",
- " # Send a request to the API.\n",
- " response = moderate_text(generated_text)\n",
- " # Show the results.\n",
- " show_text_moderation(generated_text, response)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "f615c03d6638"
- },
- "source": [
- "### Build chat applications with LLaMA2\n",
- "\n",
- "Following the [prompt template recognized by LLaMA2 models](https://huggingface.co/blog/llama2#how-to-prompt-llama-2), you can build multi-turn chat applications while having control over the system prompt.\n",
- "\n",
- "The output of the Vertex endpoint is `Prompt:\\n{{ model_prompt }}\\nOutput:\\n{{ model_output }}`. You can format this output based on the prompt template to build chat applications in the following way:\n",
- "\n",
- "In the first turn of the conversation, the user of the chat application sends `user_message_1`. With `user_message_1`, the chat application can formulate `model_prompt` as the following:\n",
- "```\n",
- "[INST] <>\n",
- "{{ system_prompt }}\n",
- "<>\n",
- "\n",
- "{{ user_message }} [/INST]\n",
- "```\n",
- "The chat application then sends `model_prompt` to the Vertex endpoint and receives an output of the form: `Prompt:\\n{{ model_prompt }}\\nOutput:\\n{{ model_output }}`. The corresponding endpoint output contains both the original `model_prompt` as well as the generated `model_output`. The chat application can then exact `model_output` to present to the user of the chat application.\n",
- "\n",
- "In the second turn of the conversation, the user of the chat application sends `user_message_2` and the chat application can formulate `model_prompt` as:\n",
- "```\n",
- "[INST] <>\n",
- "{{ system_prompt }}\n",
- "<>\n",
- "\n",
- "{{ user_message_1 }} [/INST] {{ model_answer_1 }} [INST] {{ user_message_2 }} [/INST]\n",
- "```\n",
- "where `model_answer_1` is defined to be `model_output` from the previous turn. Again, the Vertex endpoint will generate an output of the form: `Prompt:\\n{{ model_prompt }}\\nOutput:\\n{{ model_output }}`, and the chat application can present `model_output` to the user.\n",
- "\n",
- "The same approach applies to further turns of the conversation, where all prior user messages and model answers are concatenated in the prompt.\n",
- "\n",
- "Below we show how [the HuggingFace example](https://huggingface.co/blog/llama2#how-to-prompt-llama-2) can be used with Vertex endpoints. Assume that in the first turn of the conversation, the chat application receives the user message: \"There's a llama in my garden 😱 What should I do?\""
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "749d516ea1db"
- },
- "outputs": [],
- "source": [
- "instances = [\n",
- " {\n",
- " \"prompt\": \"\"\"[INST] <>\n",
- "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n",
- "\n",
- "If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\n",
- "<>\n",
- "\n",
- "There's a llama in my garden 😱 What should I do? [/INST]\"\"\",\n",
- " \"max_tokens\": 50,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- " },\n",
- "]\n",
- "response = endpoint_without_peft.predict(instances=instances)\n",
- "\n",
- "endpoint_output = response.predictions[0]\n",
- "print(f\"[Endpoint output]\\n{endpoint_output}\\n\")\n",
- "\n",
- "model_output_start_index = endpoint_output.find(\"\\nOutput:\\n\") + len(\"\\nOutput:\\n\")\n",
- "model_output = endpoint_output[model_output_start_index:]\n",
- "print(f\"[Model output (model answer) to show the user]\\n{model_output}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "YKZ4CBJ2kYaW"
- },
- "source": [
- "## Deploy Prebuilt LLaMA2 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 an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "03d504bcd60b"
- },
- "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 V100 (16G) to deploy LLaMA2 7B models.\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 1 L4 (24G) to deploy LLaMA2 7B models.\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",
- "# If A100 is not available, you may serve LLaMA2 13B models with multiple V100s\n",
- "# or L4s. Please keep in mind that the efficiency of serving with multiple\n",
- "# V100s or L4s is inferior to serving with 1 A100.\n",
- "# Sets 2 V100 (16G) to deploy LLaMA2 13B models.\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 (24G) to deploy LLaMA2 13B models.\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",
- "# Sets A100 (40G) to deploy LLaMA2 13B models.\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 8 L4 (24G) to deploy LLaMA2 70B models.\n",
- "# If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA 2 70B\n",
- "# models with 8 L4 (24G) GPUs.\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "# Sets 4 A100 (40G) to deploy LLaMA2 70B models.\n",
- "# machine_type = \"a2-highgpu-4g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 4\n",
- "\n",
- "model_without_peft_vllm, endpoint_without_peft_vllm = deploy_model_vllm(\n",
- " model_name=get_job_name_with_datetime(prefix=\"llama2-serve-vllm\"),\n",
- " model_id=model_id,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "RRR11SWykYaX"
- },
- "source": [
- "NOTE: The prebuilt model weights will be downloaded on the fly from the orginal location 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. 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": "3f5a1e1de60d"
- },
- "outputs": [],
- "source": [
- "instance = {\n",
- " \"prompt\": \"Hi, Google.\",\n",
- " \"n\": 1,\n",
- " \"max_tokens\": 50,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- "}\n",
- "response = endpoint_without_peft_vllm.predict(instances=[instance])\n",
- "print(response.predictions[0])"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e70e3519ff8b"
- },
- "source": [
- "## Finetune and deploy LLaMA2 models with PEFT\n",
- "\n",
- "This section demonstrates how to finetune and deploy LLaMA2 models with PEFT LoRA. LoRA (Low-Rank Adaptation) is one approach of PEFT (Parameter Efficient FineTuning), where pretrained model weights are frozen and rank decomposition matrices representing the change in model weights are trained during finetuning. Read more about LoRA in the following publication: [Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. *arXiv preprint arXiv:2106.09685*](https://arxiv.org/abs/2106.09685)."
- ]
- },
- {
- "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."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bf084578f5ac"
- },
- "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",
- "finetuning_precision_mode = \"float16\"\n",
- "\n",
- "# Worker pool spec.\n",
- "# Finetunes LLaMA2 7B with 1 V100 (16G).\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Finetunes LLaMA2 7B with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# If A100 is not available, you may finetune LLaMA2 13B models with multiple\n",
- "# V100s or L4s. Please keep in mind that the efficiency of finetuning with\n",
- "# multiple V100s or L4s is inferior to finetuning with 1 A100.\n",
- "# Finetunes LLaMA2 13B with 2 V100 (16G).\n",
- "# machine_type = \"n1-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Finetunes LLaMA2 13B with 2 L4 (24G).\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Finetunes LLaMA2 13B with 1 A100 (40G).\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Finetunes LLaMA2 70B with 8 L4 (24G).\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "# Finetunes LLaMA2 70B with 4 A100 (40G).\n",
- "# machine_type = \"a2-highgpu-4g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 4\n",
- "\n",
- "replica_count = 1\n",
- "\n",
- "\n",
- "# Setup training job.\n",
- "job_name = get_job_name_with_datetime(\"llama2-lora-train\")\n",
- "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "1CIt3HwL8K2w"
- },
- "source": [
- "Run the following cell to launch the training job; Or skip to the next section if you want to try hyperparameter tuning with Vertex AI."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "8gzpPPsq71yO"
- },
- "outputs": [],
- "source": [
- "# Pass training arguments and launch job.\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- ")\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}\",\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\"--precision_mode={finetuning_precision_mode}\",\n",
- " f\"--template={template}\",\n",
- " ],\n",
- " environment_variables={\"WANDB_DISABLED\": True},\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": "RnsAIEX9jlnp"
- },
- "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",
- "| Llama2-7b | 2d 10hrs | 8 | 1 | L4x1 | +0.73 | +1.61 | +5.34 | \\$49.5088 |\n",
- "| Llama2-13b | 4d 8hrs | 8 | 1 | L4x1 | +1.53 | +2.11 | +10.98 | \\$88.7744 |\n",
- "| Llama2-70b | 6d 10hrs | 8 | 2 | L4x4 | +0.77 | +0.93 | +10.63 | \\$1,312.5461 |\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": "e1289e21a9d3"
- },
- "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 LLaMA2 7B with 1 L4 (24G).\n",
- "machine_type = \"g2-standard-8\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Finetunes LLaMA2 13B with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Finetunes LLaMA2 70B with 4 L4 (24G).\n",
- "# machine_type = \"g2-standard-48\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 4"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "80a5f9508a95"
- },
- "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=`. 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": "5RbokWtvk9kS"
- },
- "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\": output_dir,\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",
- "}\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: \", output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "4ZzzbZw58bY_"
- },
- "source": [
- "Then, find the best trial from the hyperparameter tuning job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "rlvrI-W38nFV"
- },
- "outputs": [],
- "source": [
- "best_trial_id = max(\n",
- " train_hpt_job.trials, key=lambda trial: trial.final_measurement.metrics[0].value\n",
- ").id\n",
- "output_dir = os.path.join(output_dir, f\"trial_{best_trial_id}\")\n",
- "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "print(f\"Best trial {best_trial_id} saved model in:\", output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "jqmCtkGnhDmp"
- },
- "source": [
- "### Deploy fine tuned 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": "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 (16G) to deploy LLaMA2 7B models.\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 1 L4 (24G) to deploy LLaMA2 7B models.\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",
- "# If A100 is not available, you may serve LLaMA2 13B models with multiple V100s\n",
- "# or L4s. Please keep in mind that the efficiency of serving with multiple\n",
- "# V100s or L4s is inferior to serving with 1 A100.\n",
- "# Sets 2 V100 (16G) to deploy LLaMA2 13B models.\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 = 2\n",
- "\n",
- "# Sets 2 L4 (24G) to deploy LLaMA2 13B models.\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",
- "# Sets A100 (40G) to deploy LLaMA2 13B models.\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 8 L4 (24G) to deploy LLaMA2 70B models.\n",
- "# If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA 2 70B\n",
- "# models with 8 L4 (24G) GPUs.\n",
- "# Note that with the default timeout threshold of Vertex endpoints, you should\n",
- "# set a `max_length` configuration of around 1,000 tokens or fewer. If you need\n",
- "# longer generated sequences, please file a request with Vertex to allowlist\n",
- "# your project for a longer timeout threshold with Vertex endpoints.\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "# Sets 4 A100 (40G) to deploy LLaMA2 70B models.\n",
- "# machine_type = \"a2-highgpu-4g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 4\n",
- "\n",
- "# The supported precision loading types are \"4bit\", \"8bit\", \"float16\" and \"float32\".\n",
- "precision_loading_mode = finetuning_precision_mode\n",
- "\n",
- "model_with_peft, endpoint_with_peft = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"llama2-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",
- " precision_loading_mode=precision_loading_mode,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\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 $BASE_MODEL_BUCKET and LoRA model weights will be downloaded from the GCS bucket used in training above. 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": "4ab04da3ec9a"
- },
- "outputs": [],
- "source": [
- "# Loads an existing endpoint instance using the endpoint name:\n",
- "# - Using `endpoint_name = endpoint_with_peft.name` allows us to get the\n",
- "# endpoint name of the endpoint `endpoint_with_peft` 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_with_peft.name\n",
- "# # endpoint_name = \"\" # @param {type:\"string\"}\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint_with_peft = 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",
- "# If you are using L4 GPUs to serve LLaMA2 70B models, you should set\n",
- "# max_length to around 1,000 tokens or fewer. If you need longer generated\n",
- "# sequences, please file a request with Vertex to allowlist your project for a\n",
- "# longer timeout threshold with Vertex endpoints.\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_with_peft.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "yfPDd91qlSlI"
- },
- "source": [
- "Text moderation analyzes a document against a list of safety attributes, which include \"harmful categories\" and topics that may be considered sensitive."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "lEUyncyklTEE"
- },
- "outputs": [],
- "source": [
- "for generated_text in response.predictions:\n",
- " # Send a request to the API.\n",
- " response = moderate_text(generated_text)\n",
- " # Show the results.\n",
- " show_text_moderation(generated_text, response)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "YXaszFtvOk8v"
- },
- "source": [
- "## Quantize and deploy LLaMA 2 models\n",
- "\n",
- "This section demonstrates post-training quantization of LLaMA2 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": "8fZKui4v1gm3"
- },
- "source": [
- "### Deploy pre-quantized models with Google Cloud Text Moderation\n",
- "Many AWQ-quantized models are provided by TheBloke [here](https://huggingface.co/TheBloke?search_models=-awq), and GPTQ-quantized models are provided [here](https://huggingface.co/TheBloke?search_models=-gptq).\n",
- "\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.\n",
- "\n",
- "Notice that deploying a quantized model requires much less GPU.\n",
- "We can deploy a quantized 13B model with only one L4 instead of four, and\n",
- "we can deploy a quantized 70B model with only two L4s instead of eight."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "4fb99b15ee8f"
- },
- "outputs": [],
- "source": [
- "quantized_model_id = \"TheBloke/Llama-2-7B-chat-AWQ\" # @param [\"TheBloke/Llama-2-7B-chat-AWQ\", \"TheBloke/Llama-2-13B-chat-AWQ\", \"TheBloke/Llama-2-70B-chat-AWQ\", \"TheBloke/Llama-2-7B-chat-GPTQ\", \"TheBloke/Llama-2-13B-chat-GPTQ\", \"TheBloke/Llama-2-70B-chat-GPTQ\"]\n",
- "\n",
- "quantization_method = quantized_model_id.split(\"-\")[-1].lower()\n",
- "\n",
- "# 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 LLaMA2 7B and 13B models.\n",
- "machine_type = \"g2-standard-8\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets 2 L4's (24G) to deploy LLaMA2 70B models.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "model_prequantized_vllm, endpoint_prequantized_vllm = deploy_model_vllm(\n",
- " model_name=get_job_name_with_datetime(prefix=\"llama2-serve-vllm-prequantized\"),\n",
- " model_id=quantized_model_id,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " quantization_method=quantization_method,\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "qMqARZni4LXq"
- },
- "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": "2uLlv-uG4LXq"
- },
- "outputs": [],
- "source": [
- "# Loads an existing endpoint instance using the endpoint name:\n",
- "# - Using `endpoint_name = endpoint_prequantized_vllm.name` allows us to get the\n",
- "# endpoint name of the endpoint `endpoint_prequantized_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_prequantized_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_prequantized_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",
- "# If you are using L4 GPUs to serve LLaMA2 70B models, you should set\n",
- "# max_length to around 1,000 tokens or fewer. If you need longer generated\n",
- "# sequences, please file a request with Vertex to allowlist your project for a\n",
- "# longer timeout threshold with Vertex endpoints.\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_prequantized_vllm.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "VwcWPHCx4LXq"
- },
- "source": [
- "Text moderation analyzes a document against a list of safety attributes, which include \"harmful categories\" and topics that may be considered sensitive."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "uEL89BkW4LXq"
- },
- "outputs": [],
- "source": [
- "for generated_text in response.predictions:\n",
- " # Send a request to the API.\n",
- " response = moderate_text(generated_text)\n",
- " # Show the results.\n",
- " show_text_moderation(generated_text, response)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "3yHz2iiAOk8v"
- },
- "source": [
- "### Quantize LLaMA2 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). Quantizing models with AWQ will take around 0.5 hours for LLaMA2 7B, 1.5 hours for LLaMA2 13B, and 4.5 hours for LLaMA2 70B, using 1 NVIDIA_L4 GPU for 7B and 13B models and 8 NVIDIA_L4 GPUs for 70B model. Quantizing models with GPTQ will take around 1.5 hours for LLaMA2 7B, 3 hours for LLaMA2 2.5 hours for LLaMA2 13B, and 6 hours for LLaMA 70B models, using 1 NVIDIA_L4 GPU for 7B and 13B models and 8 NVIDIA_L4 GPUs for 70B model. 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": "YsJI-amIOk8v"
- },
- "outputs": [],
- "source": [
- "# Setup quantization job.\n",
- "\n",
- "# Set `finetuned_model_path` to the GCS path of the merged finetuned model from the section above, if not set, the base model will 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\"llama2-{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 7B and 13B models.\n",
- "machine_type = \"g2-standard-16\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets 8 L4 (24G) to quantize 70B models.\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "\n",
- "# Quantization parameters.\n",
- "quantization_precision_mode = \"4bit\"\n",
- "if quantization_method == \"awq\":\n",
- " awq_dataset_name = \"pileval\"\n",
- " group_size = 128\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",
- " \"--cache_examples_on_gpu=False\",\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": "bB6P-InJOk8w"
- },
- "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.\n",
- "\n",
- "Notice that deploying a quantized model requires much less GPU.\n",
- "We can deploy a quantized 13B model with only one L4 instead of four, and\n",
- "we can deploy a quantized 70B model with only two L4s instead of eight."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "7ExjQ5JAOk8w"
- },
- "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 LLaMA2 7B and 13B models.\n",
- "machine_type = \"g2-standard-8\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets 2 L4's (24G) to deploy LLaMA2 70B models.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "model_quantized_vllm, endpoint_quantized_vllm = deploy_model_vllm(\n",
- " model_name=get_job_name_with_datetime(prefix=\"llama2-serve-vllm-quantized\"),\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",
- " quantization_method=quantization_method,\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "KuywY2xnQxTK"
- },
- "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": "VwGJnNXqQzIW"
- },
- "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",
- "# If you are using L4 GPUs to serve LLaMA2 70B models, you should set\n",
- "# max_length to around 1,000 tokens or fewer. If you need longer generated\n",
- "# sequences, please file a request with Vertex to allowlist your project for a\n",
- "# longer timeout threshold with Vertex endpoints.\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": "1a0_lqNzQ25q"
- },
- "source": [
- "Text moderation analyzes a document against a list of safety attributes, which include \"harmful categories\" and topics that may be considered sensitive."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "rC4HXayZQ41C"
- },
- "outputs": [],
- "source": [
- "for generated_text in response.predictions:\n",
- " # Send a request to the API.\n",
- " response = moderate_text(generated_text)\n",
- " # Show the results.\n",
- " show_text_moderation(generated_text, response)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "baXgZzmUA9Uh"
- },
- "source": [
- "## Evaluate PEFT-finetuned LLaMA 2 models\n",
- "\n",
- "This section demonstrates evaluation of LLaMA2 models using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex Custom Job. Please reference the peak GPU memory usgaes for serving and adjust the machine type, accelerator type and accelerator count accordingly.\n",
- "\n",
- "This example uses the dataset [TruthfulQA](https://arxiv.org/abs/2109.07958).\n",
- "All the 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": "eD-3RZh7-kcn"
- },
- "outputs": [],
- "source": [
- "# Set the machine_type, accelerator_type, accelerator_count and benchmark dataset.\n",
- "eval_dataset = \"truthfulqa_mc\" # @param [\"truthfulqa_mc\", \"boolq\", \"gsm8k\", \"hellaswag\", \"natural_questions\", \"openai_humaneval\", \"openbookqa\", \"quac\", \"trivia_qa\", \"winograde\"]\n",
- "\n",
- "# Worker pool spec.\n",
- "# Find Vertex AI supported accelerators and regions in:\n",
- "# https://cloud.google.com/vertex-ai/docs/training/configure-compute\n",
- "\n",
- "# Sets V100 (16G) to evaluate LLaMA2 7B models.\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets 1 L4 (24G) to evaluate LLaMA2 7B models.\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 2 V100 (16G) to evaluate LLaMA2 13B models.\n",
- "# machine_type = \"n1-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Sets 2 L4 (24G) to evaluate LLaMA2 13B models.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Sets A100 (40G) to evaluate LLaMA2 13B models.\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 8 L4 (24G) to evaluate LLaMA2 70B models.\n",
- "# machine_type = \"g2-standard-96\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 8\n",
- "\n",
- "# Sets 4 A100 (40G) to evaluate LLaMA2 70B models.\n",
- "# machine_type = \"a2-highgpu-4g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 4\n",
- "\n",
- "# Sets A100 (80G) to evaluate LLaMA2 70B models for faster inferences.\n",
- "# machine_type = \"a2-ultragpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100_80GB\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "replica_count = 1"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "1UkYoTs0Inf9"
- },
- "outputs": [],
- "source": [
- "# Setup evaluation job.\n",
- "job_name = get_job_name_with_datetime(prefix=\"llama2-eval\")\n",
- "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "model_id_gcsfuse = model_id.replace(\"gs://\", \"/gcs/\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "g0t0RBixIw0P"
- },
- "outputs": [],
- "source": [
- "# Prepare evaluation command that runs the evaluation harness.\n",
- "# Set `trust_remote_code = True` because evaluating the model requires\n",
- "# executing code from the model repository.\n",
- "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
- "eval_command = [\n",
- " \"python\",\n",
- " \"main.py\",\n",
- " \"--model\",\n",
- " \"hf-causal-experimental\",\n",
- " \"--model_args\",\n",
- " f\"pretrained={model_id_gcsfuse},peft={output_dir_gcsfuse},trust_remote_code=True,use_accelerate=True,device_map_option=auto\",\n",
- " \"--tasks\",\n",
- " f\"{eval_dataset}\",\n",
- " \"--output_path\",\n",
- " f\"{eval_output_dir_gcsfuse}\",\n",
- "]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "_Gjzbb8pI5_V"
- },
- "outputs": [],
- "source": [
- "# Run the evaluation job.\n",
- "worker_pool_specs = [\n",
- " {\n",
- " \"machine_spec\": {\n",
- " \"machine_type\": machine_type,\n",
- " \"accelerator_type\": accelerator_type,\n",
- " \"accelerator_count\": accelerator_count,\n",
- " },\n",
- " \"replica_count\": replica_count,\n",
- " \"disk_spec\": {\n",
- " \"boot_disk_size_gb\": 500,\n",
- " },\n",
- " \"container_spec\": {\n",
- " \"image_uri\": EVAL_DOCKER_URI,\n",
- " \"command\": eval_command,\n",
- " \"args\": [],\n",
- " },\n",
- " }\n",
- "]\n",
- "\n",
- "eval_job = aiplatform.CustomJob(\n",
- " display_name=job_name,\n",
- " worker_pool_specs=worker_pool_specs,\n",
- " base_output_dir=eval_output_dir,\n",
- ")\n",
- "\n",
- "eval_job.run()\n",
- "\n",
- "print(\"Evaluation results were saved in:\", eval_output_dir)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "D6bVD7Yh-z1x"
- },
- "outputs": [],
- "source": [
- "# Fetch evaluation results.\n",
- "storage_client = storage.Client()\n",
- "BUCKET_NAME = BUCKET_URI.replace(\"gs://\", \"\")\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": [
- "if train_job._gca_resource.name:\n",
- " # Training job is submitted.\n",
- " train_job.delete()\n",
- "# Uncomment the following line to delete the hyperparameter tuning job if you\n",
- "# have created one.\n",
- "# train_hpt_job.delete()\n",
- "\n",
- "# Delete the evaluation job.\n",
- "eval_job.delete()\n",
- "\n",
- "# Delete the quantization job.\n",
- "quantize_job.delete()\n",
- "\n",
- "# Undeploy model and delete endpoint.\n",
- "endpoint_without_peft.delete(force=True)\n",
- "endpoint_with_peft.delete(force=True)\n",
- "endpoint_without_peft_vllm.delete(force=True)\n",
- "endpoint_prequantized_vllm.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_without_peft_vllm.delete()\n",
- "model_prequantized_vllm.delete()\n",
- "model_quantized_vllm.delete()\n",
- "\n",
- "# Delete Cloud Storage objects that were created\n",
- "delete_bucket = False\n",
- "if delete_bucket:\n",
- " ! gsutil -m rm -r $EXPERIMENT_BUCKET"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "fsv4jGuU89rX"
- },
- "source": [
- "# Vertex AI Model Garden - LLaMA2 (RLHF)\n",
- "\n",
- "\n",
- " \n",
- " \n",
- " Run in Colab\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " \n",
- " View on GitHub\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " Open in Vertex AI Workbench\n",
- " \n",
- " |
"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "overview:automl"
- },
- "source": [
- "## Overview\n",
- "\n",
- "This tutorial demonstrates how to use Vertex AI RLHF to tune a LLaMA2. This workflow improves a model's accuracy by fine-tuning a base model with a training dataset.\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "objective:pipelines,automl"
- },
- "source": [
- "### Objective\n",
- "\n",
- "In this tutorial, you will use `Vertex AI RLHF` to tune a large language model model, perform bulk inference and model deployment.\n",
- "\n",
- "\n",
- "This tutorial uses the following Google Cloud ML services:\n",
- "\n",
- "- `Vertex AI RLHF`\n",
- "- `Vertex AI Pipelines`\n",
- "\n",
- "\n",
- "The steps performed include:\n",
- "\n",
- "- Set the number of model tuning steps.\n",
- "- Create Vertex AI Pipeline job using a predefined template for tuning.\n",
- "- Execute the pipeline using `Vertex AI Pipelines`.\n",
- "- Perform online prediction with the tuned model."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "aef4f59195ad"
- },
- "source": [
- "### Prepare your inputs\n",
- "\n",
- "In the sample below, you will run RLHF training on LLaMA 2 70B chat model. A list of supported models, hardware and regions can be found as follows. Tuning jobs that run in us-central1 will use 8 Nvidia A100 80GB. Jobs that run in europe-west4 will use 64 v3 TPUs.\n",
- "\n",
- "**Table 1. Supported models, hardware and regions**\n",
- "\n",
- "| large_model_reference | supported accelerator_type | supported region |\n",
- "|---|---|---|\n",
- "| `llama-2-7b` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
- "| `llama-2-7b-chat` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
- "| `llama-2-13b` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
- "| `llama-2-13b-chat` | `TPU_V3`, `NVIDIA_TESLA_A100` | europe-west4, us-central1 |\n",
- "\n",
- "The workflow takes the following inputs:\n",
- "\n",
- "* **preference_dataset**: Required, `str`. Cloud storage path to a human preference dataset used to train a reward model. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
- "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
- "> * An `input_text` field containing the prompt.\n",
- "> * `candidate_0`, `candidate_1` fields containing the candidate responses.\n",
- "> * A `choice` field that specifies which candidate is preferred. This field is an integer `0` or `1`. `0` means `candidate_0` is preferred, `1` means `candidate_1` is preferred.\n",
- "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\", \"candidate_0\": \"Enjoy some fun in the sun at Gulf Shores.\", \"candidate_1\": \"A Tranquil Oasis of Natural Beauty.\", \"choice\": 0}`\n",
- "\n",
- "* **prompt_dataset**: Required, `str`. Cloud storage path to an unlabled prompt dataset used for reinforcement learning. Download the [example dataset](https://pantheon.corp.google.com/storage/browser/vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train?pageState=(%22StorageObjectListTable%22:(%22f%22:%22%255B%255D%22))&e=13802955&jsmode=O&mods=-ai_platform_fake_service&project=vertex-ai&prefix=&forceOnObjectsSortingFiltering=false), and see its content.\n",
- "> The dataset format is `jsonl`. Each line in the file contains one example:\n",
- "> * An `input_text` field containing the prompt.\n",
- "> * For example: `{\"input_text\": \"Create a description for Plantation Palms.\"}`\n",
- "\n",
- "* **eval_dataset**: Optional, `str`. Cloud storage path to the unlabeled dataset you want to do prediction on. If provided, inference will be performed on this dataset after training. Its format is exactly the same as **prompt_dataset**.\n",
- "\n",
- "* **large_model_reference**: Required, `str`. Name of the base model. In this example, we use `text-bison@001`. Valid values are listed in the table above.\n",
- "\n",
- "* **model_display_name**: Optional, `str`. Name of the fine-tuned model shown in the Model Registry. If not provided, a default name will be created.\n",
- "\n",
- "* **reward_model_train_steps**: Optional, `int`. Number of steps to use when training a reward model. Default value is 1000.\n",
- "\n",
- "* **reinforcement_learning_train_steps**: Optional, `int`. Number of reinforcement learning steps to perform when tuning a base model. Default value is 1000.\n",
- "\n",
- "* **prompt_sequence_length**: Optional, `int`. Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most `8192`. Default value is `512`.\n",
- "\n",
- "* **target_sequence_length**: Optional, `int`. Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most `1024`. Default value is `64`.\n",
- "\n",
- "* **reward_model_learning_rate_multiplier**: Optional, `float`. Constant used to adjust the base learning rate used when training a reward model. Multiply by a number > 1 to increase the magnitude of updates applied at each training step or multiply by a number < 1 to decrease the magnitude of updates. Default value is `1.0`.\n",
- "\n",
- "* **reinforcement_learning_rate_multiplier**: Optional, `float`. Constant used to adjust the base learning rate used during reinforcement learning. Multiply by a number > 1 to increase the magnitude of updates applied at each training step or multiply by a number < 1 to decrease the magnitude of updates. Default value is `1.0`.\n",
- "\n",
- "* **kl_coeff**: Optional, `float`. Coefficient for KL penalty. This regularizes the policy model and penalizes if it diverges from its initial distribution. If set to 0, the reference language model is not loaded into memory. Default value is `0.1`.\n",
- "\n",
- "* **lora_dim**: Optional, `int`. Default value is 4. If set to 0, full-model tuning is performed.\n",
- "\n",
- "* **instruction**: Optional, `str`. This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. `Classify this movie review as positive or negative` or `Translate this sentence to Danish`. See [here](https://ai.googleblog.com/2021/10/introducing-flan-more-generalizable.html) for more details on the instruction-tuned models. Do not specify this if your dataset already prepends the instruction to the inputs field.\n",
- "\n",
- "### A Note on choosing train_steps\n",
- "\n",
- "Choosing a suitable value for the number of train steps can help avoid overfitting.\n",
- "\n",
- "* **reward_model_train_steps**: This depends on the size of your \"comparison dataset\". Usually, the model should train over the comparison dataset for 20-30 times for best results.\n",
- "\n",
- "* **reinforcement_learning_train_steps**: This depends on the size of your prompt dataset. Usually, the model should train over the prompt dataset for roughly 10-20 times, but beware, if given too many training steps, the policy model may figure out a way exploit the reward and exhibit undesired behavior (i.e. \"reward hacking\").\n",
- "\n",
- "The calculator below can help you compute these numbers."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "YoCkx4iD4JEY"
- },
- "outputs": [],
- "source": [
- "# Please REPLACE with your own parameters for:\n",
- "# - N_COMPARISON_DATASET_EXAMPLES\n",
- "# - N_REWARD_MODEL_TRAINER_EPOCHS\n",
- "# - N_PROMPT_DATASET_EXAMPLES\n",
- "# - N_REINFORCER_EPOCHS\n",
- "# - BATCH_SIZE\n",
- "\n",
- "import math\n",
- "\n",
- "# This is your \"preference dataset\" size.\n",
- "N_PREFERENCE_DATASET_EXAMPLES = 3000\n",
- "# An epoch refers to one full cycle through the training dataset.\n",
- "N_REWARD_MODEL_TRAIN_EPOCHS = 30\n",
- "\n",
- "# This is your \"prompt dataset\" size.\n",
- "N_PROMPT_DATASET_EXAMPLES = 2000\n",
- "# An epoch refers to one full cycle through the training dataset.\n",
- "N_REINFORCEMENT_LEARNING_TRAIN_EPOCHS = 10\n",
- "\n",
- "BATCH_SIZE = 64 # We use fixed batch_size=64 under the hood.\n",
- "\n",
- "reward_model_train_steps = math.ceil(\n",
- " N_PREFERENCE_DATASET_EXAMPLES * N_REWARD_MODEL_TRAIN_EPOCHS / BATCH_SIZE\n",
- ") # 3000 * 30 / 64 = 1407\n",
- "\n",
- "reinforcement_learning_train_steps = math.ceil(\n",
- " N_PROMPT_DATASET_EXAMPLES * N_REINFORCEMENT_LEARNING_TRAIN_EPOCHS / BATCH_SIZE\n",
- ") # 2000 * 10 / 64 = 313\n",
- "\n",
- "print(f\"{reward_model_train_steps=}\")\n",
- "print(f\"{reinforcement_learning_train_steps=}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "costs"
- },
- "source": [
- "### Costs\n",
- "\n",
- "This tutorial uses billable components of Google Cloud:\n",
- "\n",
- "* Vertex AI\n",
- "* Cloud Storage\n",
- "\n",
- "Learn about [Vertex AI\n",
- "pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage\n",
- "pricing](https://cloud.google.com/storage/pricing), and use the [Pricing\n",
- "Calculator](https://cloud.google.com/products/calculator/)\n",
- "to generate a cost estimate based on your projected usage."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "install_aip:mbsdk"
- },
- "source": [
- "## Installation\n",
- "\n",
- "Install the following packages required to execute this notebook."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "F9dJ5Of-dORl"
- },
- "outputs": [],
- "source": [
- "! pip3 install --upgrade --quiet google-cloud-aiplatform google-cloud-pipeline-components kfp\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "restart"
- },
- "source": [
- "### Colab only: Uncomment the following cell to restart the kernel"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "HKDBjmNr9T6t"
- },
- "source": [
- "Automatically restart kernel after installs so that your environment can access the new packages"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "D-ZBOjErv5mM"
- },
- "outputs": [],
- "source": [
- "import IPython\n",
- "\n",
- "app = IPython.Application.instance()\n",
- "app.kernel.do_shutdown(True)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "before_you_begin:nogpu"
- },
- "source": [
- "## Before you begin\n",
- "\n",
- "### Set your project ID\n",
- "\n",
- "**If you don't know your project ID**, try the following:\n",
- "* Run `gcloud config list`.\n",
- "* Run `gcloud projects list`.\n",
- "* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "set_project_id"
- },
- "outputs": [],
- "source": [
- "PROJECT_ID = \"llm-experiment-393119\" # @param {type:\"string\"}\n",
- "\n",
- "# Set the project id\n",
- "! gcloud config set project {PROJECT_ID}"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "region"
- },
- "source": [
- "#### Region\n",
- "\n",
- "You can also change the `REGION` variable used by Vertex AI. For preview, only `us-central1` and `europe-west4` are supported."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "2dw8q9fdQEH5"
- },
- "outputs": [],
- "source": [
- "REGION = \"europe-west4\""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "gcp_authenticate"
- },
- "source": [
- "### Authenticate your Google Cloud account\n",
- "\n",
- "Depending on your Jupyter environment, you may need to authenticate manually. Follow the relevant instructions below.\n",
- "\n",
- "#### Vertex AI Workbench\n",
- "Do nothing as you are already authenticated.\n",
- "\n",
- "#### Local JupyterLab instance\n",
- "\n",
- "**1. Uncomment and run:**"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "ce6043da7b33"
- },
- "outputs": [],
- "source": [
- "! gcloud auth login\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "0367eac06a10"
- },
- "source": [
- "**2. Colab, uncomment and run:**"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "21ad4dbb4a61"
- },
- "outputs": [],
- "source": [
- "from google.colab import auth\n",
- "\n",
- "auth.authenticate_user()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "c13224697bfb"
- },
- "source": [
- "**3. Service account or other**\n",
- "* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bucket:mbsdk"
- },
- "source": [
- "### Create a Cloud Storage bucket\n",
- "\n",
- "Define a storage bucket to store pipeline artifacts."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "bucket"
- },
- "outputs": [],
- "source": [
- "BUCKET_URI = f\"gs://your-bucket-name-{PROJECT_ID}-unique\" # @param {type:\"string\"}"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "autoset_bucket"
- },
- "source": [
- "**Only if your bucket doesn't exist already**: Run the following cell to create your Cloud Storage bucket."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "91c46850b49b"
- },
- "outputs": [],
- "source": [
- "! gsutil mb -l $REGION -p $PROJECT_ID $BUCKET_URI\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "85c4ecfd133a"
- },
- "source": [
- "#### Service Account\n",
- "\n",
- "A service account is needed to create Vertex AI Pipeline jobs. If you don't know your service account, try to get your service account using gcloud command by executing the second cell below."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "77b01a1fdbb4"
- },
- "outputs": [],
- "source": [
- "SERVICE_ACCOUNT = \"[your-service-account]\" # @param {type:\"string\"}"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "f936bebda2d4"
- },
- "outputs": [],
- "source": [
- "import sys\n",
- "\n",
- "IS_COLAB = \"google.colab\" in sys.modules\n",
- "if (\n",
- " SERVICE_ACCOUNT == \"\"\n",
- " or SERVICE_ACCOUNT is None\n",
- " or SERVICE_ACCOUNT == \"[your-service-account]\"\n",
- "):\n",
- " # Get your service account from gcloud\n",
- " if not IS_COLAB:\n",
- " shell_output = !gcloud auth list 2>/dev/null\n",
- " SERVICE_ACCOUNT = shell_output[2].replace(\"*\", \"\").strip()\n",
- "\n",
- " else: # IS_COLAB:\n",
- " shell_output = ! gcloud projects describe $PROJECT_ID\n",
- " project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
- " SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
- "\n",
- " print(\"Service Account:\", SERVICE_ACCOUNT)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "40ef6967cad3"
- },
- "source": [
- "#### Set service account access for Vertex AI Pipelines\n",
- "\n",
- "Run the following commands to grant your service account access to read and write pipeline artifacts in the bucket that you created in the previous step. You only need to run this step once per service account."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "f88cb0488c08"
- },
- "outputs": [],
- "source": [
- "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectCreator $BUCKET_URI\n",
- "\n",
- "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.objectViewer $BUCKET_URI\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "init_aip:mbsdk"
- },
- "source": [
- "## Initialize Vertex AI SDK for Python\n",
- "\n",
- "Initialize the Vertex AI SDK for Python for your project and corresponding bucket."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "rUP-o_VB4JEY"
- },
- "outputs": [],
- "source": [
- "import google.cloud.aiplatform as aiplatform\n",
- "\n",
- "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=BUCKET_URI)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "3o6EaTXJ4JEY"
- },
- "source": [
- "## Construct the Pipeline Job and Run on Vertex AI\n",
- "\n",
- "Define a pipeline job with the following code, which will:\n",
- "\n",
- "- load the pipeline template that was compiled in the previous step, and\n",
- "- set pipeline parameters. (See [Prepare your inputs](https://colab.research.google.com/drive/1SGSTAW3dcANbU_d3g5mSGCN9LXf5cA1p?resourcekey=0-NAfp-Rrb9piBiJbvqa6bqA#scrollTo=aef4f59195ad&line=61&uniqifier=1) section above for the definition of these parameters.)\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "df55e79c"
- },
- "outputs": [],
- "source": [
- "# Create a job via the pipeline template compiled in the previous step:\n",
- "job = aiplatform.PipelineJob(\n",
- " display_name=\"tutorial-rlhf-tuning\",\n",
- " pipeline_root=f\"{BUCKET_URI}/rlhf\",\n",
- " template_path=\"https://us-kfp.pkg.dev/ml-pipeline/llm-rlhf/rlhf-llama-train-template/stable\",\n",
- " parameter_values={\n",
- " \"preference_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/summarize_from_feedback_tfds/comparisons/train/*.jsonl\",\n",
- " \"prompt_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/train/*.jsonl\",\n",
- " \"eval_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/val/*.jsonl\",\n",
- " \"large_model_reference\": \"llama-2-7b-chat\", # see Table 1 for valid values.\n",
- " \"model_display_name\": \"my_rlhf_tutorial_model\", # Optional. If omitted, a default model_display_name will be created.\n",
- " \"reward_model_train_steps\": 100, # Please remember to read \"A Note on choosing train_steps\" section.\n",
- " \"reinforcement_learning_train_steps\": 100, # Please remember to read \"A Note on choosing train_steps\" section.\n",
- " \"prompt_sequence_length\": 512,\n",
- " \"target_sequence_length\": 64,\n",
- " \"reward_model_learning_rate_multiplier\": 1.0,\n",
- " \"reinforcement_learning_rate_multiplier\": 1.0,\n",
- " \"kl_coeff\": 0.1,\n",
- " \"lora_dim\": 4,\n",
- " \"instruction\": \"Summarize in less than 50 words.\",\n",
- " },\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Lz99NXISmL0I"
- },
- "source": [
- "### Run the pipeline using Vertex AI Pipelines\n",
- "\n",
- "Execute the following cell to run the pipeline on Vertex AI."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "CtAF_q4I4JEZ"
- },
- "outputs": [],
- "source": [
- "job.run()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "SDGIpQo_4JEd"
- },
- "source": [
- "### View pipeline job in UI\n",
- "\n",
- "Go to \"Vertex AI Pipelines\" in the Google Cloud UI to view the pipeline job. When the pipeline completes, click on the **Reinforcer** step. In the output parameters section you should see the output model's GCS path.\n",
- "* If you tuned an open-source model (like the T5 models), the output model will be in your project's GCS buckets. You will be able to download the tuned model.\n",
- "* If you tuned PaLM 2 models, the output model will be in a restricted-access GCS bucket. You won't be able to download the tuned model."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "QgKNkUuWOIff"
- },
- "source": [
- "### Look at training loss and reward using Tensorboard\n",
- "\n",
- "The **RewardModelTrainer** and **Reinforcer** output the training time `rank_loss`, `reward`, `kl_loss` to tensorboard to help you examine training curves and decide how to tune hyper-parameters.\n",
- "\n",
- "Click on the **tensorboard_metrics** > `URI` attached to **RewardModelTrainer** and **Reinforcer**. If you encounter \"❗ Requested entity was not found\", this is a known issue. Hit the backarrow at the top of the page, then > `tensorboard_metrics/` > `train/` > some filename that looks like `events.out.tfevents.*-w-1.1.0.v2`. This is the tensorboard file. Download the tensorboard files from both RewardModelTrainer and Reinforcer.\n",
- "\n",
- "Visualize the tensorboards using the [tensorboard python package](https://pypi.org/project/tensorboard/). It will give you some loss curves that looks like the curves below.\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "juQ5Lk-mPayn"
- },
- "source": [
- ""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "TRfc_PDZRKF8"
- },
- "source": [
- ""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "1jSGT6VUo3FT"
- },
- "source": [
- "## Bulk Inference\n",
- "You can perform bulk inference using your tuned model. To perform bulk inference you will need the tuned model path, which can be found in the Vertex Pipelines UI under **Reinforcer** > **Output Parameters** > `output_model_path`. Then run the Bulk Inference pipeline to generate offline predictions from a tuned model checkpoint. See below for details."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "XwEw1ski4qeZ"
- },
- "source": [
- ""
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "86o4vudnkx2n"
- },
- "outputs": [],
- "source": [
- "PIPELINE_ROOT = f\"{BUCKET_URI}/tutorial-infer-eval\"\n",
- "\n",
- "# This path is the out-of-the-box llama-2-7b-chat checkpoint. If you want to\n",
- "# evaluate your own tuned model, PLEASE substitute this path with your own tuned\n",
- "# checkpoint from the RLHF training pipeline. The tuned model is stored at the\n",
- "# Reinforcer > Output Parameters > output_model_path.\n",
- "MODEL_CHECKPOINT = (\n",
- " \"gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b_chat\"\n",
- ")\n",
- "\n",
- "# Create a job via the pipeline template\n",
- "job = aiplatform.PipelineJob(\n",
- " display_name=\"tutorial-infer-eval\",\n",
- " pipeline_root=PIPELINE_ROOT,\n",
- " template_path=\"https://us-kfp.pkg.dev/ml-pipeline/llm-rlhf/llama-infer-template/stable\",\n",
- " parameter_values={\n",
- " \"large_model_reference\": \"llama-2-7b-chat\", # see Table 1 for valid values.\n",
- " \"model_checkpoint\": MODEL_CHECKPOINT,\n",
- " \"prompt_dataset\": \"gs://vertex-ai/generative-ai/rlhf/text_small/reddit_tfds/val/*.jsonl\",\n",
- " \"instruction\": \"Summarize in less than 50 words.\",\n",
- " },\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "Dudb1iFhliM9"
- },
- "source": [
- "### Run the pipeline using Vertex AI Pipelines\n",
- "\n",
- "Execute the following cell to run the pipeline on Vertex AI."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "Jp8C5W7GllQ-"
- },
- "outputs": [],
- "source": [
- "job.run()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "4tPP22Q_lsAl"
- },
- "source": [
- "## View the inference result\n",
- "\n",
- "Go to \"Vertex AI Pipelines\" in the Google Cloud UI to view the pipeline job. When the pipeline completes, click on **Bulk Inferrer** > **Output Parameters** > `output_prediction_gcs_path`. If you encounter \"❗ Requested entity was not found\", this is a known issue. Hit the backarrow at the top of the page, you will see your inference results there."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "C1Ck3ehFlw4Y"
- },
- "source": [
- ""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "NRaBADRM4JEn"
- },
- "source": [
- "## Cleaning up\n",
- "\n",
- "To clean up all Google Cloud resources used in this project, you can [delete the Google Cloud\n",
- "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
- "\n",
- "Otherwise, you can delete the individual resources you created in this tutorial."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "PAr4UWWx4JEo"
- },
- "outputs": [],
- "source": [
- "delete_bucket = False\n",
- "\n",
- "job.delete()\n",
- "\n",
- "if delete_bucket or os.getenv(\"ID_TESTING\"):\n",
- " ! gsutil rm -rf {BUCKET_URI}"
- ]
- }
- ],
- "metadata": {
- "colab": {
- "name": "model_garden_pytorch_llama2_peft.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}