Adding Phi-3-medium TPU configs to Colab deployment notebook

PiperOrigin-RevId: 687430227
This commit is contained in:
Vertex MG Team
2024-10-18 14:43:49 -07:00
committed by Copybara-Service
parent fb1871d100
commit e24c20b372
@@ -551,7 +551,7 @@
"# @markdown This section uploads prebuilt Phi-3 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
"\n",
"# @markdown Select one of the four model variations.\n",
"MODEL_ID = \"Phi-3-mini-4k-instruct\" # @param [\"Phi-3-mini-4k-instruct\", \"Phi-3-mini-128k-instruct\"] {isTemplate: true}\n",
"MODEL_ID = \"Phi-3-mini-4k-instruct\" # @param [\"Phi-3-mini-4k-instruct\", \"Phi-3-mini-128k-instruct\", \"Phi-3-medium-4k-instruct\", \"Phi-3-medium-128k-instruct\"] {isTemplate: true}\n",
"TPU_DEPLOYMENT_REGION = \"us-west1\" # @param [\"us-west1\"] {isTemplate:true}\n",
"model_id = os.path.join(MODEL_BUCKET, MODEL_ID)\n",
"hf_model_id = \"microsoft/\" + MODEL_ID\n",
@@ -564,8 +564,10 @@
"\n",
"# @markdown | Model Version | Default Max Model Length | Default TPU configuration |\n",
"# @markdown |----------------------------|------------------|-----------------------------|\n",
"# @markdown | Phi-3-mini-4k-instruct | 4096 | 1 TPU_V5e ct5lp-hightpu-1t |\n",
"# @markdown | Phi-3-mini-4k-instruct | 4096 | 1 TPU_V5e ct5lp-hightpu-1t |\n",
"# @markdown | Phi-3-mini-128k-instruct | 131072 | 4 TPU_V5e ct5lp-hightpu-4t |\n",
"# @markdown | Phi-3-medium-4k-instruct | 4096 | 4 TPU_V5e ct5lp-hightpu-4t |\n",
"# @markdown | Phi-3-medium-128k-instruct | 131072 | 4 TPU_V5e ct5lp-hightpu-4t |\n",
"\n",
"\n",
"# Note: 1 TPU V5 chip has only one core.\n",
@@ -576,7 +578,12 @@
" tpu_topo = \"1x4\"\n",
" max_model_len = 4096\n",
" machine_type = \"ct5lp-hightpu-1t\"\n",
"elif \"mini-128k\" in MODEL_ID:\n",
"elif \"medium-4k\" in MODEL_ID:\n",
" tpu_count = 4\n",
" tpu_topo = \"4x4\"\n",
" max_model_len = 4096\n",
" machine_type = \"ct5lp-hightpu-4t\"\n",
"elif \"mini-128k\" or \"medium-128k\" in MODEL_ID:\n",
" tpu_count = 4\n",
" max_model_len = 131072\n",
" tpu_topo = \"4x4\"\n",