notebooks/community/model_garden/model_garden_axolotl_finetuning.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1jg2qBjVb4yf"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YAeljAi7b4yg"
      },
      "source": [
        "# Vertex AI Model Garden - Fine-tuning with Axolotl\n",
        "\n",
        "<table><tbody><tr>\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%2Fcommunity%2Fmodel_garden%2Fmodel_garden_axolotl_finetuning.ipynb\">\n",
        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_axolotl_finetuning.ipynb\">\n",
        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</tr></tbody></table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U5NXBxyjf1xs"
      },
      "source": [
        "## Overview\n",
        "This notebook demonstrates fine-tuning using [Axolotl](https://github.com/axolotl-ai-cloud/axolotl). Axolotl streamlines AI model fine-tuning by providing a wide range of training recipes and supporting multiple configurations and architectures.\n",
        "\n",
        "We can use either Enterprise Colab runtime or Vertex AI training for fine-tuning using axolotl.\n",
        "Colab runtime has below advantages:\n",
        "- **Sanity check for flags**: Use Enterprise Colab runtime to do sanity check for Axolotl flags before running it on Vertex AI training directly.\n",
        "- **Quick experimentations**: Use Enterprise Colab runtime to do quick experimentations with Axolotl flags. The [max-steps](https://github.com/axolotl-ai-cloud/axolotl/blob/8fb72cbc0b94129141bae5fa4d84edd23b648af6/docs/config.qmd#L360) flag is useful to limit the training time.\n",
        "- **Debugging**: Use Enterprise Colab runtime to debug axolotl fine-tuning. This can be more efficient because debugging on the Vertex AI training involves waiting for resources to be provisioned, which can add delays. Also it is easier to add debug statements on Enterprise Colab runtime compared to Vertex AI training.\n",
        "\n",
        "Once the local fine-tuning is verified, the Vertex AI training is the recommended way to run the fine-tuning. Vertex AI training has several advantages, including:\n",
        "- **Running multiple training jobs in parallel**: This can be useful for hyperparameter tuning or running experiments with different datasets etc.\n",
        "- **For High End GPU**: Vertex AI training provides access to higher-end GPUs like the H100, which can be crucial if you encounter out-of-memory (OOM) errors.\n",
        "- **[DWS support](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws)**: DWS makes Vertex AI training more cost-effective, and easier to manage, especially in scenarios where GPU availability is a concern.\n",
        "Refer to [this documentation](https://cloud.google.com/vertex-ai/docs/training/overview#vertexi-ai-operationalizes-training-at-scale) for more details on Vertex AI training advantages.\n",
        "\n",
        "### Objective\n",
        "- Train model using Axolotl in local Enterprise Colab runtime.\n",
        "- Train model using Axolotl with Vertex AI Training."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K-YsE6oUoxjY"
      },
      "source": [
        "## Setup Colab Runtime\n",
        "**You need to setup the Colab Runtime with L4 GPU or A100 GPU if you want to run local finetuning. The following sections perform the setup for L4 GPU.**\n",
        "To learn more about creating runtime, you can optionally read [this](https://cloud.google.com/colab/docs/create-runtime).\n",
        "\n",
        "**Note: make sure to create a runtime with appropriate machine type and gpu type to avoid out of memory issues. [Refer this](https://huggingface.co/spaces/hf-accelerate/model-memory-usage) to decide which machine type and gpu type to select.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LGO6kScKoxjY"
      },
      "outputs": [],
      "source": [
        "# @title Create runtime\n",
        "# @markdown This cell creates a runtime template and then creates a runtime using that template.\n",
        "# @markdown **If you have already created a runtime, you can skip this cell.**\n",
        "# @markdown This cell can take up to 5 minutes to run.\n",
        "# @markdown After the cell execution finishes, you have to connect manually to the runtime by following [the instructions here](https://cloud.google.com/colab/docs/connect-to-runtime).\n",
        "\n",
        "import os\n",
        "import uuid\n",
        "import re\n",
        "\n",
        "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_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\":\n",
        "  assert RUNTIME_ACCELERATOR_COUNT in [1, 2, 4, 8], \"Only 1, 2, 4, 8 A100-80GB are supported.\"\n",
        "  RUNTIME_MACHINE_TYPE = f\"a2-ultragpu-{RUNTIME_ACCELERATOR_COUNT}g\"\n",
        "\n",
        "uuid = uuid.uuid4()\n",
        "RUNTIME_DISPLAY_NAME = f\"axolotl-{RUNTIME_ACCELERATOR_TYPE}-{RUNTIME_ACCELERATOR_COUNT}-{uuid}\"\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",
        "\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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Gq2bkOMsYFJ"
      },
      "source": [
        "## Before you begin"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "emRjTZVEsYFJ"
      },
      "outputs": [],
      "source": [
        "# @title Import utility packages for fine-tuning\n",
        "\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\n",
        "\n",
        "# Import the necessary packages.\n",
        "\n",
        "import datetime\n",
        "import importlib\n",
        "import os\n",
        "import pathlib\n",
        "import uuid\n",
        "from typing import Tuple\n",
        "\n",
        "import requests\n",
        "import yaml\n",
        "from google.cloud import aiplatform\n",
        "\n",
        "common_util = importlib.import_module(\n",
        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
        ")\n",
        "\n",
        "train_job = None\n",
        "models, endpoints = {}, {}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E0LS8jpwyUFu"
      },
      "outputs": [],
      "source": [
        "# @title Setup Google Cloud project\n",
        "\n",
        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
        "\n",
        "# @markdown 2. For finetuning, follow [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws) to use Dynamic Workload Scheduler. For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs, and [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_a100_gpus) quota for Nvidia Tesla A100 GPUs. To train using L4 gpus with default quota, check [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_l4_gpus) quota for Nvidia L4 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
        "\n",
        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
        "\n",
        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
        "\n",
        "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
        "\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
        "\n",
        "REGION = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "# Get the default cloud project id.\n",
        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
        "\n",
        "# Get the default region for launching jobs.\n",
        "if not REGION:\n",
        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
        "        raise ValueError(\n",
        "            \"REGION must be set. See\"\n",
        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
        "            \" available cloud locations.\"\n",
        "        )\n",
        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
        "\n",
        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
        "\n",
        "# Cloud Storage bucket for storing the experiment artifacts.\n",
        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
        "# prefer using your own GCS bucket, change the value yourself below.\n",
        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
        "    bucket_region = shell_output[0].strip().lower()\n",
        "    if bucket_region != REGION:\n",
        "        raise ValueError(\n",
        "            \"Bucket region %s is different from notebook region %s\"\n",
        "            % (bucket_region, REGION)\n",
        "        )\n",
        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"axolotl\")\n",
        "\n",
        "\n",
        "# Initialize Vertex AI API.\n",
        "print(\"Initializing Vertex AI API.\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "# Gets the default SERVICE_ACCOUNT.\n",
        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FcVDnCvbFZUM"
      },
      "source": [
        "## Finetune with Axolotl"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "K8uSeXw9f1xs"
      },
      "outputs": [],
      "source": [
        "# @title Set Axolotl config\n",
        "\n",
        "# @markdown You can use below axolotl configs taken from [examples directory](https://github.com/axolotl-ai-cloud/axolotl/tree/8fb72cbc0b94129141bae5fa4d84edd23b648af6/examples), which have been verified by model garden team through internal testing. Note that we have used A100 80GB and H100 80GB GPU for testing.\n",
        "# @markdown > | Model Name | Base Model | Axolotl Config |\n",
        "# @markdown | ----------- | ----------- | ----------- |\n",
        "# @markdown | code-llama | codellama/CodeLlama-7b-hf | examples/code-llama/7b/lora.yml |\n",
        "# @markdown | code-llama | codellama/CodeLlama-7b-hf | examples/code-llama/7b/qlora.yml |\n",
        "# @markdown | code-llama | codellama/CodeLlama-13b-hf | examples/code-llama/13b/lora.yml |\n",
        "# @markdown | code-llama | codellama/CodeLlama-13b-hf | examples/code-llama/13b/qlora.yml |\n",
        "# @markdown | code-llama | codellama/CodeLlama-34b-hf | examples/code-llama/34b/lora.yml |\n",
        "# @markdown | code-llama | codellama/CodeLlama-34b-hf | examples/code-llama/34b/qlora.yml |\n",
        "# @markdown | falcon | tiiuae/falcon-7b | examples/falcon/config-7b-lora.yml |\n",
        "# @markdown | falcon | tiiuae/falcon-7b | examples/falcon/config-7b.yml |\n",
        "# @markdown | gemma | google/gemma-7b | examples/gemma/qlora.yml |\n",
        "# @markdown | llama-2 | NousResearch/Llama-2-7b-hf | examples/llama-2/fft_optimized.yml |\n",
        "# @markdown | llama-2 | NousResearch/Llama-2-7b-hf | examples/llama-2/loftq.yml |\n",
        "# @markdown | llama-2 | NousResearch/Llama-2-7b-hf | examples/llama-2/lora.yml |\n",
        "# @markdown | llama-2 | NousResearch/Llama-2-7b-hf | examples/llama-2/qlora-fsdp.yml |\n",
        "# @markdown | llama-2 | NousResearch/Llama-2-7b-hf | examples/llama-2/qlora.yml |\n",
        "# @markdown | llama-3 | NousResearch/Meta-Llama-3.1-8B | examples/llama-3/fft-8b.yaml |\n",
        "# @markdown | llama-3 | NousResearch/Meta-Llama-3-8B-Instruct | examples/llama-3/instruct-lora-8b.yml |\n",
        "# @markdown | llama-3 | NousResearch/Meta-Llama-3-8B | examples/llama-3/lora-8b.yml |\n",
        "# @markdown | llama-3 | NousResearch/Llama-3.2-1B | examples/llama-3/qlora-1b.yml |\n",
        "# @markdown | llama-3 | casperhansen/llama-3-70b-fp16 | examples/llama-3/qlora-fsdp-70b.yaml |\n",
        "# @markdown | mistral | mistralai/Mistral-7B-v0.1 | examples/mistral/config.yml |\n",
        "# @markdown | mistral | mistralai/Mistral-7B-v0.1 | examples/mistral/lora-mps.yml |\n",
        "# @markdown | mistral | mistralai/Mistral-7B-v0.1 | examples/mistral/lora.yml |\n",
        "# @markdown | mistral | mistralai/Mistral-7B-v0.1 | examples/mistral/mistral-qlora-orpo.yml |\n",
        "# @markdown | mistral | mistral-community/Mixtral-8x22B-v0.1 | examples/mistral/mixtral-8x22b-qlora-fsdp.yml |\n",
        "# @markdown | mistral | mistralai/Mixtral-8x7B-v0.1 | examples/mistral/mixtral-qlora-fsdp.yml |\n",
        "# @markdown | mistral | mistralai/Mistral-7B-v0.1 | examples/mistral/qlora.yml |\n",
        "# @markdown | openllama-3b | openlm-research/open_llama_3b_v2 | examples/openllama-3b/config.yml |\n",
        "# @markdown | openllama-3b | openlm-research/open_llama_3b_v2 | examples/openllama-3b/lora.yml |\n",
        "# @markdown | openllama-3b | openlm-research/open_llama_3b_v2 | examples/openllama-3b/qlora.yml |\n",
        "# @markdown | phi | microsoft/Phi-3.5-mini-instruct | examples/phi/lora-3.5.yaml |\n",
        "# @markdown | phi | microsoft/phi-1_5 | examples/phi/phi-ft.yml |\n",
        "# @markdown | phi | microsoft/phi-1_5 | examples/phi/phi-qlora.yml |\n",
        "# @markdown | phi | microsoft/phi-2 | examples/phi/phi2-ft.yml |\n",
        "# @markdown | phi | microsoft/Phi-3-mini-4k-instruct | examples/phi/phi3-ft.yml |\n",
        "# @markdown | qwen | Qwen/Qwen1.5-MoE-A2.7B | examples/qwen/qwen2-moe-lora.yaml |\n",
        "# @markdown | qwen | Qwen/Qwen1.5-MoE-A2.7B | examples/qwen/qwen2-moe-qlora.yaml |\n",
        "# @markdown | qwen2 | Qwen/Qwen2.5-0.5B | examples/qwen2/dpo.yaml |\n",
        "# @markdown | qwen2 | Qwen/Qwen2-7B | examples/qwen2/qlora-fsdp.yaml |\n",
        "# @markdown | tiny-llama | TinyLlama/TinyLlama_v1.1 | examples/tiny-llama/lora-mps.yml |\n",
        "# @markdown | tiny-llama | TinyLlama/TinyLlama_v1.1 | examples/tiny-llama/lora.yml |\n",
        "# @markdown | tiny-llama | TinyLlama/TinyLlama-1.1B-Chat-v1.0 | examples/tiny-llama/pretrain.yml |\n",
        "# @markdown | tiny-llama | TinyLlama/TinyLlama_v1.1 | examples/tiny-llama/qlora.yml |\n",
        "\n",
        "# @markdown 1. Set Axolotl config source.<br>\n",
        "# @markdown For `GITHUB` as source, you can explore different Axolotl configurations in the [examples directory](https://github.com/axolotl-ai-cloud/axolotl/tree/8fb72cbc0b94129141bae5fa4d84edd23b648af6/examples). For `GITHUB` source, `AXOLOTL_CONFIG_PATH` should start with `examples/`. e.g. examples/tiny-llama/lora.yml.<br>\n",
        "# @markdown For `LOCAL` as source, create Axolotl config yaml file and specify correct path below. Note that, the local file will be copied to GCS bucket before running Vertex AI training job. For `LOCAL` source, `AXOLOTL_CONFIG_PATH` should be a complete path of the config file. e.g. /content/lora.yml.<br>\n",
        "# @markdown For `GCS` as source, specify the GCS URI to the Axolotl config file. Make sure the file is accessible to service account used in the notebook. For `GCS` source, `AXOLOTL_CONFIG_PATH` should be a complete GCS URI of the config file. e.g. gs://bucket/path/to/config/file.yml.\n",
        "\n",
        "AXOLOTL_SOURCE = \"GITHUB\"  # @param [\"GITHUB\", \"LOCAL\", \"GCS\"]\n",
        "\n",
        "# @markdown 2. Set the Axolotl config file path.\n",
        "AXOLOTL_CONFIG_PATH = \"examples/tiny-llama/lora.yml\"  # @param {type:\"string\"}\n",
        "\n",
        "assert AXOLOTL_CONFIG_PATH, \"AXOLOTL_CONFIG_PATH must be set.\"\n",
        "\n",
        "if AXOLOTL_SOURCE == \"GITHUB\":\n",
        "    assert AXOLOTL_CONFIG_PATH.startswith(\n",
        "        \"examples/\"\n",
        "    ), \"AXOLOTL_CONFIG_PATH must start with examples/ for GITHUB source.\"\n",
        "    github_url = f\"https://github.com/axolotl-ai-cloud/axolotl/raw/8fb72cbc0b94129141bae5fa4d84edd23b648af6/{AXOLOTL_CONFIG_PATH}\"\n",
        "    r = requests.get(github_url)\n",
        "    axolotl_config = r.content.decode(\"utf-8\")\n",
        "    axolotl_config = yaml.safe_load(axolotl_config)\n",
        "elif AXOLOTL_SOURCE == \"LOCAL\":\n",
        "    config_path = pathlib.Path(AXOLOTL_CONFIG_PATH)\n",
        "    assert config_path.exists(), \"AXOLOTL_CONFIG_PATH must exist for LOCAL source.\"\n",
        "    file_content = config_path.read_text()\n",
        "    axolotl_config = yaml.safe_load(file_content)\n",
        "elif AXOLOTL_SOURCE == \"GCS\":\n",
        "    local_path = pathlib.Path(\"/content/tmp/axolotl_config.yml\")\n",
        "    common_util.download_gcs_file_to_local(AXOLOTL_CONFIG_PATH, local_path.absolute())\n",
        "    file_content = local_path.read_text()\n",
        "    axolotl_config = yaml.safe_load(file_content)\n",
        "    AXOLOTL_CONFIG_PATH = common_util.gcs_fuse_path(AXOLOTL_CONFIG_PATH)\n",
        "else:\n",
        "    raise ValueError(f\"Unsupported AXOLOTL_SOURCE: {AXOLOTL_SOURCE}\")\n",
        "\n",
        "OUTPUT_GCS_URI = MODEL_BUCKET\n",
        "\n",
        "if not OUTPUT_GCS_URI.startswith(\"gs://\"):\n",
        "    OUTPUT_GCS_URI = f\"gs://{OUTPUT_GCS_URI}\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dySqVhK8cFoO"
      },
      "outputs": [],
      "source": [
        "# @title **[Optional]** Setup HF token\n",
        "# @markdown Some models like Gemma2, Mistral, Llama3 etc require a token to access with [gated access from huggingface](https://huggingface.co/docs/hub/en/models-gated).\n",
        "HF_TOKEN = \"\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Z3xVT_VtFZUM"
      },
      "outputs": [],
      "source": [
        "# @title **[Optional]** Setup dataset\n",
        "\n",
        "# @markdown This section configures the dataset used for fine-tuning. **Note: If you don't fill any of the dataset options given below, then the dataset used will be the one defined in the Axolotl config file.** You have two options to configure the dataset:\n",
        "\n",
        "# @markdown **1. Use a Hugging Face Dataset**\n",
        "# @markdown   - Requires specifying the dataset name and type.\n",
        "\n",
        "# @markdown **2. Load from Google Cloud Storage (GCS)**\n",
        "# @markdown   - Requires specifying the bucket name, dataset type, file type, and paths to training/test splits.\n",
        "\n",
        "# @markdown **Choose ONE of the following options:**\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown **Option 1: Hugging Face**\n",
        "\n",
        "# @markdown **Hugging Face Dataset Name:**\n",
        "HF_DATASET = \"\"  # @param {type:\"string\", placeholder: \"e.g. timdettmers/openassistant-guanaco\"}\n",
        "# @markdown **Set the dataset type:** Refer to [Axolotl config file](https://github.com/axolotl-ai-cloud/axolotl/blob/8fb72cbc0b94129141bae5fa4d84edd23b648af6/docs/config.qmd#L87) for more details.\n",
        "HF_DATASET_TYPE = \"\"  # @param {type:\"string\", placeholder: \"e.g. completion\"}\n",
        "if HF_DATASET:\n",
        "    assert HF_DATASET_TYPE, \"HF_DATASET_TYPE must be set if HF_DATASET is set.\"\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown **Option 2: GCS**\n",
        "\n",
        "# @markdown **Bucket Name:**\n",
        "DATASET_BUCKET_NAME = \"\"  # @param {type:\"string\"}\n",
        "# @markdown **Dataset Type:** Refer to the [Axolotl config file](https://github.com/axolotl-ai-cloud/axolotl/blob/8fb72cbc0b94129141bae5fa4d84edd23b648af6/docs/config.qmd#L181) for more details.\n",
        "DATASET_TYPE = \"\"  # @param {type:\"string\"}\n",
        "# @markdown **File Type**. Refer to the [Axolotl config file](https://github.com/axolotl-ai-cloud/axolotl/blob/8fb72cbc0b94129141bae5fa4d84edd23b648af6/docs/config.qmd#L178).\n",
        "FILE_TYPE = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown **Path to Training Data (relative to bucket):**\n",
        "TRAIN_DATAFILES_PATH = \"\"  # @param {type:\"string\"}\n",
        "# @markdown **[Optional] Path to Test Data (relative to bucket):**\n",
        "# @markdown To use a dedicated validation set, provide the file path. Otherwise, the training data will be split to create a validation set.\n",
        "TEST_DATAFILES_PATH = \"\"  # @param {type:\"string\"}\n",
        "\n",
        "if DATASET_BUCKET_NAME:\n",
        "    assert (\n",
        "        TRAIN_DATAFILES_PATH\n",
        "    ), \"TRAIN_DATAFILES_PATH must be set if DATASET_BUCKET_NAME is set.\"\n",
        "    assert DATASET_TYPE, \"DATASET_TYPE must be set if DATASET_BUCKET_NAME is set.\"\n",
        "    assert FILE_TYPE, \"FILE_TYPE must be set if DATASET_BUCKET_NAME is set.\"\n",
        "\n",
        "assert not (\n",
        "    HF_DATASET and DATASET_BUCKET_NAME\n",
        "), \"Only one of HF_DATASET or DATASET_BUCKET_NAME can be set.\"\n",
        "\n",
        "datasets = []\n",
        "if DATASET_BUCKET_NAME:\n",
        "    paths = TRAIN_DATAFILES_PATH.split(\",\")\n",
        "    dataset = {\n",
        "        \"path\": f\"/gcs/{DATASET_BUCKET_NAME}/\",\n",
        "        \"type\": DATASET_TYPE,\n",
        "        \"data_files\": [],\n",
        "        \"ds_type\": FILE_TYPE,\n",
        "    }\n",
        "    for path in paths:\n",
        "        if path.startswith(\"/\"):\n",
        "            path = path[1:]\n",
        "        dataset[\"data_files\"].append(f\"/gcs/{DATASET_BUCKET_NAME}/{path}\")\n",
        "        dataset[\"split\"] = \"train\"\n",
        "    datasets.append(dataset)\n",
        "\n",
        "test_datasets = []\n",
        "if TEST_DATAFILES_PATH:\n",
        "    paths = TEST_DATAFILES_PATH.split(\",\")\n",
        "    dataset = {\n",
        "        \"path\": f\"/gcs/{DATASET_BUCKET_NAME}/\",\n",
        "        \"type\": DATASET_TYPE,\n",
        "        \"data_files\": [],\n",
        "        \"ds_type\": FILE_TYPE,\n",
        "    }\n",
        "    for path in paths:\n",
        "        if path.startswith(\"/\"):\n",
        "            path = path[1:]\n",
        "        dataset[\"data_files\"].append(f\"/gcs/{DATASET_BUCKET_NAME}/{path}\")\n",
        "        dataset[\"split\"] = \"train\"\n",
        "    test_datasets.append(dataset)\n",
        "\n",
        "if HF_DATASET:\n",
        "    datasets.append({\"path\": HF_DATASET, \"type\": HF_DATASET_TYPE})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7dj8WuXRWGn8"
      },
      "outputs": [],
      "source": [
        "# @title Setup Axolotl Flags\n",
        "# @markdown This section configures additional Axolotl flags. You can explore different Axolotl flags in the [Axolotl config file](https://github.com/axolotl-ai-cloud/axolotl/blob/8fb72cbc0b94129141bae5fa4d84edd23b648af6/docs/config.qmd).\n",
        "\n",
        "# @markdown **To avoid OOM, you can reduce sequence length.** This can be done by setting `sequence_len` flag to some smaller value. But reducing sequence length will also reduce the model performance.\n",
        "# @markdown **Another alternative to avoid OOM is to use higher memory gpu.** It is recommended to use vertex ai training for Higher memory gpu like A100 and H100. Vertex AI training offers greater availability of high-end GPUs.\n",
        "\n",
        "# @markdown **Training can take a long time (20+ hours) to complete depending on the model, dataset and axololt config.** You can reduce the training time by reducing the max training steps. This can be done by setting `max_steps` flag to some smaller value. Note that this will also reduce the model performance.\n",
        "\n",
        "axolotl_flag_overrides = [\"--use-tensorboard=True\"]  # @param {type:\"raw\"}\n",
        "assert type(axolotl_flag_overrides) is list, \"axolotl_flag_overrides must be a list.\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uLMDYHuPJhOg"
      },
      "source": [
        "### Finetune with Local Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "D8kZmun8Ov3t"
      },
      "outputs": [],
      "source": [
        "# @title Install Axolotl\n",
        "! rm -rf axolotl\n",
        "! git clone https://github.com/axolotl-ai-cloud/axolotl.git\n",
        "! cd axolotl && git reset --hard 8fb72cbc0b94129141bae5fa4d84edd23b648af6\n",
        "! pip3 install packaging ninja\n",
        "! cd axolotl && pip3 install --no-build-isolation -e '.[flash-attn,deepspeed]'\n",
        "\n",
        "# This is needed because of this issue: https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1492\n",
        "! pip3 install bitsandbytes==0.45.1\n",
        "\n",
        "# @title Install GCSFUSE\n",
        "! apt-get install gcsfuse -y"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3tYJV5ScyscK"
      },
      "outputs": [],
      "source": [
        "# @title Run Local fine-tuning\n",
        "# @markdown This section runs the Axolotl training locally (i.e. colab runtime).\n",
        "# @markdown **Note: This section can take a long time to run. You can reduce the training time by reducing the max training steps as mentioned in `Setup Axolotl Flags` section.**\n",
        "# @markdown Model trained using Axolotl will be saved in the GCS bucket with the help of GCSFUSE.\n",
        "\n",
        "assert OUTPUT_GCS_URI, \"OUTPUT_GCS_URI must be set for local fine-tuning.\"\n",
        "\n",
        "# @markdown 1. Run GCSFUSE so that axolotl can store the training output in the GCS bucket.\n",
        "! mkdir -p /gcs/\n",
        "! gcsfuse /gcs\n",
        "\n",
        "# @markdown 2. Run Axolotl training.\n",
        "AXOLOTL_OUTPUT_GCS_URI = f\"{OUTPUT_GCS_URI}/axolotl_output\"\n",
        "AXOLOTL_OUTPUT_DIR = common_util.gcs_fuse_path(AXOLOTL_OUTPUT_GCS_URI)\n",
        "\n",
        "axolotl_args = f\" --output-dir={AXOLOTL_OUTPUT_DIR}\"\n",
        "if len(datasets) > 0:\n",
        "    axolotl_args += f' --datasets=\"{datasets}\"'\n",
        "if len(test_datasets) > 0:\n",
        "    axolotl_args += f' --test-datasets=\"{test_datasets}\"'\n",
        "    axolotl_args += \" --val-set-size=0\"\n",
        "additional_flags = \" \".join(axolotl_flag_overrides)\n",
        "axolotl_args += f\" {additional_flags}\"\n",
        "! accelerate launch -m axolotl.cli.train $axolotl_args /content/axolotl/$AXOLOTL_CONFIG_PATH\n",
        "\n",
        "# @markdown 3. Check the output in the bucket.\n",
        "! gsutil ls $AXOLOTL_OUTPUT_GCS_URI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HhcXl_JpyUFu"
      },
      "outputs": [],
      "source": [
        "# @title Run Local inference\n",
        "# @markdown This section performs inference using the finetuned model.\n",
        "\n",
        "# @markdown 1. Copy the finetuned model from GCS to local.\n",
        "! mkdir -p /tmp/axolotl_output\n",
        "! gsutil -m cp -r $AXOLOTL_OUTPUT_GCS_URI/* /tmp/axolotl_output/\n",
        "\n",
        "# @markdown 2. Run Axolotl inference using gradio on local finetuned model.\n",
        "! cd axolotl && axolotl inference  examples/tiny-llama/lora.yml --output-dir=/tmp/axolotl_output/ --gradio\n",
        "\n",
        "# @markdown 3. After running the cell, a public URL ([\"https://*.gradio.live\"](#)) will appear in the cell output. The playground is available in a separate browser tab when you click the URL."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "z3EvrUpxNbtN"
      },
      "outputs": [],
      "source": [
        "# @markdown This section merges the finetuned adapter with the base model.\n",
        "# @markdown **Note: This is only needed for lora and qlora. In case of full finetuning you can skip this cell.**\n",
        "\n",
        "if (\n",
        "    \"adapter\" in axolotl_config\n",
        "    and axolotl_config[\"adapter\"] != \"lora\"\n",
        "    and axolotl_config[\"adapter\"] != \"qlora\"\n",
        "):\n",
        "    raise ValueError(\"This cell is only needed for lora and qlora.\")\n",
        "\n",
        "# @markdown 1. Copy the finetuned model from GCS to local.\n",
        "! mkdir -p /tmp/axolotl_output\n",
        "! gsutil -m cp -r $AXOLOTL_OUTPUT_GCS_URI/* /tmp/axolotl_output/\n",
        "\n",
        "# @markdown 2. Run Axolotl merge.\n",
        "! cd axolotl && python3 -m axolotl.cli.merge_lora $AXOLOTL_CONFIG_PATH --output-dir=/tmp/axolotl_output/\n",
        "\n",
        "# @markdown 3. Copy the merged model to GCS.\n",
        "! gsutil -m cp -r /tmp/axolotl_output/merged /* $AXOLOTL_OUTPUT_GCS_URI/merged/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LcMgeq_0CYDJ"
      },
      "source": [
        "### Finetune with Vertex AI Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZHKMHpOGFZUM"
      },
      "outputs": [],
      "source": [
        "# @title Vertex AI fine-tuning job\n",
        "# @markdown This section runs the Axolotl training using Vertex AI training job.\n",
        "# @markdown **Note: This section can take a long time to run. You can reduce the training time by reducing the max training steps as mentioned in `Setup Axolotl Flags` section.**\n",
        "# @markdown Refer to [Axolotl config](https://axolotl-ai-cloud.github.io/axolotl/docs/config.html) to override additional Axolotl flags.\n",
        "\n",
        "from google.cloud.aiplatform.compat.types import \\\n",
        "    custom_job as gca_custom_job_compat\n",
        "\n",
        "# @markdown Acceletor type to use for training.\n",
        "training_accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_H100_80GB\"]\n",
        "\n",
        "\n",
        "replica_count = 1\n",
        "repo = \"us-docker.pkg.dev/vertex-ai\"\n",
        "per_node_accelerator_count = 1\n",
        "boot_disk_size_gb = 500\n",
        "dws_kwargs = {\n",
        "    \"max_wait_duration\": 1800,  # 30 minutes\n",
        "    \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
        "}\n",
        "is_dynamic_workload_scheduler = True\n",
        "if training_accelerator_type == \"NVIDIA_L4\":\n",
        "    training_machine_type = \"g2-standard-8\"\n",
        "    is_dynamic_workload_scheduler = False\n",
        "    dws_kwargs = {}\n",
        "elif training_accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "    training_machine_type = \"a2-highgpu-1g\"\n",
        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
        "    training_machine_type = \"a3-highgpu-8g\"\n",
        "    per_node_accelerator_count = 8\n",
        "    boot_disk_size_gb = 2000\n",
        "else:\n",
        "    raise ValueError(f\"Unsupported accelerator type: {training_accelerator_type}\")\n",
        "\n",
        "TRAIN_DOCKER_URI = (\n",
        "    f\"{repo}/vertex-vision-model-garden-dockers/axolotl-train:20250225-1800-rc0\"\n",
        ")\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=training_accelerator_type,\n",
        "    accelerator_count=per_node_accelerator_count * replica_count,\n",
        "    is_for_training=True,\n",
        "    is_restricted_image=False,\n",
        "    is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
        ")\n",
        "\n",
        "# @markdown Run Vertex AI job.\n",
        "\n",
        "# Copy the config file to the bucket.\n",
        "if AXOLOTL_SOURCE == \"LOCAL\":\n",
        "    ! gsutil -m cp $AXOLOTL_CONFIG_PATH $MODEL_BUCKET/config/\n",
        "    AXOLOTL_CONFIG_PATH = f\"{common_util.gcs_fuse_path(MODEL_BUCKET)}/config/{pathlib.Path(AXOLOTL_CONFIG_PATH).name}\"\n",
        "\n",
        "# Set axolotl flags.\n",
        "datasets = []\n",
        "if DATASET_BUCKET_NAME:\n",
        "    paths = TRAIN_DATAFILES_PATH.split(\",\")\n",
        "    dataset = {\n",
        "        \"path\": f\"/gcs/{DATASET_BUCKET_NAME}/\",\n",
        "        \"type\": DATASET_TYPE,\n",
        "        \"data_files\": [],\n",
        "        \"ds_type\": FILE_TYPE,\n",
        "    }\n",
        "    for path in paths:\n",
        "        if path.startswith(\"/\"):\n",
        "            path = path[1:]\n",
        "        dataset[\"data_files\"].append(f\"/gcs/{DATASET_BUCKET_NAME}/{path}\")\n",
        "        dataset[\"split\"] = \"train\"\n",
        "    datasets.append(dataset)\n",
        "\n",
        "test_datasets = []\n",
        "if TEST_DATAFILES_PATH:\n",
        "    paths = TEST_DATAFILES_PATH.split(\",\")\n",
        "    dataset = {\n",
        "        \"path\": f\"/gcs/{DATASET_BUCKET_NAME}/\",\n",
        "        \"type\": DATASET_TYPE,\n",
        "        \"data_files\": [],\n",
        "        \"ds_type\": FILE_TYPE,\n",
        "    }\n",
        "    for path in paths:\n",
        "        if path.startswith(\"/\"):\n",
        "            path = path[1:]\n",
        "        dataset[\"data_files\"].append(f\"/gcs/{DATASET_BUCKET_NAME}/{path}\")\n",
        "        dataset[\"split\"] = \"train\"\n",
        "    test_datasets.append(dataset)\n",
        "\n",
        "if HF_DATASET:\n",
        "    datasets.append({\"path\": HF_DATASET, \"type\": HF_DATASET_TYPE})\n",
        "\n",
        "if not OUTPUT_GCS_URI:\n",
        "    OUTPUT_GCS_URI = MODEL_BUCKET\n",
        "AXOLOTL_OUTPUT_GCS_URI = f\"{OUTPUT_GCS_URI}/axolotl_output\"\n",
        "AXOLOTL_OUTPUT_DIR = common_util.gcs_fuse_path(AXOLOTL_OUTPUT_GCS_URI)\n",
        "TRAINING_JOB_OUTPUT_DIR = f\"{AXOLOTL_OUTPUT_GCS_URI}/training_job_output\"\n",
        "\n",
        "\n",
        "axolotl_config_overwrites = []\n",
        "axolotl_config_overwrites.append(f\"--output_dir={AXOLOTL_OUTPUT_DIR}\")\n",
        "if len(datasets) > 0:\n",
        "    axolotl_config_overwrites.append(f'--datasets=\"{datasets}\"')\n",
        "if len(test_datasets) > 0:\n",
        "    axolotl_config_overwrites.append(f'--test_datasets=\"{test_datasets}\"')\n",
        "    axolotl_config_overwrites.append(\"--val_set_size=0\")\n",
        "axolotl_config_overwrites += axolotl_flag_overrides\n",
        "\n",
        "train_job_args = []\n",
        "train_job_args.append(f\"--axolotl_config_path={AXOLOTL_CONFIG_PATH}\")\n",
        "train_job_args += axolotl_config_overwrites\n",
        "\n",
        "\n",
        "train_job_envs = {}\n",
        "if HF_TOKEN:\n",
        "    train_job_envs[\"HF_TOKEN\"] = HF_TOKEN\n",
        "\n",
        "job_name = common_util.get_job_name_with_datetime(\"axolotl-train\")\n",
        "\n",
        "# Add labels for the finetuning job.\n",
        "labels = {\n",
        "    \"mg-source\": \"notebook\",\n",
        "    \"mg-notebook-name\": \"model_garden_axolotl_finetuning.ipynb\".split(\".\")[0],\n",
        "}\n",
        "\n",
        "model_name = AXOLOTL_CONFIG_PATH.split(\"/\")[1]\n",
        "publisher = axolotl_config[\"base_model\"].split(\"/\")[0]\n",
        "model_id = axolotl_config[\"base_model\"].split(\"/\")[1]\n",
        "model_id = model_id.replace(\".\", \"-\")\n",
        "labels[\"mg-tune\"] = f\"publishers-{publisher}-models-{model_name}\".lower()\n",
        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{model_id}\".lower()\n",
        "\n",
        "\n",
        "# Pass training arguments and launch job.\n",
        "train_job = aiplatform.CustomContainerTrainingJob(\n",
        "    display_name=job_name,\n",
        "    container_uri=TRAIN_DOCKER_URI,\n",
        "    labels=labels,\n",
        ")\n",
        "\n",
        "print(\"Running training job with args:\")\n",
        "print(\" \\\\\\n\".join(train_job_args))\n",
        "train_job.run(\n",
        "    args=train_job_args,\n",
        "    replica_count=replica_count,\n",
        "    machine_type=training_machine_type,\n",
        "    accelerator_type=training_accelerator_type,\n",
        "    accelerator_count=per_node_accelerator_count,\n",
        "    boot_disk_size_gb=boot_disk_size_gb,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    base_output_dir=TRAINING_JOB_OUTPUT_DIR,\n",
        "    sync=False,  # Non-blocking call to run.\n",
        "    **dws_kwargs,\n",
        ")\n",
        "\n",
        "# Wait until resource has been created.\n",
        "train_job.wait_for_resource_creation()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15BI63V2IIha"
      },
      "source": [
        "### Run TensorBoard"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RcSJinPWXy9_"
      },
      "outputs": [],
      "source": [
        "base_output_dir = AXOLOTL_OUTPUT_DIR\n",
        "\n",
        "# @markdown This section shows how to launch TensorBoard in a [Cloud Shell](https://cloud.google.com/shell/docs).\n",
        "# @markdown 1. Click the Cloud Shell icon(![terminal](https://github.com/google/material-design-icons/blob/master/png/action/terminal/materialicons/24dp/1x/baseline_terminal_black_24dp.png?raw=true)) on the top right to open the Cloud Shell.\n",
        "# @markdown 2. Copy the `tensorboard` command shown below by running this cell.\n",
        "# @markdown 3. Paste and run the command in the Cloud Shell to launch TensorBoard.\n",
        "# @markdown 4. Once the command runs (You may have to click `Authorize` if prompted), click the link starting with `http://localhost`.\n",
        "\n",
        "# @markdown Note: You may need to wait around 10 minutes after the job starts in order for the TensorBoard logs to be written to the GCS bucket.\n",
        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BO1BYNnfXy9_"
      },
      "source": [
        "## Deploy using vllm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Up326e7kXy9_"
      },
      "outputs": [],
      "source": [
        "# @markdown 1. Wait for the training job to finish.\n",
        "if train_job and 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 2. Set up VLLM docker URI and model gcs uri.\n",
        "\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241001_0916_RC00\"\n",
        "VLLM_MODEL_GCS_URI = AXOLOTL_OUTPUT_GCS_URI\n",
        "\n",
        "if \"adapter\" in axolotl_config and (\n",
        "    axolotl_config[\"adapter\"] == \"lora\" or axolotl_config[\"adapter\"] == \"qlora\"\n",
        "):\n",
        "    VLLM_MODEL_GCS_URI = f\"{AXOLOTL_OUTPUT_GCS_URI}/merged\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ulsUk-xkXy9_"
      },
      "source": [
        "### Create model endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KSW-gb6nXy9_"
      },
      "outputs": [],
      "source": [
        "# @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",
        "# @markdown 1. Set the machine type and accelerator type.\n",
        "# @markdown Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
        "machine_type = \"g2-standard-12\"  # @param {type:\"string\"}\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param {type:\"string\"}\n",
        "per_node_accelerator_count = 1  # @param {type:\"integer\"}\n",
        "\n",
        "common_util.check_quota(\n",
        "    project_id=PROJECT_ID,\n",
        "    region=REGION,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=per_node_accelerator_count,\n",
        "    is_for_training=False,\n",
        ")\n",
        "\n",
        "use_dedicated_endpoint = False\n",
        "gpu_memory_utilization = 0.95\n",
        "max_model_len = 2048\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    publisher: str,\n",
        "    publisher_model_id: str,\n",
        "    service_account: str,\n",
        "    base_model_id: str = None,\n",
        "    machine_type: str = \"g2-standard-8\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    gpu_memory_utilization: float = 0.9,\n",
        "    max_model_len: int = 4096,\n",
        "    dtype: str = \"auto\",\n",
        "    enable_trust_remote_code: bool = False,\n",
        "    enforce_eager: bool = False,\n",
        "    enable_lora: bool = False,\n",
        "    enable_chunked_prefill: bool = False,\n",
        "    enable_prefix_cache: bool = False,\n",
        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
        "    max_loras: int = 1,\n",
        "    max_cpu_loras: int = 8,\n",
        "    use_dedicated_endpoint: bool = False,\n",
        "    max_num_seqs: int = 256,\n",
        "    model_type: str = None,\n",
        "    enable_llama_tool_parser: bool = False,\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(\n",
        "        display_name=f\"{model_name}-endpoint\",\n",
        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
        "    )\n",
        "\n",
        "    if not base_model_id:\n",
        "        base_model_id = model_id\n",
        "\n",
        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
        "    vllm_args = [\n",
        "        \"python\",\n",
        "        \"-m\",\n",
        "        \"vllm.entrypoints.api_server\",\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=8080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        f\"--max-loras={max_loras}\",\n",
        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
        "        f\"--max-num-seqs={max_num_seqs}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    if enable_trust_remote_code:\n",
        "        vllm_args.append(\"--trust-remote-code\")\n",
        "\n",
        "    if enforce_eager:\n",
        "        vllm_args.append(\"--enforce-eager\")\n",
        "\n",
        "    if enable_lora:\n",
        "        vllm_args.append(\"--enable-lora\")\n",
        "\n",
        "    if enable_chunked_prefill:\n",
        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
        "\n",
        "    if enable_prefix_cache:\n",
        "        vllm_args.append(\"--enable-prefix-caching\")\n",
        "\n",
        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
        "        vllm_args.append(\n",
        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
        "        )\n",
        "\n",
        "    if model_type:\n",
        "        vllm_args.append(f\"--model-type={model_type}\")\n",
        "\n",
        "    if enable_llama_tool_parser:\n",
        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
        "    try:\n",
        "        if HF_TOKEN:\n",
        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
        "    except NameError:\n",
        "        pass\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[8080],\n",
        "        serving_container_predict_route=\"/generate\",\n",
        "        serving_container_health_route=\"/ping\",\n",
        "        serving_container_environment_variables=env_vars,\n",
        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
        "        serving_container_deployment_timeout=7200,\n",
        "        model_garden_source_model_name=(\n",
        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
        "        ),\n",
        "    )\n",
        "    print(\n",
        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
        "    )\n",
        "    model.deploy(\n",
        "        endpoint=endpoint,\n",
        "        machine_type=machine_type,\n",
        "        accelerator_type=accelerator_type,\n",
        "        accelerator_count=accelerator_count,\n",
        "        deploy_request_timeout=1800,\n",
        "        service_account=service_account,\n",
        "        system_labels={\n",
        "            \"NOTEBOOK_NAME\": \"model_garden_axolotl_finetuning.ipynb\",\n",
        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
        "        },\n",
        "    )\n",
        "    print(\"endpoint_name:\", endpoint.name)\n",
        "\n",
        "    return model, endpoint\n",
        "\n",
        "\n",
        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
        "    model_name=common_util.get_job_name_with_datetime(prefix=\"axolotl-vllm-serve\"),\n",
        "    publisher=publisher,\n",
        "    publisher_model_id=model_id,\n",
        "    model_id=VLLM_MODEL_GCS_URI,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=per_node_accelerator_count,\n",
        "    gpu_memory_utilization=gpu_memory_utilization,\n",
        "    max_model_len=max_model_len,\n",
        "    enable_lora=True,\n",
        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RO76I0p3Xy9_"
      },
      "source": [
        "### Perform Prediction"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kC9Apto7Xy9_"
      },
      "outputs": [],
      "source": [
        "def predict_vllm(\n",
        "    prompt: str,\n",
        "    max_tokens: int,\n",
        "    temperature: float,\n",
        "    top_p: float,\n",
        "    top_k: int,\n",
        "    raw_response: bool,\n",
        "    lora_weight: str = \"\",\n",
        "):\n",
        "    # Parameters for inference.\n",
        "    instance = {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "        \"raw_response\": raw_response,\n",
        "    }\n",
        "    if lora_weight:\n",
        "        instance[\"dynamic-lora\"] = lora_weight\n",
        "    instances = [instance]\n",
        "    response = endpoints[\"vllm_gpu\"].predict(\n",
        "        instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        "    )\n",
        "\n",
        "    for prediction in response.predictions:\n",
        "        print(prediction)\n",
        "\n",
        "\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
        "#   above.\n",
        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
        "#   an existing endpoint with the ID 1234567890123456789.\n",
        "# You may uncomment the code below to load an existing endpoint.\n",
        "\n",
        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoints[\"vllm_gpu\"] = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"Write a function to list n Fibonacci numbers in Python.\"  # @param {type: \"string\"}\n",
        "max_tokens = 500  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
        "top_k = 1  # @param {type:\"integer\"}\n",
        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
        "raw_response = True  # @param {type:\"boolean\"}\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "        \"raw_response\": raw_response,\n",
        "    },\n",
        "]\n",
        "response = endpoints[\"vllm_gpu\"].predict(\n",
        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
        ")\n",
        "\n",
        "# \"<|file_separator|>\" is the end of the file token.\n",
        "for prediction in response.predictions:\n",
        "    print(prediction.split(\"<|file_separator|>\")[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2grDDphYx4zI"
      },
      "source": [
        "## Clean up resources"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_Pqw3TsF2uG4"
      },
      "outputs": [],
      "source": [
        "# @markdown Delete the training job.\n",
        "\n",
        "if train_job:\n",
        "    train_job.delete()\n",
        "\n",
        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "for endpoint in endpoints.values():\n",
        "    endpoint.delete(force=True)\n",
        "\n",
        "# Delete models.\n",
        "for model in models.values():\n",
        "    model.delete()\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_NAME"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_axolotl_finetuning.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}