From 9742d29e512ff8913e2abc54cb78b3486a202a61 Mon Sep 17 00:00:00 2001 From: Rayan Dasoriya Date: Thu, 8 May 2025 09:49:43 -0700 Subject: [PATCH] Add eval harness support for Qwen2.5 finetuning notebook PiperOrigin-RevId: 756350741 --- ...el_garden_pytorch_qwen2_5_finetuning.ipynb | 160 ++++++++++++++---- 1 file changed, 127 insertions(+), 33 deletions(-) diff --git a/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb b/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb index fd755c8ab..aae56da3f 100644 --- a/notebooks/community/model_garden/model_garden_pytorch_qwen2_5_finetuning.ipynb +++ b/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", - "# @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", + "# @title Select Evaluation Checkpoint\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", - " machine_type = \"g2-standard-48\"\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", + "accelerator_type = \"NVIDIA_L4\"\n", + "if \"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",