mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
chore, refactor(egen): tpuv5e_gemma_peft_finetuning_and_serving notebook fix (#3159)
* chore: remove boilerplate and adds colab enterprise, and changes region to location * chore: lint run * chore: comments addressed and lint run
This commit is contained in:
Kaushik Koilada
@@ -32,23 +32,28 @@
|
||||
"# Vertex AI - Gemma distributed tuning with LoRA on TPUv5e, serving on L4 GPU\n",
|
||||
"\n",
|
||||
"<table align=\"left\">\n",
|
||||
" <td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb\">\n",
|
||||
" <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab (Will require the higher memory Colab pro)\n",
|
||||
" <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"><br> Run in Colab (Will require the higher memory Colab pro)\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
" <td>\n",
|
||||
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb\">\n",
|
||||
" <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
|
||||
" View on GitHub\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fofficial%2Ftraining%2Ftpuv5e_gemma_peft_finetuning_and_serving.ipynb\">\n",
|
||||
" <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
" <td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://console.cloud.google.com/vertex-ai/notebooks/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb\">\n",
|
||||
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
|
||||
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br>\n",
|
||||
"Open in Vertex AI Workbench\n",
|
||||
" </a> (An e2-standard-8 CPU w/ 250GB disk is recommended)\n",
|
||||
" </td>\n",
|
||||
" <td style=\"text-align: center\">\n",
|
||||
" <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/training/tpuv5e_gemma_peft_finetuning_and_serving.ipynb\">\n",
|
||||
" <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
|
||||
" View on GitHub\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
"</table>"
|
||||
]
|
||||
},
|
||||
@@ -89,7 +94,7 @@
|
||||
"source": [
|
||||
"### Dataset\n",
|
||||
"\n",
|
||||
"In this example, you will use the IMDB reviews dataset from TensorFlow datasets to finetune the model. Details of the dataset can be found here: https://www.tensorflow.org/datasets/catalog/imdb_reviews"
|
||||
"In this example, the IMDB reviews dataset from TensorFlow datasets is used to finetune the model. Details of the dataset can be found here: https://www.tensorflow.org/datasets/catalog/imdb_reviews"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -113,11 +118,16 @@
|
||||
"id": "i7EUnXsZhAGF"
|
||||
},
|
||||
"source": [
|
||||
"## Installation\n",
|
||||
"\n",
|
||||
"Install the following packages required to execute this notebook. \n",
|
||||
"\n",
|
||||
"Run this to install the latest google cloud platform library that supports TPUv5e"
|
||||
"## Get started"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "81711614c199"
|
||||
},
|
||||
"source": [
|
||||
"### Install Vertex AI SDK for Python and other required packages\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -143,7 +153,9 @@
|
||||
"id": "58707a750154"
|
||||
},
|
||||
"source": [
|
||||
"### Colab only: Uncomment the following cell to restart the kernel."
|
||||
"### Restart runtime (Colab only)\n",
|
||||
"\n",
|
||||
"To use the newly installed packages, you must restart the runtime on Google Colab."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -154,35 +166,53 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# import IPython\n",
|
||||
"import sys\n",
|
||||
"\n",
|
||||
"# app = IPython.Application.instance()\n",
|
||||
"# app.kernel.do_shutdown(True)"
|
||||
"if \"google.colab\" in sys.modules:\n",
|
||||
"\n",
|
||||
" import IPython\n",
|
||||
"\n",
|
||||
" app = IPython.Application.instance()\n",
|
||||
" app.kernel.do_shutdown(True)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "BF1j6f9HApxa"
|
||||
"id": "ffcde4d56c00"
|
||||
},
|
||||
"source": [
|
||||
"## Before you begin\n",
|
||||
"<div class=\"alert alert-block alert-warning\">\n",
|
||||
"<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
|
||||
"</div>\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "7176ea64999b"
|
||||
},
|
||||
"source": [
|
||||
"### Authenticate your notebook environment (Colab only)\n",
|
||||
"\n",
|
||||
"### Set up your Google Cloud project\n",
|
||||
"Authenticate your environment on Google Colab.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "7de6ef0fac42"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import sys\n",
|
||||
"\n",
|
||||
"**The following steps are required, regardless of your notebook environment.**\n",
|
||||
"if \"google.colab\" in sys.modules:\n",
|
||||
"\n",
|
||||
"1. [Select or create a Google Cloud project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
|
||||
" from google.colab import auth\n",
|
||||
"\n",
|
||||
"2. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
|
||||
"\n",
|
||||
"3. [Enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
|
||||
"\n",
|
||||
"4. If you are running this notebook locally, you need to install the [Cloud SDK](https://cloud.google.com/sdk).\n",
|
||||
"\n",
|
||||
"1. [Select or create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs.\n",
|
||||
"\n",
|
||||
"1. [Create a service account](https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console) with `Vertex AI User` and `Storage Object Admin` roles for deploying fine tuned model to Vertex AI endpoint."
|
||||
" auth.authenticate_user()"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -202,35 +232,18 @@
|
||||
"\n",
|
||||
"* Open [Kaggle settings](https://www.kaggle.com/settings)\n",
|
||||
"* Select \"Create New Token\"\n",
|
||||
"* A kaggle.json file is downloaded. It contains your Kaggle credentials. Note the username and key as you will populate this later"
|
||||
"* A kaggle.json file is downloaded. It contains your Kaggle credentials. Take note of the username and key since you'll use these later."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "WReHDGG5g0XY"
|
||||
"id": "4f872cd812d0"
|
||||
},
|
||||
"source": [
|
||||
"#### Set your project ID\n",
|
||||
"### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
|
||||
"\n",
|
||||
"**If you don't know your project ID**, try the following:\n",
|
||||
"* Run `gcloud config list`.\n",
|
||||
"* Run `gcloud projects list`.\n",
|
||||
"* See the support page: [Locate the project ID](https://support.google.com/googleapi/answer/7014113)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "oM1iC_MfAts1"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
|
||||
"\n",
|
||||
"# Set the project id\n",
|
||||
"! gcloud config set project {PROJECT_ID}"
|
||||
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -239,9 +252,9 @@
|
||||
"id": "region"
|
||||
},
|
||||
"source": [
|
||||
"#### Region\n",
|
||||
"#### Location\n",
|
||||
"\n",
|
||||
"You can also change the `REGION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
|
||||
"You can also change the `LOCATION` variable used by Vertex AI. Learn more about [Vertex AI regions](https://cloud.google.com/vertex-ai/docs/general/locations).\n",
|
||||
"\n",
|
||||
"TPUv5e is available in the [following regions listed here](https://cloud.google.com/tpu/pricing)"
|
||||
]
|
||||
@@ -254,79 +267,18 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"REGION = \"us-west1\" # @param {type: \"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "sBCra4QMA2wR"
|
||||
},
|
||||
"source": [
|
||||
"### Authenticate your Google Cloud account\n",
|
||||
"\n",
|
||||
"Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "74ccc9e52986"
|
||||
},
|
||||
"source": [
|
||||
"**1. Vertex AI Workbench**\n",
|
||||
"* Do nothing as you are already authenticated."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "de775a3773ba"
|
||||
},
|
||||
"source": [
|
||||
"**2. Local JupyterLab instance, uncomment and run:**"
|
||||
"LOCATION = \"us-west1\" # @param {type: \"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "254614fa0c46"
|
||||
"id": "oM1iC_MfAts1"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# ! gcloud auth login"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "ef21552ccea8"
|
||||
},
|
||||
"source": [
|
||||
"**3. Colab, uncomment and run:**"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "603adbbf0532"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# from google.colab import auth\n",
|
||||
"# auth.authenticate_user()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "f6b2ccc891ed"
|
||||
},
|
||||
"source": [
|
||||
"**4. Service account or other**\n",
|
||||
"* See how to grant Cloud Storage permissions to your service account at https://cloud.google.com/storage/docs/gsutil/commands/iam#ch-examples."
|
||||
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -408,7 +360,7 @@
|
||||
"id": "-EcIXiGsCePi"
|
||||
},
|
||||
"source": [
|
||||
"**Only if your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
|
||||
"**If your bucket doesn't already exist**: Run the following cell to create your Cloud Storage bucket."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -419,7 +371,7 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"! gsutil mb -l {REGION} -p {PROJECT_ID} {BUCKET_URI}"
|
||||
"! gsutil mb -l {LOCATION} -p {PROJECT_ID} {BUCKET_URI}"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -441,7 +393,7 @@
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
|
||||
"aiplatform.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -494,7 +446,7 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"image_name_train = \"gemma-lora-tuning-tpuv5e\"\n",
|
||||
"hostname = f\"{REGION}-docker.pkg.dev\"\n",
|
||||
"hostname = f\"{LOCATION}-docker.pkg.dev\"\n",
|
||||
"tag = \"latest\""
|
||||
]
|
||||
},
|
||||
@@ -507,7 +459,7 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Register gcloud as a Docker credential helper\n",
|
||||
"!gcloud auth configure-docker $REGION-docker.pkg.dev --quiet"
|
||||
"!gcloud auth configure-docker $LOCATION-docker.pkg.dev --quiet"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -520,7 +472,7 @@
|
||||
"source": [
|
||||
"# One time or use an existing repository\n",
|
||||
"!gcloud artifacts repositories create $REPOSITORY --repository-format=docker \\\n",
|
||||
"--location=$REGION --description=\"Vertex TPUv5e training repository\""
|
||||
"--location=$LOCATION --description=\"Vertex TPUv5e training repository\""
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -656,7 +608,7 @@
|
||||
"#### Kaggle credentials are required for KerasNLP training and Hex-LLM deployment with TPUs.\n",
|
||||
"Set the KAGGLE_USERNAME AND KAGGLE_KEY to pass in as an environment variable for Vertex Training to use\n",
|
||||
"Fenerate the Kaggle username and key by following [these instructions](https://github.com/Kaggle/kaggle-api?tab=readme-ov-file#api-credentials).\n",
|
||||
"You will need to review and accept the model license as mentioned earlier."
|
||||
"As mentioned earlier, you need to review and accept the model license."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1098,7 +1050,7 @@
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"job_client = aiplatform.gapic.JobServiceClient(\n",
|
||||
" client_options=dict(api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")\n",
|
||||
" client_options=dict(api_endpoint=f\"{LOCATION}-aiplatform.googleapis.com\")\n",
|
||||
")"
|
||||
]
|
||||
},
|
||||
@@ -1112,7 +1064,7 @@
|
||||
"source": [
|
||||
"create_tpuv5e_gemma_peft_job_response = job_client.create_custom_job(\n",
|
||||
" parent=\"projects/{project}/locations/{location}\".format(\n",
|
||||
" project=PROJECT_ID, location=REGION\n",
|
||||
" project=PROJECT_ID, location=LOCATION\n",
|
||||
" ),\n",
|
||||
" custom_job=tpuv5e_gemma_peft_job,\n",
|
||||
")\n",
|
||||
@@ -1313,7 +1265,7 @@
|
||||
"### Deploy fine tuned models\n",
|
||||
"This section uploads the model to Model Registry and deploys it on the Endpoint using [vLLM](https://github.com/vllm-project/vllm)\n",
|
||||
"\n",
|
||||
"The model deployment step will take 15 minutes to 1 hour to complete, depending on the model sizes."
|
||||
"The model deployment step takes 15 minutes to 1 hour to complete, depending on the model sizes."
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -1384,7 +1336,7 @@
|
||||
"startdate = startdate.strftime(\"%Y-%m-%d\")\n",
|
||||
"log_link = \"https://console.cloud.google.com/logs/query;query=resource.type=%22aiplatform.googleapis.com%2FEndpoint%22\"\n",
|
||||
"log_link += f\"%20resource.labels.endpoint_id=%22{endpoint.name}%22\"\n",
|
||||
"log_link += f\"%20resource.labels.location={REGION}\"\n",
|
||||
"log_link += f\"%20resource.labels.location={LOCATION}\"\n",
|
||||
"log_link += f\"%20timestamp%3E={startdate}\"\n",
|
||||
"print(log_link)"
|
||||
]
|
||||
@@ -1395,7 +1347,7 @@
|
||||
"id": "GI363hMzG_4U"
|
||||
},
|
||||
"source": [
|
||||
"NOTE: The overall deployment can take 30-40 minutes or more. After the deployment succeeds (15-20 minutes or so), the fine-tuned model will be downloaded from the GCS bucket used in training above. Thus, an additional ~15-20 minutes (depending on the model sizes) of waiting time is needed **after** the model deployment step above succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
|
||||
"NOTE: The overall deployment can take 30-40 minutes or more. After the deployment succeeds (15-20 minutes or so), the fine-tuned model is downloaded from the GCS bucket used in training above. Thus, an additional ~15-20 minutes (depending on the model sizes) of waiting time is needed **after** the model deployment step above succeeds and before you run the next step below. Otherwise you might see a `ServiceUnavailable: 503 502:Bad Gateway` error when you send requests to the endpoint.\n",
|
||||
"\n",
|
||||
"### Send a prediction request\n",
|
||||
"\n",
|
||||
|
||||
Reference in New Issue
Block a user