notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7d9bbf86da5e"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 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": "99c1c3fc2ca5"
      },
      "source": [
        "# Vertex AI Model Garden - Code LLaMA Deployment\n",
        "\n",
        "<table align=\"left\">\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_pytorch_codellama.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>\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
        "      View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3de7470326a2"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook showcases deploying pretrained Code LLaMA models using [vLLM](https://github.com/vllm-project/vllm). This notebook also demonstrates how to evaluate the Code LLaMA models using EleutherAI's Language Model Evaluation Harness (lm-evaluation-harness) with Vertex CustomJob.\n",
        "\n",
        "### Objective\n",
        "\n",
        "- Deploy pre-trained Code LLaMA models with [vLLM](https://github.com/vllm-project/vllm) with best serving throughput\n",
        "\n",
        "### Costs\n",
        "\n",
        "This tutorial uses billable components of Google Cloud:\n",
        "\n",
        "* Vertex AI\n",
        "* Cloud Storage\n",
        "\n",
        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bIGGC0o3vJdV"
      },
      "source": [
        "## Run the notebook"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "855d6b96f291"
      },
      "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. [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",
        "import os\n",
        "import sys\n",
        "from datetime import datetime\n",
        "from typing import Tuple\n",
        "\n",
        "from google.cloud import aiplatform\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",
        "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, please change the value yourself below.\n",
        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
        "    # Create a unique GCS bucket for this notebook if not specified\n",
        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\n",
        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
        "else:\n",
        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\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",
        "# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\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",
        "\n",
        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
        "\n",
        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_URI\n",
        "\n",
        "\n",
        "! gcloud config set project $PROJECT_ID\n",
        "\n",
        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
        "MODEL_BUCKET = os.path.join(STAGING_BUCKET, \"model\")\n",
        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
        "\n",
        "\n",
        "def get_job_name_with_datetime(prefix: str) -> str:\n",
        "    \"\"\"Gets the job name with date time when triggering training or deployment\n",
        "    jobs in Vertex AI.\n",
        "    \"\"\"\n",
        "    return prefix + datetime.now().strftime(\"_%Y%m%d_%H%M%S\")\n",
        "\n",
        "\n",
        "# The pre-built serving docker image.\n",
        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240410_0916_RC00\"\n",
        "\n",
        "\n",
        "def deploy_model_vllm(\n",
        "    model_name: str,\n",
        "    model_id: str,\n",
        "    base_model_id: str,\n",
        "    service_account: str,\n",
        "    machine_type: str = \"g2-standard-12\",\n",
        "    accelerator_type: str = \"NVIDIA_L4\",\n",
        "    accelerator_count: int = 1,\n",
        "    max_model_len: int = 8192,\n",
        "    dtype: str = \"bfloat16\",\n",
        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
        "    \"\"\"Deploys models with vLLM on GPU in Vertex AI.\"\"\"\n",
        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
        "\n",
        "    vllm_args = [\n",
        "        \"--host=0.0.0.0\",\n",
        "        \"--port=7080\",\n",
        "        f\"--model={model_id}\",\n",
        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
        "        \"--swap-space=16\",\n",
        "        \"--gpu-memory-utilization=0.9\",\n",
        "        f\"--max-model-len={max_model_len}\",\n",
        "        f\"--dtype={dtype}\",\n",
        "        \"--disable-log-stats\",\n",
        "    ]\n",
        "\n",
        "    env_vars = {\n",
        "        \"MODEL_ID\": base_model_id,\n",
        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
        "    }\n",
        "\n",
        "    model = aiplatform.Model.upload(\n",
        "        display_name=model_name,\n",
        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
        "        serving_container_command=[\"python\", \"-m\", \"vllm.entrypoints.api_server\"],\n",
        "        serving_container_args=vllm_args,\n",
        "        serving_container_ports=[7080],\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",
        "    )\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",
        "    )\n",
        "    return model, endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jwn4PcTf4EMt"
      },
      "outputs": [],
      "source": [
        "# @title Access pretrained Code LLaMA models\n",
        "\n",
        "# @markdown The original models from Meta are converted into the HuggingFace format for serving in Vertex AI.\n",
        "\n",
        "# @markdown Accept the model agreement to access the models:\n",
        "# @markdown 1. Open the [Code LLaMA model card](https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/137).\n",
        "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
        "# @markdown 3. A Cloud Storage bucket (starting with ‘gs://’) containing Code LLaMA pretrained and finetuned models will be shared under the “Documentation” section and its “Get started” subsection.\n",
        "\n",
        "# This path will be shared once click the agreement in Code LLaMA model card\n",
        "# as described in the `Access pretrained Code LLaMA models` section.\n",
        "VERTEX_AI_MODEL_GARDEN_CODE_LLAMA = \"\"  # @param {type: \"string\"}\n",
        "assert (\n",
        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA\n",
        "), \"Please click the agreement of Code LLaMA in Vertex AI Model Garden, and get the GCS path of Code LLaMA model artifacts.\"\n",
        "print(\n",
        "    \"Copying Code LLaMA model artifacts from\",\n",
        "    VERTEX_AI_MODEL_GARDEN_CODE_LLAMA,\n",
        "    \"to \",\n",
        "    MODEL_BUCKET,\n",
        ")\n",
        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_CODE_LLAMA/* $MODEL_BUCKET"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "03d504bcd60b"
      },
      "outputs": [],
      "source": [
        "# @title Deploy pretrained Code LLaMA (vLLM)\n",
        "# @markdown This section deploys prebuilt Code LLaMA models with [vLLM](https://github.com/vllm-project/vllm) on the Endpoint. Code LLaMA model weights are stored in bfloat16 precision. L4 or A100 GPUs are needed for vLLM serving at bfloat16 precision. V100 GPUs can be used with vLLM serving at float16 precision. Changing the precision from bfloat16 to float16 can result in a change in performance, and this change can be an increase and a decrease. However, the performance change should be small (within 5%).\n",
        "\n",
        "# @markdown V100 GPUs are used for demonstration. Note that V100 serving generally offers better throughput and latency performance than L4 serving, while L4 serving is generally more cost efficient than V100 serving. The serving efficiency of V100 and L4 GPUs is inferior to that of A100 GPUs, but V100 and L4 GPUs are nevertheless good serving solutions if you do not have A100 quota. The model deployment step will take 15 minutes to 1 hour to complete.\n",
        "\n",
        "# @markdown The vLLM project is an highly optimized LLM serving framework which can increase serving throughput a lot. The higher QPS you have, the more benefits you get using vLLM.\n",
        "\n",
        "\n",
        "# Finds Vertex AI prediction supported accelerators and regions in\n",
        "#  https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "\n",
        "# @markdown Set the model name.\n",
        "model_name = \"CodeLlama-7b-Instruct-hf\"  # @param [\"CodeLlama-7b-hf\", \"CodeLlama-7b-Python-hf\", \"CodeLlama-7b-Instruct-hf\", \"CodeLlama-13b-hf\", \"CodeLlama-13b-Python-hf\", \"CodeLlama-13b-Instruct-hf\", \"CodeLlama-34b-hf\", \"CodeLlama-34b-Python-hf\", \"CodeLlama-34b-Instruct-hf\", \"CodeLlama-70b-hf\", \"CodeLlama-70b-Python-hf\", \"CodeLlama-70b-Instruct-hf\"]\n",
        "model_id = os.path.join(MODEL_BUCKET, model_name)\n",
        "hf_model_id = \"meta-llama/\" + model_name\n",
        "print(model_id)\n",
        "\n",
        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
        "\n",
        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\", \"NVIDIA_TESLA_A100\"]\n",
        "\n",
        "if \"7b\" in model_name:\n",
        "    # Sets A100 (40G) to deploy 7B models.\n",
        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        machine_type = \"a2-highgpu-1g\"\n",
        "        accelerator_count = 1\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    # Sets 1 V100 (16G) to deploy 7B models..\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "        machine_type = \"n1-standard-8\"\n",
        "        accelerator_count = 1\n",
        "        vllm_precision = \"float16\"\n",
        "    # Sets 1 L4 (24G) to deploy 7B models.\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
        "        machine_type = \"g2-standard-8\"\n",
        "        accelerator_count = 1\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended GPU setting not found for: {accelerator_type} and {model_name}.\"\n",
        "        )\n",
        "elif \"13b\" in model_name:\n",
        "    # Sets A100 (40G) to deploy 13B models.\n",
        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        machine_type = \"a2-highgpu-1g\"\n",
        "        accelerator_count = 1\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    # Sets 2 V100 (16G) to deploy 13B models.\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "        machine_type = \"n1-standard-16\"\n",
        "        accelerator_count = 2\n",
        "        vllm_precision = \"float16\"\n",
        "    # Sets 2 L4 (24G) to deploy 13B models.\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
        "        machine_type = \"g2-standard-24\"\n",
        "        accelerator_count = 2\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
        "        )\n",
        "elif \"34b\" in model_name:\n",
        "    # Sets 2 A100 (40G) to deploy 34B models.\n",
        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        machine_type = \"a2-highgpu-2g\"\n",
        "        accelerator_count = 2\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    # Sets 8 V100 (16G) to deploy 34B models.\n",
        "    elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
        "        machine_type = \"n1-standard-32\"\n",
        "        accelerator_count = 8\n",
        "        vllm_precision = \"float16\"\n",
        "    # Sets 4 L4 (24G) to deploy 34B models.\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
        "        machine_type = \"g2-standard-48\"\n",
        "        accelerator_count = 4\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
        "        )\n",
        "elif \"70b\" in model_name:\n",
        "    # Sets 4 A100 (40G) to deploy 70B models.\n",
        "    if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
        "        machine_type = \"a2-highgpu-4g\"\n",
        "        accelerator_count = 4\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    # Sets 8 L4 (24G) to deploy 70B models.\n",
        "    elif accelerator_type == \"NVIDIA_L4\":\n",
        "        machine_type = \"g2-standard-96\"\n",
        "        accelerator_count = 8\n",
        "        vllm_precision = \"bfloat16\"\n",
        "    else:\n",
        "        raise ValueError(\n",
        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
        "        )\n",
        "\n",
        "\n",
        "# Note that a larger max_model_len will require more GPU memory.\n",
        "max_model_len = 2048\n",
        "\n",
        "model_vllm, endpoint_vllm = deploy_model_vllm(\n",
        "    model_name=get_job_name_with_datetime(prefix=\"code-llama-serve-vllm\"),\n",
        "    model_id=model_id,\n",
        "    base_model_id=hf_model_id,\n",
        "    service_account=SERVICE_ACCOUNT,\n",
        "    machine_type=machine_type,\n",
        "    accelerator_type=accelerator_type,\n",
        "    accelerator_count=accelerator_count,\n",
        "    max_model_len=max_model_len,\n",
        "    dtype=vllm_precision,\n",
        ")\n",
        "print(f\"Endpoint name: {endpoint_vllm.name}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E4Y1QjghSE5r"
      },
      "outputs": [],
      "source": [
        "# @title Prediction with endpoint\n",
        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Parameters supported by vLLM can be found [here](https://github.com/vllm-project/vllm/blob/2e8e49fce3775e7704d413b2f02da6d7c99525c9/vllm/sampling_params.py#L23-L64).\n",
        "\n",
        "# Loads an existing endpoint instance using the endpoint name:\n",
        "# - Using `endpoint_name = endpoint_vllm.name` allows us to get the endpoint\n",
        "#   name of the endpoint `endpoint_without_peft` created in the cell 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 = endpoint_vllm.name\n",
        "# # endpoint_name = \"\"  # @param {type:\"string\"}\n",
        "# aip_endpoint_name = (\n",
        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
        "# )\n",
        "# endpoint_vllm = aiplatform.Endpoint(aip_endpoint_name)\n",
        "\n",
        "prompt = \"import argparse\"  # @param {type:\"string\"}\n",
        "n = 1  # @param {type:\"integer\"}\n",
        "max_tokens = 200  # @param {type:\"integer\"}\n",
        "temperature = 1.0  # @param {type:\"number\"}\n",
        "top_p = 1.0  # @param {type:\"number\"}\n",
        "top_k = 10  # @param {type:\"number\"}\n",
        "\n",
        "instances = [\n",
        "    {\n",
        "        \"prompt\": prompt,\n",
        "        \"n\": n,\n",
        "        \"max_tokens\": max_tokens,\n",
        "        \"temperature\": temperature,\n",
        "        \"top_p\": top_p,\n",
        "        \"top_k\": top_k,\n",
        "    },\n",
        "]\n",
        "response = endpoint_vllm.predict(instances=instances)\n",
        "print(response.predictions[0])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "911406c1561e"
      },
      "outputs": [],
      "source": [
        "# @title Clean up resources\n",
        "# @markdown Delete the experiment models and endpoints to recycle the resources\n",
        "# @markdown and avoid unnecessary continouous charges that may incur.\n",
        "\n",
        "delete_bucket = False  # @param {type:\"boolean\"}\n",
        "if delete_bucket:\n",
        "    ! gsutil -m rm -r $BUCKET_URI\n",
        "    # Uncomment below to delete all artifacts\n",
        "    # !gsutil -m rm -r $STAGING_BUCKET $MODEL_BUCKET\n",
        "\n",
        "# Undeploy model and delete endpoint.\n",
        "endpoint_vllm.delete(force=True)\n",
        "\n",
        "# Delete model.\n",
        "model_vllm.delete()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "model_garden_pytorch_codellama.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}