xiaohei/model_garden
1
0
Fork 0
mirror of https://github.com/GoogleCloudPlatform/vertex-ai-samples.git synced 2026-09-26 14:42:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make mistral notebook tunable (#3374)

Browse Source 
* Make mistral notebook tunable

* make minor improvements

* make minor improvements

---------

Co-authored-by: Rayan Dasoriya <dasoriya@google.com>
This commit is contained in:
Rayan Dasoriya
2024-08-01 13:47:50 +00:00
committed by GitHub
co-authored by Rayan Dasoriya
parent 772881903b
commit cc2011e354
1 changed files with 325 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
notebooks/community/model_garden/model_garden_pytorch_mistral_peft_tuning.ipynb
+325 -161
View File
@@ -57,7 +57,7 @@
"\n",
"### Objective\n",
"\n",
"* Finetune and merge Mistral model with PEFT training docker image.\n",
"* Finetune and merge Mistral model using PEFT training docker image.\n",
"* Deploy the finetuned model with vLLM docker image on a Vertex AI Endpoint.\n",
"* Run inference on the deployed Vertex AI Endpoint.\n",
"\n",
@@ -68,7 +68,7 @@
"* 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."
"Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [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."
]
},
{
@@ -95,32 +95,48 @@
"\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",
"# @markdown 3. [Make sure that you have GPU quota for Vertex Training (finetuning) and Vertex Prediction (serving)](https://cloud.google.com/docs/quotas/view-manage). The quota name for Vertex Training is \"Custom model training your-gpu-type per region\" and the quota name for Vertex Prediction is \"Custom model serving your-gpu-type per region\" such as `Custom model training Nvidia L4 GPUs per region` and `Custom model serving Nvidia L4 GPUs per region` for L4 GPUs. [Submit a quota increase request](https://cloud.google.com/docs/quotas/view-manage#requesting_higher_quota) if additional quota is needed. At minimum, running this notebook requires 4 L4s for finetuning and 1 L4 for serving. More GPUs may be needed for larger models and different finetuning configurations. To secure GPUs for larger models, ask your customer engineer to get you allowlisted for a Shared Reservation or a Dynamic Workload Scheduler.\n",
"\n",
"# Import the necessary packages\n",
"! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
"\n",
"import importlib\n",
"import os\n",
"import uuid\n",
"from datetime import datetime\n",
"from typing import Tuple\n",
"\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",
"models, endpoints = {}, {}\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, change the value yourself below.\n",
"now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
"BUCKET_URI = \"gs://\" # @param {type: \"string\"}\n",
"assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
"BUCKET_URI = \"gs://\" # @param {type:\"string\"}\n",
"BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
"\n",
"# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
"assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
"if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
" BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}\"\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",
" BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\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",
@@ -128,78 +144,92 @@
" \"Bucket region %s is different from notebook region %s\"\n",
" % (bucket_region, REGION)\n",
" )\n",
"\n",
"print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
"\n",
"# @markdown Click \"Show code\" to see more details.\n",
"\n",
"! gcloud config set project $PROJECT_ID\n",
"! gcloud services enable language.googleapis.com\n",
"\n",
"STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
"EXPERIMENT_BUCKET = os.path.join(BUCKET_URI, \"peft\")\n",
"MODEL_BUCKET = os.path.join(EXPERIMENT_BUCKET, \"model\")\n",
"MODEL_BUCKET = os.path.join(BUCKET_URI, \"mistral\")\n",
"\n",
"# Gets the default BUCKET_URI and SERVICE_ACCOUNT if they were not specified by the user.\n",
"\n",
"SERVICE_ACCOUNT = None\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",
"\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",
"aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
"! gcloud config set project $PROJECT_ID\n",
"\n",
"# The pre-built training and serving docker images.\n",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240313_0916_RC00\"\n",
"TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240220_0936_RC01\"\n",
"\n",
"\n",
"def create_name_with_datetime(prefix: str) -> str:\n",
" \"\"\"Creates a 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",
"VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
"TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240724_0936_RC00\"\n",
"\n",
"\n",
"def deploy_model_vllm(\n",
" model_name: str,\n",
" 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",
" quantization_method: str = \"\",\n",
" gpu_memory_utilization: float = 0.9,\n",
" max_model_len: int = 4096,\n",
" dtype: str = \"auto\",\n",
") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
" \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
" endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
"\n",
" if not base_model_id:\n",
" base_model_id = model_id\n",
"\n",
" vllm_args = [\n",
" \"python\",\n",
" \"-m\",\n",
" \"vllm.entrypoints.api_server\",\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",
" \"--max-num-batched-tokens=4096\",\n",
" f\"--gpu-memory-utilization={gpu_memory_utilization}\",\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",
" # 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_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",
" print(\n",
" f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
" )\n",
"\n",
" model.deploy(\n",
" endpoint=endpoint,\n",
" machine_type=machine_type,\n",
@@ -208,9 +238,86 @@
" deploy_request_timeout=1800,\n",
" service_account=service_account,\n",
" )\n",
" print(\"endpoint_name:\", endpoint.name)\n",
"\n",
" return model, endpoint"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "5K169qf_udor"
},
"outputs": [],
"source": [
"# @title Set dataset\n",
"\n",
"# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
"\n",
"# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
"# @markdown You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
"\n",
"# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
"\n",
"# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
"# @markdown ```\n",
"# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
"# @markdown ```\n",
"\n",
"# @markdown The JSON object has a key `text`, which should match `instruct_column_in_dataset`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
"\n",
"# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
"\n",
"# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
"\n",
"# @markdown Optionally update the `instruct_column_in_dataset` field below if your JSON objects use a key other than the default `text`.\n",
"\n",
"# @markdown ### (Optional) Format your data with custom JSON template\n",
"\n",
"# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
"\n",
"# @markdown ```\n",
"# @markdown {\n",
"# @markdown \"description\": \"Template that accepts text-bison format.\",\n",
"# @markdown \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
"# @markdown \"prompt_input\": \"\\n\\n<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|>\\n\\n<|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
"# @markdown \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
"# @markdown \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
"# @markdown }\n",
"# @markdown ```\n",
"\n",
"\n",
"# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
"\n",
"# @markdown ```\n",
"# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
"# @markdown ```\n",
"\n",
"# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
"# @markdown\n",
"# @markdown To try such custom dataset, you can make the following changes:\n",
"# @markdown 1. Set `template` to `llama3-text-bison`\n",
"# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
"# @markdown 1. Set `train_split_name` to `train`\n",
"# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
"# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
"# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
"\n",
"# Template name or gs:// URI to a custom template.\n",
"template = \"openassistant-guanaco\" # @param {type:\"string\"}\n",
"\n",
"# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
"train_dataset_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
"train_split_name = \"train\" # @param {type:\"string\"}\n",
"eval_dataset_name = \"timdettmers/openassistant-guanaco\" # @param {type:\"string\"}\n",
"eval_split_name = \"test\" # @param {type:\"string\"}\n",
"\n",
"# Name of the dataset column containing training text input.\n",
"instruct_column_in_dataset = \"text\" # @param {type:\"string\"}"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -222,134 +329,163 @@
"source": [
"# @title Finetune\n",
"\n",
"# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
"\n",
"# @markdown **Note**:\n",
"# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
"# @markdown 1. If `max_steps > 0`, it will precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
"# @markdown 1. With the default setting, training takes between 1.5 ~ 2 hours.\n",
"\n",
"# @markdown This section demonstrates how to finetune the Mistral-7B model and merge the finetuned LoRA adapter with the base model on Vertex AI.\n",
"\n",
"# @markdown This example uses the dataset [fredmo/vertexai-qna-500](https://huggingface.co/datasets/fredmo/vertexai-qna-500). You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name.\n",
"\n",
"# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
"\n",
"# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [fredmo/vertexai-qna-500](https://huggingface.co/datasets/fredmo/vertexai-qna-500) dataset:\n",
"# @markdown ```\n",
"# @markdown {\"input_text\": \"question: What is the first step in setting up a project for Vertex AI?\", \"output_text\": \"Select or create a Google Cloud project.\"}\n",
"# @markdown ```\n",
"\n",
"# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
"\n",
"# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
"\n",
"# @markdown ### (Optional) Format your data with custom JSON template\n",
"\n",
"# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
"\n",
"# @markdown ```\n",
"# @markdown {\n",
"# @markdown \"description\": \"A short template for vertex sample dataset.\",\n",
"# @markdown \"prompt_input\": \"{input_text}{output_text}\",\n",
"# @markdown \"prompt_no_input\": \"{input_text}{output_text}\"\n",
"# @markdown }\n",
"# @markdown ```\n",
"\n",
"# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
"\n",
"# @markdown ```\n",
"# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
"# @markdown ```\n",
"\n",
"# @markdown This example template simply concatenates `input_text` with `output_text`. You can set `template` to `vertex_sample` to try out this built-in template with the dataset `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`, or build more complicated JSON templates such as [the alpaca example](https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json). To use your own JSON template, please [upload it to Google Cloud Storage](https://cloud.google.com/storage/docs/uploading-objects) and put the `gs://` URI in the `template` field below.\n",
"\n",
"# Huggingface dataset name or gs:// URI to a custom JSONL dataset.\n",
"base_model_id = \"mistralai/Mistral-7B-v0.1\"\n",
"gcs_model_id = f\"gs://vertex-model-garden-public-us/{base_model_id}\"\n",
"dataset_name = \"fredmo/vertexai-qna-500\" # @param {type:\"string\"}\n",
"pretrained_model_id = f\"gs://vertex-model-garden-public-us/{base_model_id}\"\n",
"\n",
"# Optional. Template name or gs:// URI to a custom template.\n",
"template = \"vertex_sample\" # @param {type:\"string\"}\n",
"\n",
"# Number of training steps.\n",
"max_steps = 10 # @param {type:\"integer\"}\n",
"\n",
"# LoRA parameters.\n",
"# @markdown Batch size for finetuning.\n",
"per_device_train_batch_size = 1 # @param{type:\"integer\"}\n",
"# @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
"gradient_accumulation_steps = 8 # @param{type:\"integer\"}\n",
"# @markdown Maximum sequence length.\n",
"max_seq_length = 4096 # @param{type:\"integer\"}\n",
"# @markdown Setting a positive `max_steps` here will override `num_epochs`.\n",
"max_steps = -1 # @param{type:\"integer\"}\n",
"num_epochs = 1.0 # @param{type:\"number\"}\n",
"# @markdown Precision mode for finetuning.\n",
"finetuning_precision_mode = \"4bit\" # @param [\"4bit\", \"8bit\", \"float16\"]\n",
"# @markdown Learning rate.\n",
"learning_rate = 5e-5 # @param{type:\"number\"}\n",
"# @markdown The scheduler type to use.\n",
"lr_scheduler_type = \"cosine\" # @param{type:\"string\"}\n",
"# @markdown LoRA parameters.\n",
"lora_rank = 16 # @param{type:\"integer\"}\n",
"lora_alpha = 64 # @param{type:\"integer\"}\n",
"lora_dropout = 0.1 # @param{type:\"number\"}\n",
"\n",
"# Learning rate.\n",
"learning_rate = 0.0001 # @param{type:\"number\"}\n",
"\n",
"# Precision mode for finetuning.\n",
"finetuning_precision_mode = \"float16\"\n",
"lora_alpha = 32 # @param{type:\"integer\"}\n",
"lora_dropout = 0.05 # @param{type:\"number\"}\n",
"# Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
"enable_gradient_checkpointing = True\n",
"# Attention implementation to use in the model.\n",
"attn_implementation = \"flash_attention_2\"\n",
"# The optimizer for which to schedule the learning rate.\n",
"optimizer = \"paged_adamw_32bit\"\n",
"# Define the proportion of training to be dedicated to a linear warmup where learning rate gradually increases.\n",
"warmup_ratio = \"0.01\"\n",
"# The list or string of integrations to report the results and logs to.\n",
"report_to = \"tensorboard\"\n",
"# Number of updates steps before two checkpoint saves.\n",
"save_steps = 10\n",
"# Number of update steps between two logs.\n",
"logging_steps = save_steps\n",
"# Train precision of the model.\n",
"train_precision = \"float16\"\n",
"\n",
"# Worker pool spec for 4bit finetuning.\n",
"accelerator_type = \"NVIDIA_L4\" # @param[\"NVIDIA_TESLA_V100\", \"NVIDIA_L4\", \"NVIDIA_TESLA_A100\"]\n",
"accelerator_type = \"NVIDIA_L4\" # @param[\"NVIDIA_A100_80GB\", \"NVIDIA_L4\"]\n",
"\n",
"if accelerator_type == \"NVIDIA_TESLA_V100\":\n",
" machine_type = \"n1-highmem-16\"\n",
" accelerator_count = 2\n",
"elif accelerator_type == \"NVIDIA_L4\":\n",
" machine_type = \"g2-standard-8\"\n",
" accelerator_count = 1\n",
"elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
" machine_type = \"a2-highgpu-1g\"\n",
" accelerator_count = 1\n",
"if accelerator_type == \"NVIDIA_L4\":\n",
" accelerator_count = 4\n",
" machine_type = \"g2-standard-48\"\n",
"elif accelerator_type == \"NVIDIA_A100_80GB\":\n",
" accelerator_count = 4\n",
" machine_type = \"a2-ultragpu-4g\"\n",
"else:\n",
" raise ValueError(f\"Unsupported accelerator type: {accelerator_type}\")\n",
"\n",
"replica_count = 1\n",
"\n",
"# @markdown Click \"Show code\" to see more details.\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=True,\n",
")\n",
"\n",
"# Setup training job.\n",
"job_name = create_name_with_datetime(\"mistral-lora-train\")\n",
"job_name = common_util.get_job_name_with_datetime(\"mistral-lora-train\").replace(\n",
" \"_\", \"-\"\n",
")\n",
"\n",
"base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
"# Create a GCS folder to store the LORA adapter.\n",
"lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
"# Create a GCS folder to store the merged model with the base model and the\n",
"# finetuned LORA adapter.\n",
"merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
"\n",
"\n",
"eval_args = [\n",
" f\"--eval_dataset_path={eval_dataset_name}\",\n",
" f\"--eval_column={instruct_column_in_dataset}\",\n",
" f\"--eval_template={template}\",\n",
" f\"--eval_split={eval_split_name}\",\n",
" f\"--eval_steps={save_steps}\",\n",
" \"--eval_tasks=builtin_eval\",\n",
" \"--eval_metric_name=loss\",\n",
"]\n",
"\n",
"train_job_args = [\n",
" \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
" \"--task=instruct-lora\",\n",
" \"--completion_only=False\",\n",
" f\"--pretrained_model_id={pretrained_model_id}\",\n",
" f\"--dataset_name={train_dataset_name}\",\n",
" f\"--train_split_name={train_split_name}\",\n",
" f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
" f\"--output_dir={lora_output_dir}\",\n",
" f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
" f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
" f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
" f\"--lora_rank={lora_rank}\",\n",
" f\"--lora_alpha={lora_alpha}\",\n",
" f\"--lora_dropout={lora_dropout}\",\n",
" f\"--max_steps={max_steps}\",\n",
" f\"--max_seq_length={max_seq_length}\",\n",
" f\"--learning_rate={learning_rate}\",\n",
" f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
" f\"--precision_mode={finetuning_precision_mode}\",\n",
" f\"--train_precision={train_precision}\",\n",
" f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
" f\"--num_epochs={num_epochs}\",\n",
" f\"--attn_implementation={attn_implementation}\",\n",
" f\"--optimizer={optimizer}\",\n",
" f\"--warmup_ratio={warmup_ratio}\",\n",
" f\"--report_to={report_to}\",\n",
" f\"--logging_output_dir={base_output_dir}\",\n",
" f\"--save_steps={save_steps}\",\n",
" f\"--logging_steps={logging_steps}\",\n",
" f\"--template={template}\",\n",
"] + eval_args\n",
"\n",
"\n",
"# Create TensorBoard\n",
"tensorboard = aiplatform.Tensorboard.create(job_name)\n",
"exp = aiplatform.TensorboardExperiment.create(\n",
" tensorboard_experiment_id=job_name, tensorboard_name=tensorboard.name\n",
")\n",
"\n",
"train_job = aiplatform.CustomContainerTrainingJob(\n",
" display_name=job_name,\n",
" container_uri=TRAIN_DOCKER_URI,\n",
")\n",
"\n",
"\"\"\"\n",
"In the below code, the finetuned LoRA adapter will be saved to a GCS bucket\n",
"specified by the variable lora_output_dir below; and you merge the\n",
"LoRa adapter with the base model, and save it to a separate GCS bucket\n",
"specified by merged_model_output_dir below.\n",
"\"\"\"\n",
"\n",
"# Create a GCS folder to store the LORA adapter.\n",
"lora_adapter_dir = create_name_with_datetime(\"mistral-lora-adapter\")\n",
"lora_output_dir = os.path.join(MODEL_BUCKET, lora_adapter_dir)\n",
"\n",
"# Create a GCS folder to store the merged model with the base model and the\n",
"# finetuned LORA adapter.\n",
"merged_model_dir = create_name_with_datetime(\"mistral-merged-model\")\n",
"merged_model_output_dir = os.path.join(MODEL_BUCKET, merged_model_dir)\n",
"\n",
"# Pass training arguments and launch job.\n",
"train_job.run(\n",
" args=[\n",
" \"--task=causal-language-modeling-lora\",\n",
" f\"--pretrained_model_id={gcs_model_id}\",\n",
" f\"--dataset_name={dataset_name}\",\n",
" f\"--output_dir={lora_output_dir}\",\n",
" f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
" f\"--lora_rank={lora_rank}\",\n",
" f\"--lora_alpha={lora_alpha}\",\n",
" f\"--lora_dropout={lora_dropout}\",\n",
" \"--warmup_steps=10\",\n",
" f\"--max_steps={max_steps}\",\n",
" f\"--learning_rate={learning_rate}\",\n",
" f\"--precision_mode={finetuning_precision_mode}\",\n",
" f\"--template={template}\",\n",
" ],\n",
" args=train_job_args,\n",
" environment_variables={\"WANDB_DISABLED\": True},\n",
" replica_count=replica_count,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" boot_disk_size_gb=500,\n",
" service_account=SERVICE_ACCOUNT,\n",
" tensorboard=tensorboard.resource_name,\n",
" base_output_dir=base_output_dir,\n",
")\n",
"\n",
"print(\"The finetuned Lora adapter can be found at: \", lora_output_dir)\n",
"print(\n",
" \"The finetuned Lora adapter merged with the base model can be found at: \",\n",
" merged_model_output_dir,\n",
")"
"print(\"LoRA adapter was saved in: \", lora_output_dir)\n",
"print(\"Trained and merged models were saved in: \", merged_model_output_dir)\n",
"\n",
"# @markdown Click \"Show Code\" to see more details."
]
},
{
@@ -362,25 +498,39 @@
"outputs": [],
"source": [
"# @title Deploy\n",
"# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. The model deployment step will take ~15 minutes to complete.\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 depending on the size of model.\n",
"\n",
"# @markdown Click \"Show code\" to see more details.\n",
"print(\"Deploying models in: \", merged_model_output_dir)\n",
"\n",
"# Finds Vertex AI prediction supported accelerators and regions in\n",
"# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
"# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
"\n",
"model, endpoint = deploy_model_vllm(\n",
" model_name=create_name_with_datetime(prefix=\"mistral-peft-serve-vllm\"),\n",
"common_util.check_quota(\n",
" project_id=PROJECT_ID,\n",
" region=REGION,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" is_for_training=False,\n",
")\n",
"\n",
"gpu_memory_utilization = 0.85\n",
"max_model_len = 8192 # Maximum context length.\n",
"\n",
"# Ensure max_model_len does not exceed the limit\n",
"if max_model_len > 8192:\n",
" raise ValueError(\"max_model_len cannot exceed 8192\")\n",
"\n",
"models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
" model_name=common_util.get_job_name_with_datetime(prefix=\"mistral-vllm-serve\"),\n",
" model_id=merged_model_output_dir,\n",
" service_account=SERVICE_ACCOUNT,\n",
" machine_type=machine_type,\n",
" accelerator_type=accelerator_type,\n",
" accelerator_count=accelerator_count,\n",
" gpu_memory_utilization=gpu_memory_utilization,\n",
" max_model_len=max_model_len,\n",
")\n",
"\n",
"print(\"endpoint_name:\", endpoint.name)\n",
"print(\"model_name:\", model.display_name)\n",
"print(\"model_id:\", model.resource_name)"
"# @markdown Click \"Show code\" to see more details."
]
},
{
@@ -393,8 +543,8 @@
"outputs": [],
"source": [
"# @title Predict\n",
"# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
"# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\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",
"# @markdown Example:\n",
"\n",
@@ -402,6 +552,7 @@
"# @markdown Human: What is a car?\n",
"# @markdown Assistant: A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
"# @markdown ```\n",
"# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
"\n",
"# Loads an existing endpoint instance using the endpoint name:\n",
"# - Using `endpoint_name = endpoint.name` allows us to get the\n",
@@ -417,24 +568,31 @@
"# )\n",
"# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
"\n",
"prompt = \"What is Vertex AI?\" # @param {type: \"string\"}\n",
"max_tokens = 100 # @param {type:\"integer\"}\n",
"prompt = \"What is a car?\" # @param {type: \"string\"}\n",
"# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
"max_tokens = 50 # @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",
"raw_response = False # @param {type:\"boolean\"}\n",
"\n",
"instance = {\n",
" \"prompt\": prompt,\n",
" \"max_tokens\": max_tokens,\n",
" \"temperature\": temperature,\n",
" \"top_p\": top_p,\n",
" \"top_k\": top_k,\n",
"}\n",
"\n",
"response = endpoint.predict(instances=[instance])\n",
"# Overrides parameters for inferences.\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(instances=instances)\n",
"\n",
"for prediction in response.predictions:\n",
" print(prediction)"
" print(prediction)\n",
"\n",
"# @markdown Click \"Show Code\" to see more details."
]
},
{
@@ -446,16 +604,22 @@
},
"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",
"# @title Delete the model and endpoint\n",
"\n",
"endpoint.delete(force=True)\n",
"model.delete()\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_URI"
" ! gsutil -m rm -r $BUCKET_NAME"
]
}
],