Use separate regions for training, evaluation, and deployment in Llama 3.1 finetuning notebook

PiperOrigin-RevId: 835030738

Author: vertex-mg-bot

notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb

@@ -139,7 +139,7 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# Import the necessary packages\n",
+        "# Import the necessary packages.\n",
         "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "! cd vertex-ai-samples && git reset --hard 7ae13b346a72ee2a2dc8152dd40c6ddd72d6c810\n",
         "\n",
@@ -266,9 +266,7 @@
         "        VERTEX_AI_MODEL_GARDEN_LLAMA3_1\n",
         "    ), \"Click the agreement of Llama3.1 in Vertex AI Model Garden, and get the GCS path of the model artifacts.\"\n",
         "\n",
-        "MODEL_BUCKET = VERTEX_AI_MODEL_GARDEN_LLAMA3_1\n",
-        "\n",
-        "# @markdown ---"
+        "MODEL_BUCKET = VERTEX_AI_MODEL_GARDEN_LLAMA3_1"
       ]
     },
     {
@@ -451,6 +449,13 @@
         "# @markdown Acceletor type to use for training.\n",
         "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
         "\n",
+        "# @markdown Set the Training Region. If not set, it will be set to default region.\n",
+        "TRAINING_REGION = \"\"  # @param {type: \"string\"}\n",
+        "if not TRAINING_REGION:\n",
+        "    TRAINING_REGION = REGION\n",
+        "\n",
+        "aiplatform.init(location=TRAINING_REGION)\n",
+        "\n",
         "# The pre-built training docker image.\n",
         "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
         "    repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
@@ -544,7 +549,7 @@
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
+        "    region=TRAINING_REGION,\n",
         "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count * replica_count,\n",
         "    is_for_training=True,\n",
@@ -701,6 +706,13 @@
         "# @markdown Set `RUN_EVALUATION` to False to skip the evaluation job.\n",
         "RUN_EVALUATION = True  # @param {type:\"boolean\"}\n",
         "\n",
+        "# @markdown Set the Evaluation Region. If not set, it will be set to default region.\n",
+        "EVAL_REGION = \"\"  # @param {type: \"string\"}\n",
+        "if not EVAL_REGION:\n",
+        "    EVAL_REGION = REGION\n",
+        "\n",
+        "aiplatform.init(location=EVAL_REGION)\n",
+        "\n",
         "if \"8b\" in base_model_id.lower():\n",
         "    eval_machine_type = \"g2-standard-24\"\n",
         "    eval_accelerator_type = \"NVIDIA_L4\"\n",
@@ -759,7 +771,7 @@
         "    ]\n",
         "    common_util.check_quota(\n",
         "        project_id=PROJECT_ID,\n",
-        "        region=REGION,\n",
+        "        region=EVAL_REGION,\n",
         "        accelerator_type=eval_accelerator_type,\n",
         "        accelerator_count=eval_accelerator_count,\n",
         "        is_for_training=True,\n",
@@ -804,6 +816,13 @@
         "# The pre-built serving docker image for vLLM.\n",
         "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
         "\n",
+        "# @markdown Set the Deployment Region. If not set, it will be set to default region.\n",
+        "DEPLOY_REGION = \"\"  # @param {type: \"string\"}\n",
+        "if not DEPLOY_REGION:\n",
+        "    DEPLOY_REGION = REGION\n",
+        "\n",
+        "aiplatform.init(location=DEPLOY_REGION)\n",
+        "\n",
         "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
         "if \"8b\" in base_model_id.lower():\n",
         "    machine_type = \"g2-standard-12\"\n",
@@ -822,7 +841,7 @@
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
+        "    region=DEPLOY_REGION,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count,\n",
         "    is_for_training=False,\n",