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

Update Mixtral vLLM container and support OpenAI vLLM server. (#2630)

Browse Source
This commit is contained in:
Kathy Yu
2024-01-18 19:06:48 +00:00
committed by GitHub
parent b679b3d72d
commit c8f95d2121
1 changed files with 112 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/model_garden/model_garden_pytorch_mistral.ipynb
+112 -63
View File
@@ -243,15 +243,14 @@
},
{
"cell_type": "code",
"execution_count": 1,
"execution_count": null,
"metadata": {
"id": "b42bd4fa2b2d"
},
"outputs": [],
"source": [
"# The pre-built serving docker images with vLLM\n",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20231002_0916_RC00\"\n",
"MIXTRAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-mixtral-serve:20231214\""
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240112_0916_RC00\""
]
},
{
@@ -265,7 +264,7 @@
},
{
"cell_type": "code",
"execution_count": 2,
"execution_count": 1,
"metadata": {
"id": "354da31189dc"
},
@@ -292,8 +291,41 @@
" machine_type: str = \"g2-standard-8\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
" accelerator_count: int = 1,\n",
" max_model_len: int = 4096,\n",
" use_openai_server: bool = False,\n",
" use_chat_completions_if_openai_server: bool = False,\n",
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
" \"\"\"Deploys Mistral models with vLLM on Vertex AI.\"\"\"\n",
" \"\"\"Deploys Mistral models with vLLM on Vertex AI.\n",
"\n",
" Args:\n",
" model_name: Display name of the model.\n",
" model_id: Model ID or path to model weights.\n",
" service_account: Service account for model uploading and deployment.\n",
" machine_type: Deployment machine type.\n",
" accelerator_type: Deployment accelerator type.\n",
" accelerator_count: Number of accelerators to use.\n",
" max_model_len: Maximum model length.\n",
" use_openai_server: Whether to use the OpenAI-format vLLM model server.\n",
" use_chat_completions_if_openai_server: If the OpenAI model server is\n",
" used, whether to use the chat completion API as opposed to the text\n",
" completion API. The vLLM text completion API mimics the OpenAI text\n",
" completion API:\n",
" https://platform.openai.com/docs/api-reference/completions/create.\n",
" It has two required parameters: the model ID to direct requests to\n",
" and the prompt. The response includes a \"choices\" field that\n",
" contains the generated text and a \"usage\" field that contains token\n",
" counts. The vLLM chat completion API mimics the OpenAI chat\n",
" completion API:\n",
" https://platform.openai.com/docs/api-reference/chat/create. It has\n",
" two required parameters: the model ID to direct requests to and\n",
" \"messages\" which is a sequence of system/user/assistant/tool\n",
" messages that can represent a multi-turn chat conversation. The\n",
" response includes a \"choices\" field that contains the generated\n",
" message from a role and a \"usage\" field that contains token counts.\n",
"\n",
" Returns:\n",
" Model instance and endpoint instance.\n",
" \"\"\"\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
"\n",
" dtype = \"bfloat16\"\n",
@@ -308,60 +340,31 @@
" \"--swap-space=16\",\n",
" f\"--dtype={dtype}\",\n",
" \"--gpu-memory-utilization=0.9\",\n",
" \"--max-num-batched-tokens=4096\",\n",
" f\"--max-model-len={max_model_len}\",\n",
" \"--disable-log-stats\",\n",
" ]\n",
" if use_openai_server:\n",
" if use_chat_completions_if_openai_server:\n",
" serving_container_predict_route = \"/v1/chat/completions\"\n",
" else:\n",
" serving_container_predict_route = \"/v1/completions\"\n",
" else:\n",
" serving_container_predict_route = \"/generate\"\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_command=[\n",
" \"python\",\n",
" \"-m\",\n",
" (\n",
" \"vllm.entrypoints.api_server\"\n",
" if not use_openai_server\n",
" else \"vllm.entrypoints.openai.api_server\"\n",
" ),\n",
" ],\n",
" serving_container_args=vllm_args,\n",
" serving_container_ports=[7080],\n",
" serving_container_predict_route=\"/generate\",\n",
" serving_container_health_route=\"/ping\",\n",
" )\n",
"\n",
" model.deploy(\n",
" endpoint=endpoint,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" return model, endpoint\n",
"\n",
"\n",
"def deploy_model_mixtral(\n",
" model_name: str,\n",
" model_id: str,\n",
" service_account: str,\n",
" machine_type: str = \"g2-standard-96\",\n",
" accelerator_type: str = \"NVIDIA_L4\",\n",
" accelerator_count: str = 8,\n",
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
" \"\"\"Deploys Mixtral 8x7B with vLLM on Vertex AI.\"\"\"\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
"\n",
" dtype = \"bfloat16\"\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",
" f\"--dtype={dtype}\",\n",
" \"--gpu-memory-utilization=0.8\",\n",
" \"--max-model-len=4096\",\n",
" \"--download-dir=/model_cache\",\n",
" \"--disable-log-stats\",\n",
" ]\n",
" model = aiplatform.Model.upload(\n",
" display_name=model_name,\n",
" serving_container_image_uri=MIXTRAL_DOCKER_URI,\n",
" serving_container_args=vllm_args,\n",
" serving_container_ports=[7080],\n",
" serving_container_predict_route=\"/generate\",\n",
" serving_container_predict_route=serving_container_predict_route,\n",
" serving_container_health_route=\"/health\",\n",
" )\n",
"\n",
@@ -370,7 +373,7 @@
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" deploy_request_timeout=3600,\n",
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" return model, endpoint"
@@ -498,6 +501,10 @@
"# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
"# accelerator_count = 1\n",
"\n",
"# Larger setting of `max-model-len` can lead to higher requirements on\n",
"# `gpu-memory-utilization` and GPU configuration.\n",
"max_model_len = 4096\n",
"\n",
"model, endpoint = deploy_model_vllm(\n",
" model_name=get_job_name_with_datetime(prefix=\"mistral-serve-vllm\"),\n",
" model_id=prebuilt_model_id,\n",
@@ -505,6 +512,9 @@
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" max_model_len=max_model_len,\n",
" use_openai_server=False,\n",
" use_chat_completions_if_openai_server=False,\n",
")"
]
},
@@ -565,13 +575,29 @@
"response = endpoint.predict(instances=instances)\n",
"\n",
"for prediction in response.predictions:\n",
" print(prediction)"
" print(prediction)\n",
"\n",
"# Reference the following code for using the OpenAI vLLM server.\n",
"# import json\n",
"# response = endpoint.raw_predict(\n",
"# body=json.dumps({\n",
"# \"model\": prebuilt_model_id,\n",
"# \"prompt\": \"My favourite condiment is\",\n",
"# \"n\": 1,\n",
"# \"max_tokens\": 200,\n",
"# \"temperature\": 1.0,\n",
"# \"top_p\": 1.0,\n",
"# \"top_k\": 10,\n",
"# }),\n",
"# headers={\"Content-Type\": \"application/json\"},\n",
"# )\n",
"# print(response.json())"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "YKZ4CBJ2kYaW"
"id": "wOh9irbqJ-MM"
},
"source": [
"## Deploy Prebuilt Mixtral 8x7B model with vLLM\n",
@@ -584,7 +610,7 @@
{
"cell_type": "markdown",
"metadata": {
"id": "25b5b3a44cf8"
"id": "h2uCSnoaJ-MM"
},
"source": [
"Set the prebuilt model id."
@@ -594,7 +620,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "10547af949fc"
"id": "-X42gkGYJ-MM"
},
"outputs": [],
"source": [
@@ -605,7 +631,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "03d504bcd60b"
"id": "M-YiJXT3J-MM"
},
"outputs": [],
"source": [
@@ -622,20 +648,27 @@
"# accelerator_type = \"NVIDIA_TESLA_A100\"\n",
"# accelerator_count = 4\n",
"\n",
"model, endpoint = deploy_model_mixtral(\n",
"# Larger setting of `max-model-len` can lead to higher requirements on\n",
"# `gpu-memory-utilization` and GPU configuration.\n",
"max_model_len = 4096\n",
"\n",
"model, endpoint = deploy_model_vllm(\n",
" model_name=get_job_name_with_datetime(prefix=\"mixtral-serve-vllm\"),\n",
" model_id=prebuilt_model_id,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" max_model_len=max_model_len,\n",
" use_openai_server=False,\n",
" use_chat_completions_if_openai_server=False,\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "RRR11SWykYaX"
"id": "agDw0_7JJ-MM"
},
"source": [
"NOTE: If you see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint, the model server is likely still initializing. Please retry later.\n",
@@ -648,7 +681,7 @@
{
"cell_type": "markdown",
"metadata": {
"id": "3a7948c56e3d"
"id": "zLyWJK5aJ-MM"
},
"source": [
"### Run sample prompt"
@@ -658,7 +691,7 @@
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3f5a1e1de60d"
"id": "NYN1Z49SJ-MM"
},
"outputs": [],
"source": [
@@ -688,7 +721,23 @@
"response = endpoint.predict(instances=instances)\n",
"\n",
"for prediction in response.predictions:\n",
" print(prediction)"
" print(prediction)\n",
"\n",
"# Reference the following code for using the OpenAI vLLM server.\n",
"# import json\n",
"# response = endpoint.raw_predict(\n",
"# body=json.dumps({\n",
"# \"model\": prebuilt_model_id,\n",
"# \"prompt\": \"My favourite condiment is\",\n",
"# \"n\": 1,\n",
"# \"max_tokens\": 200,\n",
"# \"temperature\": 1.0,\n",
"# \"top_p\": 1.0,\n",
"# \"top_k\": 10,\n",
"# }),\n",
"# headers={\"Content-Type\": \"application/json\"},\n",
"# )\n",
"# print(response.json())"
]
},
{