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

refactor: refactor gemma notebooks (#3255)

Browse Source 
Co-authored-by: Rayan Dasoriya <dasoriya@google.com>
This commit is contained in:
Rayan Dasoriya
2024-07-11 19:23:15 +00:00
committed by GitHub
co-authored by Rayan Dasoriya
parent 9387236c01
commit 671e9f84fa
5 changed files with 385 additions and 622 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb
+23 -18
View File
@@ -31,19 +31,18 @@
"source": [
"# Gemma deployment to GKE using TGI on GPU\n",
"\n",
"\u003ctable align=\"left\"\u003e\u003ctbody\u003e\u003ctr\u003e\n",
" \u003ctd\u003e\n",
" \u003ca 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_gemma_deployment_on_gke.ipynb\"\u003e\n",
" \u003cimg alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"\u003e\u003cbr\u003e Run in Colab Enterprise\n",
" \u003c/a\u003e\n",
" \u003c/td\u003e\n",
" \u003ctd\u003e\n",
" \u003ca href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma_deployment_on_gke.ipynb\"\u003e\n",
" \u003cimg src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"\u003e\u003cbr\u003e\n",
" View on GitHub\n",
" \u003c/a\u003e\n",
" \u003c/td\u003e\n",
"\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e"
"<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_gemma_deployment_on_gke.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_gemma_deployment_on_gke.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>"
]
},
{
@@ -157,7 +156,7 @@
"# Create Kubernetes secret for Hugging Face credentials\n",
"! kubectl create secret generic hf-secret \\\n",
" --from-literal=hf_api_token={HF_TOKEN} \\\n",
" --dry-run=client -o yaml \u003e hf-secret.yaml\n",
" --dry-run=client -o yaml > hf-secret.yaml\n",
"\n",
"! kubectl apply -f hf-secret.yaml"
]
@@ -340,13 +339,22 @@
"command = f\"\"\"kubectl exec -t $( kubectl get pod -l app=gemma-server -o jsonpath=\"{{.items[0].metadata.name}}\" ) -c inference-server -- curl -X POST http://localhost:8000/generate \\\n",
" -H \"Content-Type: application/json\" \\\n",
" -d '{json.dumps(request)}' \\\n",
" 2\u003e /dev/null\"\"\"\n",
" 2> /dev/null\"\"\"\n",
"\n",
"output = !{command}\n",
"print(\"Output:\")\n",
"print(json.loads(output[0])[\"generated_text\"])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "wbRmgoOZF6es"
},
"source": [
"## Clean up resources"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -356,8 +364,6 @@
},
"outputs": [],
"source": [
"# @title Clean up resources\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",
@@ -377,7 +383,6 @@
"metadata": {
"colab": {
"name": "model_garden_gemma_deployment_on_gke.ipynb",
"provenance": [],
"toc_visible": true
},
"kernelspec": {
notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb
+127 -183
View File
@@ -82,20 +82,6 @@
"## Before you begin"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "C_wC61dhpWXj"
},
"outputs": [],
"source": [
"# @title Request for TPU 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)."
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -109,28 +95,44 @@
"\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 **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\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",
"import json\n",
"# @markdown 3. 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",
"# Import the necessary packages\n",
"\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"import importlib\n",
"import os\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, please change the value yourself below.\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",
" # Create a unique GCS bucket for this notebook if not specified\n",
" BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
" ! gsutil mb -l {REGION} {BUCKET_URI}\n",
"else:\n",
@@ -147,6 +149,10 @@
"\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\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",
@@ -155,13 +161,12 @@
"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",
"# Enable Vertex AI and Cloud Compute APIs.\n",
"! gcloud config set project $PROJECT_ID\n",
"! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
"\n",
"# @markdown ## Access Gemma Models\n",
"# @markdown Choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
@@ -171,7 +176,9 @@
"# @markdown Alternatively, you can also load the original Gemma models for serving from Vertex AI after accepting the agreement.\n",
"\n",
"# @markdown **Please only select and fill one of the two following sections.**\n",
"LOAD_MODEL_FROM = \"Hugging Face\" # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
"LOAD_MODEL_FROM = (\n",
" \"Hugging Face\" # @param [\"Hugging Face\", \"Google Cloud\"] {isTemplate:true}\n",
")\n",
"\n",
"# @markdown ---\n",
"\n",
@@ -219,20 +226,14 @@
"HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240508_0916_RC02\"\n",
"\n",
"SERVICE_ENDPOINT = \"aiplatform.googleapis.com\"\n",
"\n",
"\n",
"def get_job_name_with_datetime(prefix: str) -> str:\n",
" \"\"\"Gets the job name with date time when triggering deployment jobs.\"\"\"\n",
" return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\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-1t\",\n",
" base_model_id: str = None,\n",
" tensor_parallel_size: int = 1,\n",
" machine_type: str = \"ct5lp-hightpu-1t\",\n",
" hbm_utilization_factor: float = 0.6,\n",
" max_running_seqs: int = 256,\n",
" endpoint_id: str = \"\",\n",
@@ -246,28 +247,42 @@
" )\n",
" endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
" else:\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
" endpoint = aiplatform.Endpoint.create(\n",
" display_name=f\"{model_name}-endpoint\",\n",
" location=TPU_DEPLOYMENT_REGION,\n",
" )\n",
"\n",
" if not base_model_id:\n",
" base_model_id = model_id\n",
"\n",
" if not tensor_parallel_size:\n",
" tensor_parallel_size = int(machine_type[-2])\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",
" \"--enable_jit\",\n",
" \"--load_format=auto\",\n",
" f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
" f\"--max_running_seqs={max_running_seqs}\",\n",
" ]\n",
" hexllm_envs = {\n",
"\n",
" env_vars = {\n",
" \"MODEL_ID\": base_model_id,\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",
" if HF_TOKEN:\n",
" hexllm_envs.update({\"HF_TOKEN\": HF_TOKEN})\n",
"\n",
" try:\n",
" if HF_TOKEN:\n",
" env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
" except:\n",
" pass\n",
"\n",
" model = aiplatform.Model.upload(\n",
" display_name=model_name,\n",
@@ -277,9 +292,10 @@
" 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_environment_variables=env_vars,\n",
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
" serving_container_deployment_timeout=7200,\n",
" location=TPU_DEPLOYMENT_REGION,\n",
" )\n",
"\n",
" model.deploy(\n",
@@ -297,33 +313,42 @@
" model_name: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"g2-standard-12\",\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",
" max_model_len: int = 8192,\n",
" dtype: str = \"bfloat16\",\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 models with vLLM on GPU in Vertex AI.\"\"\"\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",
" \"--gpu-memory-utilization=0.9\",\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\": model_id,\n",
" \"MODEL_ID\": base_model_id,\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
" if HF_TOKEN:\n",
" env_vars[\"HF_TOKEN\"] = HF_TOKEN\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",
@@ -337,7 +362,9 @@
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
" serving_container_deployment_timeout=7200,\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",
@@ -346,103 +373,9 @@
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" return model, endpoint\n",
" print(\"endpoint_name:\", endpoint.name)\n",
"\n",
"\n",
"def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
" \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
" quota_list_output = !gcloud alpha services quota list --service=$SERVICE_ENDPOINT --consumer=projects/$project_id --filter=\"$SERVICE_ENDPOINT/$resource_id\" --format=json\n",
" # Use '.s' on the command output because it is an SList type.\n",
" quota_data = json.loads(quota_list_output.s)\n",
" if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
" return -1\n",
" if (\n",
" len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
" or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
" ):\n",
" return -1\n",
" all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
" for region_data in all_regions_data:\n",
" if (\n",
" region_data.get(\"dimensions\")\n",
" and region_data[\"dimensions\"][\"region\"] == region\n",
" ):\n",
" if \"effectiveLimit\" in region_data:\n",
" return int(region_data[\"effectiveLimit\"])\n",
" else:\n",
" return 0\n",
" return -1\n",
"\n",
"\n",
"def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
" \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
" Args:\n",
" accelerator_type: The accelerator type.\n",
" is_for_training: Whether the resource is used for training. Set false\n",
" for serving use case.\n",
" Returns:\n",
" The resource id.\n",
" \"\"\"\n",
" training_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
" \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
" }\n",
" serving_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
" }\n",
" if is_for_training:\n",
" if accelerator_type in training_accelerator_map:\n",
" return training_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
" )\n",
" else:\n",
" if accelerator_type in serving_accelerator_map:\n",
" return serving_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
" )\n",
"\n",
"\n",
"def check_quota(\n",
" project_id: str,\n",
" region: str,\n",
" accelerator_type: str,\n",
" accelerator_count: int,\n",
" is_for_training: bool,\n",
"):\n",
" \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
" resource_id = get_resource_id(accelerator_type, is_for_training)\n",
" quota = get_quota(project_id, region, resource_id)\n",
" quota_request_instruction = (\n",
" \"Either use \"\n",
" \"a different region or request additional quota. Follow \"\n",
" \"instructions here \"\n",
" \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
" \" to check quota in a region or request additional quota for \"\n",
" \"your project.\"\n",
" )\n",
" if quota == -1:\n",
" raise ValueError(\n",
" f\"\"\"Quota not found for: {resource_id} in {region}.\n",
" {quota_request_instruction}\"\"\"\n",
" )\n",
" if quota < accelerator_count:\n",
" raise ValueError(\n",
" f\"\"\"Quota not enough for {resource_id} in {region}:\n",
" {quota} < {accelerator_count}.\n",
" {quota_request_instruction}\"\"\"\n",
" )"
" return model, endpoint"
]
},
{
@@ -472,6 +405,7 @@
"\n",
"# @markdown Select one of the six model variations.\n",
"MODEL_ID = \"gemma-1.1-2b-it\" # @param [\"gemma-2b\", \"gemma-2b-it\", \"gemma-7b\", \"gemma-7b-it\", \"gemma-1.1-2b-it\", \"gemma-1.1-7b-it\"] {allow-input: true, isTemplate: true}\n",
"TPU_DEPLOYMENT_REGION = \"us-west1\" # @param [\"us-west1\"] {isTemplate:true}\n",
"model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
"\n",
"# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
@@ -489,7 +423,7 @@
" # Note: 1 TPU V5 chip has only one core.\n",
" accelerator_count = 4\n",
"\n",
"check_quota(\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -498,20 +432,18 @@
")\n",
"\n",
"# Server parameters.\n",
"tensor_parallel_size = accelerator_count\n",
"hbm_utilization_factor = 0.6 # 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",
"hbm_utilization_factor = 0.6 # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
"max_running_seqs = 256\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=get_job_name_with_datetime(prefix=MODEL_ID),\n",
"models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\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",
" min_replica_count=min_replica_count,\n",
@@ -530,7 +462,9 @@
"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",
"# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts based on your `template`. Note that the first few prompts will take longer to execute.\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",
"# @markdown Example:\n",
"\n",
@@ -542,8 +476,8 @@
"# @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",
"# - Using `endpoint_name = endpoint.name` allows us to get the endpoint\n",
"# name of the endpoint `endpoint` 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",
@@ -552,7 +486,7 @@
"# aip_endpoint_name = (\n",
"# f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
"# )\n",
"# endpoint_hexllm = aiplatform.Endpoint(aip_endpoint_name)\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",
@@ -568,10 +502,10 @@
" \"top_k\": top_k,\n",
" },\n",
"]\n",
"response = endpoint_hexllm.predict(instances=instances)\n",
"response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
"\n",
"prediction = response.predictions[0]\n",
"print(prediction)"
"for prediction in response.predictions:\n",
" print(prediction)"
]
},
{
@@ -639,7 +573,7 @@
" \"top_k\": 1,\n",
" },\n",
"]\n",
"response = endpoint_hexllm.predict(instances=instances)\n",
"response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
"\n",
"prediction = response.predictions[0]\n",
"print(prediction)"
@@ -726,7 +660,7 @@
" % accelerator_type\n",
" )\n",
"\n",
"check_quota(\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -737,8 +671,8 @@
"# Note that a larger max_model_len will require more GPU memory.\n",
"max_model_len = 2048\n",
"\n",
"model_vllm, endpoint_vllm = deploy_model_vllm(\n",
" model_name=get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
"models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
" model_id=model_id,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
@@ -769,26 +703,37 @@
"source": [
"# @title Predict\n",
"\n",
"# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Note that the first few prompts will take longer to execute.\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_vllm.name` allows us to get the endpoint\n",
"# name of the endpoint `endpoint_vllm` created in the cell above.\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",
"# endpoint_name = endpoint_vllm.name\n",
"# # endpoint_name = \"\" # @param {type:\"string\"}\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_vllm = aiplatform.Endpoint(aip_endpoint_name)\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 = 10 # @param {type: \"integer\"}\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",
"instances = [\n",
" {\n",
" \"prompt\": prompt,\n",
@@ -796,12 +741,13 @@
" \"temperature\": temperature,\n",
" \"top_p\": top_p,\n",
" \"top_k\": top_k,\n",
" \"raw_response\": raw_response,\n",
" },\n",
"]\n",
"response = endpoint_vllm.predict(instances=instances)\n",
"response = endpoints[\"vllm_gpu\"].predict(instances=instances)\n",
"\n",
"prediction = response.predictions[0]\n",
"print(prediction)"
"for prediction in response.predictions:\n",
" print(prediction)"
]
},
{
@@ -833,20 +779,18 @@
},
"outputs": [],
"source": [
"# @title Delete the models and endpoints\n",
"\n",
"# @markdown Delete the experiment models and endpoints to recycle the resources\n",
"# @markdown and avoid unnecessary continouous charges that may incur.\n",
"# Undeploy models and delete endpoints.\n",
"endpoint_hexllm.delete(force=True)\n",
"endpoint_vllm.delete(force=True)\n",
"\n",
"# Undeploy model and delete endpoint.\n",
"for endpoint in endpoints.values():\n",
" endpoint.delete(force=True)\n",
"\n",
"# Delete models.\n",
"model_hexllm.delete()\n",
"model_vllm.delete()\n",
"for model in models.values():\n",
" model.delete()\n",
"\n",
"# Delete Cloud Storage objects.\n",
"delete_bucket = False # @param {type:\"boolean\", isTemplate: true}\n",
"delete_bucket = False # @param {type:\"boolean\"}\n",
"if delete_bucket:\n",
" ! gsutil -m rm -r $BUCKET_URI"
]
notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb
+52 -118
View File
@@ -3,6 +3,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"id": "B8S-yo8qTIcO"
},
@@ -91,6 +92,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"cellView": "form",
"id": "81CC3tL1T_TL"
@@ -98,18 +100,26 @@
"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",
"# 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 json\n",
"from datetime import datetime\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",
@@ -122,16 +132,16 @@
"\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, please change the value yourself below.\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",
"assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
"\n",
"if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
" # Create a unique GCS bucket for this notebook, if not specified by the user\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",
@@ -139,127 +149,37 @@
" \"Bucket region %s is different from notebook region %s\"\n",
" % (bucket_region, REGION)\n",
" )\n",
"\n",
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
"\n",
"# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
"SERVICE_ACCOUNT = None\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\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",
"# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
"! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
"\n",
"# Initialize Vertex AI API.\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"staging\")\n",
"MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"model\")\n",
"print(\"Initializing Vertex AI API.\")\n",
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\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",
"# The evaluation docker image.\n",
"EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\"\n",
"\n",
"# Define common functions\n",
"\n",
"\n",
"def get_job_name_with_datetime(prefix: str) -> str:\n",
" \"\"\"Gets the job name with date time when triggering training or deployment\n",
" jobs in Vertex AI.\n",
" \"\"\"\n",
" return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
"\n",
"\n",
"def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
" \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
" service_endpoint = \"aiplatform.googleapis.com\"\n",
" quota_list_output = !gcloud alpha services quota list --service=$service_endpoint --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
" # Use '.s' on the command output because it is an SList type.\n",
" quota_data = json.loads(quota_list_output.s)\n",
" if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
" return -1\n",
" if len(quota_data[0][\"consumerQuotaLimits\"]) == 0 or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]:\n",
" return -1\n",
" all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
" for region_data in all_regions_data:\n",
" if region_data.get('dimensions') and region_data['dimensions']['region'] == region:\n",
" if 'effectiveLimit' in region_data:\n",
" return int(region_data['effectiveLimit'])\n",
" else:\n",
" return 0\n",
" return -1\n",
"\n",
"\n",
"def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
" \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
" Args:\n",
" accelerator_type: The accelerator type.\n",
" is_for_training: Whether the resource is used for training. Set false\n",
" for serving use case.\n",
" Returns:\n",
" The resource id.\n",
" \"\"\"\n",
" training_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
" \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
" }\n",
" serving_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
" }\n",
" if is_for_training:\n",
" if accelerator_type in training_accelerator_map:\n",
" return training_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
" )\n",
" else:\n",
" if accelerator_type in serving_accelerator_map:\n",
" return serving_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
" )\n",
"\n",
"\n",
"def check_quota(project_id:str, region: str, accelerator_type: str,\n",
" accelerator_count: int, is_for_training: bool):\n",
" \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
" resource_id = get_resource_id(accelerator_type, is_for_training)\n",
" quota = get_quota(project_id, region, resource_id)\n",
" quota_request_instruction = (\"Either use \"\n",
" \"a different region or request additional quota. Follow \"\n",
" \"instructions here \"\n",
" \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
" \" to check quota in a region or request additional quota for \"\n",
" \"your project.\")\n",
" if quota == -1:\n",
" raise ValueError(\n",
" f\"\"\"Quota not found for: {resource_id} in {region}.\n",
" {quota_request_instruction}\"\"\"\n",
" )\n",
" if quota < accelerator_count:\n",
" raise ValueError(\n",
" f\"\"\"Quota not enough for {resource_id} in {region}:\n",
" {quota} < {accelerator_count}.\n",
" {quota_request_instruction}\"\"\"\n",
" )"
"EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20240320_0655_RC00\""
]
},
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"cellView": "form",
"id": "pNHMbjr0UjrK"
@@ -285,12 +205,14 @@
"# Setup evaluation job.\n",
"# @markdown Set the base model id.\n",
"base_model_id = \"google/gemma-1.1-2b-it\" # @param[\"google/gemma-2b\", \"google/gemma-2b-it\", \"google/gemma-7b\", \"google/gemma-7b-it\", \"google/gemma-1.1-2b-it\", \"google/gemma-1.1-7b-it\"] {isTemplate:true}\n",
"job_name = get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
"job_name = common_util.get_job_name_with_datetime(prefix=\"gemma-eval\")\n",
"eval_output_dir = os.path.join(MODEL_BUCKET, job_name)\n",
"eval_output_dir_gcsfuse = eval_output_dir.replace(\"gs://\", \"/gcs/\")\n",
"\n",
"# @markdown Set the accelerator type.\n",
"accelerator_type = \"NVIDIA_L4\" # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
"accelerator_type = (\n",
" \"NVIDIA_L4\" # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
")\n",
"\n",
"# @markdown To evaluate a PEFT-finetuned model, enter the PEFT output directory to the LoRA adapter below.\n",
"# @markdown Otherwise, leave it empty.\n",
@@ -313,11 +235,13 @@
"\n",
"replica_count = 1\n",
"\n",
"check_quota(project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=True)\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",
"# Prepare evaluation command that runs the evaluation harness.\n",
"# Set `trust_remote_code = True` because evaluating the model requires\n",
@@ -409,6 +333,17 @@
"print(f\"Evaluation result:\\n{result_formatted}\")"
]
},
{
"cell_type": "markdown",
"execution_count": null,
"metadata": {
"id": "unjukbcjEBOd"
},
"outputs": [],
"source": [
"## Clean up resources"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -418,7 +353,6 @@
},
"outputs": [],
"source": [
"# @title Clean up resources\n",
"# Delete evaluation job.\n",
"\n",
"delete_bucket = False # @param {type:\"boolean\"}\n",
notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb
+124 -167
View File
@@ -3,6 +3,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"id": "7d9bbf86da5e"
},
@@ -25,6 +26,7 @@
},
{
"cell_type": "markdown",
"language": "markdown",
"metadata": {
"id": "99c1c3fc2ca5"
},
@@ -87,6 +89,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"cellView": "form",
"id": "855d6b96f291"
@@ -94,51 +97,76 @@
"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 **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\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",
"import json\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, please change the value yourself below.\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",
"if not BUCKET_URI.strip() or BUCKET_URI == \"gs://\":\n",
" # Create a unique GCS bucket for this notebook if not specified\n",
" BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\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",
"assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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, \"gemma\")\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",
"\n",
"print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\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",
"# Enable Vertex AI and Cloud Compute APIs.\n",
"! gcloud config set project $PROJECT_ID\n",
"! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
"\n",
"# @markdown ## Access Gemma Models\n",
"# @markdown For GPU based finetuning and serving, choose between accessing Gemma models on [Hugging Face](https://huggingface.co/)\n",
@@ -174,10 +202,6 @@
"\n",
"VERTEX_AI_MODEL_GARDEN_GEMMA = \"\" # @param {type:\"string\", isTemplate:true}\n",
"\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma\")\n",
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
"\n",
"if LOAD_MODEL_FROM == \"Google Cloud\":\n",
" assert (\n",
" VERTEX_AI_MODEL_GARDEN_GEMMA\n",
@@ -204,7 +228,7 @@
"# and Hex-LLM serving.\n",
"KERAS_TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-train-tpu:20240422_0939_RC00\"\n",
"KERAS_MODEL_CONVERSION_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-keras-model-conversion:20240422_0949_RC00\"\n",
"HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20240220_0936_RC01\"\n",
"HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
"conversion_job = None\n",
"\n",
"# @markdown *--- Or ---*\n",
@@ -221,34 +245,31 @@
"# @markdown ---\n",
"\n",
"\n",
"def get_job_name_with_datetime(prefix: str) -> str:\n",
" \"\"\"Gets the job name with date time when triggering training or deployment\n",
" jobs in Vertex AI.\n",
" \"\"\"\n",
" return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
"\n",
"\n",
"def deploy_model_vllm(\n",
" model_name: str,\n",
" base_model_id: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"g2-standard-12\",\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",
" max_model_len: int = 8192,\n",
" dtype: str = \"bfloat16\",\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 models with vLLM on GPU in Vertex AI.\"\"\"\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",
" \"--gpu-memory-utilization=0.95\",\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",
@@ -258,8 +279,12 @@
" \"MODEL_ID\": base_model_id,\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
" if HF_TOKEN:\n",
" env_vars[\"HF_TOKEN\"] = HF_TOKEN\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",
@@ -273,7 +298,9 @@
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
" serving_container_deployment_timeout=7200,\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",
@@ -282,38 +309,52 @@
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" print(\"endpoint_name:\", endpoint.name)\n",
"\n",
" return model, endpoint\n",
"\n",
"\n",
"def deploy_model_hexllm(\n",
" model_name: str,\n",
" base_model_id: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" base_model_id: str = None,\n",
" tensor_parallel_size: int = 1,\n",
" machine_type: str = \"ct5lp-hightpu-1t\",\n",
" max_num_batched_tokens: int = 11264,\n",
" tokens_pad_multiple: int = 1024,\n",
" seqs_pad_multiple: int = 32,\n",
" hbm_utilization_factor: float = 0.6,\n",
" max_running_seqs: int = 256,\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",
" endpoint = aiplatform.Endpoint.create(\n",
" display_name=f\"{model_name}-endpoint\",\n",
" location=TPU_DEPLOYMENT_REGION,\n",
" )\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(\n",
" display_name=f\"{model_name}-endpoint\",\n",
" location=TPU_DEPLOYMENT_REGION,\n",
" )\n",
"\n",
" if not base_model_id:\n",
" base_model_id = model_id\n",
"\n",
" if not tensor_parallel_size:\n",
" tensor_parallel_size = int(machine_type[-2])\n",
"\n",
" num_tpu_chips = int(machine_type[-2])\n",
" hexllm_args = [\n",
" \"--host=0.0.0.0\",\n",
" \"--port=7080\",\n",
" \"--log_level=INFO\",\n",
" f\"--model={model_id}\",\n",
" f\"--tensor_parallel_size={num_tpu_chips}\",\n",
" \"--num_nodes=1\",\n",
" \"--use_ray\",\n",
" \"--batch_mode=continuous\",\n",
" f\"--max_num_batched_tokens={max_num_batched_tokens}\",\n",
" f\"--tokens_pad_multiple={tokens_pad_multiple}\",\n",
" f\"--seqs_pad_multiple={seqs_pad_multiple}\",\n",
" f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
" \"--enable_jit\",\n",
" \"--load_format=auto\",\n",
" f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
" f\"--max_running_seqs={max_running_seqs}\",\n",
" ]\n",
"\n",
" env_vars = {\n",
@@ -323,14 +364,17 @@
" \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
" \"DEPLOY_SOURCE\": \"notebook\",\n",
" }\n",
" if KAGGLE_USERNAME and KAGGLE_KEY:\n",
" env_vars[\"KAGGLE_USERNAME\"] = KAGGLE_USERNAME\n",
" env_vars[\"KAGGLE_KEY\"] = KAGGLE_KEY\n",
"\n",
" try:\n",
" if HF_TOKEN:\n",
" env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
" except:\n",
" pass\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.entrypoints.api_server\"],\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",
@@ -346,105 +390,10 @@
" 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 get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
" \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
" service_endpoint = \"aiplatform.googleapis.com\" # noqa: F841\n",
" quota_list_output = !gcloud alpha services quota list --service=$service_endpoint --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
" # Use '.s' on the command output because it is an SList type.\n",
" quota_data = json.loads(quota_list_output.s)\n",
" if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
" return -1\n",
" if (\n",
" len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
" or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
" ):\n",
" return -1\n",
" all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
" for region_data in all_regions_data:\n",
" if (\n",
" region_data.get(\"dimensions\")\n",
" and region_data[\"dimensions\"][\"region\"] == region\n",
" ):\n",
" if \"effectiveLimit\" in region_data:\n",
" return int(region_data[\"effectiveLimit\"])\n",
" else:\n",
" return 0\n",
" return -1\n",
"\n",
"\n",
"def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
" \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
" Args:\n",
" accelerator_type: The accelerator type.\n",
" is_for_training: Whether the resource is used for training. Set false\n",
" for serving use case.\n",
" Returns:\n",
" The resource id.\n",
" \"\"\"\n",
" training_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
" \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
" }\n",
" serving_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
" }\n",
" if is_for_training:\n",
" if accelerator_type in training_accelerator_map:\n",
" return training_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
" )\n",
" else:\n",
" if accelerator_type in serving_accelerator_map:\n",
" return serving_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
" )\n",
"\n",
"\n",
"def check_quota(\n",
" project_id: str,\n",
" region: str,\n",
" accelerator_type: str,\n",
" accelerator_count: int,\n",
" is_for_training: bool,\n",
"):\n",
" \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
" resource_id = get_resource_id(accelerator_type, is_for_training)\n",
" quota = get_quota(project_id, region, resource_id)\n",
" quota_request_instruction = (\n",
" \"Either use \"\n",
" \"a different region or request additional quota. Follow \"\n",
" \"instructions here \"\n",
" \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
" \" to check quota in a region or request additional quota for \"\n",
" \"your project.\"\n",
" )\n",
" if quota == -1:\n",
" raise ValueError(\n",
" f\"\"\"Quota not found for: {resource_id} in {region}.\n",
" {quota_request_instruction}\"\"\"\n",
" )\n",
" if quota < accelerator_count:\n",
" raise ValueError(\n",
" f\"\"\"Quota not enough for {resource_id} in {region}:\n",
" {quota} < {accelerator_count}.\n",
" {quota_request_instruction}\"\"\"\n",
" )"
" return model, endpoint"
]
},
{
@@ -582,7 +531,7 @@
"\n",
" replica_count = 1\n",
"\n",
" check_quota(\n",
" common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=ACCELERATOR_TYPE,\n",
@@ -591,7 +540,7 @@
" )\n",
"\n",
" # Setup training job.\n",
" job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
" job_name = common_util.get_job_name_with_datetime(\"gemma-lora-train\")\n",
"\n",
" # Pass training arguments and launch job.\n",
" train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -600,12 +549,12 @@
" )\n",
"\n",
" # Create a GCS folder to store the LORA adapter.\n",
" lora_adapter_dir = get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
" lora_adapter_dir = common_util.get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
" lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\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 = get_job_name_with_datetime(\"gemma-merged-model\")\n",
" merged_model_dir = common_util.get_job_name_with_datetime(\"gemma-merged-model\")\n",
" merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
"\n",
" train_job.run(\n",
@@ -666,7 +615,7 @@
" accelerator_type = \"NVIDIA_L4\"\n",
" accelerator_count = 1\n",
"\n",
" check_quota(\n",
" common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -678,7 +627,7 @@
" max_model_len = 2048\n",
"\n",
" model, endpoint = deploy_model_vllm(\n",
" model_name=get_job_name_with_datetime(prefix=\"gemma-vllm-serve\"),\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-vllm-serve\"),\n",
" base_model_id=f\"google/{MODEL_ID}\",\n",
" model_id=merged_model_output_dir,\n",
" service_account=SERVICE_ACCOUNT,\n",
@@ -700,6 +649,8 @@
"outputs": [],
"source": [
"# @title Predict\n",
"\n",
"\n",
"# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
"\n",
"# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the finetuning outcome:\n",
@@ -905,7 +856,7 @@
"\n",
" replica_count = 1\n",
"\n",
" check_quota(\n",
" common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -914,7 +865,7 @@
" )\n",
"\n",
" # Setup training job.\n",
" job_name = get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
" job_name = common_util.get_job_name_with_datetime(\"gemma-keras-lora-train\")\n",
"\n",
" # Pass training arguments and launch job.\n",
" train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -1004,7 +955,7 @@
" replica_count = 1\n",
"\n",
" # Setup training job.\n",
" job_name = get_job_name_with_datetime(\"gemma-keras-model-conversion\")\n",
" job_name = common_util.get_job_name_with_datetime(\"gemma-keras-model-conversion\")\n",
"\n",
" # Pass training arguments and launch job.\n",
" conversion_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -1065,7 +1016,7 @@
" # Note: 1 TPU-V5e chip has only 1 core.\n",
" accelerator_count = 4\n",
"\n",
" check_quota(\n",
" common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -1084,7 +1035,7 @@
"\n",
" print(\"Using model from: \", output_folder)\n",
" model, endpoint = deploy_model_hexllm(\n",
" model_name=get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-hexllm\"),\n",
" base_model_id=f\"google/{KAGGLE_MODEL_ID}\",\n",
" model_id=output_folder,\n",
" service_account=SERVICE_ACCOUNT,\n",
@@ -1099,6 +1050,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"cellView": "form",
"id": "lgJPw3e9yL5a"
@@ -1183,14 +1135,19 @@
"if conversion_job:\n",
" conversion_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",
"endpoint.delete(force=True)\n",
"for endpoint in endpoints.values():\n",
" endpoint.delete(force=True)\n",
"\n",
"# Delete models.\n",
"model.delete()\n",
"for model in models.values():\n",
" model.delete()\n",
"\n",
"# Delete Cloud Storage bucket that was created.\n",
"if BUCKET_URI == f\"gs://{PROJECT_ID}-tmp-{now}\":\n",
"delete_bucket = False # @param {type:\"boolean\"}\n",
"if delete_bucket:\n",
" ! gsutil -m rm -r $BUCKET_URI"
]
}
notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb
+59 -136
View File
@@ -3,6 +3,7 @@
{
"cell_type": "code",
"execution_count": null,
"language": "python",
"metadata": {
"id": "7d9bbf86da5e"
},
@@ -94,55 +95,77 @@
"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 **[Optional]** Set the GCS BUCKET_URI to store the experiment artifacts, if you want to use your own bucket. **If not set, a unique GCS bucket will be created automatically on your behalf**.\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",
"import json\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"import importlib\n",
"import os\n",
"import sys\n",
"import uuid\n",
"from datetime import datetime\n",
"from typing import Tuple\n",
"\n",
"from google.cloud import aiplatform, language\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, please change the value yourself below.\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",
"if not BUCKET_URI.strip() or BUCKET_URI == \"gs://\":\n",
" # Create a unique GCS bucket for this notebook if not specified\n",
" BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\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",
"assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\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, \"gemma\")\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",
"\n",
"print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\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",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
"\n",
"# Enable Vertex AI, Cloud Compute, and Cloud Language APIs.\n",
"! gcloud config set project $PROJECT_ID\n",
"! gcloud services enable aiplatform.googleapis.com compute.googleapis.com language.googleapis.com\n",
"\n",
"# @markdown ## Access Gemma Models\n",
"\n",
@@ -161,13 +184,6 @@
" auth.authenticate_user(project_id=PROJECT_ID)\n",
"\n",
"\n",
"def get_job_name_with_datetime(prefix: str) -> str:\n",
" \"\"\"Gets the job name with date time when triggering training or deployment\n",
" jobs in Vertex AI.\n",
" \"\"\"\n",
" return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
"\n",
"\n",
"def moderate_text(text: str) -> language.ModerateTextResponse:\n",
" \"\"\"Calls Vertex AI APIs to analyze text moderations.\"\"\"\n",
" client = language.LanguageServiceClient()\n",
@@ -236,104 +252,7 @@
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" return model, endpoint\n",
"\n",
"\n",
"def get_quota(project_id: str, region: str, resource_id: str) -> int:\n",
" \"\"\"Returns the quota for a resource in a region. Returns -1 if can not figure out the quota.\"\"\"\n",
" service_endpoint = \"aiplatform.googleapis.com\" # noqa: F841\n",
" quota_list_output = !gcloud alpha services quota list --service=$service_endpoint --consumer=projects/$project_id --filter=\"$service_endpoint/$resource_id\" --format=json\n",
" # Use '.s' on the command output because it is an SList type.\n",
" quota_data = json.loads(quota_list_output.s)\n",
" if len(quota_data) == 0 or \"consumerQuotaLimits\" not in quota_data[0]:\n",
" return -1\n",
" if (\n",
" len(quota_data[0][\"consumerQuotaLimits\"]) == 0\n",
" or \"quotaBuckets\" not in quota_data[0][\"consumerQuotaLimits\"][0]\n",
" ):\n",
" return -1\n",
" all_regions_data = quota_data[0][\"consumerQuotaLimits\"][0][\"quotaBuckets\"]\n",
" for region_data in all_regions_data:\n",
" if (\n",
" region_data.get(\"dimensions\")\n",
" and region_data[\"dimensions\"][\"region\"] == region\n",
" ):\n",
" if \"effectiveLimit\" in region_data:\n",
" return int(region_data[\"effectiveLimit\"])\n",
" else:\n",
" return 0\n",
" return -1\n",
"\n",
"\n",
"def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:\n",
" \"\"\"Returns the resource id for a given accelerator type and the use case.\n",
" Args:\n",
" accelerator_type: The accelerator type.\n",
" is_for_training: Whether the resource is used for training. Set false\n",
" for serving use case.\n",
" Returns:\n",
" The resource id.\n",
" \"\"\"\n",
" training_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_training_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_training_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_training_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_training_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_training_tpu_v5e\",\n",
" \"TPU_V3\": \"custom_model_training_tpu_v3\",\n",
" }\n",
" serving_accelerator_map = {\n",
" \"NVIDIA_TESLA_V100\": \"custom_model_serving_nvidia_v100_gpus\",\n",
" \"NVIDIA_L4\": \"custom_model_serving_nvidia_l4_gpus\",\n",
" \"NVIDIA_TESLA_A100\": \"custom_model_serving_nvidia_a100_gpus\",\n",
" \"NVIDIA_TESLA_T4\": \"custom_model_serving_nvidia_t4_gpus\",\n",
" \"TPU_V5e\": \"custom_model_serving_tpu_v5e\",\n",
" }\n",
" if is_for_training:\n",
" if accelerator_type in training_accelerator_map:\n",
" return training_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for training.\"\n",
" )\n",
" else:\n",
" if accelerator_type in serving_accelerator_map:\n",
" return serving_accelerator_map[accelerator_type]\n",
" else:\n",
" raise ValueError(\n",
" f\"Could not find accelerator type: {accelerator_type} for serving.\"\n",
" )\n",
"\n",
"\n",
"def check_quota(\n",
" project_id: str,\n",
" region: str,\n",
" accelerator_type: str,\n",
" accelerator_count: int,\n",
" is_for_training: bool,\n",
"):\n",
" \"\"\"Checks if the project and the region has the required quota.\"\"\"\n",
" resource_id = get_resource_id(accelerator_type, is_for_training)\n",
" quota = get_quota(project_id, region, resource_id)\n",
" quota_request_instruction = (\n",
" \"Either use \"\n",
" \"a different region or request additional quota. Follow \"\n",
" \"instructions here \"\n",
" \"https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota\"\n",
" \" to check quota in a region or request additional quota for \"\n",
" \"your project.\"\n",
" )\n",
" if quota == -1:\n",
" raise ValueError(\n",
" f\"\"\"Quota not found for: {resource_id} in {region}.\n",
" {quota_request_instruction}\"\"\"\n",
" )\n",
" if quota < accelerator_count:\n",
" raise ValueError(\n",
" f\"\"\"Quota not enough for {resource_id} in {region}:\n",
" {quota} < {accelerator_count}.\n",
" {quota_request_instruction}\"\"\"\n",
" )"
" return model, endpoint"
]
},
{
@@ -480,7 +399,7 @@
"\n",
"replica_count = 1\n",
"\n",
"check_quota(\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=ACCELERATOR_TYPE,\n",
@@ -489,7 +408,7 @@
")\n",
"\n",
"# Setup training job.\n",
"job_name = get_job_name_with_datetime(\"gemma-lora-train\")\n",
"job_name = common_util.get_job_name_with_datetime(\"gemma-lora-train\")\n",
"\n",
"# Pass training arguments and launch job.\n",
"train_job = aiplatform.CustomContainerTrainingJob(\n",
@@ -498,12 +417,12 @@
")\n",
"\n",
"# Create a GCS folder to store the LORA adapter.\n",
"lora_adapter_dir = get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
"lora_adapter_dir = common_util.get_job_name_with_datetime(\"gemma-lora-adapter\")\n",
"lora_output_dir = os.path.join(STAGING_BUCKET, lora_adapter_dir)\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 = get_job_name_with_datetime(\"gemma-merged-model\")\n",
"merged_model_dir = common_util.get_job_name_with_datetime(\"gemma-merged-model\")\n",
"merged_model_output_dir = os.path.join(STAGING_BUCKET, merged_model_dir)\n",
"\n",
"train_job.run(\n",
@@ -560,7 +479,7 @@
"accelerator_type = \"NVIDIA_L4\"\n",
"accelerator_count = 1\n",
"\n",
"check_quota(\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
@@ -573,8 +492,8 @@
"max_total_tokens = 1024\n",
"max_batch_prefill_tokens = 2048\n",
"\n",
"model, endpoint = deploy_model_tgi(\n",
" model_name=get_job_name_with_datetime(prefix=\"gemma-tgi-serve\"),\n",
"models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-tgi-serve\"),\n",
" model_id=merged_model_output_dir,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
@@ -583,8 +502,7 @@
" max_input_length=max_input_length,\n",
" max_total_tokens=max_total_tokens,\n",
" max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
")\n",
"print(\"endpoint_name:\", endpoint.name)"
")"
]
},
{
@@ -641,7 +559,7 @@
" },\n",
" },\n",
"]\n",
"response = endpoint.predict(instances=instances)\n",
"response = endpoints[\"tgi\"].predict(instances=instances)\n",
"\n",
"for prediction in response.predictions:\n",
" print(prediction)"
@@ -686,15 +604,20 @@
"# Delete the train job.\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",
"endpoint.delete(force=True)\n",
"for endpoint in endpoints.values():\n",
" endpoint.delete(force=True)\n",
"\n",
"# Delete models.\n",
"model.delete()\n",
"for model in models.values():\n",
" model.delete()\n",
"\n",
"# Delete Cloud Storage bucket that was created.\n",
"if BUCKET_URI == f\"gs://{PROJECT_ID}-tmp-{now}\":\n",
" ! gsutil -m rm -r $STAGING_BUCKET"
"delete_bucket = False # @param {type:\"boolean\"}\n",
"if delete_bucket:\n",
" ! gsutil -m rm -r $BUCKET_URI"
]
}
],