Add eval harness support for Qwen2.5 finetuning notebook

PiperOrigin-RevId: 756350741

notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb

@@ -140,7 +140,7 @@
         "\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 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
+        "! cd vertex-ai-samples && git reset --hard f147b50332ec996f1c309653c35b91a5eed7824e\n",
         "\n",
         "import datetime\n",
         "import importlib\n",
@@ -330,7 +330,7 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "iu1YAu8315sG"
+        "id": "oHvYSPr7JdZq"
       },
       "outputs": [],
       "source": [
@@ -346,7 +346,7 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "_mNcaofv4zpv"
+        "id": "X3N8O9_0JdZq"
       },
       "outputs": [],
       "source": [
@@ -400,7 +400,7 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "ivVGS9dHXPOz"
+        "id": "VvQLIpjhJdZq"
       },
       "outputs": [],
       "source": [
@@ -602,7 +602,7 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "x93f7805YwJg"
+        "id": "6jQZVMzcJdZq"
       },
       "outputs": [],
       "source": [
@@ -622,18 +622,129 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "qmHW6m8xG_4U"
+        "id": "Ys3FzqDaJdZq"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select Evaluation Checkpoint\n",
-        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
         "\n",
         "if train_job.end_time is None:\n",
         "    print(\"Waiting for the training job to finish...\")\n",
         "    train_job.wait()\n",
         "    print(\"The training job has finished.\")\n",
         "\n",
+        "# @markdown The following checkpoints are available for evaluation:\n",
+        "! gcloud storage ls \"{lora_output_dir}/node-0\" | grep \"checkpoint-\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "dztN2-JdJdZq"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run Evaluation Job\n",
+        "# @markdown This section runs the evaluation using [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) on the finetuned model. The evaluation takes approximately 20 mins to finish.\n",
+        "\n",
+        "# The pre-built evaluation docker image for LM Evaluation Harness.\n",
+        "LM_EVAL_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-lm-evaluation-harness:20250410_1035_RC00\"\n",
+        "\n",
+        "# @markdown Set `RUN_EVALUATION` to False to skip the evaluation job.\n",
+        "RUN_EVALUATION = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "eval_accelerator_type = \"NVIDIA_L4\"\n",
+        "gpu_memory_utilization = 0.85\n",
+        "\n",
+        "if \"7B\" in base_model_id or \"14B\" in base_model_id:\n",
+        "    eval_machine_type = \"g2-standard-48\"\n",
+        "    eval_accelerator_count = 4\n",
+        "elif \"32B\" in base_model_id:\n",
+        "    eval_machine_type = \"g2-standard-96\"\n",
+        "    eval_accelerator_count = 8\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        \"Recommended machine settings not found for model: %s\" % base_model_id\n",
+        "    )\n",
+        "\n",
+        "# @markdown Set `evaluation_checkpoint_dir` to an intermediate checkpoint from the above training job. If not set, the evaluation job will use the merged model.\n",
+        "evaluation_checkpoint_dir = \"\"  # @param {type:\"string\"}\n",
+        "if evaluation_checkpoint_dir:\n",
+        "    pretrained = pretrained_model_id\n",
+        "else:\n",
+        "    pretrained = merged_model_output_dir\n",
+        "\n",
+        "# @markdown Evaluation tasks to run.\n",
+        "eval_tasks = \"coqa\"  # @param {type:\"string\"}\n",
+        "# @markdown Model to use for evaluation.\n",
+        "model = \"vllm\"  # @param {type:\"string\"}\n",
+        "# @markdown Batch size for evaluation.\n",
+        "batch_size = \"auto\"  # @param {type:\"string\"}\n",
+        "apply_chat_template = True\n",
+        "max_model_len = 4096  # Maximum context length.\n",
+        "\n",
+        "model_args = f\"tensor_parallel_size={eval_accelerator_count},max_model_len={max_model_len},gpu_memory_utilization={gpu_memory_utilization},enforce_eager=True\"\n",
+        "eval_output_dir = os.path.join(base_output_dir, \"lm_eval\")\n",
+        "\n",
+        "lm_eval_job_args = [\n",
+        "    \"--task=lm_eval\",\n",
+        "    f\"--model={model}\",\n",
+        "    f\"--eval_tasks={eval_tasks}\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained}\",\n",
+        "    f\"--model_args={model_args}\",\n",
+        "    f\"--output_dir={eval_output_dir}\",\n",
+        "    f\"--apply_chat_template={apply_chat_template}\",\n",
+        "    f\"--batch_size={batch_size}\",\n",
+        "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
+        "]\n",
+        "\n",
+        "if evaluation_checkpoint_dir:\n",
+        "    lm_eval_job_args.append(f\"--lora_path={evaluation_checkpoint_dir}\")\n",
+        "\n",
+        "if RUN_EVALUATION:\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=eval_accelerator_type,\n",
+        "        accelerator_count=eval_accelerator_count,\n",
+        "        is_for_training=True,\n",
+        "    )\n",
+        "    lm_eval_job = aiplatform.CustomContainerTrainingJob(\n",
+        "        display_name=common_util.get_job_name_with_datetime(\"qwen2_5-lm-eval\"),\n",
+        "        container_uri=LM_EVAL_DOCKER_URI,\n",
+        "        labels=labels,\n",
+        "    )\n",
+        "\n",
+        "    print(\"Running evaluation job with args:\")\n",
+        "    print(\" \\\\\\n\".join(lm_eval_job_args))\n",
+        "    lm_eval_job.run(\n",
+        "        args=lm_eval_job_args,\n",
+        "        replica_count=1,\n",
+        "        machine_type=eval_machine_type,\n",
+        "        accelerator_type=eval_accelerator_type,\n",
+        "        accelerator_count=eval_accelerator_count,\n",
+        "        boot_disk_size_gb=boot_disk_size_gb,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        base_output_dir=base_output_dir,\n",
+        "    )\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qWpHMUt3JdZq"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
+        "\n",
         "print(\"Deploying models in:\", merged_model_output_dir)\n",
         "\n",
         "# The pre-built serving docker image for vLLM.\n",
@@ -643,20 +754,15 @@
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
-        "if \"32b\" in base_model_id.lower():\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
-        "    machine_type = \"g2-standard-48\"\n",
+        "if \"7b\" in base_model_id.lower():\n",
-        "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    per_node_accelerator_count = 4\n",
-        "elif \"14b\" in base_model_id.lower():\n",
-        "    machine_type = \"g2-standard-48\"\n",
-        "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    per_node_accelerator_count = 4\n",
-        "elif \"7b\" in base_model_id.lower():\n",
         "    machine_type = \"g2-standard-24\"\n",
-        "    accelerator_type = \"NVIDIA_L4\"\n",
         "    per_node_accelerator_count = 2\n",
+        "elif \"14b\" in base_model_id.lower() or \"32b\" in base_model_id.lower():\n",
+        "    machine_type = \"g2-standard-48\"\n",
+        "    per_node_accelerator_count = 4\n",
         "else:\n",
-        "    raise ValueError(f\"Unsupported model ID or GCS path: {base_model_id}.\")\n",
+        "    raise ValueError(f\"Unsupported model ID: {base_model_id}.\")\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
@@ -674,18 +780,6 @@
         "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
         "\n",
         "\n",
-        "def get_deploy_source() -> str:\n",
-        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
-        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
-        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
-        "        return \"notebook_colab_enterprise\"\n",
-        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
-        "        return \"notebook_workbench\"\n",
-        "    else:\n",
-        "        # Legacy workbench, legacy colab, or other custom environments.\n",
-        "        return \"notebook_environment_unspecified\"\n",
-        "\n",
-        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
@@ -805,7 +899,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_qwen2_5_finetuning.ipynb\",\n",
-        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -816,7 +910,7 @@
         "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"qwen2_5-vllm-serve\"),\n",
         "    model_id=merged_model_output_dir,\n",
-        "    publisher=\"alibaba\",\n",
+        "    publisher=\"qwen\",\n",
         "    publisher_model_id=\"qwen2_5\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",