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

Add YaRN scaling to QwQ deployment notebook for 128k context length. Fix chatCompletions format.

Browse Source 
PiperOrigin-RevId: 738871259
This commit is contained in:
Dustin Luong
2025-03-20 11:09:46 -07:00
committed by Copybara-Service
parent 8ea6932dc2
commit dc59488638
1 changed files with 22 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/model_garden/model_garden_pytorch_qwq_deployment.ipynb
+22 -7
View File
@@ -224,17 +224,25 @@
"use_dedicated_endpoint = True # @param {type:\"boolean\"}\n",
"\n",
"accelerator_type = \"NVIDIA_L4\" # @param [\"NVIDIA_L4\", \"NVIDIA_H100_80GB\"] {isTemplate: true}\n",
"max_model_len = 32768 # Maximum context length without YaRN.\n",
"gpu_memory_utilization = 0.9\n",
"\n",
"# @markdown For inputs exceeding 8,192 tokens, enable [YaRN](https://arxiv.org/abs/2309.00071) on the deployment to improve the model's ability to capture long-sequence information effectively.\n",
"# @markdown Enabling YaRN will allow for 128k context-length but may require more GPUs.\n",
"enable_yarn_scaling = False # @param {type:\"boolean\"}\n",
"\n",
"if enable_yarn_scaling:\n",
" max_model_len = 131072\n",
"else:\n",
" max_model_len = 32768\n",
"\n",
"if accelerator_type == \"NVIDIA_L4\":\n",
" accelerator_count = 4\n",
" # Sets machine type to g2-standard-48 for 4 L4s\n",
" machine_type = \"g2-standard-48\"\n",
" if enable_yarn_scaling:\n",
" accelerator_count = 8\n",
" machine_type = \"g2-standard-96\"\n",
" else:\n",
" accelerator_count = 4\n",
" machine_type = \"g2-standard-48\"\n",
"elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
" accelerator_count = 2\n",
" # Sets machine type to a3-highgpu-2g for 2 NVIDIA H100 80GBs.\n",
" machine_type = \"a3-highgpu-2g\"\n",
"else:\n",
" raise ValueError(\n",
@@ -276,6 +284,7 @@
" use_dedicated_endpoint: bool = False,\n",
" max_num_seqs: int = 256,\n",
" model_type: str = None,\n",
" enable_yarn_scaling: bool = False,\n",
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
" \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
" endpoint = aiplatform.Endpoint.create(\n",
@@ -324,6 +333,10 @@
" vllm_args.append(\n",
" f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
" )\n",
" if enable_yarn_scaling:\n",
" vllm_args.append(\n",
" '--rope-scaling=\\'{\"factor\":4.0,\"original_max_position_embeddings\":32768,\"type\":\"yarn\"}\\''\n",
" )\n",
"\n",
" if model_type:\n",
" vllm_args.append(f\"--model-type={model_type}\")\n",
@@ -383,9 +396,11 @@
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" max_model_len=max_model_len,\n",
" gpu_memory_utilization=gpu_memory_utilization,\n",
" gpu_memory_utilization=0.9,\n",
" enable_chunked_prefill=True,\n",
" use_dedicated_endpoint=use_dedicated_endpoint,\n",
" max_num_seqs=128,\n",
" enable_yarn_scaling=enable_yarn_scaling,\n",
")\n",
"\n",
"# @markdown Click \"Show Code\" to see more details."