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

change default gpu type as empty to avoid unintended gpu runtime creation

Browse Source 
PiperOrigin-RevId: 756199290
This commit is contained in:
Vertex MG Team
2025-05-08 01:17:09 -07:00
committed by Copybara-Service
parent 6888728be5
commit 30fda48397
1 changed files with 40 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/model_garden/model_garden_axolotl_finetuning.ipynb
+40 -37
View File
@@ -107,52 +107,55 @@
"RUNTIME_PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
"RUNTIME_REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
"\n",
"RUNTIME_ACCELERATOR_TYPE = \"NVIDIA_L4\" # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\"]\n",
"RUNTIME_ACCELERATOR_TYPE = \"\" # @param [\"\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\"]\n",
"RUNTIME_ACCELERATOR_COUNT = \"1\" # @param [1, 2, 4, 8, 16]\n",
"RUNTIME_ACCELERATOR_COUNT = int(RUNTIME_ACCELERATOR_COUNT)\n",
"\n",
"if RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 1:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-8\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 2:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-24\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 4:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-48\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 8:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-96\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_TESLA_A100\" and RUNTIME_ACCELERATOR_COUNT != 16:\n",
" RUNTIME_MACHINE_TYPE = f\"a2-highgpu-{RUNTIME_ACCELERATOR_COUNT}g\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_TESLA_A100\" and RUNTIME_ACCELERATOR_COUNT == 16:\n",
" RUNTIME_MACHINE_TYPE = \"a2-megagpu-16g\"\n",
"elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_A100_80GB\" and RUNTIME_ACCELERATOR_COUNT in [1, 2, 4, 8]:\n",
" RUNTIME_MACHINE_TYPE = f\"a2-ultragpu-{RUNTIME_ACCELERATOR_COUNT}g\"\n",
"if not RUNTIME_ACCELERATOR_TYPE:\n",
" print(\"No accelerator type specified. Skipping runtime creation.\")\n",
"else:\n",
" raise ValueError(f\"Invalid GPU type {RUNTIME_ACCELERATOR_TYPE}, and count {RUNTIME_ACCELERATOR_COUNT} combination.\")\n",
" if RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 1:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-8\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 2:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-24\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 4:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-48\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_L4\" and RUNTIME_ACCELERATOR_COUNT == 8:\n",
" RUNTIME_MACHINE_TYPE = \"g2-standard-96\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_TESLA_A100\" and RUNTIME_ACCELERATOR_COUNT != 16:\n",
" RUNTIME_MACHINE_TYPE = f\"a2-highgpu-{RUNTIME_ACCELERATOR_COUNT}g\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_TESLA_A100\" and RUNTIME_ACCELERATOR_COUNT == 16:\n",
" RUNTIME_MACHINE_TYPE = \"a2-megagpu-16g\"\n",
" elif RUNTIME_ACCELERATOR_TYPE == \"NVIDIA_A100_80GB\":\n",
" RUNTIME_MACHINE_TYPE = f\"a2-ultragpu-{RUNTIME_ACCELERATOR_COUNT}g\"\n",
" else:\n",
" raise ValueError(f\"Invalid GPU type {RUNTIME_ACCELERATOR_TYPE}, and count {RUNTIME_ACCELERATOR_COUNT} combination.\")\n",
"\n",
"print(f\"Machine type: {RUNTIME_MACHINE_TYPE}\")\n",
" print(f\"Machine type: {RUNTIME_MACHINE_TYPE}\")\n",
"\n",
"uuid = uuid.uuid4()\n",
"RUNTIME_DISPLAY_NAME = f\"axolotl-{RUNTIME_ACCELERATOR_TYPE}-{RUNTIME_ACCELERATOR_COUNT}-{uuid}\"\n",
"print(f\"Creating runtime with display name: {RUNTIME_DISPLAY_NAME}\")\n",
" uuid = uuid.uuid4()\n",
" RUNTIME_DISPLAY_NAME = f\"axolotl-{RUNTIME_ACCELERATOR_TYPE}-{RUNTIME_ACCELERATOR_COUNT}-{uuid}\"\n",
" print(f\"Creating runtime with display name: {RUNTIME_DISPLAY_NAME}\")\n",
"\n",
"# create runtime template\n",
"shell_output = ! gcloud colab runtime-templates create --display-name=$RUNTIME_DISPLAY_NAME \\\n",
" --project=$RUNTIME_PROJECT_ID --region=$RUNTIME_REGION \\\n",
" --machine-type=$RUNTIME_MACHINE_TYPE --accelerator-type=$RUNTIME_ACCELERATOR_TYPE \\\n",
" --accelerator-count=$RUNTIME_ACCELERATOR_COUNT --disk-type=PD_BALANCED\n",
"shell_output = \"\\n\".join(shell_output)\n",
"print(shell_output)\n",
"RUNTIME_TEMPLATE_ID = re.search(r\"projects/.*/locations/.*/notebookRuntimeTemplates/(\\d+)\", shell_output).group(1)\n",
" # create runtime template\n",
" shell_output = ! gcloud colab runtime-templates create --display-name=$RUNTIME_DISPLAY_NAME \\\n",
" --project=$RUNTIME_PROJECT_ID --region=$RUNTIME_REGION \\\n",
" --machine-type=$RUNTIME_MACHINE_TYPE --accelerator-type=$RUNTIME_ACCELERATOR_TYPE \\\n",
" --accelerator-count=$RUNTIME_ACCELERATOR_COUNT --disk-type=PD_BALANCED\n",
" shell_output = \"\\n\".join(shell_output)\n",
" print(shell_output)\n",
" RUNTIME_TEMPLATE_ID = re.search(r\"projects/.*/locations/.*/notebookRuntimeTemplates/(\\d+)\", shell_output).group(1)\n",
"\n",
"# create runtime\n",
"shell_output = ! gcloud colab runtimes create --display-name=$RUNTIME_DISPLAY_NAME \\\n",
" --runtime-template=$RUNTIME_TEMPLATE_ID --project=$RUNTIME_PROJECT_ID \\\n",
" --region=$RUNTIME_REGION\n",
"shell_output = \"\\n\".join(shell_output)\n",
"print(shell_output)\n",
"RUNTIME_ID = re.search(r\"projects/.*/locations/.*/notebookRuntimes/(\\d+)\", shell_output).group(1)\n",
" # create runtime\n",
" shell_output = ! gcloud colab runtimes create --display-name=$RUNTIME_DISPLAY_NAME \\\n",
" --runtime-template=$RUNTIME_TEMPLATE_ID --project=$RUNTIME_PROJECT_ID \\\n",
" --region=$RUNTIME_REGION\n",
" shell_output = \"\\n\".join(shell_output)\n",
" print(shell_output)\n",
" RUNTIME_ID = re.search(r\"projects/.*/locations/.*/notebookRuntimes/(\\d+)\", shell_output).group(1)\n",
"\n",
"# start runtime\n",
"! gcloud colab runtimes start $RUNTIME_ID --project=$RUNTIME_PROJECT_ID --region=$RUNTIME_REGION"
" # start runtime\n",
" ! gcloud colab runtimes start $RUNTIME_ID --project=$RUNTIME_PROJECT_ID --region=$RUNTIME_REGION"
]
},
{