mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
Fix format issues
PiperOrigin-RevId: 681739757
This commit is contained in:
Vertex MG Team
committed by
Copybara-Service
Copybara-Service
parent
fe42990c9d
commit
575025a1cd
@@ -138,8 +138,6 @@
|
||||
")\n",
|
||||
"\n",
|
||||
"models, endpoints = {}, {}\n",
|
||||
"# Set use_dedicated_endpoint to False if don't want to use dedicated endpoint.\n",
|
||||
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
|
||||
"\n",
|
||||
"# Get the default cloud project id.\n",
|
||||
"PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
|
||||
@@ -204,11 +202,83 @@
|
||||
" HF_TOKEN\n",
|
||||
"), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
|
||||
"\n",
|
||||
"model_path_prefix = \"google/\"\n",
|
||||
"model_path_prefix = \"google/\""
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "8neJc8CnDDpu"
|
||||
},
|
||||
"source": [
|
||||
"## Deploy Gemma 2 models with Hex-LLM on TPU\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": {
|
||||
"cellView": "form",
|
||||
"id": "E8OiHHNNE_wj"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# @title Deploy\n",
|
||||
"# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS 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:gemma2\"\n",
|
||||
"TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
|
||||
"\n",
|
||||
"# @markdown Select one of the four model variations.\n",
|
||||
"MODEL_ID = \"gemma-2-2b-it\" # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-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",
|
||||
"# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
|
||||
"if \"2b\" in model_id:\n",
|
||||
" # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2 2B models.\n",
|
||||
" machine_type = \"ct5lp-hightpu-1t\"\n",
|
||||
" accelerator_type = \"TPU_V5e\"\n",
|
||||
" # Note: 1 TPU V5 chip has only one core.\n",
|
||||
" accelerator_count = 1\n",
|
||||
"elif \"9b\" in model_id:\n",
|
||||
" # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 2 9B 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",
|
||||
"else:\n",
|
||||
" # Sets ct5lp-hightpu-8t (8 TPU chips) to deploy Gemma 2 27B models.\n",
|
||||
" machine_type = \"ct5lp-hightpu-8t\"\n",
|
||||
" accelerator_type = \"TPU_V5e\"\n",
|
||||
" # Note: 1 TPU V5 chip has only one core.\n",
|
||||
" accelerator_count = 8\n",
|
||||
"\n",
|
||||
"common_util.check_quota(\n",
|
||||
" project_id=PROJECT_ID,\n",
|
||||
" region=TPU_DEPLOYMENT_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.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",
|
||||
"\n",
|
||||
"# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
|
||||
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
|
||||
"\n",
|
||||
"# Endpoint configurations.\n",
|
||||
"min_replica_count = 1\n",
|
||||
"max_replica_count = 1\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def deploy_model_hexllm(\n",
|
||||
" model_name: str,\n",
|
||||
@@ -297,127 +367,6 @@
|
||||
" return model, endpoint\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def deploy_model_tgi(\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",
|
||||
" max_input_length: int = 2047,\n",
|
||||
" max_total_tokens: int = 2048,\n",
|
||||
" max_batch_prefill_tokens: int = 2048,\n",
|
||||
" use_dedicated_endpoint: bool = False,\n",
|
||||
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
|
||||
" \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
|
||||
" endpoint = aiplatform.Endpoint.create(\n",
|
||||
" display_name=f\"{model_name}-endpoint\",\n",
|
||||
" dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" env_vars = {\n",
|
||||
" \"MODEL_ID\": model_id,\n",
|
||||
" \"NUM_SHARD\": f\"{accelerator_count}\",\n",
|
||||
" \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
|
||||
" \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
|
||||
" \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
|
||||
" \"DEPLOY_SOURCE\": \"notebook\",\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" # HF_TOKEN is not a compulsory field and may not be defined.\n",
|
||||
" try:\n",
|
||||
" if HF_TOKEN:\n",
|
||||
" env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
|
||||
" except NameError:\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
" model = aiplatform.Model.upload(\n",
|
||||
" display_name=model_name,\n",
|
||||
" serving_container_image_uri=TGI_DOCKER_URI,\n",
|
||||
" serving_container_ports=[8080],\n",
|
||||
" serving_container_environment_variables=env_vars,\n",
|
||||
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" model.deploy(\n",
|
||||
" endpoint=endpoint,\n",
|
||||
" machine_type=machine_type,\n",
|
||||
" accelerator_type=accelerator_type,\n",
|
||||
" accelerator_count=accelerator_count,\n",
|
||||
" deploy_request_timeout=1800,\n",
|
||||
" service_account=service_account,\n",
|
||||
" )\n",
|
||||
" return model, endpoint"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "8neJc8CnDDpu"
|
||||
},
|
||||
"source": [
|
||||
"## Deploy Gemma 2 models with Hex-LLM on TPU\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": {
|
||||
"cellView": "form",
|
||||
"id": "E8OiHHNNE_wj"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# @title Deploy\n",
|
||||
"# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
|
||||
"\n",
|
||||
"# @markdown Select one of the four model variations.\n",
|
||||
"MODEL_ID = \"gemma-2-2b-it\" # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-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",
|
||||
"# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
|
||||
"if \"2b\" in model_id:\n",
|
||||
" # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy Gemma 2 2B models.\n",
|
||||
" machine_type = \"ct5lp-hightpu-1t\"\n",
|
||||
" accelerator_type = \"TPU_V5e\"\n",
|
||||
" # Note: 1 TPU V5 chip has only one core.\n",
|
||||
" accelerator_count = 1\n",
|
||||
"elif \"9b\" in model_id:\n",
|
||||
" # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy Gemma 2 9B 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",
|
||||
"else:\n",
|
||||
" # Sets ct5lp-hightpu-8t (8 TPU chips) to deploy Gemma 2 27B models.\n",
|
||||
" machine_type = \"ct5lp-hightpu-8t\"\n",
|
||||
" accelerator_type = \"TPU_V5e\"\n",
|
||||
" # Note: 1 TPU V5 chip has only one core.\n",
|
||||
" accelerator_count = 8\n",
|
||||
"\n",
|
||||
"common_util.check_quota(\n",
|
||||
" project_id=PROJECT_ID,\n",
|
||||
" region=TPU_DEPLOYMENT_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.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",
|
||||
"\n",
|
||||
"# Endpoint configurations.\n",
|
||||
"min_replica_count = 1\n",
|
||||
"max_replica_count = 1\n",
|
||||
"\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",
|
||||
@@ -519,6 +468,10 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# @title Deploy\n",
|
||||
"\n",
|
||||
"# The pre-built serving docker image.\n",
|
||||
"TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
|
||||
"\n",
|
||||
"MODEL_ID = \"gemma-2-2b\" # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
|
||||
"model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
|
||||
"\n",
|
||||
@@ -575,6 +528,60 @@
|
||||
"max_total_tokens = 2048\n",
|
||||
"max_batch_prefill_tokens = 2048\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def deploy_model_tgi(\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",
|
||||
" max_input_length: int = 2047,\n",
|
||||
" max_total_tokens: int = 2048,\n",
|
||||
" max_batch_prefill_tokens: int = 2048,\n",
|
||||
" use_dedicated_endpoint: bool = False,\n",
|
||||
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
|
||||
" \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
|
||||
" endpoint = aiplatform.Endpoint.create(\n",
|
||||
" display_name=f\"{model_name}-endpoint\",\n",
|
||||
" dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" env_vars = {\n",
|
||||
" \"MODEL_ID\": model_id,\n",
|
||||
" \"NUM_SHARD\": f\"{accelerator_count}\",\n",
|
||||
" \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
|
||||
" \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
|
||||
" \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\n",
|
||||
" \"DEPLOY_SOURCE\": \"notebook\",\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" # HF_TOKEN is not a compulsory field and may not be defined.\n",
|
||||
" try:\n",
|
||||
" if HF_TOKEN:\n",
|
||||
" env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
|
||||
" except NameError:\n",
|
||||
" pass\n",
|
||||
"\n",
|
||||
" model = aiplatform.Model.upload(\n",
|
||||
" display_name=model_name,\n",
|
||||
" serving_container_image_uri=TGI_DOCKER_URI,\n",
|
||||
" serving_container_ports=[8080],\n",
|
||||
" serving_container_environment_variables=env_vars,\n",
|
||||
" serving_container_shared_memory_size_mb=(16 * 1024), # 16 GB\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
" model.deploy(\n",
|
||||
" endpoint=endpoint,\n",
|
||||
" machine_type=machine_type,\n",
|
||||
" accelerator_type=accelerator_type,\n",
|
||||
" accelerator_count=accelerator_count,\n",
|
||||
" deploy_request_timeout=1800,\n",
|
||||
" service_account=service_account,\n",
|
||||
" )\n",
|
||||
" return model, endpoint\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
|
||||
" model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
|
||||
" model_id=model_id,\n",
|
||||
|
||||
@@ -125,8 +125,6 @@
|
||||
")\n",
|
||||
"\n",
|
||||
"models, endpoints = {}, {}\n",
|
||||
"# Set use_dedicated_endpoint to False if don't want to use dedicated endpoint.\n",
|
||||
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
|
||||
"\n",
|
||||
"# Get the default cloud project id.\n",
|
||||
"PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
|
||||
@@ -296,6 +294,9 @@
|
||||
"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",
|
||||
"# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
|
||||
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
|
||||
"\n",
|
||||
"# Endpoint configurations.\n",
|
||||
"min_replica_count = 1\n",
|
||||
"max_replica_count = 1\n",
|
||||
@@ -387,6 +388,7 @@
|
||||
" )\n",
|
||||
" return model, endpoint\n",
|
||||
"\n",
|
||||
"\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",
|
||||
@@ -733,6 +735,7 @@
|
||||
"\n",
|
||||
" return model, endpoint\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
|
||||
" model_name=common_util.get_job_name_with_datetime(prefix=\"gemma-serve-vllm\"),\n",
|
||||
" model_id=model_id,\n",
|
||||
|
||||
Reference in New Issue
Block a user