\n",
- " \n",
- " \n",
- " ![\"Colab](\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\") Run in Colab\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " ![\"GitHub](\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\") \n",
- " View on GitHub\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- " ![\"Vertex](\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\") \n",
- "Open in Vertex AI Workbench\n",
- " (A Python-3 GPU notebook is recommended)\n",
- " | \n",
- "
"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "3de7470326a2"
- },
- "source": [
- "## Overview\n",
- "\n",
- "This notebook demonstrates running local inference with prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA, deploying prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm), finetuning and deploying OpenLLaMA with performance efficient finetuning libraries ([PEFT](https://github.com/huggingface/peft)), quantizing and deploying OpenLLaMA with AWQ or GPTQ, and evaluating PEFT-finetuned OpenLLaMA in Vertex AI.\n",
- "\n",
- "### Objective\n",
- "\n",
- "- Run local inference with prebuilt OpenLLaMA\n",
- "- Deploy prebuilt OpenLLaMA\n",
- "- Deploy prebuilt OpenLLaMA with [vLLM](https://github.com/vllm-project/vllm) to improve serving throughput\n",
- "- Finetune and deploy OpenLLaMA with PEFT\n",
- "- Quantize and deploy OpenLLaMA models with AWQ or GPTQ\n",
- "- Evaluate finetuned OpenLLaMA with PEFT\n",
- "\n",
- "| Models | LoRA |\n",
- "| :- | :- |\n",
- "| [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) | Y |\n",
- "| [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) | Y |\n",
- "| [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) | Y |\n",
- "\n",
- "### Costs\n",
- "\n",
- "This tutorial uses billable components of Google Cloud:\n",
- "\n",
- "* Vertex AI\n",
- "* Cloud Storage\n",
- "\n",
- "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "264c07757582"
- },
- "source": [
- "## Before you begin\n",
- "\n",
- "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands.\n",
- "\n",
- "Running local inference with OpenLLaMA requires a GPU."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ioensNKM8ned"
- },
- "source": [
- "### Colab only\n",
- "Run the following commands for Colab and skip this section if you are using Workbench."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "2707b02ef5df"
- },
- "outputs": [],
- "source": [
- "import sys\n",
- "\n",
- "if \"google.colab\" in sys.modules:\n",
- " ! pip3 install --upgrade google-cloud-aiplatform\n",
- " from google.colab import auth as google_auth\n",
- "\n",
- " google_auth.authenticate_user()\n",
- "\n",
- " # Restart the notebook kernel after installs.\n",
- " import IPython\n",
- "\n",
- " app = IPython.Application.instance()\n",
- " app.kernel.do_shutdown(True)\n",
- "! pip3 install transformers==4.31.0\n",
- "! pip3 install sentencepiece==0.1.99\n",
- "! pip3 install accelerate==0.21.0"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bb7adab99e41"
- },
- "source": [
- "### Setup Google Cloud project\n",
- "\n",
- "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
- "\n",
- "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
- "\n",
- "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
- "\n",
- "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
- "\n",
- "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "6c460088b873"
- },
- "source": [
- "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "855d6b96f291"
- },
- "outputs": [],
- "source": [
- "# Cloud project id.\n",
- "PROJECT_ID = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# The region you want to launch jobs in.\n",
- "REGION = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# The Cloud Storage bucket for storing experiments output.\n",
- "# Start with gs:// prefix, e.g. gs://foo_bucket.\n",
- "BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n",
- "\n",
- "! gcloud config set project $PROJECT_ID\n",
- "\n",
- "import os\n",
- "\n",
- "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
- "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
- "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
- "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
- "\n",
- "# The service account looks like:\n",
- "# '@.iam.gserviceaccount.com'\n",
- "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
- "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
- "# The service account for deploying fine tuned model.\n",
- "SERVICE_ACCOUNT = \"\" # @param {type:\"string\"}"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e828eb320337"
- },
- "source": [
- "### Initialize Vertex AI API"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "12cd25839741"
- },
- "outputs": [],
- "source": [
- "from google.cloud import aiplatform\n",
- "\n",
- "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "2cc825514deb"
- },
- "source": [
- "### Define constants"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "b42bd4fa2b2d"
- },
- "outputs": [],
- "source": [
- "# The pre-built training and serving docker images.\n",
- "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00\"\n",
- "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:20231130_0948_RC00\"\n",
- "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20231127_0916_RC00\"\n",
- "VLLM_GPTQ_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:gptq\"\n",
- "EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20231011_0934_RC00\""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "0c250872074f"
- },
- "source": [
- "### Define common functions"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "354da31189dc"
- },
- "outputs": [],
- "source": [
- "from datetime import datetime\n",
- "from typing import Tuple\n",
- "\n",
- "from google.cloud import aiplatform\n",
- "\n",
- "\n",
- "def create_name_with_datetime(prefix: str) -> str:\n",
- " \"\"\"Creates a name with date time when triggering training or deployment\n",
- " jobs in Vertex AI.\n",
- " \"\"\"\n",
- " return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
- "\n",
- "\n",
- "def deploy_model(\n",
- " model_name: str,\n",
- " model_id: str,\n",
- " finetuned_lora_model_path: str,\n",
- " service_account: str,\n",
- " task: str,\n",
- " precision_loading_mode: str = \"float16\",\n",
- " machine_type: str = \"n1-standard-8\",\n",
- " accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
- " accelerator_count: int = 1,\n",
- ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
- " \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
- " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
- " serving_env = {\n",
- " \"MODEL_ID\": model_id,\n",
- " \"PRECISION_LOADING_MODE\": precision_loading_mode,\n",
- " \"TASK\": task,\n",
- " \"DEPLOY_SOURCE\": \"notebook\",\n",
- " }\n",
- " if finetuned_lora_model_path:\n",
- " serving_env[\"FINETUNED_LORA_MODEL_PATH\"] = finetuned_lora_model_path\n",
- " model = aiplatform.Model.upload(\n",
- " display_name=model_name,\n",
- " serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
- " serving_container_ports=[7080],\n",
- " serving_container_predict_route=\"/predictions/peft_serving\",\n",
- " serving_container_health_route=\"/ping\",\n",
- " serving_container_environment_variables=serving_env,\n",
- " model_garden_source_model_name=\"publishers/openlm-research/models/openllama\"\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",
- " system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_openllama_peft.ipynb\"},\n",
- " )\n",
- " return model, endpoint\n",
- "\n",
- "\n",
- "def deploy_model_vllm(\n",
- " model_name: str,\n",
- " model_id: str,\n",
- " service_account: str,\n",
- " machine_type: str = \"n1-standard-8\",\n",
- " accelerator_type: str = \"NVIDIA_TESLA_V100\",\n",
- " accelerator_count: int = 1,\n",
- " quantization_method: str = \"\",\n",
- ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
- " \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
- " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
- "\n",
- " vllm_args = [\n",
- " \"--host=0.0.0.0\",\n",
- " \"--port=7080\",\n",
- " f\"--model={model_id}\",\n",
- " f\"--tensor-parallel-size={accelerator_count}\",\n",
- " \"--swap-space=16\",\n",
- " \"--gpu-memory-utilization=0.9\",\n",
- " \"--max-num-batched-tokens=4096\",\n",
- " \"--disable-log-stats\",\n",
- " ]\n",
- " if quantization_method:\n",
- " vllm_args.append(f\"--quantization={quantization_method}\")\n",
- " if quantization_method == \"gptq\":\n",
- " vllm_docker_uri = VLLM_GPTQ_DOCKER_URI\n",
- " else:\n",
- " vllm_docker_uri = VLLM_DOCKER_URI\n",
- "\n",
- " serving_env = {\n",
- " \"MODEL_ID\": \"openlm-research/open_llama\",\n",
- " \"DEPLOY_SOURCE\": \"notebook\",\n",
- " }\n",
- " model = aiplatform.Model.upload(\n",
- " display_name=model_name,\n",
- " serving_container_image_uri=vllm_docker_uri,\n",
- " serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
- " serving_container_args=vllm_args,\n",
- " serving_container_ports=[7080],\n",
- " serving_container_predict_route=\"/generate\",\n",
- " serving_container_health_route=\"/ping\",\n",
- " serving_container_environment_variables=serving_env,\n",
- " model_garden_source_model_name=\"publishers/openlm-research/models/openllama\"\n",
- " )\n",
- "\n",
- " model.deploy(\n",
- " endpoint=endpoint,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " deploy_request_timeout=1800,\n",
- " service_account=service_account,\n",
- " )\n",
- " return model, endpoint"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "65eaa62632d1"
- },
- "source": [
- "## Run inferences locally with prebuilt OpenLLaMA"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "339601a9500b"
- },
- "outputs": [],
- "source": [
- "import torch\n",
- "from transformers import LlamaForCausalLM, LlamaTokenizer\n",
- "\n",
- "model_path = \"openlm-research/open_llama_3b\"\n",
- "\n",
- "tokenizer = LlamaTokenizer.from_pretrained(model_path)\n",
- "\n",
- "model = LlamaForCausalLM.from_pretrained(\n",
- " model_path,\n",
- " torch_dtype=torch.float16,\n",
- " device_map=\"auto\",\n",
- ")\n",
- "\n",
- "prompt = \"Q: What is the largest animal?\\nA:\"\n",
- "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
- "input_ids = input_ids.to(\"cuda\")\n",
- "generation_output = model.generate(input_ids=input_ids, max_new_tokens=32)\n",
- "print(tokenizer.decode(generation_output[0]))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "V7VOhhHGpUrj"
- },
- "source": [
- "## Deploy Prebuilt OpenLLaMA with vLLM\n",
- "\n",
- "This section deploys prebuilt OpenLLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. The model deployment step will take ~15 minutes to complete.\n",
- "\n",
- "vLLM is a highly optimized LLM serving framework that can significantly increase serving throughput. The higher QPS you have, the more performance benefits you get from using vLLM."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "4GTNnnuYqrW_"
- },
- "source": [
- "Set the prebuilt model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "kLsRoc4Kqrkx"
- },
- "outputs": [],
- "source": [
- "prebuilt_model_id = \"openlm-research/open_llama_7b\" # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "YI0vaDi6p2fi"
- },
- "outputs": [],
- "source": [
- "# Finds Vertex AI prediction supported accelerators and regions in\n",
- "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
- "\n",
- "# Sets V100 to deploy open_llama_3b and open_llama_7b.\n",
- "# V100 serving has better throughput and latency performance than L4 serving.\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets L4 to deploy open_llama_3b and open_llama_7b.\n",
- "# L4 serving is more cost efficient than V100 serving.\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 2 V100 to deploy open_llama_13b.\n",
- "# V100 serving has better throughput and latency performance than L4 serving.\n",
- "# machine_type = \"n1-standard-16\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Sets 2 L4 to deploy open_llama_13b.\n",
- "# L4 serving is more cost efficient than V100 serving.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
- " # vLLM currently does not support OpenLLaMA 3B.\n",
- " precision_loading_mode = \"float16\"\n",
- " model_without_peft, endpoint_without_peft = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"openllama-serve\"),\n",
- " model_id=model_id,\n",
- " finetuned_lora_model_path=\"\", # This will avoid override finetuning models.\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"causal-language-modeling-lora\",\n",
- " precision_loading_mode=precision_loading_mode,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )\n",
- "else:\n",
- " model_without_peft, endpoint_without_peft = deploy_model_vllm(\n",
- " model_name=create_name_with_datetime(prefix=\"openllama-serve-vllm\"),\n",
- " model_id=prebuilt_model_id,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "dWYmYWoqqBuZ"
- },
- "source": [
- "NOTE: The prebuilt model weights will be downloaded on the fly from the original location after the deployment succeeds. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts. If you are interested in additional serving parameters, please refer to the vLLM GitHub [examples/api_client.py](https://github.com/vllm-project/vllm/blob/main/examples/api_client.py) for more details."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "fjO4z3qAp3pK"
- },
- "outputs": [],
- "source": [
- "instance = {\n",
- " \"prompt\": \"Hi, Google. How are you doing?\",\n",
- " \"n\": 1,\n",
- " \"max_tokens\": 32,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- "}\n",
- "response = endpoint_without_peft.predict(instances=[instance])\n",
- "print(response.predictions[0])"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e70e3519ff8b"
- },
- "source": [
- "## Finetune and deploy OpenLLaMA with PEFT\n",
- "\n",
- "This section demonstrates how to finetune the OpenLLaMA-7b model, merge the finetuned LoRA adapter with the base model, and serve using vLLM."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "5qCrm_kJH5cz"
- },
- "source": [
- "Set the base model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "N3UBLiYrM3sU"
- },
- "outputs": [],
- "source": [
- "model_id = \"openlm-research/open_llama_7b\" # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "iWGwJHqI7LMs"
- },
- "source": [
- "### Finetune"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "KKEYoRfiHDVv"
- },
- "source": [
- "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
- "\n",
- "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes). You can either use a [dataset from huggingface](https://huggingface.co/datasets) or a custom JSONL dataset in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) stored in Cloud Storage. The `template` parameter is optional.\n",
- "\n",
- "In order to make the finetuning efficiently, we enabled quantization for loading pretrained models for finetuning LoRA models. Precision options include `\"4bit\"`, `\"8bit\"`, `\"float16\"` (default) and `\"float32\"`, and the precision can be set via `\"--precision_mode\"`. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. `open_llama_3b` and `open_llama_7b` can be finetuned on **1 V100 (16G)** and **1 L4 (24G)**, and `open_llama_13b` can be finetuned on **1 L4 (24G)**.\n",
- "\n",
- "In this section, the finetuned LoRA adapter will be saved to a GCS bucket specified by the variable `lora_adapter_dir` below; and we merge the LoRa adapter with the base model, and save it to a separate GCS bucket specified by `merged_model_output_dir` below.\n"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "0810ef72dd9f"
- },
- "source": [
- "#### [Optional] Finetune with a custom dataset\n",
- "\n",
- "To use a custom dataset, you should supply a `gs://` URI to a JSONL file in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) in the `dataset_name` below.\n",
- "\n",
- "For example, here is one data point from the sample dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`:\n",
- "\n",
- "```json\n",
- "{\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
- "```\n",
- "\n",
- "To use this sample dataset that contains `input_text` and `output_text` fields, set `dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl` and `template` to `vertex_sample`. For advanced usage with custom datatset fields, see [the template example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json) and supply your own JSON template as `gs://` URIs."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "65467b361315"
- },
- "outputs": [],
- "source": [
- "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
- "dataset_name = \"Abirate/english_quotes\" # @param {type:\"string\"}\n",
- "# Optional. Template name or gs:// URI to a custom template.\n",
- "template = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# Worker pool spec.\n",
- "# Finetunes open_llama_3b and open_llama_7b with 1 V100 (16G).\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Finetunes open_llama_3b and open_llama_7b with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Finetunes open_llama_13b with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Finetunes open_llama_13b with 1 A100 (40G).\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "replica_count = 1\n",
- "\n",
- "\n",
- "# Setup training job.\n",
- "job_name = create_name_with_datetime(\"openllama-lora-train\")\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- ")\n",
- "\n",
- "# Create a GCS folder to store the LORA adapter.\n",
- "lora_adapter_dir = create_name_with_datetime(\"openllama-lora-adapter\")\n",
- "lora_output_dir = os.path.join(MODEL_BUCKET, lora_adapter_dir)\n",
- "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Create a GCS folder to store the merged model with the base model and the\n",
- "# finetuned LORA adapter.\n",
- "merged_model_dir = create_name_with_datetime(\"openllama-merged-model\")\n",
- "merged_model_output_dir = os.path.join(MODEL_BUCKET, merged_model_dir)\n",
- "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Pass training arguments and launch job.\n",
- "train_job.run(\n",
- " args=[\n",
- " \"--task=causal-language-modeling-lora\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--dataset_name={dataset_name}\",\n",
- " f\"--output_dir={lora_output_dir_gcsfuse}\",\n",
- " f\"--merge_base_and_lora_output_dir={merged_model_output_dir_gcsfuse}\",\n",
- " \"--lora_rank=16\",\n",
- " \"--lora_alpha=32\",\n",
- " \"--lora_dropout=0.05\",\n",
- " \"--warmup_steps=10\",\n",
- " \"--max_steps=10\",\n",
- " \"--learning_rate=2e-4\",\n",
- " f\"--template={template}\",\n",
- " ],\n",
- " replica_count=replica_count,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " boot_disk_size_gb=500,\n",
- ")\n",
- "\n",
- "print(\"The finetuned Lora adapter can be found at: \", lora_output_dir)\n",
- "print(\n",
- " \"The finetuned Lora adapter merged with the base model can be found at: \",\n",
- " merged_model_output_dir,\n",
- ")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "53b8a1ad6def"
- },
- "source": [
- "### [Optional] Hyperparameter tuning\n",
- "\n",
- "You can use the Vertex AI SDK to create and run the [hyperparameter tuning job](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) to obtain a better performance by experimenting with different hyperparameters such as learning rates.\n",
- "\n",
- "Define the following specifications:\n",
- "\n",
- "- `worker_pool_specs`: Dictionary specifying the machine type and Docker image.\n",
- "\n",
- "- `parameter_spec`: Dictionary specifying the parameters to optimize. The dictionary key is the string assigned to the command line argument for each hyperparameter in your training application code, and the dictionary value is the parameter specification. The parameter specification includes the type, min/max values, and scale for the hyperparameter.\n",
- "\n",
- "- `metric_spec`: Dictionary specifying the metric to optimize. The dictionary key is the hyperparameter_metric_tag that you set in your training application code, and the value is the optimization goal.\n",
- "\n",
- "The following 4bit QLoRA experiment results show the effectiveness of hyperparameter tuning evaluated on the ARC Challenge dataset (for reference only):\n",
- "\n",
- "| Model | Training time | Trials | Parallel Trials | GPU | ∆arc challenge | ∆hellaswag | ∆truthfulqa_mc | cost |\n",
- "|---------------|---------------|--------|-----------------|------|----------------|------------|----------------|----------|\n",
- "| Openllama-3b | 2d 10hrs | 8 | 1 | L4x1 | +1.62 | +7.32 | +3.34 | \\$29.0232 |\n",
- "| Openllama-7b | 1d 4hrs | 8 | 2 | L4x1 | +2.82 | +3.55 | +6.68 | \\$47.8016 |\n",
- "| Openllama-13b | 6d 10hrs | 8 | 2 | L4x1 | +1.01 | +3.67 | +6.19 | \\$87.9208 |\n",
- "\n",
- "The following example runs 8 trials on `timdettmers/openassistant-guanaco` with different learning rates, and evaluates the model on `arc_challenge` dataset. You can customize the search space by extending the range of learning rates, adding other parameters such as LoRA rank, etc. Please refer to the [hyperparameter tuning documentation](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview) for more information."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "a81402f70641"
- },
- "outputs": [],
- "source": [
- "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
- "dataset_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
- "# Optional. Template name or gs:// URI to a custom template.\n",
- "template = \"\" # @param {type:\"string\"}\n",
- "\n",
- "hpt_precision_mode = \"4bit\"\n",
- "\n",
- "# Worker pool spec for 4bit finetuning.\n",
- "\n",
- "# Finetunes Openllama 3B / 7B / 13B with 1 L4 (24G).\n",
- "machine_type = \"g2-standard-8\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "55ed2a2a9d54"
- },
- "source": [
- "### [Optional] Custom evaluation dataset\n",
- "\n",
- "To obtain a model with better performance on some specific tasks, you might want to run hyperparameter tuning with a custom evaluation dataset. The hyperparameter tuning service will pick the model according to the evaluation dataset and the metrics you selected. You can use any of the following tasks as the `eval_task` in the code cell below:\n",
- "\n",
- "1. The name of a [lm-evaluation-harness task](https://github.com/EleutherAI/lm-evaluation-harness/tree/big-refactor/lm_eval/tasks).\n",
- "\n",
- "2. `custom_likelihood`. Then, add a flag `--eval_dataset_path=`. The JSONL file must be in the format in Vertex AI language model's [prepare evaluation dataset](https://cloud.google.com/vertex-ai/docs/generative-ai/models/evaluate-models#classification) page.\n",
- "\n",
- "3. `builtin_eval`. The built-in evaluation loop of the trainer will be used to evaluate the model instead of the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) library. You can supply any eval dataset in the same format as the training dataset by specifying `--eval_dataset_path`, `--eval_split`, `--eval_template`, and `--eval_column`."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "7864cff27197"
- },
- "outputs": [],
- "source": [
- "from google.cloud.aiplatform import hyperparameter_tuning as hpt\n",
- "\n",
- "eval_task = \"arc_challenge\" # @param {type:\"string\"}\n",
- "eval_metric_name = \"acc_norm\" # @param {type:\"string\"}\n",
- "\n",
- "# Runs 10 training steps as a minimal example. Use 1000 to reproduce the experiment results.\n",
- "max_steps = 10 # @param {type:\"integer\"}\n",
- "# Evaluates the model on 10 examples. Use 10000 to reproduce the experiment results.\n",
- "eval_limit = 10 # @param {type:\"integer\"}\n",
- "\n",
- "flags = {\n",
- " \"learning_rate\": 1e-5,\n",
- " \"precision_mode\": hpt_precision_mode,\n",
- " \"task\": \"instruct-lora\",\n",
- " \"pretrained_model_id\": model_id,\n",
- " \"output_dir\": lora_output_dir_gcsfuse,\n",
- " \"warmup_steps\": 10,\n",
- " \"max_steps\": max_steps,\n",
- " \"lora_rank\": 32,\n",
- " \"lora_alpha\": 64,\n",
- " \"lora_dropout\": 0.05,\n",
- " \"dataset_name\": dataset_name,\n",
- " \"eval_steps\": max_steps + 1, # Only evaluates in the end.\n",
- " \"eval_tasks\": eval_task,\n",
- " \"eval_limit\": eval_limit,\n",
- " \"eval_metric_name\": eval_metric_name,\n",
- " \"merge_base_and_lora_output_dir\": merged_model_output_dir_gcsfuse,\n",
- "}\n",
- "\n",
- "# Add labels for the finetuning job.\n",
- "labels = {\n",
- " \"mg-source\": \"notebook\",\n",
- " \"mg-notebook-name\": \"model_garden_pytorch_openllama_peft.ipynb\".split(\".\")[0],\n",
- "}\n",
- "\n",
- "labels[\"mg-tune\"] = \"publishers-openlm-research-models-openllama\"\n",
- "versioned_model_id = prebuilt_model_id.split(\"/\")[1].replace(\"_\", \"-\")\n",
- "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\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",
- " labels=labels,\n",
- ")\n",
- "\n",
- "train_hpt_job = aiplatform.HyperparameterTuningJob(\n",
- " display_name=f\"{job_name}_hpt\",\n",
- " custom_job=train_job,\n",
- " metric_spec=metric_spec,\n",
- " parameter_spec=parameter_spec,\n",
- " max_trial_count=8,\n",
- " parallel_trial_count=2,\n",
- ")\n",
- "\n",
- "train_hpt_job.run()\n",
- "\n",
- "print(\"Trained models were saved in: \", lora_output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "afad2431f37f"
- },
- "source": [
- "Then, find the best trial from the hyperparameter tuning job."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "138b09ee6313"
- },
- "outputs": [],
- "source": [
- "best_trial_id = max(\n",
- " train_hpt_job.trials, key=lambda trial: trial.final_measurement.metrics[0].value\n",
- ").id\n",
- "lora_output_dir = os.path.join(lora_output_dir, f\"trial_{best_trial_id}\")\n",
- "lora_output_dir_gcsfuse = lora_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "print(f\"Best trial {best_trial_id} saved model in:\", lora_output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "jqmCtkGnhDmp"
- },
- "source": [
- "### Deploy with vLLM\n",
- "This section uploads the model to Model Registry and deploys it on the Endpoint. vLLM currently does not support serving finetuned [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b).\n",
- "\n",
- "The model deployment step will take ~15 minutes to complete.\n",
- "\n",
- "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~5.3G, ~8.7G and ~15.2G respectively with the default settings.\n",
- "\n",
- "NOTE: vLLM requires a merged model with the base model and the finetuned LoRA adapter. Based on your business need, if you need the base model and the finetuned LoRA weight to be served separately, please consider using the regular Vertex serving instead.\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "bf55e38815dc"
- },
- "outputs": [],
- "source": [
- "# Finds Vertex AI prediction supported accelerators and regions in\n",
- "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
- "\n",
- "# Sets V100 to deploy open_llama_3b and open_llama_7b.\n",
- "# V100 serving has better throughput and latency performance than L4 serving.\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets L4 to deploy open_llama_3b and open_llama_7b.\n",
- "# L4 serving is more cost efficient than V100 serving.\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 2 V100 to deploy open_llama_13b.\n",
- "# V100 serving has better throughput and latency performance than L4 serving.\n",
- "# machine_type = \"n1-standard-16\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Sets 2 L4 to deploy open_llama_13b.\n",
- "# L4 serving is more cost efficient than V100 serving.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
- " # vLLM currently does not support OpenLLaMA 3B.\n",
- " precision_loading_mode = \"float16\"\n",
- " model_with_peft, endpoint_with_peft = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
- " model_id=model_id,\n",
- " finetuned_lora_model_path=lora_output_dir, # This will avoid override finetuning models.\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"causal-language-modeling-lora\",\n",
- " precision_loading_mode=precision_loading_mode,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )\n",
- "else:\n",
- " model_with_peft, endpoint_with_peft = deploy_model_vllm(\n",
- " model_name=create_name_with_datetime(prefix=\"openllama-peft-serve-vllm\"),\n",
- " model_id=merged_model_output_dir,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )\n",
- "\n",
- "print(\"endpoint_name:\", endpoint_with_peft.name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "80b3fd2ace09"
- },
- "source": [
- "NOTE: After the deployment succeeds, the base model weights will be downloaded on the fly from the original location and LoRA model weights will be downloaded from the GCS bucket used in training above. Thus, an additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you can run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts. Parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "4ab04da3ec9a"
- },
- "outputs": [],
- "source": [
- "instance = {\n",
- " \"prompt\": \"Hi, Google. How are you doing?\",\n",
- " \"n\": 1,\n",
- " \"max_tokens\": 32,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- "}\n",
- "response = endpoint_with_peft.predict(instances=[instance])\n",
- "print(response.predictions[0])"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "vhDf9dNNn4bP"
- },
- "source": [
- "### [Optional] Merge a previously trained LoRA adapter with the base model\n",
- "\n",
- "This section demonstrates how to merge a previously trained LoRA adapter with a base model, and save the merged model to a GCS bucket. Please be aware that the LoRA adapter should be trained on the same base model."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "YHdru1aRqRFF"
- },
- "outputs": [],
- "source": [
- "merge_job_name = create_name_with_datetime(prefix=\"openllama-peft-merge\")\n",
- "\n",
- "# The base model to be merged upon. It can be a huggingface model id, or a GCS\n",
- "# path where the base model was stored.\n",
- "base_model_dir = \"gs://\" # @param {type:\"string\"}\n",
- "# The previously trained LoRA adapter. It needs to be stored in a GCS path.\n",
- "finetuned_lora_adapter_dir = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# The GCS path to save the merged model\n",
- "merged_model_output_dir = os.path.join(MODEL_BUCKET, merge_job_name)\n",
- "merged_model_output_dir_gcsfuse = merged_model_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Worker pool spec.\n",
- "# Merges open_llama_3b and open_llama_7b with 1 V100 (16G).\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "\n",
- "# Merges open_llama_3b and open_llama_7b with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "\n",
- "# Merges open_llama_13b with 1 L4 (24G).\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "\n",
- "# Merges open_llama_13b with 1 A100 (40G).\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "\n",
- "worker_pool_specs = [\n",
- " {\n",
- " \"machine_spec\": {\n",
- " \"machine_type\": machine_type,\n",
- " \"accelerator_type\": accelerator_type,\n",
- " \"accelerator_count\": 1,\n",
- " },\n",
- " \"replica_count\": 1,\n",
- " \"container_spec\": {\n",
- " \"image_uri\": TRAIN_DOCKER_URI,\n",
- " \"command\": [],\n",
- " \"args\": [\n",
- " \"--task=merge-causal-language-model-lora\",\n",
- " \"--merge_model_precision_mode=float16\",\n",
- " \"--pretrained_model_id=%s\" % base_model_dir,\n",
- " \"--finetuned_lora_model_dir=%s\" % finetuned_lora_adapter_dir,\n",
- " \"--merge_base_and_lora_output_dir=%s\" % merged_model_output_dir_gcsfuse,\n",
- " ],\n",
- " },\n",
- " }\n",
- "]\n",
- "\n",
- "merge_custom_job = aiplatform.CustomJob(\n",
- " display_name=merge_job_name,\n",
- " project=PROJECT_ID,\n",
- " worker_pool_specs=worker_pool_specs,\n",
- " staging_bucket=STAGING_BUCKET,\n",
- ")\n",
- "\n",
- "merge_custom_job.run()\n",
- "\n",
- "print(\"The merged model is stored at: \", merged_model_output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "wi6UZnldXpdi"
- },
- "source": [
- "## Quantize and deploy OpenLLaMA 2 models\n",
- "\n",
- "This section demonstrates post-training quantization of OpenLLaMA models with Vertex Custom Job. Quantization reduces the memory required by a model while attempting to retain the same performance. Two such algorithms to do so are AWQ and GPTQ. Read more about AWQ in the following publication: [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://arxiv.org/abs/2306.00978). Read more about GPTQ in the following publication: [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers\n",
- "](https://arxiv.org/abs/2210.17323)."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "HTUo8dWSo8t9"
- },
- "source": [
- "### Quantize OpenLLaMA models\n",
- "\n",
- "Quantization reduces the amount of GPU required to serve a model by reducing the bit precision of the weights while minimizing drop in performance. Serving quantized models on VLLM requires models to be quantized to 4 bits. It is recommended to first search if a model has already been quantized and made publicly available: [AWQ](https://huggingface.co/TheBloke?search_models=-awq) and [GPTQ](https://huggingface.co/TheBloke?search_models=-gptq).\n",
- "\n",
- "Quantizing models with AWQ with 1 NVIDIA_L4 GPU will take around\n",
- "20 minutes for OpenLLaMA 3B, 30 minutes for OpenLLaMA 7B, and 1 hour for OpenLLaMA 13B.\n",
- "\n",
- "Quantizing models with GPTQ with 1 NVIDIA_L4 GPU will take around 30 minutes for OpenLLaMA 3B, 45 minutes for OpenLLaMA 7B, and 1.5 hour for OpenLLaMA 13B. Finetuned models can also be quantized, so long as the LoRA weights are merged with the base model."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "_YB31uAEpusN"
- },
- "outputs": [],
- "source": [
- "# Setup quantization job.\n",
- "\n",
- "# Set `finetuned_model_path` to `merged_model_output_dir` from the previous\n",
- "# section above to quantize the finetuned model, if not set the base model will\n",
- "# be quantized.\n",
- "finetuned_model_path = \"\" # @param {type:\"string\"}\n",
- "if finetuned_model_path:\n",
- " prequantized_model_path = finetuned_model_path\n",
- "else:\n",
- " prequantized_model_path = model_id\n",
- "\n",
- "quantization_method = \"awq\" # @param [\"awq\", \"gptq\"]\n",
- "quantization_job_name = get_job_name_with_datetime(\n",
- " f\"openllama-{quantization_method}-quantize\"\n",
- ")\n",
- "\n",
- "quantization_output_dir = os.path.join(MODEL_BUCKET, quantization_job_name)\n",
- "quantization_output_dir_gcsfuse = quantization_output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Worker pool spec.\n",
- "\n",
- "# Sets 1 L4 (24G) to quantize OpenLLaMA model.\n",
- "machine_type = \"g2-standard-16\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "\n",
- "# Quantization parameters.\n",
- "quantization_precision_mode = \"4bit\"\n",
- "if quantization_method == \"awq\":\n",
- " awq_dataset_name = \"pileval\"\n",
- " group_size = 64\n",
- " quantization_args = [\n",
- " \"--task=quantize-model\",\n",
- " f\"--quantization_method={quantization_method}\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--quantization_precision_mode={quantization_precision_mode}\",\n",
- " f\"--quantization_output_dir={quantization_output_dir_gcsfuse}\",\n",
- " f\"--quantization_dataset_name={awq_dataset_name}\",\n",
- " f\"--group_size={group_size}\",\n",
- " ]\n",
- "else:\n",
- " # The original datasets used in GPTQ paper [\"wikitext2\",\"c4\",\"c4-new\",\"ptb\",\"ptb-new\"].\n",
- " gptq_dataset_name = \"c4\" # @param {type:\"string\"}\n",
- " gptq_precision_mode = \"4bit\"\n",
- " group_size = -1\n",
- " damp_percent = 0.1\n",
- " desc_act = True\n",
- " quantization_args = [\n",
- " \"--task=quantize-model\",\n",
- " f\"--quantization_method={quantization_method}\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--quantization_precision_mode={quantization_precision_mode}\",\n",
- " f\"--quantization_output_dir={quantization_output_dir_gcsfuse}\",\n",
- " f\"--quantization_dataset_name={gptq_dataset_name}\",\n",
- " f\"--group_size={group_size}\",\n",
- " f\"--damp_percent={damp_percent}\",\n",
- " f\"--desc_act={desc_act}\",\n",
- " ]\n",
- "\n",
- "# Pass quantization arguments and launch job.\n",
- "worker_pool_specs = [\n",
- " {\n",
- " \"machine_spec\": {\n",
- " \"machine_type\": machine_type,\n",
- " \"accelerator_type\": accelerator_type,\n",
- " \"accelerator_count\": accelerator_count,\n",
- " },\n",
- " \"replica_count\": 1,\n",
- " \"disk_spec\": {\n",
- " \"boot_disk_type\": \"pd-ssd\",\n",
- " \"boot_disk_size_gb\": 500,\n",
- " },\n",
- " \"container_spec\": {\n",
- " \"image_uri\": TRAIN_DOCKER_URI,\n",
- " \"env\": [\n",
- " {\n",
- " \"name\": \"PYTORCH_CUDA_ALLOC_CONF\",\n",
- " \"value\": \"max_split_size_mb:32\",\n",
- " },\n",
- " ],\n",
- " \"command\": [],\n",
- " \"args\": quantization_args,\n",
- " },\n",
- " }\n",
- "]\n",
- "\n",
- "print(f\"Quantizing {prequantized_model_path}.\")\n",
- "quantize_job = aiplatform.CustomJob(\n",
- " display_name=quantization_job_name,\n",
- " project=PROJECT_ID,\n",
- " worker_pool_specs=worker_pool_specs,\n",
- " staging_bucket=STAGING_BUCKET,\n",
- ")\n",
- "quantize_job.run()\n",
- "\n",
- "print(\"Quantized models were saved in: \", quantization_output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ocm3m2Hzr4ln"
- },
- "source": [
- "### Deploy quantized models with Google Cloud Text Moderation\n",
- "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
- "\n",
- "The model deployment step will take 15 minutes to 1 hour to complete, depending on the model sizes."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "DdINEB5Ur9FX"
- },
- "outputs": [],
- "source": [
- "# Finds Vertex AI prediction supported accelerators and regions in\n",
- "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
- "\n",
- "# Sets 1 L4 (24G) to deploy OpenLLaMA models.\n",
- "machine_type = \"g2-standard-8\"\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "accelerator_count = 1\n",
- "\n",
- "\n",
- "if prebuilt_model_id == \"openlm-research/open_llama_3b\":\n",
- " # vLLM currently does not support OpenLLaMA 3B.\n",
- " precision_loading_mode = \"float16\"\n",
- " model_quantized_vllm, endpoint_quantized_vllm = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"openllama-quantized-serve\"),\n",
- " model_id=quantization_output_dir,\n",
- " finetuned_lora_model_path=\"\", # This will avoid override finetuning models.\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"causal-language-modeling-lora\",\n",
- " precision_loading_mode=precision_loading_mode,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )\n",
- "else:\n",
- " model_quantized_vllm, endpoint_quantized_vllm = deploy_model_vllm(\n",
- " model_name=create_name_with_datetime(prefix=\"openllama-quantized-serve-vllm\"),\n",
- " model_id=quantization_output_dir,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " )"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "y3Iy3q7CsEGe"
- },
- "source": [
- "NOTE: After the deployment succeeds, the model weights will be downloaded on the fly. Thus additional 10 ~ 40 minutes (depending on the model sizes) of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
- "\n",
- "Example:\n",
- "\n",
- "```\n",
- "Human: What is a car?\n",
- "Assistant: A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
- "```"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "eGP9sCq9sF0V"
- },
- "outputs": [],
- "source": [
- "# Loads an existing endpoint instance using the endpoint name:\n",
- "# - Using `endpoint_name = endpoint_quantized_vllm.name` allows us to get the\n",
- "# endpoint name of the endpoint `endpoint_quantized_vllm` created in the cell\n",
- "# above.\n",
- "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
- "# an existing endpoint with the ID 1234567890123456789.\n",
- "# You may uncomment the code below to load an existing endpoint.\n",
- "\n",
- "# endpoint_name = endpoint_quantized_vllm.name\n",
- "# # endpoint_name = \"\" # @param {type:\"string\"}\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint_quantized_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
- "\n",
- "\n",
- "# Overides max_length and top_k parameters during inferences.\n",
- "# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
- "# you can reduce the max length, such as set max_length as 20.\n",
- "instances = [\n",
- " {\n",
- " \"prompt\": \"What is a car?\",\n",
- " \"max_tokens\": 50,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- " },\n",
- "]\n",
- "response = endpoint_quantized_vllm.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "JmuUk3l1DoEo"
- },
- "source": [
- "## Evaluate PEFT-finetuned OpenLLaMA\n",
- "\n",
- "This section demonstrates how to evaluate the OpenLLaMA model fintuned with PEFT LoRA using EleutherAI's [Language Model Evaluation Harness (lm-evaluation-harness)](https://github.com/EleutherAI/lm-evaluation-harness) with Vertex CustomJob.\n",
- "\n",
- "This example uses the dataset [TruthfulQA](https://arxiv.org/abs/2109.07958). All supported tasks are listed in [this task table](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/docs/task_table.md)."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "gM4SXaquDoEo"
- },
- "outputs": [],
- "source": [
- "eval_dataset = \"truthfulqa_mc\" # @param {type:\"string\"}\n",
- "\n",
- "# Worker pool spec.\n",
- "# Sets L4 to evaluate open_llama_3b and open_llama_7b.\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "accelerator_count = 1\n",
- "\n",
- "# Sets L4 to evaluate open_llama_3b and open_llama_7b.\n",
- "# machine_type = \"g2-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 1\n",
- "\n",
- "# Sets 2 V100 to evaluate open_llama_13b.\n",
- "# machine_type = \"n1-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "# Sets 2 L4 to evaluate open_llama_13b.\n",
- "# machine_type = \"g2-standard-24\"\n",
- "# accelerator_type = \"NVIDIA_L4\"\n",
- "# accelerator_count = 2\n",
- "\n",
- "replica_count = 1\n",
- "\n",
- "# Setup evaluation job.\n",
- "job_name = create_name_with_datetime(prefix=\"openllama-peft-eval\")\n",
- "eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "g0t0RBixIw0P"
- },
- "outputs": [],
- "source": [
- "# Prepare evaluation command that runs the evaluation harness.\n",
- "# Set `use_accelerate = True` to enable evaluation across multiple GPUs.\n",
- "eval_command = [\n",
- " \"python\",\n",
- " \"main.py\",\n",
- " \"--model\",\n",
- " \"hf-causal-experimental\",\n",
- " \"--model_args\",\n",
- " f\"pretrained={merged_model_output_dir_gcsfuse},use_accelerate=True,device_map_option=auto\",\n",
- " \"--tasks\",\n",
- " f\"{eval_dataset}\",\n",
- " \"--output_path\",\n",
- " f\"{eval_output_dir_gcsfuse}\",\n",
- "]"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ItWB0WS__CX-"
- },
- "source": [
- "### Submit evaluation CustomJob"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "BbyIk99bDoEo"
- },
- "outputs": [],
- "source": [
- "# Pass evaluation arguments and launch job.\n",
- "worker_pool_specs = [\n",
- " {\n",
- " \"machine_spec\": {\n",
- " \"machine_type\": machine_type,\n",
- " \"accelerator_type\": accelerator_type,\n",
- " \"accelerator_count\": accelerator_count,\n",
- " },\n",
- " \"replica_count\": replica_count,\n",
- " \"disk_spec\": {\n",
- " \"boot_disk_size_gb\": 500,\n",
- " },\n",
- " \"container_spec\": {\n",
- " \"image_uri\": EVAL_DOCKER_URI,\n",
- " \"command\": eval_command,\n",
- " \"args\": [],\n",
- " },\n",
- " }\n",
- "]\n",
- "\n",
- "eval_job = aiplatform.CustomJob(\n",
- " display_name=job_name,\n",
- " worker_pool_specs=worker_pool_specs,\n",
- " base_output_dir=eval_output_dir,\n",
- ")\n",
- "\n",
- "eval_job.run()\n",
- "\n",
- "print(\"Evaluation results were saved in:\", eval_output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "kN0lE2iu_NXN"
- },
- "source": [
- "### Fetch and print evaluation results"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "927oRxoADoEp"
- },
- "outputs": [],
- "source": [
- "import json\n",
- "\n",
- "from google.cloud import storage\n",
- "\n",
- "# Fetch evaluation results.\n",
- "storage_client = storage.Client()\n",
- "BUCKET_NAME = BUCKET_URI.split(\"gs://\")[1]\n",
- "bucket = storage_client.get_bucket(BUCKET_NAME)\n",
- "RESULT_FILE_PATH = eval_output_dir[len(BUCKET_URI) + 1 :]\n",
- "blob = bucket.blob(RESULT_FILE_PATH)\n",
- "raw_result = blob.download_as_string()\n",
- "\n",
- "# Print evaluation results.\n",
- "result = json.loads(raw_result)\n",
- "result_formatted = json.dumps(result, indent=2)\n",
- "print(f\"Evaluation result:\\n{result_formatted}\")"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "af21a3cff1e0"
- },
- "source": [
- "## Clean up resources"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "911406c1561e"
- },
- "outputs": [],
- "source": [
- "# Delete custom train and evaluation jobs.\n",
- "train_job.delete()\n",
- "eval_job.delete()\n",
- "quantize_job.delete()\n",
- "\n",
- "# Undeploy models and delete endpoints.\n",
- "endpoint_without_peft.delete(force=True)\n",
- "endpoint_with_peft.delete(force=True)\n",
- "endpoint_quantized_vllm.delete(force=True)\n",
- "\n",
- "# Delete models.\n",
- "model_without_peft.delete()\n",
- "model_with_peft.delete()\n",
- "model_quantized_vllm.delete()"
- ]
- }
- ],
- "metadata": {
- "colab": {
- "name": "model_garden_pytorch_openllama_peft.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/notebooks/community/model_garden/model_garden_pytorch_peft.ipynb b/notebooks/community/model_garden/model_garden_pytorch_peft.ipynb
deleted file mode 100644
index 2d3d974ad..000000000
--- a/notebooks/community/model_garden/model_garden_pytorch_peft.ipynb
+++ /dev/null
@@ -1,1047 +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 - PEFT [DEPRECATED]\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 finetuning models with performance efficient finetuning libaries ([PEFT](https://github.com/huggingface/peft)), and running inferences with containers in Vertex AI. There are various models supported in [PEFT](https://github.com/huggingface/peft). This notebook shows examples with some models, such as OpenLLaMA, Falcon-instruct, BERT, RoBERTa-large, and XLM-RoBERTa-large, etc.\n",
- "\n",
- "### Objective\n",
- "\n",
- "- Finetune causal language models with PEFT and run inferences with Vertex AI, supporting\n",
- "\n",
- "| Models | LoRA |\n",
- "| :- | :- |\n",
- "| [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b) | Y |\n",
- "| [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b) | Y |\n",
- "| [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) | Y |\n",
- "\n",
- "- Finetune instruct models with PEFT and run inferences with Vertex AI, supporting\n",
- "\n",
- "| Models | LoRA |\n",
- "| :- | :- |\n",
- "| [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct) | Y |\n",
- "| [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) | Y |\n",
- "\n",
- "\n",
- "- Finetune sequence classification models with PEFT and run inferences with Vertex AI, supporting\n",
- "\n",
- "| Models | LoRA |\n",
- "| :- | :- |\n",
- "| [bert-base-uncase](https://huggingface.co/bert-base-uncased) | Y |\n",
- "| [RoBERTa-large](https://huggingface.co/roberta-large) | Y |\n",
- "| [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) | Y |\n",
- "\n",
- "### Costs\n",
- "\n",
- "This tutorial uses billable components of Google Cloud:\n",
- "\n",
- "* Vertex AI\n",
- "* Cloud Storage\n",
- "\n",
- "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "264c07757582"
- },
- "source": [
- "## Before you begin\n",
- "\n",
- "**NOTE**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ioensNKM8ned"
- },
- "source": [
- "### Colab only\n",
- "Run the following commands for Colab and skip this section if you are using Workbench."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "2707b02ef5df"
- },
- "outputs": [],
- "source": [
- "import sys\n",
- "\n",
- "if \"google.colab\" in sys.modules:\n",
- " ! pip3 install --upgrade google-cloud-aiplatform\n",
- " from google.colab import auth as google_auth\n",
- "\n",
- " google_auth.authenticate_user()\n",
- " # Install gdown for downloading example training images.\n",
- " ! pip3 install gdown\n",
- "\n",
- " # Restart the notebook kernel after installs.\n",
- " import IPython\n",
- "\n",
- " app = IPython.Application.instance()\n",
- " app.kernel.do_shutdown(True)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bb7adab99e41"
- },
- "source": [
- "### Setup Google Cloud project\n",
- "\n",
- "1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
- "\n",
- "1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
- "\n",
- "1. [Enable the Vertex AI API and Compute Engine API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,compute_component).\n",
- "\n",
- "1. [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
- "\n",
- "1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "6c460088b873"
- },
- "source": [
- "Set the following variables for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the specified region (`REGION`). Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\")."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "855d6b96f291"
- },
- "outputs": [],
- "source": [
- "# Cloud project id.\n",
- "PROJECT_ID = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# The region you want to launch jobs in.\n",
- "REGION = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# The Cloud Storage bucket for storing experiments output.\n",
- "BUCKET_URI = \"\" # @param {type:\"string\"}\n",
- "\n",
- "! gcloud config set project $PROJECT_ID\n",
- "\n",
- "import os\n",
- "\n",
- "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
- "EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
- "DATA_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"data\")\n",
- "MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
- "\n",
- "# The service account looks like:\n",
- "# '@.iam.gserviceaccount.com'\n",
- "# Please go to https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console\n",
- "# and create service account with `Vertex AI User` and `Storage Object Admin` roles.\n",
- "# The service account for deploying fine tuned model.\n",
- "SERVICE_ACCOUNT = \"\" # @param {type:\"string\"}"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e828eb320337"
- },
- "source": [
- "### Initialize Vertex AI API"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "12cd25839741"
- },
- "outputs": [],
- "source": [
- "from google.cloud import aiplatform\n",
- "\n",
- "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "2cc825514deb"
- },
- "source": [
- "### Define constants"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "b42bd4fa2b2d"
- },
- "outputs": [],
- "source": [
- "# The pre-built training and serving docker images.\n",
- "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231020_0936_RC00\"\n",
- "PREDICTION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:20231108_1540_RC00\""
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "0c250872074f"
- },
- "source": [
- "### Define common functions"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "354da31189dc"
- },
- "outputs": [],
- "source": [
- "import os\n",
- "from datetime import datetime\n",
- "\n",
- "from google.cloud import aiplatform\n",
- "\n",
- "\n",
- "def get_job_name_with_datetime(prefix: str):\n",
- " \"\"\"Gets the job name with date time when triggering training or deployment\n",
- " jobs in Vertex AI.\n",
- " \"\"\"\n",
- " return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
- "\n",
- "\n",
- "def deploy_model(\n",
- " model_name,\n",
- " model_id,\n",
- " publisher,\n",
- " publisher_model_id,\n",
- " finetuned_lora_model_path,\n",
- " service_account,\n",
- " task,\n",
- " machine_type=\"n1-standard-8\",\n",
- " accelerator_type=\"NVIDIA_TESLA_V100\",\n",
- "):\n",
- " \"\"\"Deploys trained models into Vertex AI.\"\"\"\n",
- " endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
- " serving_env = {\n",
- " \"MODEL_ID\": model_id,\n",
- " \"FINETUNED_LORA_MODEL_PATH\": finetuned_lora_model_path,\n",
- " \"TASK\": task,\n",
- " \"DEPLOY_SOURCE\": \"notebook\",\n",
- " }\n",
- " model = aiplatform.Model.upload(\n",
- " display_name=model_name,\n",
- " serving_container_image_uri=PREDICTION_DOCKER_URI,\n",
- " serving_container_ports=[7080],\n",
- " serving_container_predict_route=\"/predictions/peft_serving\",\n",
- " serving_container_health_route=\"/ping\",\n",
- " serving_container_environment_variables=serving_env,\n",
- " model_garden_source_model_name=(\n",
- " f\"publishers/{publisher}/models/{publisher_model_id}\",\n",
- " ),\n",
- " )\n",
- " model.deploy(\n",
- " endpoint=endpoint,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=1,\n",
- " deploy_request_timeout=1800,\n",
- " service_account=service_account,\n",
- " system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_peft.ipynb\"},\n",
- " )\n",
- " return model, endpoint"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "e70e3519ff8b"
- },
- "source": [
- "## Casual Language Modeling + PEFT\n",
- "\n",
- "This section demonstrates how to finetune OpenLLaMA with PEFT LoRA, including [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b)."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "iWGwJHqI7LMs"
- },
- "source": [
- "### Finetune with LoRA"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "5qCrm_kJH5cz"
- },
- "source": [
- "Set the base model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "N3UBLiYrM3sU"
- },
- "outputs": [],
- "source": [
- "model_id = \"openlm-research/open_llama_3b\" # @param [\"openlm-research/open_llama_3b\", \"openlm-research/open_llama_7b\", \"openlm-research/open_llama_13b\"]"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "KKEYoRfiHDVv"
- },
- "source": [
- "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
- "\n",
- "This example uses the dataset [Abirate/english_quotes](https://huggingface.co/datasets/Abirate/english_quotes). You can either use a [dataset from huggingface](https://huggingface.co/datasets) or a custom JSONL dataset in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) stored in Cloud Storage.\n",
- "\n",
- "In order to make the finetunig efficiently, we enabled quantization (8bits) when loading pretrained models for finetuning LoRA models. The peak GPU memory usages are ~7G, ~10G and ~16G for finetuning LoRA models for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) separately with default training parameters and the example dataset. In theory, open_llama_3b and open_llama_7b can be finetuned on 1 V100, and open_llama_13b can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "bd995e7aa529"
- },
- "source": [
- "#### [Optional] Finetune with a custom dataset\n",
- "\n",
- "To use a custom dataset, you should supply a `gs://` URI to a JSONL file in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) in the `dataset_name` below. The `template` parameter is optional.\n",
- "\n",
- "For example, here is one data point from the sample dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`:\n",
- "\n",
- "```json\n",
- "{\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
- "```\n",
- "\n",
- "To use this sample dataset that contains `input_text` and `output_text` fields, set `dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl` and `template` to `vertex_sample`. For advanced usage with custom datatset fields, see [the template example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json) and supply your own JSON template as `gs://` URIs."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "65467b361315"
- },
- "outputs": [],
- "source": [
- "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
- "dataset_name = \"Abirate/english_quotes\" # @param {type:\"string\"}\n",
- "# Optional. Template name or gs:// URI to a custom template.\n",
- "template = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# Worker pool spec.\n",
- "# machine_type = \"n1-standard-8\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "machine_type = \"a2-highgpu-1g\"\n",
- "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "replica_count = 1\n",
- "accelerator_count = 1\n",
- "\n",
- "# Add labels for the finetuning job.\n",
- "labels = {\n",
- " \"mg-source\": \"notebook\",\n",
- " \"mg-notebook-name\": \"model_garden_pytorch_peft.ipynb\".split(\".\")[0],\n",
- "}\n",
- "\n",
- "labels[\"mg-tune\"] = \"publishers-openlm-research-models-openllama\"\n",
- "versioned_model_id = model_id.split(\"/\")[1].replace(\"_\", \"-\")\n",
- "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
- "\n",
- "# Setup training job.\n",
- "job_name = get_job_name_with_datetime(\"openllama-lora-train\")\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- " labels=labels,\n",
- ")\n",
- "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Pass training arguments and launch job.\n",
- "train_job.run(\n",
- " args=[\n",
- " \"--task=causal-language-modeling-lora\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--dataset_name={dataset_name}\",\n",
- " f\"--output_dir={output_dir_gcsfuse}\",\n",
- " \"--lora_rank=16\",\n",
- " \"--lora_alpha=32\",\n",
- " \"--lora_dropout=0.05\",\n",
- " \"--warmup_steps=10\",\n",
- " \"--max_steps=10\",\n",
- " \"--learning_rate=2e-4\",\n",
- " f\"--template={template}\",\n",
- " ],\n",
- " replica_count=replica_count,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " boot_disk_size_gb=500,\n",
- ")\n",
- "\n",
- "print(\"Trained models were saved in: \", output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "jqmCtkGnhDmp"
- },
- "source": [
- "### Run inferences with serving images\n",
- "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
- "\n",
- "The model deployment step will take ~15 minutes to complete.\n",
- "\n",
- "The peak GPU memory usages for [openlm-research/open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b), [openlm-research/open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b), and [openlm-research/open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b) with LoRA weights are ~5.3G, ~8.7G and ~15.2G separately with the default settings. We use V100 in deployments for simplicity."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "bf55e38815dc"
- },
- "outputs": [],
- "source": [
- "model, endpoint = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"openllama-peft-serve\"),\n",
- " model_id=model_id,\n",
- " publisher=\"openlm-research\",\n",
- " publisher_model_id=\"openllama\",\n",
- " finetuned_lora_model_path=output_dir,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"causal-language-modeling-lora\",\n",
- ")\n",
- "print(\"endpoint_name:\", endpoint.name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "80b3fd2ace09"
- },
- "source": [
- "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed **after** the above model deployment step succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "4ab04da3ec9a"
- },
- "outputs": [],
- "source": [
- "# # Loads an existing endpoint as below.\n",
- "# endpoint_name = endpoint.name\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
- "instances = [\n",
- " {\n",
- " \"prompt\": \"Hi Google.\",\n",
- " \"max_tokens\": 50,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- " },\n",
- "]\n",
- "response = endpoint.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "af21a3cff1e0"
- },
- "source": [
- "### Clean up resources"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "911406c1561e"
- },
- "outputs": [],
- "source": [
- "# Delete custom train jobs.\n",
- "train_job.delete()\n",
- "\n",
- "# Undeploy model and delete endpoint.\n",
- "endpoint.delete(force=True)\n",
- "\n",
- "# Delete models.\n",
- "model.delete()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "W2USaHtQbE-l"
- },
- "source": [
- "## Instruct + PEFT\n",
- "\n",
- "This section demonstrates how to finetune Falcon-instruct with PEFT LoRA, including [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct)."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "PwluMqvYbLu9"
- },
- "source": [
- "### Finetune with LoRA"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "ap1pQd-ckAAb"
- },
- "source": [
- "Set the base model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "vVTfG2bpbR90"
- },
- "outputs": [],
- "source": [
- "model_id = \"tiiuae/falcon-7b-instruct\" # @param [\"tiiuae/falcon-7b-instruct\", \"tiiuae/falcon-40b-instruct\"]"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "iJe_us8mkZEV"
- },
- "source": [
- "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
- "\n",
- "This example uses the dataset [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco). You can either use a [dataset from huggingface](https://huggingface.co/datasets) or a custom JSONL dataset in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) stored in Cloud Storage. The `template` parameter is optional.\n",
- "\n",
- "The peak GPU memory usages are ~11G and ~34G for finetuning LoRA models for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) separately with default training parameters and the example dataset. In theory, falcon-7b-instruct can be finetuned on 1 P100/V100, and falcon-40b-instruct can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "4c128f63be2a"
- },
- "source": [
- "#### [Optional] Finetune with a custom dataset\n",
- "\n",
- "To use a custom dataset, you should supply a `gs://` URI to a JSONL file in [Vertex text model dataset format](https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-text-models-supervised#dataset-format) in the `dataset_name` below.\n",
- "\n",
- "For example, here is one data point from the sample dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`:\n",
- "\n",
- "```json\n",
- "{\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
- "```\n",
- "\n",
- "To use this sample dataset that contains `input_text` and `output_text` fields, set `dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl` and `template` to `vertex_sample`. For advanced usage with custom datatset fields, see [the template example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json) and supply your own JSON template as `gs://` URIs."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "-i4M7mWPbV8s"
- },
- "outputs": [],
- "source": [
- "# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
- "dataset_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
- "# Optional. Template name or gs:// URI to a custom template.\n",
- "template = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# Worker pool spec.\n",
- "machine_type = \"a2-highgpu-1g\"\n",
- "accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "replica_count = 1\n",
- "accelerator_count = 1\n",
- "\n",
- "# Setup training job.\n",
- "job_name = get_job_name_with_datetime(\"falcon-instruct-lora-train\")\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- ")\n",
- "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Pass training arguments and launch job.\n",
- "max_steps = 10\n",
- "train_job.run(\n",
- " args=[\n",
- " \"--task=instruct-lora\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--dataset_name={dataset_name}\",\n",
- " f\"--output_dir={output_dir_gcsfuse}\",\n",
- " \"--lora_rank=64\",\n",
- " \"--lora_alpha=16\",\n",
- " \"--lora_dropout=0.1\",\n",
- " \"--warmup_ratio=0.03\",\n",
- " f\"--max_steps={max_steps}\",\n",
- " \"--max_seq_length=512\",\n",
- " \"--learning_rate=2e-4\",\n",
- " f\"--template={template}\",\n",
- " ],\n",
- " replica_count=replica_count,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " boot_disk_size_gb=500,\n",
- ")\n",
- "\n",
- "print(\"Trained models were saved in: \", output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "oCgivsdllZ1x"
- },
- "source": [
- "### Run inferences with serving images"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "_lNIFvielcyV"
- },
- "source": [
- "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
- "\n",
- "The model deployment step will take ~15 minutes to complete.\n",
- "\n",
- "The peak GPU memory usages for [tiiuae/falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), and [tiiuae/falcon-40b-instruct](https://huggingface.co/tiiuae/falcon-40b-instruct) with LoRA weights are ~15.5G and ~38.2G separately with the default settings. We use V100 in deployments for simplicity."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "inrBXzoxk53t"
- },
- "outputs": [],
- "source": [
- "# # If deploy finetuned falcon-40b-instruct models, please set\n",
- "# machine_type = \"a2-highgpu-1g\",\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "\n",
- "model, endpoint = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"falcon-peft-serve\"),\n",
- " model_id=model_id,\n",
- " publisher=\"tiiuae\",\n",
- " publisher_model_id=\"falcon-instruct-7b-peft\",\n",
- " finetuned_lora_model_path=os.path.join(output_dir, \"checkpoint-\" + str(max_steps)),\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"instruct-lora\",\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- ")\n",
- "print(\"endpoint_name: \", endpoint.name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "_dmO3XooliAs"
- },
- "source": [
- "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
- "\n",
- "Example:\n",
- "\n",
- "```\n",
- "Human: What is a car?\n",
- "Assistant: A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
- "```"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "JrKsrffGl14T"
- },
- "outputs": [],
- "source": [
- "# # Loads an existing endpoint as below.\n",
- "# endpoint_name = endpoint.name\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
- "instances = [\n",
- " {\n",
- " \"prompt\": \"What is a car?\",\n",
- " \"max_tokens\": 50,\n",
- " \"temperature\": 1.0,\n",
- " \"top_p\": 1.0,\n",
- " \"top_k\": 10,\n",
- " },\n",
- "]\n",
- "response = endpoint.predict(instances=instances)\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "uk-t69nFl9rr"
- },
- "source": [
- "### Clean up resources"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "CcIX2WAJmAI7"
- },
- "outputs": [],
- "source": [
- "# Delete custom train jobs.\n",
- "train_job.delete()\n",
- "\n",
- "# Undeploy model and delete endpoint.\n",
- "endpoint.delete(force=True)\n",
- "\n",
- "# Delete models.\n",
- "model.delete()"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "og_s64QVmJDb"
- },
- "source": [
- "## Sequence Classification + PEFT\n",
- "\n",
- "This section demonstrates how to finetune sequence classification models with PEFT LoRA, including [bert-base-uncased](https://huggingface.co/bert-base-uncased), [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large)."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "QqJcvCb2mn8l"
- },
- "source": [
- "### Finetune with PEFT"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "c0E7dfqloeBB"
- },
- "source": [
- "Set the base model id."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "eru79ot6mtNG"
- },
- "outputs": [],
- "source": [
- "model_id = \"xlm-roberta-large\" # @param [\"bert-base-uncased\", \"roberta-large\", \"xlm-roberta-large\"]"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "y9n-qheaotS9"
- },
- "source": [
- "Use the Vertex AI SDK to create and run the custom training jobs with Vertex AI Model Garden training images.\n",
- "\n",
- "This example uses the dataset [glue](https://huggingface.co/datasets/glue).\n",
- "\n",
- "The peak GPU memory usages are ~8.8G and ~15G for finetuning LoRA models for [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) separately with default training parameters and the example dataset. In theory, RoBERTa-large can be finetuned on 1 P100/V100, and XLM-RoBERTa-large can be finetuned on 1 A100 (40G). We choose to use 1 A100 (40G) by default to support all these models in this notebook for simplicity."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "OHeSYEKinfP3"
- },
- "outputs": [],
- "source": [
- "dataset_name = \"glue\" # @param {type:\"string\"}\n",
- "\n",
- "# Worker pool spec.\n",
- "# # Please switch to A100 or other powerful machines if you run out of memory.\n",
- "# machine_type = \"a2-highgpu-1g\"\n",
- "# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
- "machine_type = \"n1-standard-8\"\n",
- "accelerator_type = \"NVIDIA_TESLA_V100\"\n",
- "replica_count = 1\n",
- "accelerator_count = 1\n",
- "\n",
- "# Setup training job.\n",
- "job_name = get_job_name_with_datetime(os.path.basename(model_id) + \"-lora-train\")\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- ")\n",
- "output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
- "output_dir_gcsfuse = output_dir.replace(\"gs://\", \"/gcs/\")\n",
- "\n",
- "# Pass training arguments and launch job.\n",
- "train_job.run(\n",
- " args=[\n",
- " \"--task=sequence-classification-lora\",\n",
- " f\"--pretrained_model_id={model_id}\",\n",
- " f\"--dataset_name={dataset_name}\",\n",
- " f\"--output_dir={output_dir_gcsfuse}\",\n",
- " \"--lora_rank=8\",\n",
- " \"--lora_alpha=16\",\n",
- " \"--lora_dropout=0.1\",\n",
- " \"--num_epochs=20\",\n",
- " \"--batch_size=32\",\n",
- " \"--learning_rate=3e-4\",\n",
- " ],\n",
- " replica_count=replica_count,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=accelerator_count,\n",
- " boot_disk_size_gb=500,\n",
- ")\n",
- "\n",
- "print(\"Trained models were saved in: \", output_dir)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "kJuHaBCqoacg"
- },
- "source": [
- "### Run inferences with serving images"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "jw1GFXydoft0"
- },
- "source": [
- "This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
- "\n",
- "The model deployment step will take ~15 minutes to complete.\n",
- "\n",
- "The peak GPU memory usages for [RoBERTa-large](https://huggingface.co/roberta-large), and [XLM-RoBERTa-large](https://huggingface.co/xlm-roberta-large) with LoRA weights are ~2G and ~2.4G separately with the default settings. We use V100 in deployments for simplicity."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "A5j0a3AZoj3h"
- },
- "outputs": [],
- "source": [
- "if model_id == \"bert-base-uncased\":\n",
- " publisher = \"google\"\n",
- " publisher_model_id = \"bert-base-uncased\"\n",
- "elif model_id == \"roberta-large\":\n",
- " publisher = \"meta\"\n",
- " publisher_model_id = \"roberta-large\"\n",
- "elif model_id == \"xlm-roberta-large\":\n",
- " publisher = \"meta\"\n",
- " publisher_model_id = \"xlm-roberta-large\"\n",
- "else:\n",
- " raise ValueError(f\"Unsupported model id: {model_id}\")\n",
- "\n",
- "model, endpoint = deploy_model(\n",
- " model_name=get_job_name_with_datetime(prefix=\"sequence-classification-peft-serve\"),\n",
- " model_id=model_id,\n",
- " publisher=publisher,\n",
- " publisher_model_id=publisher_model_id,\n",
- " finetuned_lora_model_path=output_dir,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " task=\"sequence-classification-lora\",\n",
- ")\n",
- "print(\"endpoint_name:\", endpoint.name)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "6dR3-Hc7oqu-"
- },
- "source": [
- "NOTE: The model weights will be downloaded after the deployment succeeds. Thus additional 5 minutes of waiting time is needed after the above model deployment step succeeds and before you run the next step below. Otherwise you might see a ServiceUnavailable: 503 502:Bad Gateway error when you send requests to the endpoint.\n",
- "\n",
- "Once deployment succeeds, you can send requests to the endpoint with text prompts."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "xAmULPegoyXl"
- },
- "outputs": [],
- "source": [
- "# # Loads an existing endpoint as below.\n",
- "# endpoint_name = endpoint.name\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
- "instances = [\n",
- " {\"prompt\": \"The cat sat on the mat.\"},\n",
- "]\n",
- "response = endpoint.predict(instances=instances)\n",
- "labels = [int(item) for item in response.predictions]\n",
- "print(labels)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "1Vg_b46Jo6C1"
- },
- "source": [
- "### Clean up resources"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "id": "vNiTDL_VpAlp"
- },
- "outputs": [],
- "source": [
- "# Delete custom train jobs.\n",
- "train_job.delete()\n",
- "\n",
- "# Undeploy model and delete endpoint.\n",
- "endpoint.delete(force=True)\n",
- "\n",
- "# Delete models.\n",
- "model.delete()"
- ]
- }
- ],
- "metadata": {
- "colab": {
- "name": "model_garden_pytorch_peft.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb b/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb
deleted file mode 100644
index 24bf834a8..000000000
--- a/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb
+++ /dev/null
@@ -1,1050 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "7d9bbf86da5e"
- },
- "outputs": [],
- "source": [
- "# Copyright 2026 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 - Qwen 2.5 Finetuning\n",
- "\n",
- "\n",
- " \n",
- " \n",
- " ![\"Google](\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\")
Run in Colab Enterprise\n",
- " \n",
- " | \n",
- " \n",
- " \n",
- "
View on GitHub\n",
- " \n",
- " | \n",
- "
"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "3de7470326a2"
- },
- "source": [
- "## Overview\n",
- "\n",
- "This notebook demonstrates finetuning and deploying Qwen 2.5 models with Vertex AI. All of the examples in this notebook use parameter efficient finetuning methods [PEFT (LoRA)](https://github.com/huggingface/peft) to reduce training and storage costs. LoRA (Low-Rank Adaptation) is one approach of Parameter Efficient FineTuning (PEFT), 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).\n",
- "\n",
- "After finetuning, we can deploy models on Vertex with GPU.\n",
- "\n",
- "\n",
- "### Objective\n",
- "\n",
- "- Finetune Qwen 2.5 models with Vertex AI Custom Training Jobs.\n",
- "- Deploy finetuned Qwen 2.5 models on Vertex AI Prediction.\n",
- "- Send prediction requests to your finetuned Qwen 2.5 models.\n",
- "\n",
- "### File a bug\n",
- "\n",
- "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
- "\n",
- "### Costs\n",
- "\n",
- "This tutorial uses billable components of Google Cloud:\n",
- "\n",
- "* Vertex AI\n",
- "* Cloud Storage\n",
- "\n",
- "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "264c07757582"
- },
- "source": [
- "## Before you begin"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "855d6b96f291"
- },
- "outputs": [],
- "source": [
- "# @title Install Python Packages for Finetuning\n",
- "\n",
- "# @markdown 1. Install packages to validate dataset with template.\n",
- "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
- "! pip install --upgrade --quiet accelerate==0.34.2\n",
- "! pip install --upgrade --quiet transformers==4.47.1\n",
- "! pip install --upgrade --quiet datasets==2.20.0"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "50273xHFJi5T"
- },
- "outputs": [],
- "source": [
- "# @title Setup Google Cloud project\n",
- "\n",
- "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
- "\n",
- "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
- "\n",
- "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
- "\n",
- "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
- "# @markdown | ----------- | ----------- | ----------- |\n",
- "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
- "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
- "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
- "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
- "\n",
- "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
- "\n",
- "BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n",
- "\n",
- "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
- "\n",
- "REGION = \"\" # @param {type:\"string\"}\n",
- "\n",
- "# Import the necessary packages\n",
- "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
- "! cd vertex-ai-samples && git reset --hard 7ae13b346a72ee2a2dc8152dd40c6ddd72d6c810\n",
- "\n",
- "import datetime\n",
- "import importlib\n",
- "import os\n",
- "import uuid\n",
- "from typing import Tuple\n",
- "\n",
- "from google.cloud import aiplatform\n",
- "from google.cloud.aiplatform.compat.types import \\\n",
- " custom_job as gca_custom_job_compat\n",
- "\n",
- "common_util = importlib.import_module(\n",
- " \"vertex-ai-samples.notebooks.community.model_garden.docker_source_codes.notebook_util.common_util\"\n",
- ")\n",
- "\n",
- "models, endpoints = {}, {}\n",
- "\n",
- "# Get the default cloud project id.\n",
- "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
- "\n",
- "# Get the default region for launching jobs.\n",
- "if not REGION:\n",
- " if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
- " raise ValueError(\n",
- " \"REGION must be set. See\"\n",
- " \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
- " \" available cloud locations.\"\n",
- " )\n",
- " REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
- "\n",
- "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
- "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
- "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
- "\n",
- "# Cloud Storage bucket for storing the experiment artifacts.\n",
- "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
- "# prefer using your own GCS bucket, change the value yourself below.\n",
- "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
- "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
- "\n",
- "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
- " BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
- " BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
- " ! gsutil mb -l {REGION} {BUCKET_URI}\n",
- "else:\n",
- " assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
- " shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
- " bucket_region = shell_output[0].strip().lower()\n",
- " if bucket_region != REGION:\n",
- " raise ValueError(\n",
- " \"Bucket region %s is different from notebook region %s\"\n",
- " % (bucket_region, REGION)\n",
- " )\n",
- "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
- "\n",
- "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
- "MODEL_BUCKET = os.path.join(BUCKET_URI, \"qwen2_5\")\n",
- "\n",
- "\n",
- "# Initialize Vertex AI API.\n",
- "print(\"Initializing Vertex AI API.\")\n",
- "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
- "\n",
- "# Gets the default SERVICE_ACCOUNT.\n",
- "shell_output = ! gcloud projects describe $PROJECT_ID\n",
- "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
- "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
- "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
- "\n",
- "\n",
- "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
- "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
- "\n",
- "! gcloud config set project $PROJECT_ID\n",
- "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
- "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "36c21f10355f"
- },
- "outputs": [],
- "source": [
- "# @title Access Qwen 2.5 models\n",
- "\n",
- "# @markdown ### Access Qwen 2.5 models on Hugging Face for GPU based finetuning and serving\n",
- "# @markdown You must provide a Hugging Face User Access Token (with read access) to access the Qwen 2.5 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
- "\n",
- "HF_TOKEN = \"\" # @param {type:\"string\", isTemplate:true}\n",
- "assert (\n",
- " HF_TOKEN\n",
- "), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
- "# @markdown ---"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "cb56d402e84a"
- },
- "source": [
- "## Finetune with HuggingFace PEFT and deploy with vLLM on GPUs"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "KwAW99YZHTdy"
- },
- "outputs": [],
- "source": [
- "# @title Set dataset\n",
- "\n",
- "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
- "\n",
- "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
- "# @markdown You can set `train_dataset` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `train_column` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `train_column` to `text` in this notebook.\n",
- "\n",
- "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
- "\n",
- "# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
- "# @markdown ```\n",
- "# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
- "# @markdown ```\n",
- "\n",
- "# @markdown The JSON object has a key `text`, which should match `train_column`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
- "\n",
- "# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
- "\n",
- "# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
- "\n",
- "# @markdown Optionally update the `train_column` field below if your JSON objects use a key other than the default `text`.\n",
- "\n",
- "# @markdown ### (Optional) Format your data with custom JSON template\n",
- "\n",
- "# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
- "\n",
- "# @markdown ```\n",
- "# @markdown {\n",
- "# @markdown \"description\": \"Template used by Llama 3.1, accepting text-bison format.\",\n",
- "# @markdown \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
- "# @markdown \"prompt_input\": \"<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
- "# @markdown \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
- "# @markdown \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
- "# @markdown }\n",
- "# @markdown ```\n",
- "\n",
- "# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
- "\n",
- "# @markdown ```\n",
- "# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
- "# @markdown ```\n",
- "\n",
- "# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
- "# @markdown\n",
- "# @markdown To try such custom dataset, you can make the following changes:\n",
- "# @markdown 1. Set `template` to `llama3-text-bison`\n",
- "# @markdown 1. Set `train_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
- "# @markdown 1. Set `train_split` to `train`\n",
- "# @markdown 1. Set `eval_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
- "# @markdown 1. Set `eval_split` to `train` (**NOT** `test`)\n",
- "# @markdown 1. Set `train_column` as `input_text`.\n",
- "\n",
- "# Template name or gs:// URI to a custom template.\n",
- "template = \"openassistant-guanaco\" # @param {type:\"string\"}\n",
- "\n",
- "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
- "train_dataset = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
- "train_split = \"train\" # @param {type:\"string\"}\n",
- "eval_dataset = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
- "eval_split = \"test\" # @param {type:\"string\"}\n",
- "\n",
- "# Name of the dataset column containing training text input.\n",
- "train_column = \"text\" # @param {type:\"string\"}\n",
- "# Maximum sequence length.\n",
- "max_seq_length = 4096 # @param{type:\"integer\"}"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "oHvYSPr7JdZq"
- },
- "outputs": [],
- "source": [
- "# @title Set model\n",
- "\n",
- "# @markdown Select a model variant of Qwen 2.5.\n",
- "base_model_id = \"Qwen/Qwen2.5-7B-Instruct\" # @param [\"Qwen/Qwen2.5-7B-Instruct\", \"Qwen/Qwen2.5-14B-Instruct\", \"Qwen/Qwen2.5-32B-Instruct\"] {isTemplate:true}\n",
- "pretrained_model_id = base_model_id"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "X3N8O9_0JdZq"
- },
- "outputs": [],
- "source": [
- "# @title Validate Dataset with Template\n",
- "\n",
- "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
- "\n",
- "dataset_validation_util = importlib.import_module(\n",
- " \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
- ")\n",
- "\n",
- "if dataset_validation_util.is_gcs_path(pretrained_model_id):\n",
- " # Download tokenizer.\n",
- " ! mkdir tokenizer\n",
- " ! gsutil cp {pretrained_model_id}/tokenizer.json ./tokenizer\n",
- " ! gsutil cp {pretrained_model_id}/config.json ./tokenizer\n",
- " tokenizer_path = \"./tokenizer\"\n",
- " access_token = \"\"\n",
- "else:\n",
- " tokenizer_path = pretrained_model_id\n",
- " access_token = HF_TOKEN\n",
- "\n",
- "tokenizer = dataset_validation_util.load_tokenizer(tokenizer_path, None, access_token)\n",
- "\n",
- "# Validate the train dataset.\n",
- "dataset_validation_util.validate_dataset_with_template(\n",
- " dataset_name=train_dataset,\n",
- " split=train_split,\n",
- " input_column=train_column,\n",
- " template=template,\n",
- " max_seq_length=max_seq_length,\n",
- " use_multiprocessing=False,\n",
- " tokenizer=tokenizer,\n",
- ")\n",
- "\n",
- "# Validate the eval dataset if it exists.\n",
- "if eval_dataset:\n",
- " dataset_validation_util.validate_dataset_with_template(\n",
- " dataset_name=eval_dataset,\n",
- " split=eval_split,\n",
- " input_column=train_column,\n",
- " template=template,\n",
- " max_seq_length=max_seq_length,\n",
- " use_multiprocessing=False,\n",
- " tokenizer=tokenizer,\n",
- " )"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "VvQLIpjhJdZq"
- },
- "outputs": [],
- "source": [
- "# @title Finetune\n",
- "\n",
- "# @markdown This section demonstrates how to finetune the Qwen 2.5 model on Vertex AI. It uses the Vertex AI SDK to create and run the custom training jobs.\n",
- "\n",
- "# @markdown The training job takes approximately between 10 to 20 mins to set-up. Once done, the training job is expected to take around 2.5 hours with the default configurations. To find the training time, throughput, and memory usage of your training job, you can go to the training logs and check the log line of the last training epoch.\n",
- "\n",
- "# @markdown **Note**:\n",
- "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
- "# @markdown 1. If `max_steps > 0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
- "\n",
- "# @markdown Acceletor type to use for training.\n",
- "# fmt: off\n",
- "training_accelerator_type = \"NVIDIA_A100_80GB\" # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
- "# fmt: on\n",
- "\n",
- "# The pre-built training docker image.\n",
- "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
- " repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
- " is_restricted_image = True\n",
- " is_dynamic_workload_scheduler = False\n",
- " dws_kwargs = {}\n",
- "else:\n",
- " repo = \"us-docker.pkg.dev/vertex-ai\"\n",
- " is_restricted_image = False\n",
- " is_dynamic_workload_scheduler = True\n",
- " dws_kwargs = {\n",
- " \"max_wait_duration\": 1800, # 30 minutes\n",
- " \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
- " }\n",
- "\n",
- "TRAIN_DOCKER_URI = (\n",
- " f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250705\"\n",
- ")\n",
- "\n",
- "# Worker pool spec.\n",
- "boot_disk_size_gb = 500\n",
- "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
- " per_node_accelerator_count = 8\n",
- " training_machine_type = \"a2-ultragpu-8g\"\n",
- "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
- " per_node_accelerator_count = 8\n",
- " training_machine_type = \"a3-highgpu-8g\"\n",
- "else:\n",
- " raise ValueError(\n",
- " f\"Recommended machine settings not found for: {training_accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` by clicking `Show Code` and then modifying the code.\"\n",
- " )\n",
- "\n",
- "# The number of nodes to use for this worker pool in distributed training.\n",
- "replica_count = 1\n",
- "\n",
- "# Set config file.\n",
- "config_file = \"vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml\"\n",
- "\n",
- "# @markdown Batch size for finetuning.\n",
- "per_device_train_batch_size = 1 # @param{type:\"integer\"}\n",
- "# @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
- "gradient_accumulation_steps = 4 # @param{type:\"integer\"}\n",
- "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
- "max_steps = -1 # @param{type:\"integer\"}\n",
- "num_train_epochs = 1.0 # @param{type:\"number\"}\n",
- "# @markdown Precision mode for finetuning.\n",
- "finetuning_precision_mode = \"4bit\" # @param [\"4bit\", \"float16\"]\n",
- "# @markdown Learning rate.\n",
- "learning_rate = 5e-5 # @param{type:\"number\"}\n",
- "# @markdown The scheduler type to use.\n",
- "lr_scheduler_type = \"cosine\" # @param{type:\"string\"}\n",
- "# @markdown LoRA parameters.\n",
- "lora_rank = 16 # @param{type:\"integer\"}\n",
- "lora_alpha = 32 # @param{type:\"integer\"}\n",
- "lora_dropout = 0.05 # @param{type:\"number\"}\n",
- "# Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
- "gradient_checkpointing = True\n",
- "# Attention implementation to use in the model.\n",
- "attn_implementation = \"flash_attention_2\"\n",
- "# The optimizer for which to schedule the learning rate.\n",
- "optimizer = \"adamw_torch\"\n",
- "# Define the proportion of training to be dedicated to a linear warmup where learning rate gradually increases.\n",
- "warmup_ratio = \"0.01\"\n",
- "# The list or string of integrations to report the results and logs to.\n",
- "report_to = \"tensorboard\"\n",
- "# Number of updates steps before two checkpoint saves.\n",
- "save_steps = 10\n",
- "# Number of update steps between two logs.\n",
- "logging_steps = save_steps\n",
- "\n",
- "# @markdown Evaluation metrics to compute. Supported eval metrics: loss, perplexity, bleu, google_bleu, rouge1, rouge2, rougeL, rougeLsum.\n",
- "eval_metric_name = \"loss,perplexity,bleu\" # @param{type:\"string\"}\n",
- "# @markdown Metric to use for best model selection. This will save the best checkpoint based on the eval metric.\n",
- "metric_for_best_model = \"perplexity\" # @param{type:\"string\"}\n",
- "\n",
- "common_util.check_quota(\n",
- " project_id=PROJECT_ID,\n",
- " region=REGION,\n",
- " accelerator_type=training_accelerator_type,\n",
- " accelerator_count=per_node_accelerator_count * replica_count,\n",
- " is_for_training=True,\n",
- " is_restricted_image=is_restricted_image,\n",
- " is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
- ")\n",
- "\n",
- "job_name = common_util.get_job_name_with_datetime(\"qwen2_5-lora-train\")\n",
- "\n",
- "base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
- "# Create a GCS folder to store the LORA adapter.\n",
- "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
- "# Create a GCS folder to store the merged model with the base model and the\n",
- "# finetuned LORA adapter.\n",
- "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
- "\n",
- "# Add labels for the finetuning job.\n",
- "labels = {\n",
- " \"mg-source\": \"notebook\",\n",
- " \"mg-notebook-name\": \"model_garden_pytorch_qwen2_5_finetuning.ipynb\".split(\".\")[0],\n",
- "}\n",
- "\n",
- "labels[\"mg-tune\"] = \"publishers-qwen-models-qwen2-5\"\n",
- "versioned_model_id = base_model_id.split(\"/\")[1].lower().replace(\".\", \"-\")\n",
- "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
- "\n",
- "eval_args = [\n",
- " f\"--eval_dataset={eval_dataset}\",\n",
- " f\"--eval_column={train_column}\",\n",
- " f\"--eval_template={template}\",\n",
- " f\"--eval_split={eval_split}\",\n",
- " f\"--eval_steps={save_steps}\",\n",
- " f\"--eval_metric_name={eval_metric_name}\",\n",
- " f\"--metric_for_best_model={metric_for_best_model}\",\n",
- "]\n",
- "\n",
- "train_job_args = [\n",
- " f\"--config_file={config_file}\",\n",
- " \"--task=instruct-lora\",\n",
- " \"--input_masking=True\",\n",
- " f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
- " f\"--train_dataset={train_dataset}\",\n",
- " f\"--train_split={train_split}\",\n",
- " f\"--train_column={train_column}\",\n",
- " f\"--output_dir={lora_output_dir}\",\n",
- " f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
- " f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
- " f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
- " f\"--lora_rank={lora_rank}\",\n",
- " f\"--lora_alpha={lora_alpha}\",\n",
- " f\"--lora_dropout={lora_dropout}\",\n",
- " f\"--max_steps={max_steps}\",\n",
- " f\"--max_seq_length={max_seq_length}\",\n",
- " f\"--learning_rate={learning_rate}\",\n",
- " f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
- " f\"--precision_mode={finetuning_precision_mode}\",\n",
- " f\"--gradient_checkpointing={gradient_checkpointing}\",\n",
- " f\"--num_train_epochs={num_train_epochs}\",\n",
- " f\"--attn_implementation={attn_implementation}\",\n",
- " f\"--optimizer={optimizer}\",\n",
- " f\"--warmup_ratio={warmup_ratio}\",\n",
- " f\"--report_to={report_to}\",\n",
- " f\"--logging_output_dir={base_output_dir}\",\n",
- " f\"--save_steps={save_steps}\",\n",
- " f\"--logging_steps={logging_steps}\",\n",
- " f\"--train_template={template}\",\n",
- " f\"--huggingface_access_token={HF_TOKEN}\",\n",
- "] + eval_args\n",
- "\n",
- "# Pass training arguments and launch job.\n",
- "train_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=job_name,\n",
- " container_uri=TRAIN_DOCKER_URI,\n",
- " labels=labels,\n",
- ")\n",
- "\n",
- "print(\"Running training job with args:\")\n",
- "print(\" \\\\\\n\".join(train_job_args))\n",
- "train_job.run(\n",
- " args=train_job_args,\n",
- " replica_count=replica_count,\n",
- " machine_type=training_machine_type,\n",
- " accelerator_type=training_accelerator_type,\n",
- " accelerator_count=per_node_accelerator_count,\n",
- " boot_disk_size_gb=boot_disk_size_gb,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " base_output_dir=base_output_dir,\n",
- " sync=False, # Non-blocking call to run.\n",
- " **dws_kwargs,\n",
- ")\n",
- "\n",
- "# Wait until resource has been created.\n",
- "train_job.wait_for_resource_creation()\n",
- "\n",
- "print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
- "print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
- "\n",
- "# @markdown Click \"Show Code\" to see more details."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "6jQZVMzcJdZq"
- },
- "outputs": [],
- "source": [
- "# @title Run TensorBoard\n",
- "# @markdown This section shows how to launch TensorBoard in a [Cloud Shell](https://cloud.google.com/shell/docs).\n",
- "# @markdown 1. Click the Cloud Shell icon() on the top right to open the Cloud Shell.\n",
- "# @markdown 2. Copy the `tensorboard` command shown below by running this cell.\n",
- "# @markdown 3. Paste and run the command in the Cloud Shell to launch TensorBoard.\n",
- "# @markdown 4. Once the command runs (You may have to click `Authorize` if prompted), click the link starting with `http://localhost`.\n",
- "\n",
- "# @markdown Note: You may need to wait around 10 minutes after the job starts in order for the TensorBoard logs to be written to the GCS bucket.\n",
- "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "Ys3FzqDaJdZq"
- },
- "outputs": [],
- "source": [
- "# @title Select Evaluation Checkpoint\n",
- "\n",
- "if train_job.end_time is None:\n",
- " print(\"Waiting for the training job to finish...\")\n",
- " train_job.wait()\n",
- " print(\"The training job has finished.\")\n",
- "\n",
- "# @markdown The following checkpoints are available for evaluation:\n",
- "! gcloud storage ls \"{lora_output_dir}/node-0\" | grep \"checkpoint-\""
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "dztN2-JdJdZq"
- },
- "outputs": [],
- "source": [
- "# @title Run Evaluation Job\n",
- "# @markdown This section runs the evaluation using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) on the finetuned model. The evaluation takes approximately 20 mins to finish.\n",
- "\n",
- "# The pre-built evaluation docker image for LM Evaluation Harness.\n",
- "LM_EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20250410_1035_RC00\"\n",
- "\n",
- "# @markdown Set `RUN_EVALUATION` to False to skip the evaluation job.\n",
- "RUN_EVALUATION = True # @param {type:\"boolean\"}\n",
- "\n",
- "eval_accelerator_type = \"NVIDIA_L4\"\n",
- "gpu_memory_utilization = 0.85\n",
- "\n",
- "if \"7B\" in base_model_id or \"14B\" in base_model_id:\n",
- " eval_machine_type = \"g2-standard-48\"\n",
- " eval_accelerator_count = 4\n",
- "elif \"32B\" in base_model_id:\n",
- " eval_machine_type = \"g2-standard-96\"\n",
- " eval_accelerator_count = 8\n",
- "else:\n",
- " raise ValueError(\n",
- " \"Recommended machine settings not found for model: %s\" % base_model_id\n",
- " )\n",
- "\n",
- "# @markdown Set `evaluation_checkpoint_dir` to an intermediate checkpoint from the above training job. If not set, the evaluation job will use the merged model.\n",
- "evaluation_checkpoint_dir = \"\" # @param {type:\"string\"}\n",
- "if evaluation_checkpoint_dir:\n",
- " pretrained = pretrained_model_id\n",
- "else:\n",
- " pretrained = merged_model_output_dir\n",
- "\n",
- "# @markdown Evaluation tasks to run.\n",
- "eval_tasks = \"coqa\" # @param {type:\"string\"}\n",
- "# @markdown Model to use for evaluation.\n",
- "model = \"vllm\" # @param {type:\"string\"}\n",
- "# @markdown Batch size for evaluation.\n",
- "batch_size = \"auto\" # @param {type:\"string\"}\n",
- "apply_chat_template = True\n",
- "max_model_len = 4096 # Maximum context length.\n",
- "\n",
- "model_args = f\"tensor_parallel_size={eval_accelerator_count},max_model_len={max_model_len},gpu_memory_utilization={gpu_memory_utilization},enforce_eager=True\"\n",
- "eval_output_dir = os.path.join(base_output_dir, \"lm_eval\")\n",
- "\n",
- "lm_eval_job_args = [\n",
- " \"--task=lm_eval\",\n",
- " f\"--model={model}\",\n",
- " f\"--eval_tasks={eval_tasks}\",\n",
- " f\"--pretrained_model_name_or_path={pretrained}\",\n",
- " f\"--model_args={model_args}\",\n",
- " f\"--output_dir={eval_output_dir}\",\n",
- " f\"--apply_chat_template={apply_chat_template}\",\n",
- " f\"--batch_size={batch_size}\",\n",
- " f\"--huggingface_access_token={HF_TOKEN}\",\n",
- "]\n",
- "\n",
- "if evaluation_checkpoint_dir:\n",
- " lm_eval_job_args.append(f\"--lora_path={evaluation_checkpoint_dir}\")\n",
- "\n",
- "if RUN_EVALUATION:\n",
- " common_util.check_quota(\n",
- " project_id=PROJECT_ID,\n",
- " region=REGION,\n",
- " accelerator_type=eval_accelerator_type,\n",
- " accelerator_count=eval_accelerator_count,\n",
- " is_for_training=True,\n",
- " )\n",
- " lm_eval_job = aiplatform.CustomContainerTrainingJob(\n",
- " display_name=common_util.get_job_name_with_datetime(\"qwen2_5-lm-eval\"),\n",
- " container_uri=LM_EVAL_DOCKER_URI,\n",
- " labels=labels,\n",
- " )\n",
- "\n",
- " print(\"Running evaluation job with args:\")\n",
- " print(\" \\\\\\n\".join(lm_eval_job_args))\n",
- " lm_eval_job.run(\n",
- " args=lm_eval_job_args,\n",
- " replica_count=1,\n",
- " machine_type=eval_machine_type,\n",
- " accelerator_type=eval_accelerator_type,\n",
- " accelerator_count=eval_accelerator_count,\n",
- " boot_disk_size_gb=boot_disk_size_gb,\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " base_output_dir=base_output_dir,\n",
- " )\n",
- "\n",
- "# @markdown Click \"Show Code\" to see more details."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "qWpHMUt3JdZq"
- },
- "outputs": [],
- "source": [
- "# @title Deploy\n",
- "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
- "\n",
- "print(\"Deploying models in:\", merged_model_output_dir)\n",
- "\n",
- "# The pre-built serving docker image for vLLM.\n",
- "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
- "\n",
- "# @markdown Set `use_dedicated_endpoint` to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
- "use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
- "\n",
- "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
- "accelerator_type = \"NVIDIA_L4\"\n",
- "if \"7b\" in base_model_id.lower():\n",
- " machine_type = \"g2-standard-24\"\n",
- " per_node_accelerator_count = 2\n",
- "elif \"14b\" in base_model_id.lower() or \"32b\" in base_model_id.lower():\n",
- " machine_type = \"g2-standard-48\"\n",
- " per_node_accelerator_count = 4\n",
- "else:\n",
- " raise ValueError(f\"Unsupported model ID: {base_model_id}.\")\n",
- "\n",
- "common_util.check_quota(\n",
- " project_id=PROJECT_ID,\n",
- " region=REGION,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=per_node_accelerator_count,\n",
- " is_for_training=False,\n",
- ")\n",
- "\n",
- "gpu_memory_utilization = 0.95\n",
- "max_model_len = 8192 # Maximum context length.\n",
- "\n",
- "# Ensure max_model_len does not exceed the limit\n",
- "if max_model_len > 8192:\n",
- " raise ValueError(\"max_model_len cannot exceed 8192\")\n",
- "\n",
- "\n",
- "def deploy_model_vllm(\n",
- " model_name: str,\n",
- " model_id: str,\n",
- " publisher: str,\n",
- " publisher_model_id: str,\n",
- " service_account: str,\n",
- " base_model_id: str = None,\n",
- " machine_type: str = \"g2-standard-8\",\n",
- " accelerator_type: str = \"NVIDIA_L4\",\n",
- " accelerator_count: int = 1,\n",
- " gpu_memory_utilization: float = 0.9,\n",
- " max_model_len: int = 4096,\n",
- " dtype: str = \"auto\",\n",
- " enable_trust_remote_code: bool = False,\n",
- " enforce_eager: bool = False,\n",
- " enable_lora: bool = False,\n",
- " enable_chunked_prefill: bool = False,\n",
- " enable_prefix_cache: bool = False,\n",
- " host_prefix_kv_cache_utilization_target: float = 0.0,\n",
- " max_loras: int = 1,\n",
- " max_cpu_loras: int = 8,\n",
- " use_dedicated_endpoint: bool = False,\n",
- " max_num_seqs: int = 256,\n",
- " model_type: str = None,\n",
- " enable_llama_tool_parser: bool = False,\n",
- ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
- " \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
- " endpoint = aiplatform.Endpoint.create(\n",
- " display_name=f\"{model_name}-endpoint\",\n",
- " dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
- " )\n",
- "\n",
- " if not base_model_id:\n",
- " base_model_id = model_id\n",
- "\n",
- " # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
- " vllm_args = [\n",
- " \"python\",\n",
- " \"-m\",\n",
- " \"vllm.entrypoints.api_server\",\n",
- " \"--host=0.0.0.0\",\n",
- " \"--port=8080\",\n",
- " f\"--model={model_id}\",\n",
- " f\"--tensor-parallel-size={accelerator_count}\",\n",
- " \"--swap-space=16\",\n",
- " f\"--max-model-len={max_model_len}\",\n",
- " f\"--dtype={dtype}\",\n",
- " f\"--max-loras={max_loras}\",\n",
- " f\"--max-cpu-loras={max_cpu_loras}\",\n",
- " f\"--max-num-seqs={max_num_seqs}\",\n",
- " \"--disable-log-stats\",\n",
- " ]\n",
- "\n",
- " if gpu_memory_utilization:\n",
- " vllm_args.append(f\"--gpu-memory-utilization={gpu_memory_utilization}\")\n",
- "\n",
- " if enable_trust_remote_code:\n",
- " vllm_args.append(\"--trust-remote-code\")\n",
- "\n",
- " if enforce_eager:\n",
- " vllm_args.append(\"--enforce-eager\")\n",
- "\n",
- " if enable_lora:\n",
- " vllm_args.append(\"--enable-lora\")\n",
- "\n",
- " if enable_chunked_prefill:\n",
- " vllm_args.append(\"--enable-chunked-prefill\")\n",
- "\n",
- " if enable_prefix_cache:\n",
- " vllm_args.append(\"--enable-prefix-caching\")\n",
- "\n",
- " if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
- " vllm_args.append(\n",
- " f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
- " )\n",
- "\n",
- " if model_type:\n",
- " vllm_args.append(f\"--model-type={model_type}\")\n",
- "\n",
- " if enable_llama_tool_parser:\n",
- " vllm_args.append(\"--enable-auto-tool-choice\")\n",
- " vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
- "\n",
- " env_vars = {\n",
- " \"MODEL_ID\": base_model_id,\n",
- " \"DEPLOY_SOURCE\": \"notebook\",\n",
- " }\n",
- "\n",
- " # HF_TOKEN is not a compulsory field and may not be defined.\n",
- " try:\n",
- " if HF_TOKEN:\n",
- " env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
- " except NameError:\n",
- " pass\n",
- "\n",
- " model = aiplatform.Model.upload(\n",
- " display_name=model_name,\n",
- " serving_container_image_uri=VLLM_DOCKER_URI,\n",
- " serving_container_args=vllm_args,\n",
- " serving_container_ports=[8080],\n",
- " serving_container_predict_route=\"/generate\",\n",
- " serving_container_health_route=\"/ping\",\n",
- " serving_container_environment_variables=env_vars,\n",
- " serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
- " serving_container_deployment_timeout=7200,\n",
- " model_garden_source_model_name=(\n",
- " f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
- " ),\n",
- " )\n",
- " print(\n",
- " f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\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",
- " system_labels={\n",
- " \"NOTEBOOK_NAME\": \"model_garden_pytorch_qwen2_5_finetuning.ipynb\",\n",
- " \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
- " },\n",
- " )\n",
- " print(\"endpoint_name:\", endpoint.name)\n",
- "\n",
- " return model, endpoint\n",
- "\n",
- "\n",
- "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
- " model_name=common_util.get_job_name_with_datetime(prefix=\"qwen2_5-vllm-serve\"),\n",
- " model_id=merged_model_output_dir,\n",
- " publisher=\"qwen\",\n",
- " publisher_model_id=\"qwen2_5\",\n",
- " service_account=SERVICE_ACCOUNT,\n",
- " machine_type=machine_type,\n",
- " accelerator_type=accelerator_type,\n",
- " accelerator_count=per_node_accelerator_count,\n",
- " gpu_memory_utilization=gpu_memory_utilization,\n",
- " max_model_len=max_model_len,\n",
- " enable_lora=True,\n",
- " use_dedicated_endpoint=use_dedicated_endpoint,\n",
- ")\n",
- "\n",
- "# @markdown Click \"Show Code\" to see more details."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "2UYUNn60G_4U"
- },
- "outputs": [],
- "source": [
- "# @title Predict\n",
- "\n",
- "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
- "\n",
- "# @markdown Example:\n",
- "\n",
- "# @markdown ```\n",
- "# @markdown Human: What is a car?\n",
- "# @markdown 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",
- "# @markdown ```\n",
- "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
- "\n",
- "# Loads an existing endpoint instance using the endpoint name:\n",
- "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
- "# endpoint name of the endpoint `endpoint` created in the cell\n",
- "# above.\n",
- "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
- "# an existing endpoint with the ID 1234567890123456789.\n",
- "# You may uncomment the code below to load an existing endpoint.\n",
- "\n",
- "# endpoint_name = \"\" # @param {type:\"string\"}\n",
- "# aip_endpoint_name = (\n",
- "# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
- "# )\n",
- "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
- "\n",
- "prompt = \"What is a car?\" # @param {type: \"string\"}\n",
- "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\n",
- "max_tokens = 50 # @param {type:\"integer\"}\n",
- "temperature = 1.0 # @param {type:\"number\"}\n",
- "top_p = 1.0 # @param {type:\"number\"}\n",
- "top_k = 1 # @param {type:\"integer\"}\n",
- "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
- "raw_response = False # @param {type:\"boolean\"}\n",
- "\n",
- "# Overrides parameters for inferences.\n",
- "instances = [\n",
- " {\n",
- " \"prompt\": prompt,\n",
- " \"max_tokens\": max_tokens,\n",
- " \"temperature\": temperature,\n",
- " \"top_p\": top_p,\n",
- " \"top_k\": top_k,\n",
- " \"raw_response\": raw_response,\n",
- " },\n",
- "]\n",
- "response = endpoints[\"vllm_gpu\"].predict(\n",
- " instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
- ")\n",
- "\n",
- "for prediction in response.predictions:\n",
- " print(prediction)\n",
- "\n",
- "# @markdown Click \"Show Code\" to see more details."
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "af21a3cff1e0"
- },
- "source": [
- "## Clean up resources"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "cellView": "form",
- "id": "911406c1561e"
- },
- "outputs": [],
- "source": [
- "# @title Delete the model and endpoint\n",
- "\n",
- "if train_job:\n",
- " train_job.delete()\n",
- "if RUN_EVALUATION and lm_eval_job:\n",
- " lm_eval_job.delete()\n",
- "\n",
- "\n",
- "# @markdown Delete the experiment models and endpoints to recycle the resources\n",
- "# @markdown and avoid unnecessary continuous charges that may incur.\n",
- "\n",
- "# Undeploy model and delete endpoint.\n",
- "for endpoint in endpoints.values():\n",
- " endpoint.delete(force=True)\n",
- "\n",
- "# Delete models.\n",
- "for model in models.values():\n",
- " model.delete()\n",
- "\n",
- "delete_bucket = False # @param {type:\"boolean\"}\n",
- "if delete_bucket:\n",
- " ! gsutil -m rm -r $BUCKET_NAME"
- ]
- }
- ],
- "metadata": {
- "colab": {
- "name": "model_garden_pytorch_qwen2_5_finetuning.ipynb",
- "toc_visible": true
- },
- "kernelspec": {
- "display_name": "Python 3",
- "name": "python3"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}