xiaohei/model_garden
1
0
Fork 0
mirror of https://github.com/GoogleCloudPlatform/vertex-ai-samples.git synced 2026-09-26 14:42:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add llama3.1 finetune and deploy notebooks (#3309)

Browse Source
This commit is contained in:
Xiang Xu
2024-07-23 15:15:47 +00:00
committed by GitHub
parent 81d900051b
commit 27f903deb0
3 changed files with 1393 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/CODEOWNERS
+2
View File
@@ -125,6 +125,8 @@
/notebooks/community/model_garden/model_garden_pytorch_llama2_rlhf_tuning.ipynb @genquan9
/notebooks/community/model_garden/model_garden_pytorch_llama3_deployment.ipynb @kathyyu-google
/notebooks/community/model_garden/model_garden_pytorch_llama3_finetuning.ipynb @kathyyu-google
/notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb @xiangxu-google
/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb @wrzhao-work
/notebooks/community/model_garden/model_garden_pytorch_wizard_coder.ipynb @KCFindstr
/notebooks/community/model_registry/get_started_with_vertex_ai_deployer.ipynb angelmontero@ @inardini
/notebooks/community/model_garden/model_garden_pytorch_wizard_lm.ipynb @KCFindstr
notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb
+673
View File
@@ -0,0 +1,673 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "7d9bbf86da5e"
},
"outputs": [],
"source": [
"# Copyright 2024 Google LLC\n",
"#\n",
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "99c1c3fc2ca5"
},
"source": [
"# Vertex AI Model Garden - Llama 3.1 (Deployment)\n",
"\n",
"<table><tbody><tr>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_deployment.ipynb\">\n",
" <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb\">\n",
" <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</tr></tbody></table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3de7470326a2"
},
"source": [
"## Overview\n",
"\n",
"This notebook demonstrates downloading, deploying, and serving prebuilt Llama 3.1 models with [Hex-LLM](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm) or [vLLM](https://github.com/vllm-project/vllm).\n",
"\n",
"\n",
"### Objective\n",
"\n",
"- Deploy Llama 3.1 8B with Hex-LLM on TPU.\n",
"- Deploy Llama 3.1 70B and 405B with vLLM on GPU, optionally with dynamic LoRA adapters.\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": "ax7zWynUDcjk"
},
"outputs": [],
"source": [
"# @title Request for quota\n",
"\n",
"# @markdown By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
"\n",
"# @markdown By default, the quota for H100 deployment `Custom model serving per region` is 0. You need to request for H100 quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "YXFGIp1l-qtT"
},
"outputs": [],
"source": [
"# @title Setup Google Cloud project\n",
"# Import the necessary packages\n",
"\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"import importlib\n",
"import os\n",
"import re\n",
"from datetime import datetime\n",
"from typing import Tuple\n",
"\n",
"from google.cloud import aiplatform\n",
"\n",
"common_util = importlib.import_module(\n",
" \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
")\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. [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",
"# Get the default cloud project id.\n",
"PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
"\n",
"# Get the default region for launching jobs.\n",
"REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
"\n",
"# 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.now().strftime(\"%Y%m%d%H%M%S\")\n",
"BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n",
"\n",
"if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
" BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
" ! gsutil mb -l {REGION} {BUCKET_URI}\n",
"else:\n",
" assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
" BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
" shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
" bucket_region = shell_output[0].strip().lower()\n",
" if bucket_region != REGION:\n",
" raise ValueError(\n",
" \"Bucket region %s is different from notebook region %s\"\n",
" % (bucket_region, REGION)\n",
" )\n",
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
"\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama_3_1\")\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",
"BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
"! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
"\n",
"! gcloud config set project $PROJECT_ID\n",
"\n",
"# @markdown # Access Llama 3.1 models on Vertex AI for serving\n",
"# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
"# @markdown Accept the model agreement to access the models:\n",
"# @markdown 1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
"# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
"# @markdown 3. After accepting the agreement of Llama 3.1, a `gs://` URI containing Llama 3.1 pretrained and finetuned models will be shared.\n",
"# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA_3_1` field below.\n",
"# @markdown 5. The Llama 3.1 models will be copied into `BUCKET_URI`.\n",
"\n",
"\n",
"VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = \"\" # @param {type:\"string\", isTemplate:true}\n",
"assert (\n",
" VERTEX_AI_MODEL_GARDEN_LLAMA_3_1\n",
"), \"Please click the agreement of Llama 3.1 in Vertex AI Model Garden, and get the GCS path of Llama 3.1 model artifacts.\"\n",
"parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_3_1)\n",
"if parsed_gcs_url:\n",
" VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = parsed_gcs_url.group()\n",
"assert VERTEX_AI_MODEL_GARDEN_LLAMA_3_1.startswith(\n",
" \"gs://\"\n",
"), \"VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 is expected to be a GCS URI and must start with `gs://`.\"\n",
"print(\n",
" \"Copying LLaMA3 model artifacts from\",\n",
" VERTEX_AI_MODEL_GARDEN_LLAMA_3_1,\n",
" \"to \",\n",
" MODEL_BUCKET,\n",
")\n",
"\n",
"! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_3_1/* $MODEL_BUCKET\n",
"\n",
"# The pre-built serving docker images.\n",
"HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:llama3.1\"\n",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240722_2146_RC00\"\n",
"\n",
"SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
"\n",
"\n",
"def deploy_model_hexllm(\n",
" model_name: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"ct5lp-hightpu-4t\",\n",
" tensor_parallel_size: int = 4,\n",
" hbm_utilization_factor: float = 0.8,\n",
" max_running_seqs: int = 256,\n",
" max_model_len: int = 8192,\n",
" endpoint_id: str = \"\",\n",
" min_replica_count: int = 1,\n",
" max_replica_count: int = 1,\n",
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
" \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
" if endpoint_id:\n",
" aip_endpoint_name = (\n",
" f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
" )\n",
" endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
" else:\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
"\n",
" hexllm_args = [\n",
" \"--host=0.0.0.0\",\n",
" \"--port=7080\",\n",
" \"--log_level=INFO\",\n",
" \"--enable_jit\",\n",
" f\"--model={model_id}\",\n",
" \"--load_format=auto\",\n",
" f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
" f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
" f\"--max_running_seqs={max_running_seqs}\",\n",
" f\"--max_model_len={max_model_len}\",\n",
" ]\n",
" hexllm_envs = {\n",
" \"PJRT_DEVICE\": \"TPU\",\n",
" \"RAY_DEDUP_LOGS\": \"0\",\n",
" \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
" \"MODEL_ID\": model_id,\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
"\n",
" model = aiplatform.Model.upload(\n",
" display_name=model_name,\n",
" serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
" serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
" serving_container_args=hexllm_args,\n",
" serving_container_ports=[7080],\n",
" serving_container_predict_route=\"/generate\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=hexllm_envs,\n",
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
" serving_container_deployment_timeout=7200,\n",
" )\n",
"\n",
" model.deploy(\n",
" endpoint=endpoint,\n",
" machine_type=machine_type,\n",
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" min_replica_count=min_replica_count,\n",
" max_replica_count=max_replica_count,\n",
" )\n",
" return model, endpoint\n",
"\n",
"\n",
"def deploy_model_vllm(\n",
" model_name: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"g2-standard-8\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
" accelerator_count: int = 1,\n",
" gpu_memory_utilization: float = 0.9,\n",
" max_model_len: int = 8192,\n",
" max_loras: int = 1,\n",
" max_cpu_loras: int = 16,\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",
" \"python\",\n",
" \"-m\",\n",
" \"vllm.entrypoints.api_server\",\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",
" f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
" f\"--max-model-len={max_model_len}\",\n",
" \"--enable-lora\",\n",
" \"--disable-custom-all-reduce\",\n",
" f\"--max-loras={max_loras}\",\n",
" f\"--max-cpu-loras={max_cpu_loras}\",\n",
" \"--disable-log-stats\",\n",
" ]\n",
"\n",
" env_vars = {\"MODEL_ID\": model_id, \"DEPLOY_SOURCE\": \"notebook\"}\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=[7080],\n",
" serving_container_predict_route=\"/generate\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=env_vars,\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",
" )\n",
" print(\"endpoint_name:\", endpoint.name)\n",
"\n",
" return model, endpoint"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "mAt6NcA5Dcjl"
},
"source": [
"## Deploy prebuilt Llama 3.1 8B with Hex-LLM\n",
"\n",
"**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
"\n",
"Refer to the \"Request for TPU quota\" section for TPU quota."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9-5obzXZDcjl"
},
"outputs": [],
"source": [
"# @title Deploy\n",
"\n",
"# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
"\n",
"# @markdown Select one of the four model variations. More model variants will be supported by Hex-LLM in the future.\n",
"MODEL_ID = \"Meta-Llama-3.1-8B\" # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\"] {allow-input: true, isTemplate: true}\n",
"model_id = os.path.join(MODEL_BUCKET, MODEL_ID)\n",
"\n",
"# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
"# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
"\n",
"# Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Llama 3.1 8B models.\n",
"machine_type = \"ct5lp-hightpu-4t\"\n",
"accelerator_type = \"TPU_V5e\"\n",
"# Note: 1 TPU V5 chip has only one core.\n",
"accelerator_count = 4\n",
"\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=False,\n",
")\n",
"\n",
"# Server parameters.\n",
"tensor_parallel_size = accelerator_count\n",
"hbm_utilization_factor = 0.8 # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
"max_running_seqs = 256 # Maximum number of running sequences in a continuous batch.\n",
"max_model_len = 8192\n",
"\n",
"# Endpoint configurations.\n",
"min_replica_count = 1\n",
"max_replica_count = 1\n",
"\n",
"model_hexllm, endpoint_hexllm = deploy_model_hexllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"llama_3_1-hexllm-serve\"),\n",
" model_id=model_id,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
" tensor_parallel_size=tensor_parallel_size,\n",
" hbm_utilization_factor=hbm_utilization_factor,\n",
" max_running_seqs=max_running_seqs,\n",
" max_model_len=max_model_len,\n",
" min_replica_count=min_replica_count,\n",
" max_replica_count=max_replica_count,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cEq8oadxDcjl"
},
"outputs": [],
"source": [
"# @title Predict\n",
"\n",
"# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The first few requests may have high latency. This is because the server needs to warm up with the initial requests. The following requests should not have the same delay.\n",
"\n",
"# @markdown Example:\n",
"\n",
"# @markdown ```\n",
"# @markdown > What is a car?\n",
"# @markdown > A car is a four-wheeled vehicle designed for the transportation of passengers and their belongings.\n",
"# @markdown ```\n",
"\n",
"# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
"\n",
"# Loads an existing endpoint instance using the endpoint name:\n",
"# - Using `endpoint_name = endpoint_hexllm.name` allows us to get the endpoint\n",
"# name of the endpoint `endpoint_hexllm` created in the cell above.\n",
"# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
"# an existing endpoint with the ID 1234567890123456789.\n",
"# You may uncomment the code below to load an existing endpoint:\n",
"# endpoint_name = endpoint_without_peft.name\n",
"# # endpoint_name = \"\" # @param {type:\"string\"}\n",
"# aip_endpoint_name = (\n",
"# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
"# )\n",
"# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\n",
"\n",
"prompt = \"What is a car?\" # @param {type: \"string\"}\n",
"max_tokens = 50 # @param {type: \"integer\"}\n",
"temperature = 1.0 # @param {type: \"number\"}\n",
"top_p = 1.0 # @param {type: \"number\"}\n",
"top_k = 1 # @param {type: \"integer\"}\n",
"instances = [\n",
" {\n",
" \"prompt\": prompt,\n",
" \"max_tokens\": max_tokens,\n",
" \"temperature\": temperature,\n",
" \"top_p\": top_p,\n",
" \"top_k\": top_k,\n",
" },\n",
"]\n",
"response = endpoint_hexllm.predict(instances=instances)\n",
"\n",
"prediction = response.predictions[0]\n",
"print(prediction)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "z-XybZjtgF9M"
},
"source": [
"## Deploy prebuilt Llama 3.1 70B and 405B with vLLM"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "E8OiHHNNE_wj"
},
"outputs": [],
"source": [
"# @title Deploy\n",
"\n",
"# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
"\n",
"# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of H100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have H100 quota.\n",
"\n",
"# @markdown H100 is hard to get for now. It's recommended to use the deployment button in the model card. You can still try to deploy H100 endpoint through the notebook, but there is a chance that resource is not available.\n",
"\n",
"# @markdown Set the model to deploy.\n",
"\n",
"base_model_name = \"Meta-Llama-3.1-70B\" # @param [\"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\", \"Meta-Llama-3.1-405B-Instruct-FP8\"] {isTemplate:true}\n",
"model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
"\n",
"# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
"\n",
"if \"70\" in base_model_name:\n",
" accelerator_type = \"NVIDIA_L4\"\n",
" machine_type = \"g2-standard-8\"\n",
" accelerator_count = 8\n",
"elif \"405\" in base_model_name:\n",
" accelerator_type = \"NVIDIA_H100_80GB\"\n",
" machine_type = \"a3-highgpu-8g\"\n",
" accelerator_count = 8\n",
"else:\n",
" raise ValueError(\n",
" f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
" )\n",
"\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=False,\n",
")\n",
"\n",
"gpu_memory_utilization = 0.8\n",
"max_model_len = 8192 # Maximum context length.\n",
"\n",
"model, endpoint = deploy_model_vllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"llama_3_1-vllm-serve\"),\n",
" model_id=model_id,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" gpu_memory_utilization=gpu_memory_utilization,\n",
" max_model_len=max_model_len,\n",
")\n",
"# @markdown Click \"Show Code\" to see more details."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "rDHsCOqvFYBi"
},
"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",
"# @markdown Optionally, you can apply LoRA weights to prediction. Set `lora_weight` to be either a GCS URI or a HuggingFace repo containing the LoRA weight.\n",
"\n",
"prompt = \"What is a car?\" # @param {type: \"string\"}\n",
"max_tokens = 50 # @param {type:\"integer\"}\n",
"temperature = 1.0 # @param {type:\"number\"}\n",
"top_p = 1.0 # @param {type:\"number\"}\n",
"top_k = 1 # @param {type:\"integer\"}\n",
"raw_response = False # @param {type:\"boolean\"}\n",
"lora_weight = \"\" # @param {type:\"string\", isTemplate: true}\n",
"\n",
"# Overides parameters for inferences.\n",
"# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
"# you can reduce the maximum number of output tokens, such as set max_tokens as 20.\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",
" \"dynamic-lora\": lora_weight,\n",
" },\n",
"]\n",
"response = endpoint.predict(instances=instances)\n",
"\n",
"for prediction in response.predictions:\n",
" print(prediction)\n",
"\n",
"# @markdown You can also use the `@requestFormat` parameter to send the OpenAI chat completions request.\n",
"\n",
"message_role = \"user\" # @param {type: \"string\"}\n",
"message_content = \"Why is a car?\" # @param {type: \"string\"}\n",
"\n",
"messages = [\n",
" {\n",
" \"role\": message_role,\n",
" \"content\": message_content,\n",
" }\n",
"]\n",
"\n",
"instances = [\n",
" {\n",
" \"messages\": messages,\n",
" \"@requestFormat\": \"chatCompletions\",\n",
" },\n",
"]\n",
"response = endpoint.predict(instances=instances)\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": "vfK-kZmwV7Bb"
},
"source": [
"## Use Llama guard model\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Empf8V2GWJJI"
},
"source": [
"You can use the Llama Guard model together with the Llama 3.1 405B Instruct API. See the [Llama Guard model](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama-guard) for details."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "JETd33jIDcjm"
},
"source": [
"## Clean up resources"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "911406c1561e"
},
"outputs": [],
"source": [
"# @title Delete the models and endpoints\n",
"# @markdown Delete the experiment models and endpoints to recycle the resources\n",
"# @markdown and avoid unnecessary continouous charges that may incur.\n",
"\n",
"# Undeploy model and delete endpoint.\n",
"endpoint.delete(force=True)\n",
"\n",
"# Delete models.\n",
"model.delete()\n",
"\n",
"delete_bucket = False # @param {type:\"boolean\"}\n",
"if delete_bucket:\n",
" ! gsutil -m rm -r $BUCKET_URI"
]
}
],
"metadata": {
"colab": {
"name": "model_garden_pytorch_llama3_1_deployment.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}
notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb
+718
View File
@@ -0,0 +1,718 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "7d9bbf86da5e"
},
"outputs": [],
"source": [
"# Copyright 2024 Google LLC\n",
"#\n",
"# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
"# you may not use this file except in compliance with the License.\n",
"# You may obtain a copy of the License at\n",
"#\n",
"# https://www.apache.org/licenses/LICENSE-2.0\n",
"#\n",
"# Unless required by applicable law or agreed to in writing, software\n",
"# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
"# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
"# See the License for the specific language governing permissions and\n",
"# limitations under the License."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "99c1c3fc2ca5"
},
"source": [
"# Vertex AI Model Garden - Llama 3.1 Finetuning\n",
"\n",
"<table><tbody><tr>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
" <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
" </a>\n",
" </td>\n",
" <td style=\"text-align: center\">\n",
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb\">\n",
" <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
" </a>\n",
" </td>\n",
"</tr></tbody></table>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "3de7470326a2"
},
"source": [
"## Overview\n",
"\n",
"This notebook demonstrates finetuning and deploying Llama 3.1 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 Llama 3.1 models with Vertex AI Custom Training Jobs.\n",
"- Deploy finetuned Llama 3.1 models on Vertex AI Prediction.\n",
"- Send prediction requests to your finetuned Llama 3.1 models.\n",
"\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 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. [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",
"# @markdown 3. [Make sure that you have GPU quota for Vertex Training (finetuing) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 4 L4s for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations. To secure GPUs for larger models, ask your customer engineer to get you allowlisted for a Shared Reservation or a Dynamic Workload Scheduler.\n",
"\n",
"# Import the necessary packages\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"import importlib\n",
"import os\n",
"import uuid\n",
"from datetime import datetime\n",
"from typing import Tuple\n",
"\n",
"from google.cloud import aiplatform\n",
"\n",
"common_util = importlib.import_module(\n",
" \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.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",
"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.now().strftime(\"%Y%m%d%H%M%S\")\n",
"BUCKET_URI = \"gs://\" # @param {type:\"string\"}\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, \"llama3_1\")\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"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "36c21f10355f"
},
"outputs": [],
"source": [
"# @title Access Llama 3.1 models\n",
"\n",
"# @markdown For GPU based finetuning and serving, choose between accessing Llama 3.1 models on [Hugging Face](https://huggingface.co/)\n",
"# @markdown or Vertex AI as described below.\n",
"\n",
"# @markdown If you already obtained access to Llama 3.1 models on [Hugging Face](https://huggingface.co/), you can load models from there.\n",
"# @markdown Alternatively, you can also load the original Llama 3.1 models for finetuning and serving from Vertex AI after accepting the agreement.\n",
"\n",
"# @markdown **Only select and fill one of the following sections.**\n",
"# fmt: off\n",
"LOAD_MODEL_FROM = \"Hugging Face\" # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
"# fmt: on\n",
"\n",
"# @markdown ---\n",
"\n",
"# @markdown ### Access Llama 3.1 models on Hugging Face for GPU based finetuning and serving\n",
"# @markdown You must provide a Hugging Face User Access Token (read) to access the Llama 3.1 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
"\n",
"HF_TOKEN = \"\" # @param {type:\"string\", isTemplate:true}\n",
"if LOAD_MODEL_FROM == \"Hugging Face\":\n",
" assert (\n",
" HF_TOKEN\n",
" ), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
"\n",
"# @markdown *--- Or ---*\n",
"# @markdown ### Access Llama 3.1 models on Vertex AI for GPU based serving\n",
"# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
"# @markdown Accept the model agreement to access the models:\n",
"# @markdown 1. Open the [Llama 3.1 model card](https://console.cloud.google.com/vertex-ai/publishers/meta/model-garden/llama3_1) from [Vertex AI Model Garden](https://cloud.google.com/model-garden).\n",
"# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
"# @markdown 3. After accepting the agreement of Llama 3.1, a `gs://` URI containing Llama 3.1 pretrained and finetuned models will be shared.\n",
"# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA3_1` field below.\n",
"\n",
"VERTEX_AI_MODEL_GARDEN_LLAMA3_1 = \"\" # @param {type:\"string\", isTemplate:true}\n",
"MODEL_BUCKET = VERTEX_AI_MODEL_GARDEN_LLAMA3_1\n",
"\n",
"# @markdown ---\n",
"\n",
"\n",
"# The pre-built serving docker image.\n",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
"\n",
"\n",
"def deploy_model_vllm(\n",
" model_name: str,\n",
" 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",
") -> 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",
" if not base_model_id:\n",
" base_model_id = model_id\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",
" f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
" f\"--max-model-len={max_model_len}\",\n",
" f\"--dtype={dtype}\",\n",
" \"--disable-log-stats\",\n",
" ]\n",
"\n",
" env_vars = {\n",
" \"MODEL_ID\": base_model_id,\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
"\n",
" try:\n",
" if HF_TOKEN:\n",
" env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
" except:\n",
" pass\n",
"\n",
" model = aiplatform.Model.upload(\n",
" display_name=model_name,\n",
" serving_container_image_uri=VLLM_DOCKER_URI,\n",
" serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
" serving_container_args=vllm_args,\n",
" serving_container_ports=[7080],\n",
" serving_container_predict_route=\"/generate\",\n",
" serving_container_health_route=\"/ping\",\n",
" serving_container_environment_variables=env_vars,\n",
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
" serving_container_deployment_timeout=7200,\n",
" )\n",
" 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",
" )\n",
" print(\"endpoint_name:\", endpoint.name)\n",
"\n",
" return model, endpoint"
]
},
{
"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 `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` 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 `instruct_column_in_dataset` 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 `instruct_column_in_dataset`; 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 `instruct_column_in_dataset` 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_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
"# @markdown 1. Set `train_split_name` to `train`\n",
"# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
"# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
"# @markdown 1. Set `instruct_column_in_dataset` 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_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
"train_split_name = \"train\" # @param {type:\"string\"}\n",
"eval_dataset_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
"eval_split_name = \"test\" # @param {type:\"string\"}\n",
"\n",
"# Name of the dataset column containing training text input.\n",
"instruct_column_in_dataset = \"text\" # @param {type:\"string\"}"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "ivVGS9dHXPOz"
},
"outputs": [],
"source": [
"# @title Finetune\n",
"# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\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. We recommend using NVIDIA_L4 for 8B models and NVIDIA_A100_80GB for 70B models.\n",
"# @markdown 1. If `max_steps>0`, it will precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
"# @markdown 1. With the default setting, training takes between 1.5 ~ 2 hours.\n",
"\n",
"TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240722_0936_RC01\"\n",
"\n",
"\n",
"# The Llama 3.1 base model.\n",
"MODEL_ID = \"meta-llama/Meta-Llama-3.1-8B-Instruct\" # @param [\"meta-llama/Meta-Llama-3.1-8B\", \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"meta-llama/Meta-Llama-3.1-70B\", \"meta-llama/Meta-Llama-3.1-70B-Instruct\"] {isTemplate:true}\n",
"if LOAD_MODEL_FROM == \"Google Cloud\":\n",
" base_model_id = os.path.join(MODEL_BUCKET, MODEL_ID.split(\"/\")[-1])\n",
"else:\n",
" base_model_id = MODEL_ID\n",
"\n",
"# The accelerator to use.\n",
"accelerator_type = \"NVIDIA_L4\" # @param [\"NVIDIA_L4\", \"NVIDIA_A100_80GB\"]\n",
"\n",
"# Batch size for finetuning.\n",
"per_device_train_batch_size = 1 # @param{type:\"integer\"}\n",
"gradient_accumulation_steps = 8 # @param{type:\"integer\"}\n",
"# Maximum sequence length.\n",
"max_seq_length = 4096 # @param{type:\"integer\"}\n",
"# Setting a positive `max_steps` here will override `num_epochs`\n",
"max_steps = -1 # @param{type:\"integer\"}\n",
"num_epochs = 1.0 # @param{type:\"number\"}\n",
"# Precision mode for finetuning.\n",
"finetuning_precision_mode = \"4bit\" # @param [\"4bit\", \"8bit\", \"float16\"]\n",
"# Learning rate.\n",
"learning_rate = 5e-5 # @param{type:\"number\"}\n",
"lr_scheduler_type = \"cosine\" # @param{type:\"string\"}\n",
"# 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",
"enable_gradient_checkpointing = True\n",
"attn_implementation = \"flash_attention_2\"\n",
"optimizer = \"paged_adamw_32bit\"\n",
"warmup_ratio = \"0.01\"\n",
"report_to = \"tensorboard\"\n",
"save_steps = 10\n",
"logging_steps = save_steps\n",
"\n",
"# Worker pool spec.\n",
"machine_type = None\n",
"if \"8b\" in MODEL_ID.lower():\n",
" if accelerator_type == \"NVIDIA_L4\":\n",
" accelerator_count = 4\n",
" machine_type = \"g2-standard-48\"\n",
" else:\n",
" raise ValueError(\n",
" f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
" )\n",
"elif \"70b\" in MODEL_ID.lower():\n",
" if accelerator_type == \"NVIDIA_A100_80GB\":\n",
" accelerator_count = 4\n",
" machine_type = \"a2-ultragpu-4g\"\n",
" else:\n",
" raise ValueError(\n",
" f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
" )\n",
"else:\n",
" raise ValueError(f\"Unsupported model ID or GCS path: {MODEL_ID}.\")\n",
"\n",
"replica_count = 1\n",
"\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=True,\n",
")\n",
"\n",
"job_name = common_util.get_job_name_with_datetime(\"llama3_1-lora-train\").replace(\n",
" \"_\", \"-\"\n",
")\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",
"eval_args = [\n",
" f\"--eval_dataset_path={eval_dataset_name}\",\n",
" f\"--eval_column={instruct_column_in_dataset}\",\n",
" f\"--eval_template={template}\",\n",
" f\"--eval_split={eval_split_name}\",\n",
" f\"--eval_steps={save_steps}\",\n",
" \"--eval_tasks=builtin_eval\",\n",
" \"--eval_metric_name=loss\",\n",
"]\n",
"\n",
"train_job_args = [\n",
" \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
" \"--task=instruct-lora\",\n",
" \"--completion_only=True\",\n",
" f\"--pretrained_model_id={base_model_id}\",\n",
" f\"--dataset_name={train_dataset_name}\",\n",
" f\"--train_split_name={train_split_name}\",\n",
" f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
" f\"--output_dir={lora_output_dir}\",\n",
" f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
" f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
" f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\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\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
" f\"--num_epochs={num_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\"--template={template}\",\n",
" f\"--huggingface_access_token={HF_TOKEN}\",\n",
"] + eval_args\n",
"\n",
"# Create TensorBoard\n",
"tensorboard = aiplatform.Tensorboard.create(job_name)\n",
"exp = aiplatform.TensorboardExperiment.create(\n",
" tensorboard_experiment_id=job_name, tensorboard_name=tensorboard.name\n",
")\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",
")\n",
"\n",
"train_job.run(\n",
" args=train_job_args,\n",
" environment_variables={\"WANDB_DISABLED\": True},\n",
" replica_count=replica_count,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" boot_disk_size_gb=500,\n",
" service_account=SERVICE_ACCOUNT,\n",
" tensorboard=tensorboard.resource_name,\n",
" base_output_dir=base_output_dir,\n",
")\n",
"\n",
"print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
"print(\"Trained and merged models were 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": "qmHW6m8xG_4U"
},
"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",
"# Find Vertex AI prediction supported accelerators and regions in [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
"if \"8b\" in MODEL_ID.lower():\n",
" machine_type = \"g2-standard-12\"\n",
" accelerator_type = \"NVIDIA_L4\"\n",
" accelerator_count = 1\n",
"else:\n",
" machine_type = \"g2-standard-96\"\n",
" accelerator_type = \"NVIDIA_L4\"\n",
" accelerator_count = 8\n",
"\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=False,\n",
")\n",
"\n",
"gpu_memory_utilization = 0.85\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",
"models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_1-vllm-serve\"),\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",
" gpu_memory_utilization=gpu_memory_utilization,\n",
" max_model_len=max_model_len,\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",
"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",
"raw_response = False # @param {type:\"boolean\"}\n",
"\n",
"# Overrides parameters for inferences.\n",
"# If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`,\n",
"# you can reduce the maximum number of output tokens, such as set max_tokens as 20.\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(instances=instances)\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",
"train_job.delete()\n",
"\n",
"# @markdown Delete the experiment models and endpoints to recycle the resources\n",
"# @markdown and avoid unnecessary continouous 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_llama3_1_finetuning.ipynb",
"toc_visible": true
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
}
},
"nbformat": 4,
"nbformat_minor": 0
}