feat: Update the model id to a default value, this is not used when the deployed endpoint is passed in

PiperOrigin-RevId: 725755135

.github/dependabot.yml

@@ -0,0 +1,10 @@
+version: 2
+updates:
+
+  # Ignore model garden dockerfiles:
+  - package-ecosystem: "npm"
+    directory: "/community-content/vertex_model_garden"
+    schedule:
+      interval: "monthly"
+    ignore:
+      - dependency-name: "*"

.github/workflows/linter/Dockerfile

@@ -4,7 +4,7 @@
 # 2. To lint specific notebooks:
 # docker run -v ${PWD}:/setup/app gcr.io/python-docs-samples-tests/notebook_linter:latest notebooks/1.ipynb notebooks/2.ipynb
 
-FROM python:3.12
+FROM python:3.13
 
 WORKDIR setup
 

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==24.8.0
-pyupgrade==3.17.0
-isort==5.13.2
-flake8==7.1.0
-nbqa==1.9.0
+black==25.1.0
+pyupgrade==3.19.1
+isort==6.0.0
+flake8==7.1.1
+nbqa==1.9.1
 

community-content/CODEOWNERS

@@ -20,6 +20,7 @@
 /vertex_model_garden/model_oss/movinet @KCFindstr
 /vertex_model_garden/model_oss/data_converter @KCFindstr
 /vertex_model_garden/model_oss/peft @weigary
+/vertex_model_garden/model_oss/peft/templates @rayandasoriya
 /vertex_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
 /vertex_model_garden/model_oss/tfvision @dstnluong-google
 /vertex_model_garden/model_oss/fvlm @minwoo33park

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile

@@ -1,16 +1,40 @@
-FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
+# Stage 1: Build Environment
+FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime AS builder
 
+# Install necessary tools and dependencies
 RUN apt-get update && \
     apt-get install -y curl gnupg && \
     echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
-    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg  add - && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
     apt-get update -y && \
-    apt-get install google-cloud-sdk -y
+    apt-get install -y google-cloud-sdk
 
+# Copy application code
 COPY . /trainer
 
+# Set working directory
 WORKDIR /trainer
 
-RUN pip install -r requirements.txt
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
 
-ENTRYPOINT ["python", "-m", "task"]
+# Stage 2: Runtime Environment
+FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
+
+# Install Google Cloud SDK
+RUN apt-get update && \
+    apt-get install -y curl gnupg && \
+    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
+    apt-get update -y && \
+    apt-get install -y google-cloud-sdk && \
+    apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Copy from the builder stage
+COPY --from=builder /trainer /trainer
+
+# Set working directory
+WORKDIR /trainer
+
+# Set the entry point
+ENTRYPOINT ["python", "-m", "task"]

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/src/simulator/requirements.txt

@@ -1,5 +1,5 @@
 dataclasses==0.6
 google-cloud-aiplatform==1.8.1
-tensorflow==2.7.2
+tensorflow==2.12.1
 pillow==10.3.0
 tf-agents==0.8.0

community-content/vertex_cpr_samples/sklearn/predictor_Classifier.py

@@ -0,0 +1,33 @@
+import numpy as np
+import os
+import pickle
+
+from google.cloud.aiplatform.constants import prediction
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from sklearn.datasets import load_breast_cancer
+from sklearn.linear_model import RidgeClassifier
+
+class LinearRegressionPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        if os.path.exists(prediction.MODEL_FILENAME_PKL):
+            self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
+        else:
+            self._model = RidgeClassifier()
+        X, y = load_breast_cancer(return_X_y=True)
+        self._model.fit(X, y)
+
+    def preprocess(self, prediction_input: dict) -> np.ndarray:
+        instances = prediction_input["instances"]
+        return np.asarray(instances)
+
+    def predict(self, instances: np.ndarray) -> np.ndarray:
+        return self._model.predict(instances)
+
+    def postprocess(self, prediction_results: np.ndarray) -> dict:
+        return {"predictions": prediction_results.tolist()}

community-content/vertex_cpr_samples/sklearn/predictor_SGDClassifier.py

@@ -0,0 +1,33 @@
+import numpy as np
+import os
+import pickle
+
+from google.cloud.aiplatform.constants import prediction
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from sklearn.linear_model import SGDClassifier
+
+class SGDClassifierPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        if os.path.exists(prediction.MODEL_FILENAME_PKL):
+            self._model = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
+        else:
+            self._model = SGDClassifier(max_iter=5)
+        X = [[0., 0.], [1., 1.]]
+        y = [0, 1]
+        self._model.fit(X, y)
+
+    def preprocess(self, prediction_input: dict) -> np.ndarray:
+        instances = prediction_input["instances"]
+        return np.asarray(instances)
+
+    def predict(self, instances: np.ndarray) -> np.ndarray:
+        return self._model.predict(instances)
+
+    def postprocess(self, prediction_results: np.ndarray) -> dict:
+        return {"predictions": prediction_results.tolist()}

community-content/vertex_cpr_samples/torch/predictor_resnet.py

@@ -0,0 +1,34 @@
+import os
+import torch
+
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from torchvision.models import detection, resnet50, ResNet50_Weights
+from typing import Dict, List
+
+class ResNetPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        if os.path.exists("model.pth.tar"):
+            self.model = detection.fasterrcnn_resnet50_fpn(pretrained=True)
+            stat_dic = torch.load("model.pth.tar")
+            self.model.load_state_dict(stat_dic['state_dict'])
+        else:
+            weights = ResNet50_Weights.DEFAULT
+            self.model = resnet50(weights=weights)
+        self.model.eval()
+
+    def preprocess(self, prediction_input: dict) -> torch.Tensor:
+        instances = prediction_input["instances"]
+        return torch.Tensor(instances)
+
+    @torch.inference_mode()
+    def predict(self, instances: torch.Tensor) -> List[str]:
+        return self._model(instances)
+
+    def postprocess(self, prediction_results: List[str]) -> Dict:
+        return {"predictions": prediction_results}

community-content/vertex_cpr_samples/xgboost/predictor_Classifier.py

@@ -0,0 +1,37 @@
+import os
+import numpy as np
+import pickle
+import xgboost as xgb
+
+from google.cloud.aiplatform.constants import prediction
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+from sklearn.datasets import make_blobs
+from xgboost import XGBClassifier
+
+
+class ClassifierPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        if os.path.exists(prediction.MODEL_FILENAME_PKL):
+            booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
+        else:
+            X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=1)
+            model = XGBClassifier()
+            model.fit(X, y)
+            booster = model.get_booster()
+        self._booster = booster
+
+    def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
+        instances = prediction_input["instances"]
+        return xgb.DMatrix(instances)
+
+    def predict(self, instances: xgb.DMatrix) -> np.ndarray:
+        return self._booster.predict(instances)
+
+    def postprocess(self, prediction_results: np.ndarray) -> dict:
+        return {"predictions": prediction_results.tolist()}

community-content/vertex_cpr_samples/xgboost/predictor_xgbranker.py

@@ -0,0 +1,41 @@
+import os
+import numpy as np
+import pandas as pd
+import pickle
+import xgboost as xgb
+
+from google.cloud.aiplatform.constants import prediction
+from google.cloud.aiplatform.utils import prediction_utils
+from google.cloud.aiplatform.prediction.predictor import Predictor
+
+class XGBRankerPredictor(Predictor):
+
+    def __init__(self):
+        return
+
+    def load(self, artifacts_uri: str) -> None:
+        prediction_utils.download_model_artifacts(artifacts_uri)
+        if os.path.exists(prediction.MODEL_FILENAME_PKL):
+            booster = pickle.load(open(prediction.MODEL_FILENAME_PKL, "rb"))
+            self._booster = booster
+        else:
+            N = 500
+            dates = pd.date_range(start='2023-01-01', end='2023-01-12', periods=N)
+            X = pd.DataFrame(np.random.randn(N, 5), columns=list('ABCDE'), index=dates)
+            y = pd.Series(np.random.randint(0, 10, size=N), index=dates, name='label')
+            group = X.groupby(dates + pd.offsets.MonthEnd(0)).size()
+            sample_weight = pd.Series(np.arange(len(group)), index=group.index)
+            model = xgb.XGBRanker(objective='rank:pairwise', max_depth=3, learning_rate=0.1, booster='gbtree', tree_method='hist', n_jobs=4, n_estimators=50, enable_categorical=False, random_state=42)
+            model.fit(X=X, y=y, group=group, sample_weight=sample_weight, verbose=True)
+            booster = model.get_booster()
+            self._booster = booster
+
+    def preprocess(self, prediction_input: dict) -> xgb.DMatrix:
+        instances = prediction_input["instances"]
+        return xgb.DMatrix(instances)
+
+    def predict(self, instances: xgb.DMatrix) -> np.ndarray:
+        return self._booster.predict(instances, output_margin=False, ntree_limit=0)
+
+    def postprocess(self, prediction_results: np.ndarray) -> dict:
+        return {"predictions": prediction_results.tolist()}

community-content/vertex_model_garden/model_oss/notebook_util/common_util.py

@@ -230,7 +230,7 @@ def download_image(url: str) -> str:
     base64 encoded image.
   """
   response = requests.get(url)
-  return Image.open(io.BytesIO(response.content))
+  return Image.open(io.BytesIO(response.content))  # pytype: disable=bad-return-type  # pillow-102-upgrade
 
 
 def resize_image(image: Any, new_width: int = 1000) -> Any:
@@ -422,6 +422,41 @@ def detect_predict(
   return response.predictions[0].get("response")
 
 
+def copy_model_artifacts(
+    model_id: str,
+    model_source: str,
+    model_destination: str,
+) -> None:
+  """Copies model artifacts from model_source to model_destination.
+
+  model_source and model_destination should be GCS path.
+
+  Args:
+    model_id: The model id.
+    model_source: The source of the model artifact.
+    model_destination: The destination of the model artifact.
+  """
+  if not model_source.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{model_source} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  if not model_destination.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{model_destination} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  model_source = f"{model_source}/{model_id}"
+  model_destination = f"{model_destination}/{model_id}"
+  print("Copying model artifact from ", model_source, " to ", model_destination)
+  subprocess.check_output([
+      "gcloud",
+      "storage",
+      "cp",
+      "-r",
+      model_source,
+      model_destination,
+  ])
+
+
 def get_quota(project_id: str, region: str, resource_id: str) -> int:
   """Returns the quota for a resource in a region.
 
@@ -476,6 +511,7 @@ def get_resource_id(
     accelerator_type: str,
     is_for_training: bool,
     is_restricted_image: bool = False,
+    is_dynamic_workload_scheduler: bool = False,
 ) -> str:
   """Returns the resource id for a given accelerator type and the use case.
 
@@ -484,45 +520,62 @@ def get_resource_id(
     is_for_training: Whether the resource is used for training. Set false for
       serving use case.
     is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
+    is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
+      Workload Scheduler.
 
   Returns:
     The resource id.
   """
+  accelerator_suffix_map = {
+      "NVIDIA_TESLA_V100": "nvidia_v100_gpus",
+      "NVIDIA_TESLA_P100": "nvidia_p100_gpus",
+      "NVIDIA_L4": "nvidia_l4_gpus",
+      "NVIDIA_TESLA_A100": "nvidia_a100_gpus",
+      "NVIDIA_A100_80GB": "nvidia_a100_80gb_gpus",
+      "NVIDIA_H100_80GB": "nvidia_h100_gpus",
+      "NVIDIA_TESLA_T4": "nvidia_t4_gpus",
+      "TPU_V5e": "tpu_v5e",
+      "TPU_V3": "tpu_v3",
+  }
   default_training_accelerator_map = {
-      "NVIDIA_TESLA_V100": "custom_model_training_nvidia_v100_gpus",
-      "NVIDIA_L4": "custom_model_training_nvidia_l4_gpus",
-      "NVIDIA_TESLA_A100": "custom_model_training_nvidia_a100_gpus",
-      "NVIDIA_A100_80GB": "custom_model_training_nvidia_a100_80gb_gpus",
-      "NVIDIA_H100_80GB": "custom_model_training_nvidia_h100_gpus",
-      "NVIDIA_TESLA_T4": "custom_model_training_nvidia_t4_gpus",
-      "TPU_V5e": "custom_model_training_tpu_v5e",
-      "TPU_V3": "custom_model_training_tpu_v3",
+      key: f"custom_model_training_{accelerator_suffix_map[key]}"
+      for key in accelerator_suffix_map
+  }
+  dws_training_accelerator_map = {
+      key: f"custom_model_training_preemptible_{accelerator_suffix_map[key]}"
+      for key in accelerator_suffix_map
   }
   restricted_image_training_accelerator_map = {
       "NVIDIA_A100_80GB": "restricted_image_training_nvidia_a100_80gb_gpus",
   }
   serving_accelerator_map = {
-      "NVIDIA_TESLA_V100": "custom_model_serving_nvidia_v100_gpus",
-      "NVIDIA_L4": "custom_model_serving_nvidia_l4_gpus",
-      "NVIDIA_TESLA_A100": "custom_model_serving_nvidia_a100_gpus",
-      "NVIDIA_A100_80GB": "custom_model_serving_nvidia_a100_80gb_gpus",
-      "NVIDIA_H100_80GB": "custom_model_serving_nvidia_h100_gpus",
-      "NVIDIA_TESLA_T4": "custom_model_serving_nvidia_t4_gpus",
-      "TPU_V5e": "custom_model_serving_tpu_v5e",
+      key: f"custom_model_serving_{accelerator_suffix_map[key]}"
+      for key in accelerator_suffix_map
   }
+
   if is_for_training:
+    if is_restricted_image and is_dynamic_workload_scheduler:
+      raise ValueError(
+          "Dynamic Workload Scheduler does not work for restricted image"
+          " training."
+      )
     training_accelerator_map = (
         restricted_image_training_accelerator_map
         if is_restricted_image
         else default_training_accelerator_map
     )
     if accelerator_type in training_accelerator_map:
-      return training_accelerator_map[accelerator_type]
+      if is_dynamic_workload_scheduler:
+        return dws_training_accelerator_map[accelerator_type]
+      else:
+        return training_accelerator_map[accelerator_type]
     else:
       raise ValueError(
           f"Could not find accelerator type: {accelerator_type} for training."
       )
   else:
+    if is_dynamic_workload_scheduler:
+      raise ValueError("Dynamic Workload Scheduler does not work for serving.")
     if accelerator_type in serving_accelerator_map:
       return serving_accelerator_map[accelerator_type]
     else:
@@ -538,10 +591,14 @@ def check_quota(
     accelerator_count: int,
     is_for_training: bool,
     is_restricted_image: bool = False,
+    is_dynamic_workload_scheduler: bool = False,
 ):
   """Checks if the project and the region has the required quota."""
   resource_id = get_resource_id(
-      accelerator_type, is_for_training, is_restricted_image
+      accelerator_type,
+      is_for_training=is_for_training,
+      is_restricted_image=is_restricted_image,
+      is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,
   )
   quota = get_quota(project_id, region, resource_id)
   quota_request_instruction = (
@@ -562,3 +619,4 @@ def check_quota(
         f"Quota not enough for {resource_id} in {region}: {quota} <"
         f" {accelerator_count}. {quota_request_instruction}"
     )
+

community-content/vertex_model_garden/model_oss/notebook_util/dataset_validation_util.py

@@ -0,0 +1,544 @@
+"""Functions for dataset validation.
+
+This tool is used to validate the dataset against the given template.
+"""
+
+import json
+import multiprocessing
+import os
+import subprocess
+from typing import Any, Callable, Dict, Union
+from absl import logging
+import accelerate
+import datasets
+import transformers
+
+GCS_URI_PREFIX = "gs://"
+GCSFUSE_URI_PREFIX = "/gcs/"
+LOCAL_BASE_MODEL_DIR = "/tmp/base_model_dir"
+LOCAL_TEMPLATE_DIR = "/tmp/template_dir"
+_TEMPLATE_DIRNAME = "templates"
+_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME = "vertex-ai-samples"
+_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR = (
+    "community-content/vertex_model_garden/model_oss/peft/train/vmg/templates"
+)
+_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
+_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
+_DESCRIPTION_KEY = "description"
+_SOURCE_KEY = "source"
+_PROMPT_INPUT_KEY = "prompt_input"
+_PROMPT_NO_INPUT_KEY = "prompt_no_input"
+_RESPONSE_SEPARATOR = "response_separator"
+_INSTRUCTION_SEPARATOR = "instruction_separator"
+_CHAT_TEMPLATE_KEY = "chat_template"
+_KNOWN_KEYS = (
+    _DESCRIPTION_KEY,
+    _SOURCE_KEY,
+    _PROMPT_INPUT_KEY,
+    _PROMPT_NO_INPUT_KEY,
+    _RESPONSE_SEPARATOR,
+    _INSTRUCTION_SEPARATOR,
+    _CHAT_TEMPLATE_KEY,
+)
+
+
+def is_gcs_path(input_path: str) -> bool:
+  """Checks if the input path is a Google Cloud Storage (GCS) path.
+
+  Args:
+      input_path: The input path to be checked.
+
+  Returns:
+      True if the input path is a GCS path, False otherwise.
+  """
+  return input_path is not None and input_path.startswith(GCS_URI_PREFIX)
+
+
+def force_gcs_fuse_path(gcs_uri: str) -> str:
+  """Converts gs:// uris to their /gcs/ equivalents. No-op for other uris.
+
+  Args:
+    gcs_uri: The GCS URI to convert.
+
+  Returns:
+    The converted GCS URI.
+  """
+  if is_gcs_path(gcs_uri):
+    return GCSFUSE_URI_PREFIX + gcs_uri[len(GCS_URI_PREFIX) :]
+  else:
+    return gcs_uri
+
+
+def download_gcs_uri_to_local(
+    gcs_uri: str, destination_dir: str = LOCAL_BASE_MODEL_DIR
+) -> str:
+  """Downloads GCS URI to local.
+
+  If GCS URI is a directory, gs://some/folder is downloaded to
+  /destination_dir/folder. If GCS URI is a file, gs://some/file is downloaded to
+  /destination_dir/file.
+
+  Args:
+    gcs_uri: GCS URI to download.
+    destination_dir: Local directory directory.
+
+  Returns:
+    Local path to target folder/file.
+  """
+  target = os.path.join(
+      destination_dir,
+      os.path.basename(os.path.normpath(gcs_uri)),
+  )
+  if os.path.exists(target):
+    logging.info("File %s already exists.", target)
+    return target
+  if accelerate.PartialState().is_local_main_process:
+    logging.info(
+        "Downloading file(s) from %s to %s...", gcs_uri, destination_dir
+    )
+    if not os.path.exists(destination_dir):
+      os.mkdir(destination_dir)
+    subprocess.check_output([
+        "gsutil",
+        "-m",
+        "cp",
+        "-r",
+        gcs_uri,
+        destination_dir,
+    ])
+    logging.info("Downloaded file(s) from %s to %s.", gcs_uri, destination_dir)
+  # Make sure ALL processes process to next step after data downloading is done.
+  # It matters for the main process to wait for other processes as well.
+  accelerate.PartialState().wait_for_everyone()
+  return target
+
+
+def get_template(template_path: str) -> Dict[str, str]:
+  """Gets the template dictionary given the file path.
+
+  Args:
+    template_path: Path to the template file.
+
+  Returns:
+    A dictionary of the template.
+
+  Raises:
+    ValueError: If the template file does not exist or contains unknown keys.
+  """
+  if is_gcs_path(template_path):
+    template_path = force_gcs_fuse_path(template_path)
+  elif not os.path.isfile(template_path):
+    template_path = os.path.join(
+        os.path.dirname(__file__),
+        _TEMPLATE_DIRNAME,
+        template_path + ".json",
+    )
+  if not os.path.isfile(template_path):
+    raise ValueError(f"Template file {template_path} does not exist.")
+  with open(template_path, "r") as f:
+    template_json: dict[str, str] = json.load(f)
+    for key in template_json:
+      if key not in _KNOWN_KEYS:
+        raise ValueError(f"Unknown key {key} in template {template_path}.")
+    return template_json
+
+
+def get_response_separator(template_json: Dict[str, str]) -> Union[str, None]:
+  return template_json.get(_RESPONSE_SEPARATOR, None)
+
+
+def get_instruction_separator(
+    template_json: Dict[str, str],
+) -> Union[str, None]:
+  return template_json.get(_INSTRUCTION_SEPARATOR, None)
+
+
+def _format_template_fn(
+    template: str,
+    input_column: str,
+    tokenizer: transformers.PreTrainedTokenizer | None = None,
+) -> Callable[[Dict[str, str]], Dict[str, str]]:
+  """Formats a dataset example according to a template.
+
+  Args:
+    template: Name of the JSON template file under `templates/` or GCS path to
+      the template file.
+    input_column: The input column in the dataset to be used or updated by the
+      template. If it does not exist, the template's `prompt_no_input` will be
+      used, and the input_column will be created.
+    tokenizer: The tokenizer to use for chat_template templates.
+
+  Returns:
+    A function that formats data according to the template.
+  """
+  template_json = get_template(template)
+
+  if _CHAT_TEMPLATE_KEY not in template_json:
+
+    def format_fn(example: Dict[str, str]) -> Dict[str, str]:
+      format_dict = {key: value for key, value in example.items()}
+      if format_dict.get(input_column):
+        format_str = template_json[_PROMPT_INPUT_KEY]
+      elif _PROMPT_NO_INPUT_KEY in template_json:
+        format_str = template_json[_PROMPT_NO_INPUT_KEY]
+      else:
+        raise KeyError(
+            f"The template {os.path.basename(template)} does not contain"
+            f" {_PROMPT_INPUT_KEY} or {_PROMPT_NO_INPUT_KEY} key."
+        )
+      try:
+        return {input_column: format_str.format(**format_dict)}
+      except KeyError as e:
+        raise KeyError(
+            f"The template {os.path.basename(template)} contains a key {e} in"
+            f" {_PROMPT_INPUT_KEY} or {_PROMPT_NO_INPUT_KEY} that does not"
+            " exist in the dataset example. The dataset example looks like"
+            f" {format_dict}."
+        ) from e
+
+    return format_fn
+  elif (
+      _PROMPT_INPUT_KEY in template_json
+      or _PROMPT_NO_INPUT_KEY in template_json
+  ):
+    raise ValueError(
+        f"chat_template templates do not support {_PROMPT_INPUT_KEY} or"
+        f" {_PROMPT_NO_INPUT_KEY} templates."
+    )
+  else:
+    if tokenizer is None:
+      raise ValueError("A tokenizer is required for chat_template templates.")
+    # Assign HuggingFace jinja template.
+    tokenizer.chat_template = template_json[_CHAT_TEMPLATE_KEY]
+
+    def format_fn(example: Dict[str, str]) -> Dict[str, str]:
+      try:
+        return {
+            input_column: tokenizer.apply_chat_template(
+                example[input_column],
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+        }
+      except KeyError as e:
+        raise KeyError(
+            f"The template {os.path.basename(template)} contains a key {e} in"
+            f" {_CHAT_TEMPLATE_KEY} that does not exist in the dataset example."
+        ) from e
+
+    return format_fn
+
+
+def _get_split_string(
+    split: str,
+    dataset_percent: int | None = None,
+    dataset_k_rows: int | None = None,
+) -> str:
+  """Gets the formatted split string for the dataset.
+
+  This is used to format the split string as per
+  https://huggingface.co/docs/datasets/v2.21.0/loading#slice-splits. Also, this
+  function will only be used to load the partial dataset for validating the
+  dataset against the template.
+
+  Args:
+    split: Split of the dataset.
+    dataset_percent: The percentage of the dataset to load.
+    dataset_k_rows: The top k sequences to load from the dataset.
+
+  Returns:
+    A formatted split string.
+  """
+  # Validate the dataset_percent and dataset_k_rows values.
+  if dataset_percent and dataset_k_rows:
+    raise ValueError(
+        "You can set either validate_percentage_of_dataset or"
+        " validate_k_rows_of_dataset, but not both."
+    )
+
+  if dataset_percent:
+    logging.info("Loading %d percent of the dataset...", dataset_percent)
+    return f"{split}[:{dataset_percent}%]"
+
+  if dataset_k_rows:
+    logging.info("Loading top %d rows of the dataset...", dataset_k_rows)
+    return f"{split}[:{dataset_k_rows}]"
+
+  return split
+
+
+def _github_template_path(template: str) -> str:
+  """Generates the path to the template in the Vertex AI Samples GitHub repo.
+
+  Args:
+    template: Name of the template.
+
+  Returns:
+    The path to the template in the Vertex AI Samples GitHub repo.
+  """
+  # vertex-ai-samples directory may lie under separate directory depending on
+  # the scratch_dir parameter in the notebook execution environment.
+  vertex_ai_samples_abs_path = os.getcwd().split(
+      _VERTEX_AI_SAMPLES_GITHUB_REPO_NAME
+  )[0]
+  return os.path.join(
+      vertex_ai_samples_abs_path,
+      _VERTEX_AI_SAMPLES_GITHUB_REPO_NAME,
+      _VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR,
+      template + ".json",
+  )
+
+
+def _get_dataset(
+    dataset_name: str,
+    split: str,
+    num_proc: int | None = None,
+) -> datasets.DatasetDict:
+  """Gets a dataset.
+
+  Args:
+    dataset_name: Name of the dataset or path to a custom dataset.
+    split: Split of the dataset.
+    num_proc: Number of processors to use.
+
+  Returns:
+    A dataset.
+  """
+  dataset_name = force_gcs_fuse_path(dataset_name)
+  if os.path.isfile(dataset_name):
+    # Custom dataset.
+    return datasets.load_dataset(
+        "json",
+        data_files=[dataset_name],
+        split=split,
+        num_proc=num_proc,
+    )
+  # HF dataset.
+  return datasets.load_dataset(dataset_name, split=split, num_proc=num_proc)
+
+
+def should_add_pad_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special pad token.
+
+  Args:
+    model_id: The name of the model.
+
+  Returns:
+    True if the model requires adding a special pad token, False otherwise.
+  """
+  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
+
+
+def should_add_eos_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special eos token.
+
+  Args:
+    model_id: The name of the model.
+
+  Returns:
+    True if the model requires adding a special eos token, False otherwise.
+  """
+  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
+
+
+def load_tokenizer(
+    pretrained_model_id: str,
+    padding_side: str | None = None,
+    access_token: str | None = None,
+) -> transformers.AutoTokenizer:
+  """Loads tokenizer based on `pretrained_model_id`.
+
+  Args:
+    pretrained_model_id: The name of the pretrained model.
+    padding_side: The side to pad the input on.
+    access_token: The access token to use for the tokenizer.
+
+  Returns:
+    The tokenizer.
+  """
+  tokenizer_kwargs = {}
+  if should_add_eos_token(pretrained_model_id):
+    tokenizer_kwargs["add_eos_token"] = True
+  if padding_side:
+    tokenizer_kwargs["padding_side"] = padding_side
+
+  with accelerate.PartialState().local_main_process_first():
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        pretrained_model_id,
+        trust_remote_code=False,
+        use_fast=True,
+        token=access_token,
+        **tokenizer_kwargs,
+    )
+
+  if should_add_pad_token(pretrained_model_id):
+    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+  return tokenizer
+
+
+def get_filtered_dataset(
+    dataset: Any,
+    input_column: str,
+    max_seq_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Any:
+  """Returns the dataset by removing examples that are longer than max_seq_length.
+
+  Args:
+    dataset: The dataset to filter.
+    input_column: The input column in the dataset to be used.
+    max_seq_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+  """
+  actual_dataset_length = len(dataset)
+  filtered_dataset = dataset.filter(
+      lambda x: len(tokenizer(x[input_column])["input_ids"]) <= max_seq_length
+  )
+  filtered_dataset_length = len(filtered_dataset)
+  if actual_dataset_length != filtered_dataset_length:
+    examples_removed_percent = (
+        (actual_dataset_length - filtered_dataset_length)
+        * 100
+        / actual_dataset_length
+    )
+    logging.info(
+        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
+        " max-seq-length. Filtering out %d example(s) which are longer than"
+        " max-seq-length.",
+        100 - examples_removed_percent,
+        max_seq_length,
+        examples_removed_percent,
+        actual_dataset_length - filtered_dataset_length,
+    )
+
+  return filtered_dataset
+
+
+def load_dataset_with_template(
+    dataset_name: str,
+    split: str,
+    input_column: str,
+    template: str = None,
+    tokenizer: transformers.PreTrainedTokenizer | None = None,
+) -> Any:
+  """Loads dataset with templates.
+
+  Args:
+    dataset_name: Name of the dataset or path to a custom dataset.
+    split: Split of the dataset.
+    input_column: The input column in the dataset to be used or updaded by the
+      template. If it does not exist, the template's `prompt_no_input` will be
+      used, and the input_column will be created.
+    template: Name of the JSON template file under `templates/` or GCS path to
+      the template file.
+    tokenizer: The tokenizer to use for chat_template templates.
+
+  Returns:
+    A dataset compatible with the template.
+  """
+  dataset = _get_dataset(dataset_name, split=split)
+  if template:
+    dataset = dataset.map(
+        _format_template_fn(
+            template,
+            input_column=input_column,
+            tokenizer=tokenizer,
+        )
+    )
+
+  return dataset
+
+
+def validate_dataset_with_template(
+    dataset_name: str,
+    split: str,
+    input_column: str,
+    template: str,
+    tokenizer: transformers.PreTrainedTokenizer | None = None,
+    max_seq_length: int | None = None,
+    use_multiprocessing: bool = False,
+    validate_percentage_of_dataset: int | None = None,
+    validate_k_rows_of_dataset: int | None = None,
+) -> Any:
+  """Validates dataset with templates.
+
+  This function will be used to load the dataset and validate it against the
+  template. In case of validation, we also allow the users to load the dataset
+  partially by allowing them to read x% or top k rows of the dataset. To
+  validate the dataset, the template file must be available in the GCS bucket
+  and the dataset must be available either in the GCS bucket or Hugging Face.
+
+  Args:
+    dataset_name: Name of the dataset or path to a custom dataset.
+    split: Split of the dataset.
+    input_column: The input column in the dataset to be used or updaded by the
+      template. If it does not exist, the template's `prompt_no_input` will be
+      used, and the input_column will be created.
+    template: Name of the JSON template file under `templates/` or GCS path to
+      the template file.
+    tokenizer: The tokenizer to use for chat_template templates.
+    max_seq_length: The maximum sequence length.
+    use_multiprocessing: If True, it will use multiprocessing to load the
+      dataset.
+    validate_percentage_of_dataset: The percentage of the dataset to load.
+    validate_k_rows_of_dataset: The top k sequences to load from the dataset.
+
+  Returns:
+    None if the validation is successful, otherwise returns the error message.
+  """
+  if not template:
+    raise ValueError("template is required for validate_dataset.")
+
+  if not dataset_name:
+    raise ValueError("dataset_name is empty.")
+
+  if not split:
+    raise ValueError("split is empty.")
+
+  split = _get_split_string(
+      split,
+      validate_percentage_of_dataset,
+      validate_k_rows_of_dataset,
+  )
+
+  num_proc = multiprocessing.cpu_count() if use_multiprocessing else 1
+
+  # gcsfuse cannot be used from the notebook runtime env. Hence, we have
+  # to download dataset and template from gcs to local.
+  if is_gcs_path(dataset_name):
+    dataset_name = download_gcs_uri_to_local(dataset_name, LOCAL_BASE_MODEL_DIR)
+
+  if is_gcs_path(template):
+    template_path = download_gcs_uri_to_local(template, LOCAL_TEMPLATE_DIR)
+  elif os.path.isfile(_github_template_path(template)):
+    template_path = _github_template_path(template)
+  else:
+    raise ValueError(
+        f"Template file {template} does not exist. To validate the"
+        " dataset, please provide a valid GCS path for the template or a valid"
+        " template name from"
+        f" https://github.com/GoogleCloudPlatform/{_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME}/tree/main/{_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR}."
+    )
+
+  dataset = _get_dataset(dataset_name, split, num_proc).map(
+      _format_template_fn(
+          template_path,
+          input_column=input_column,
+          tokenizer=tokenizer,
+      )
+  )
+
+  if tokenizer is not None:
+    get_filtered_dataset(
+        dataset=dataset,
+        input_column=input_column,
+        max_seq_length=max_seq_length,
+        tokenizer=tokenizer,
+    )
+  print(
+      "Dataset {} is compatible with the {} template.".format(
+          os.path.basename(dataset_name), os.path.basename(template)
+      )
+  )
+

community-content/vertex_model_garden/model_oss/peft/causal_language_modeling_lora.py

@@ -1,144 +0,0 @@
-"""Causal language modeling with LoRA models."""
-
-# pylint: disable=g-importing-member
-
-from datasets import load_dataset
-from peft import get_peft_model
-from peft import LoraConfig
-import torch
-from torch import nn
-import transformers
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
-from transformers import BitsAndBytesConfig
-from transformers import TrainingArguments
-from typing import List
-from util import constants
-
-
-def finetune_causal_language_modeling(
-    pretrained_model_id: str,
-    dataset_name: str,
-    output_dir: str,
-    precision_mode: str = None,
-    lora_rank: int = 16,
-    lora_alpha: int = 32,
-    lora_dropout: float = 0.05,
-    target_modules: List[str] = constants.CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES,
-    warmup_steps: int = 10,
-    max_steps: int = 10,
-    learning_rate: float = 2e-4,
-    local_pretrained_model_id: str = None,
-) -> None:
-  """Finetunes causal language modelings."""
-  if precision_mode == constants.PRECISION_MODE_32:
-    model = AutoModelForCausalLM.from_pretrained(
-        local_pretrained_model_id
-        if local_pretrained_model_id
-        else pretrained_model_id,
-        torch_dtype=torch.float32,
-        device_map="auto",
-    )
-  elif precision_mode == constants.PRECISION_MODE_16:
-    model = AutoModelForCausalLM.from_pretrained(
-        local_pretrained_model_id
-        if local_pretrained_model_id
-        else pretrained_model_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-    )
-  elif precision_mode == constants.PRECISION_MODE_8:
-    quantization_config = BitsAndBytesConfig(
-        load_in_8bit=True, int8_threshold=0
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        local_pretrained_model_id
-        if local_pretrained_model_id
-        else pretrained_model_id,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        quantization_config=quantization_config,
-    )
-  else:
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_quant_type="nf4",
-        bnb_4bit_compute_dtype=torch.bfloat16,
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        local_pretrained_model_id
-        if local_pretrained_model_id
-        else pretrained_model_id,
-        device_map="auto",
-        torch_dtype=torch.bfloat16,
-        quantization_config=quantization_config,
-    )
-
-  tokenizer = AutoTokenizer.from_pretrained(
-      local_pretrained_model_id
-      if local_pretrained_model_id
-      else pretrained_model_id
-  )
-  if "llama" in pretrained_model_id:
-    tokenizer.pad_token = "[PAD]"
-
-  for param in model.parameters():
-    # Freezes the model - train adapters later.
-    param.requires_grad = False
-    if param.ndim == 1:
-      # Casts the small parameters (e.g. layernorm) to fp32 for stability.
-      param.data = param.data.to(torch.float32)
-
-  # Reduces the number of stored activations.
-  model.gradient_checkpointing_enable()
-  model.enable_input_require_grads()
-
-  class CastOutputToFloat(nn.Sequential):
-
-    def forward(self, x):
-      return super().forward(x).to(torch.float32)
-
-  model.lm_head = CastOutputToFloat(model.lm_head)
-
-  config = LoraConfig(
-      r=lora_rank,
-      lora_alpha=lora_alpha,
-      target_modules=target_modules,
-      lora_dropout=lora_dropout,
-      bias="none",
-      task_type="CAUSAL_LM",
-  )
-
-  model = get_peft_model(model, config)
-  model.print_trainable_parameters()
-
-  data = load_dataset(dataset_name)
-  data = data.map(
-      lambda samples: tokenizer(samples["quote"]),
-      batched=True,
-  )
-
-  trainer = transformers.Trainer(
-      model=model,
-      train_dataset=data["train"],
-      args=TrainingArguments(
-          per_device_train_batch_size=4,
-          gradient_accumulation_steps=4,
-          warmup_steps=warmup_steps,
-          max_steps=max_steps,
-          learning_rate=learning_rate,
-          fp16=True,
-          logging_steps=1,
-          output_dir=output_dir,
-          ddp_find_unused_parameters=False,
-      ),
-      data_collator=transformers.DataCollatorForLanguageModeling(
-          tokenizer,
-          mlm=False,
-      ),
-  )
-  # Silence the warnings. Please re-enable for inference!
-  model.config.use_cache = False
-  trainer.train()
-
-  model.save_pretrained(output_dir)

community-content/vertex_model_garden/model_oss/peft/dockerfile/code_llama_convert.Dockerfile

@@ -0,0 +1,28 @@
+# Base on pytorch-cuda image.
+FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-devel
+
+# Install tools.
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update
+RUN apt-get install -y --no-install-recommends apt-utils
+RUN apt-get install -y --no-install-recommends curl
+RUN apt-get install -y --no-install-recommends wget
+RUN apt-get install -y --no-install-recommends git
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN python3 -m pip install --upgrade pip
+RUN pip install tokenizers==0.13.3
+RUN pip install accelerate==0.21.0
+RUN pip install sentencepiece==0.1.99
+RUN pip install datasets==2.14.4
+RUN pip install protobuf==4.24.1
+
+# Install transformers
+RUN git clone https://github.com/huggingface/transformers.git
+WORKDIR transformers
+# Pin the commit to add-code-llama 08/25/2023
+RUN git reset --hard 015f8e110d270a0ad42de4ae5b98198d69eb1964
+RUN pip install -e .
+
+ENTRYPOINT ["python","src/transformers/models/llama/convert_llama_weights_to_hf.py"]

community-content/vertex_model_garden/model_oss/peft/dockerfile/conversion.Dockerfile

@@ -0,0 +1,22 @@
+# Dockerfile for Language Model Conversion.
+#
+# To build:
+# docker build -f model_oss/peft/dockerfile/conversion.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM tensorflow/build:2.14-python3.8
+
+RUN git clone https://github.com/facebookresearch/llama-recipes.git && \
+    cd llama-recipes && \
+    pip install -r requirements.txt && \
+    pip freeze | grep transformers && \
+    git clone https://github.com/huggingface/transformers.git && \
+    cd transformers && \
+    pip install protobuf
+
+WORKDIR /llama-recipes/transformers
+
+ENTRYPOINT ["python","src/transformers/models/llama/convert_llama_weights_to_hf.py"]

community-content/vertex_model_garden/model_oss/peft/dockerfile/serve.Dockerfile

@@ -7,39 +7,40 @@
 # docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
 # docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
 
-FROM pytorch/torchserve:0.7.0-gpu
+FROM pytorch/torchserve:0.11.0-gpu
 
 USER root
 
-ENV infer_port=7080
-ENV mng_port=7081
-ENV model_name="peft_serving"
+ENV INFER_PORT=7080
+ENV MNG_PORT=7081
+ENV MODEL="peft_serving"
 ENV PATH="/home/model-server/:${PATH}"
 
-RUN apt-get update && apt-get install -y --no-install-recommends \
+RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
         curl \
         wget \
         vim \
         git \
         git-lfs
 RUN git lfs install
+RUN apt-get autoremove -y
 
 # Install libraries.
 ENV PIP_ROOT_USER_ACTION=ignore
 RUN python3 -m pip install --upgrade pip
-RUN pip install --upgrade torch==2.0.1
+RUN pip install --upgrade torch==2.0.1 --index-url https://download.pytorch.org/whl/cu118
 RUN pip install torchvision==0.15.2
 RUN pip install tokenizers==0.13.3
 RUN pip install accelerate==0.21.0
 RUN pip install sentencepiece==0.1.99
 RUN pip install grpcio-status==1.33.2
 RUN pip install protobuf==3.19.6
-RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
+RUN pip install peft==0.5.0
 RUN pip install datasets==2.14.4
-RUN pip install triton==2.0.0.dev20221120
+RUN pip install triton==3.0.0
 RUN pip install xformers==0.0.20
-RUN pip install google-cloud-storage==2.7.0
-RUN pip install absl-py==1.4.0
+RUN pip install google-cloud-storage
+RUN pip install absl-py
 RUN pip install scipy==1.10.1
 RUN pip install evaluate==0.4.0
 RUN pip install scikit-learn==1.2.2
@@ -47,52 +48,43 @@ RUN pip install loralib==0.1.1
 RUN pip install bitsandbytes==0.39.0
 RUN pip install trl==0.4.4
 RUN pip install einops==0.6.1
-
-# Install diffusers from source.
-RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
-WORKDIR diffusers
-RUN pip install -e .
-WORKDIR /home/model-server
-
-# Install transformers from source.
-RUN git clone --depth 1 --branch v4.31.0 https://github.com/huggingface/transformers.git
-# The patch is used to change the transformers loading model behavior:
-# 1) For models on Huggingface hub: if the model has multiple shards, each shard
-#    will be downloaded separately and get deleted after loading to GPU.
-# 2) For models on local disk: if a model bin file is actually a text file
-#    recording a GCS path, the model file will be downloaded and get deleted
-#    after loading to GPU.
-COPY model_oss/peft/hf_transformers_lazy_download.patch /home/model-server/hf_transformers_lazy_download.patch
-WORKDIR transformers
-RUN git apply /home/model-server/hf_transformers_lazy_download.patch
-RUN pip install -e .
-WORKDIR /home/model-server
+RUN pip install optimum==1.13.2
+RUN pip install auto-gptq==0.4.2
+RUN pip install https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.7/autoawq-0.1.7+cu118-cp39-cp39-linux_x86_64.whl
+RUN pip install diffusers==0.27.2
+RUN pip install tiktoken==0.6.0
+RUn pip install git+https://github.com/huggingface/transformers.git@76fa17c1663a0efeca7208c20579833365584889
+RUN pip install pynvml==11.4.0
+RUN pip install -i https://test.pypi.org/simple/ bitsandbytes
 
 # Copy license.
+WORKDIR /home/model-server
 RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
 
 # Copy model artifacts.
 COPY model_oss/peft/handler.py /home/model-server/handler.py
 COPY model_oss/peft/config.properties /home/model-server/config.properties
 COPY model_oss/util/ /home/model-server/util/
+COPY model_oss/util/pytorch_startup_prober.sh /model_garden/scripts/pytorch_startup_prober.sh
 ENV PYTHONPATH /home/model-server/
 
 # Expose ports.
-EXPOSE ${infer_port}
-EXPOSE ${mng_port}
+EXPOSE ${INFER_PORT}
+EXPOSE ${MNG_PORT}
 
 # Set environments.
 ENV TASK "causal-language-modeling-lora"
-ENV MODEL_ID "openlm-research/open_llama_7b"
+ENV BASE_MODEL_ID ""
+ENV MODEL_ID ""
 ENV PRECISION_LOADING_MODE "float16"
 ENV FINETUNED_LORA_MODEL_PATH ""
-
+ENV TRUST_REMOTE_CODE ""
 
 # Archive model artifacts and dependencies.
 # Do not set --model-file and --serialized-file because model and checkpoint
 # will be dynamically loaded in handler.py.
 RUN torch-model-archiver \
---model-name=${model_name} \
+--model-name=${MODEL} \
 --version=1.0 \
 --handler=/home/model-server/handler.py \
 --runtime=python3 \
@@ -103,5 +95,5 @@ RUN torch-model-archiver \
 # Run Torchserve HTTP serve to respond to prediction requests.
 CMD ["torchserve", "--start", \
      "--ts-config", "/home/model-server/config.properties", \
-     "--models", "${model_name}=${model_name}.mar", \
+     "--models", "${MODEL}=${MODEL}.mar", \
      "--model-store", "/home/model-server/model-store"]

community-content/vertex_model_garden/model_oss/peft/dockerfile/train.Dockerfile

@@ -1,111 +0,0 @@
-# Dockerfile for PEFT Training.
-#
-# To build:
-# docker build -f model_oss/peft/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
-#
-# To push to gcr:
-# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
-# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
-
-# Builds GPU docker image of PyTorch
-# Uses multi-staged approach to reduce size
-# Stage 1
-# Use base conda image to reduce time
-FROM continuumio/miniconda3:latest AS compile-image
-# Specify py version
-ENV PYTHON_VERSION=3.8
-# Install apt libs - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
-RUN apt-get update && \
-    apt-get install -y curl git wget software-properties-common git-lfs && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists*
-
-# Install audio-related libraries
-RUN apt-get update && \
-    apt install -y ffmpeg
-
-RUN apt install -y libsndfile1-dev
-RUN git lfs install
-
-# Create our conda env - copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
-RUN conda create --name peft python=${PYTHON_VERSION} ipython jupyter pip
-RUN python3 -m pip install --no-cache-dir --upgrade pip
-
-# Below is copied from https://github.com/huggingface/accelerate/blob/main/docker/accelerate-gpu/Dockerfile
-# We don't install pytorch here yet since CUDA isn't available
-# instead we use the direct torch wheel
-ENV PATH /opt/conda/envs/peft/bin:$PATH
-# Activate our bash shell
-RUN chsh -s /bin/bash
-SHELL ["/bin/bash", "-c"]
-# Activate the conda env and install transformers + accelerate from source
-RUN source activate peft
-RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/transformers
-RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/accelerate
-RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft#egg=peft[test]
-RUN python3 -m pip install --no-cache-dir bitsandbytes
-
-# Stage 2
-FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
-COPY --from=compile-image /opt/conda /opt/conda
-ENV PATH /opt/conda/bin:$PATH
-
-# Install apt libs
-RUN apt-get update && \
-    apt-get install -y curl git wget vim && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists*
-
-# Copy license.
-RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
-
-RUN echo "source activate peft" >> ~/.profile
-
-# Install libraries.
-RUN pip install --upgrade torch==2.0.1
-RUN pip install torchvision==0.15.2
-RUN pip install git+https://github.com/huggingface/transformers@de9255de27abfcae4a1f816b904915f0b1e23cd9
-RUN pip install transformers -U
-RUN pip install accelerate==0.21.0
-RUN pip install sentencepiece==0.1.99
-RUN pip install grpcio-status==1.33.2
-RUN pip install protobuf==3.19.6
-RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
-RUN pip install datasets==2.9.0
-RUN pip install triton==2.0.0.dev20221120
-RUN pip install xformers==0.0.20
-RUN pip install Jinja2==3.1.2
-RUN pip install ftfy==6.1.1
-RUN pip install cloudml-hypertune==0.1.0.dev6
-RUN pip install tensorboard==2.12.0
-RUN pip install scipy==1.10.1
-RUN pip install evaluate==0.4.0
-RUN pip install scikit-learn==1.2.2
-RUN pip install loralib==0.1.1
-RUN pip install bitsandbytes==0.39.0
-RUN pip install trl==0.4.4
-RUN pip install einops==0.6.1
-RUN pip install google-cloud-storage==2.7.0
-
-RUN git clone --depth 1 --branch v0.16.1 https://github.com/huggingface/diffusers.git
-WORKDIR diffusers
-RUN pip install -e .
-
-# Switch to diffusers examples folder.
-WORKDIR examples
-
-# NOTE: use 'sed' to modify train_text_to_image_lora.py to
-#       fix the bug for accelerator.
-RUN sed -i \
-"s#logging_dir=logging_dir#project_dir=logging_dir#g" \
-text_to_image/train_text_to_image_lora.py
-
-# Config accelerate.
-RUN mkdir -p ./vertex_vision_model_garden_peft/
-COPY model_oss/peft/train.sh ./vertex_vision_model_garden_peft/train.sh
-COPY model_oss/peft/*.py     ./vertex_vision_model_garden_peft/
-COPY model_oss/util /diffusers/examples/util
-ENV PYTHONPATH /diffusers/examples/
-
-# Generate accelerate config at the beginning of docker run.
-ENTRYPOINT ["python3", "vertex_vision_model_garden_peft/main.py"]

community-content/vertex_model_garden/model_oss/peft/handler.py

@@ -72,15 +72,39 @@ class PeftHandler(BaseHandler):
         "PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
     )
     self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
-    self.base_model_id = os.environ.get("BASE_MODEL_ID", None)
-    self.model_id = self.base_model_id
-    if not self.base_model_id:
-      self.model_id = os.environ.get("MODEL_ID", "")
+    trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", None)
+    if trust_remote_code == "false":
+      self.trust_remote_code = False
+    else:
+      self.trust_remote_code = True
+
+    # If present, the path of the model in the container.
+    aip_storage_dir = os.environ.get("AIP_STORAGE_DIR", None)
+
+    # If present, the URI of the model in a google owned GCS bucket.
+    aip_storage_uri = os.environ.get("AIP_STORAGE_URI", None)
+
+    model_id = os.environ.get("MODEL_ID", None)
+    base_model_id = os.environ.get("BASE_MODEL_ID", None)
+
+    self.model_id = None
+    if aip_storage_dir:
+      self.model_id = aip_storage_dir
+      logging.info(f"Loaded base model from AIP_STORAGE_DIR: {self.model_id}.")
+    elif aip_storage_uri:
+      self.model_id = aip_storage_uri
+      logging.info(f"Loaded base model from AIP_STORAGE_URI: {self.model_id}.")
+    elif model_id:
+      self.model_id = model_id
+      logging.info(f"Loaded base model from MODEL_ID: {self.model_id}.")
+    elif base_model_id:
+      # Note: BASE_MODEL_ID has been unified with MODEL_ID.
+      # MODEL_ID should be used whenever possible.
+      self.model_id = base_model_id
+      logging.info(f"Loaded base model from BASE_MODEL_ID: {self.model_id}.")
+
     self.quantization = os.environ.get("QUANTIZATION", None)
-    logging.info(f"Load base model id from MODEL_ID:{self.model_id}.")
-    if not self.model_id:
-      self.model_id = os.environ.get("AIP_STORAGE_URI", "")
-      logging.info(f"Load base model id from AIP_STORAGE_URI: {self.model_id}.")
+
     if not self.model_id:
       raise ValueError("Base model id is must be set.")
     if fileutils.is_gcs_path(self.model_id):
@@ -101,8 +125,7 @@ class PeftHandler(BaseHandler):
 
     logging.info(
         f"Using task:{self.task}, base model:{self.model_id}, lora model:"
-        f" {self.finetuned_lora_model_path}, and precision"
-        f" {self.precision_mode}."
+        f" {self.finetuned_lora_model_path}, precision {self.precision_mode}."
     )
 
     self.pipeline = None
@@ -145,11 +168,18 @@ class PeftHandler(BaseHandler):
     elif (
         self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
     ):
-      tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+      tokenizer = AutoTokenizer.from_pretrained(
+          self.model_id,
+          trust_remote_code=self.trust_remote_code,
+      )
+
       logging.debug("Initialized the tokenizer.")
       if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
         if self.quantization == constants.AWQ:
-          model = AutoAWQForCausalLM.from_quantized(self.model_id)
+          model = AutoAWQForCausalLM.from_quantized(
+              self.model_id,
+              trust_remote_code=self.trust_remote_code,
+          )
         elif self.quantization == constants.GPTQ or not self.quantization:
           if self.precision_mode == constants.PRECISION_MODE_32:
             model = AutoModelForCausalLM.from_pretrained(
@@ -157,6 +187,7 @@ class PeftHandler(BaseHandler):
                 return_dict=True,
                 torch_dtype=torch.float32,
                 device_map="auto",
+                trust_remote_code=self.trust_remote_code,
             )
           elif self.precision_mode == constants.PRECISION_MODE_16B:
             model = AutoModelForCausalLM.from_pretrained(
@@ -164,6 +195,7 @@ class PeftHandler(BaseHandler):
                 return_dict=True,
                 torch_dtype=torch.bfloat16,
                 device_map="auto",
+                trust_remote_code=self.trust_remote_code,
             )
           elif self.precision_mode == constants.PRECISION_MODE_16:
             model = AutoModelForCausalLM.from_pretrained(
@@ -171,6 +203,7 @@ class PeftHandler(BaseHandler):
                 return_dict=True,
                 torch_dtype=torch.float16,
                 device_map="auto",
+                trust_remote_code=self.trust_remote_code,
             )
           elif self.precision_mode == constants.PRECISION_MODE_8:
             quantization_config = BitsAndBytesConfig(
@@ -182,6 +215,7 @@ class PeftHandler(BaseHandler):
                 torch_dtype=torch.float16,
                 device_map="auto",
                 quantization_config=quantization_config,
+                trust_remote_code=self.trust_remote_code,
             )
           else:
             quantization_config = BitsAndBytesConfig(
@@ -195,6 +229,7 @@ class PeftHandler(BaseHandler):
                 device_map="auto",
                 torch_dtype=torch.bfloat16,
                 quantization_config=quantization_config,
+                trust_remote_code=self.trust_remote_code,
             )
         else:
           raise ValueError(f"Invalid QUANTIZATION value: {self.quantization}")
@@ -203,14 +238,14 @@ class PeftHandler(BaseHandler):
           model = AutoModelForCausalLM.from_pretrained(
               self.model_id,
               torch_dtype=torch.bfloat16,
-              trust_remote_code=True,
+              trust_remote_code=self.trust_remote_code,
               device_map="auto",
           )
         except:  # pylint: disable=bare-except
           model = AutoModelForCausalLM.from_pretrained(
               self.model_id,
               torch_dtype=torch.bfloat16,
-              trust_remote_code=True,
+              trust_remote_code=self.trust_remote_code,
               device_map="auto",
           )
       logging.debug("Initialized the base model.")
@@ -329,4 +364,4 @@ class PeftHandler(BaseHandler):
     return f"Prompt:\n{prompt.strip()}\nOutput:\n{output}"
 
 
-# pylint: enable=logging-fstring-interpolation
+# pylint: enable=logging-fstring-interpolation

community-content/vertex_model_garden/model_oss/peft/hf_transformers_lazy_download.patch

@@ -1,131 +0,0 @@
-diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
-index 45459ed..32527f4 100644
---- a/src/transformers/modeling_utils.py
-+++ b/src/transformers/modeling_utils.py
-@@ -32,6 +32,8 @@ import torch
- from packaging import version
- from torch import Tensor, nn
- from torch.nn import CrossEntropyLoss
-+from huggingface_hub import hf_hub_download
-+from google.cloud import storage
-
- from .activations import get_activation
- from .configuration_utils import PretrainedConfig
-@@ -442,6 +444,29 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
-     """
-     Reads a PyTorch checkpoint file, returning properly formatted errors if they arise.
-     """
-+    delete_download = False
-+    tmp_dir = "/tmp/model"
-+    os.makedirs(tmp_dir, exist_ok=True)
-+    if isinstance(checkpoint_file, dict):
-+        # Download model file from huggingface
-+        print(f"==> Download model from HF: {checkpoint_file}")
-+        checkpoint_file = hf_hub_download(
-+            local_dir=tmp_dir, local_dir_use_symlinks=False, force_download=True, resume_download=True, **checkpoint_file)
-+        delete_download = True
-+    else:
-+        with open(checkpoint_file, "rb") as f:
-+            is_gcs_file = (f.read(2) == b"gs")
-+        if is_gcs_file:
-+            # Download model file from GCS
-+            with open(checkpoint_file, "r") as f:
-+                gcs_file = f.read()
-+            checkpoint_file = os.path.join(tmp_dir, gcs_file.split("/")[-1])
-+            print(f"==> Download model from GCS: {gcs_file} to: {checkpoint_file}")
-+            client = storage.Client()
-+            with open(checkpoint_file, 'wb') as f:
-+                client.download_blob_to_file(gcs_file, f)
-+            delete_download = True
-+
-     if checkpoint_file.endswith(".safetensors") and is_safetensors_available():
-         # Check format of the archive
-         with safe_open(checkpoint_file, framework="pt") as f:
-@@ -455,9 +480,9 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
-             raise NotImplementedError(
-                 f"Conversion from a {metadata['format']} safetensors archive to PyTorch is not implemented yet."
-             )
--        return safe_load_file(checkpoint_file)
-+        state_dict = safe_load_file(checkpoint_file)
-     try:
--        return torch.load(checkpoint_file, map_location="cpu")
-+        state_dict = torch.load(checkpoint_file, map_location="cpu")
-     except Exception as e:
-         try:
-             with open(checkpoint_file) as f:
-@@ -478,6 +503,10 @@ def load_state_dict(checkpoint_file: Union[str, os.PathLike]):
-                 f"at '{checkpoint_file}'. "
-                 "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
-             )
-+    if delete_download:
-+        print(f"==> Delete downloaded model: {checkpoint_file}")
-+        os.remove(checkpoint_file)
-+    return state_dict
-
-
- def set_initialized_submodules(model, state_dict_keys):
-@@ -3179,7 +3208,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
-             return mismatched_keys
-
-         if resolved_archive_file is not None:
--            folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
-+            if isinstance(resolved_archive_file, str):
-+                folder = os.path.sep.join(resolved_archive_file[0].split(os.path.sep)[:-1])
-+            else:
-+                folder = None
-         else:
-             folder = None
-         if device_map is not None and is_safetensors:
-diff --git a/src/transformers/utils/hub.py b/src/transformers/utils/hub.py
-index ffed743..4b15770 100644
---- a/src/transformers/utils/hub.py
-+++ b/src/transformers/utils/hub.py
-@@ -414,20 +414,34 @@ def cached_file(
-     user_agent = http_user_agent(user_agent)
-     try:
-         # Load from URL or cache if already cached
--        resolved_file = hf_hub_download(
--            path_or_repo_id,
--            filename,
--            subfolder=None if len(subfolder) == 0 else subfolder,
--            repo_type=repo_type,
--            revision=revision,
--            cache_dir=cache_dir,
--            user_agent=user_agent,
--            force_download=force_download,
--            proxies=proxies,
--            resume_download=resume_download,
--            use_auth_token=use_auth_token,
--            local_files_only=local_files_only,
--        )
-+        if filename.endswith(".bin"):
-+            # NOTE: To save disk we do not download bin file eagerly. Do not support safetensors.
-+            resolved_file = dict(
-+                repo_id=path_or_repo_id,
-+                filename=filename,
-+                subfolder=None if len(subfolder) == 0 else subfolder,
-+                repo_type=repo_type,
-+                revision=revision,
-+                user_agent=user_agent,
-+                proxies=proxies,
-+                use_auth_token=use_auth_token,
-+            )
-+            print(f"--> Apply lazy download to bin file: {resolved_file}")
-+        else:
-+            resolved_file = hf_hub_download(
-+                path_or_repo_id,
-+                filename,
-+                subfolder=None if len(subfolder) == 0 else subfolder,
-+                repo_type=repo_type,
-+                revision=revision,
-+                cache_dir=cache_dir,
-+                user_agent=user_agent,
-+                force_download=force_download,
-+                proxies=proxies,
-+                resume_download=resume_download,
-+                use_auth_token=use_auth_token,
-+                local_files_only=local_files_only,
-+            )
-
-     except RepositoryNotFoundError:
-         raise EnvironmentError(

community-content/vertex_model_garden/model_oss/peft/instruct_lora.py

@@ -1,95 +0,0 @@
-"""Instruct/Chat with LoRA models."""
-
-# pylint: disable=g-importing-member
-from datasets import load_dataset
-from peft import LoraConfig
-import torch
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
-from transformers import BitsAndBytesConfig
-from transformers import TrainingArguments
-from trl import SFTTrainer
-from typing import List
-from util import constants
-
-
-def finetune_instruct(
-    pretrained_model_id: str,
-    dataset_name: str,
-    output_dir: str,
-    lora_rank: int = 64,
-    lora_alpha: int = 16,
-    lora_dropout: float = 0.1,
-    target_modules: List[str] = constants.INSTRUCT_LORA_TARGET_MODULES,
-    warmup_ratio: int = 0.03,
-    max_steps: int = 10,
-    max_seq_length: int = 512,
-    learning_rate: float = 2e-4,
-) -> None:
-  """Finetunes instruct."""
-  dataset = load_dataset(dataset_name, split="train")
-
-  bnb_config = BitsAndBytesConfig(
-      load_in_4bit=True,
-      bnb_4bit_quant_type="nf4",
-      bnb_4bit_compute_dtype=torch.float16,
-  )
-
-  model = AutoModelForCausalLM.from_pretrained(
-      pretrained_model_id,
-      quantization_config=bnb_config,
-      trust_remote_code=True,
-  )
-  model.config.use_cache = False
-
-  tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_id, trust_remote_code=True
-  )
-  tokenizer.pad_token = tokenizer.eos_token
-
-  peft_config = LoraConfig(
-      lora_alpha=lora_alpha,
-      lora_dropout=lora_dropout,
-      r=lora_rank,
-      bias="none",
-      task_type="CAUSAL_LM",
-      target_modules=target_modules,
-  )
-
-  per_device_train_batch_size = 4
-  gradient_accumulation_steps = 4
-  optim = "paged_adamw_32bit"
-  save_steps = 10
-  logging_steps = 10
-  max_grad_norm = 0.3
-  lr_scheduler_type = "constant"
-
-  training_arguments = TrainingArguments(
-      output_dir=output_dir,
-      per_device_train_batch_size=per_device_train_batch_size,
-      gradient_accumulation_steps=gradient_accumulation_steps,
-      optim=optim,
-      save_steps=save_steps,
-      logging_steps=logging_steps,
-      learning_rate=learning_rate,
-      fp16=True,
-      max_grad_norm=max_grad_norm,
-      max_steps=max_steps,
-      warmup_ratio=warmup_ratio,
-      group_by_length=True,
-      lr_scheduler_type=lr_scheduler_type,
-  )
-
-  trainer = SFTTrainer(
-      model=model,
-      train_dataset=dataset,
-      peft_config=peft_config,
-      dataset_text_field="text",
-      max_seq_length=max_seq_length,
-      tokenizer=tokenizer,
-      args=training_arguments,
-  )
-  for name, module in trainer.model.named_modules():
-    if "norm" in name:
-      module = module.to(torch.float32)
-  trainer.train()

community-content/vertex_model_garden/model_oss/peft/main.py

@@ -1,185 +0,0 @@
-"""Main function to start PEFT finetuning."""
-import subprocess
-
-from absl import app
-from absl import flags
-from absl import logging
-
-from peft import causal_language_modeling_lora
-from peft import instruct_lora
-from peft import sequence_classification_lora
-from util import constants
-from util import fileutils
-
-_TASK = flags.DEFINE_string(
-    'task',
-    constants.CAUSAL_LANGUAGE_MODELING_LORA,
-    'The supported PEFT tasks.',
-)
-
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
-    None,
-    'The pretrained model id. Supported models can be causal language modeling'
-    ' models from https://github.com/huggingface/peft/tree/main.',
-    required=True,
-)
-
-_DATASET_NAME = flags.DEFINE_string(
-    'dataset_name',
-    None,
-    'The dataset name in huggingface.',
-    required=True,
-)
-
-_OUTPUT_DIR = flags.DEFINE_string(
-    'output_dir',
-    None,
-    'The output directory.',
-    required=True,
-)
-
-_PRECISION_MODE = flags.DEFINE_string(
-    'precision_mode',
-    constants.PRECISION_MODE_16,
-    'Supported finetuning precision_modes are `{}` and `{}`.'.format(
-        constants.PRECISION_MODE_8, constants.PRECISION_MODE_16
-    ),
-)
-
-_LORA_RANK = flags.DEFINE_integer(
-    'lora_rank',
-    16,
-    'The rank of the update matrices, expressed in int. Lower rank results in'
-    ' smaller update matrices with fewer trainable parameters, referring to'
-    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
-)
-
-_LORA_ALPHA = flags.DEFINE_integer(
-    'lora_alpha',
-    32,
-    'LoRA scaling factor, referring to'
-    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
-)
-
-_LORA_DROPOUT = flags.DEFINE_float(
-    'lora_dropout',
-    0.05,
-    'dropout probability of the LoRA layers, referring to'
-    ' https://huggingface.co/docs/peft/task_guides/token-classification-lora.',
-)
-
-_TARGET_MODULES = flags.DEFINE_list(
-    'target_modules',
-    constants.CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES,
-    'The comma separated list of target modules for LoRa training.',
-)
-
-_WARMUP_STEPS = flags.DEFINE_integer(
-    'warmup_steps',
-    10,
-    'Number of steps for the warmup in the learning rate scheduler.',
-)
-
-_WARMUP_RATIO = flags.DEFINE_float(
-    'warmup_ratio',
-    0.03,
-    'The warmup ratio in the learning rate scheduler.',
-)
-
-_MAX_STEPS = flags.DEFINE_integer(
-    'max_steps',
-    10,
-    'Total number of training steps.',
-)
-
-_MAX_SEQ_LENGTH = flags.DEFINE_integer(
-    'max_seq_length',
-    512,
-    'The maximum sequence length.',
-)
-
-_NUM_EPOCHS = flags.DEFINE_integer(
-    'num_epochs',
-    20,
-    'The number of training epochs.',
-)
-
-_BATCH_SIZE = flags.DEFINE_integer(
-    'batch_size',
-    32,
-    'The batch size.',
-)
-
-_LEARNING_RATE = flags.DEFINE_float(
-    'learning_rate',
-    2e-4,
-    'The learning rate after the potential warmup period.',
-)
-
-
-def main(_) -> None:
-  task = _TASK.value
-  pretrained_model_id = _PRETRAINED_MODEL_ID.value
-  local_pretrained_model_id = None
-  if pretrained_model_id.startswith(constants.GCS_URI_PREFIX):
-    logging.info(
-        'Start to copy pretrained models locally: %s.', pretrained_model_id
-    )
-    fileutils.download_gcs_dir_to_local(
-        pretrained_model_id, constants.LOCAL_BASE_MODEL_DIR
-    )
-    local_pretrained_model_id = constants.LOCAL_BASE_MODEL_DIR
-    logging.info(
-        'Finished copying pretrained models locally to: %s.',
-        local_pretrained_model_id,
-    )
-  if task == constants.TEXT_TO_IMAGE_LORA:
-    subprocess.run(['/bin/bash', 'train.sh'], check=True)
-  elif task == constants.SEQUENCE_CLASSIFICATION_LORA:
-    sequence_classification_lora.finetune_sequence_classification(
-        pretrained_model_id=pretrained_model_id,
-        dataset_name=_DATASET_NAME.value,
-        output_dir=_OUTPUT_DIR.value,
-        lora_rank=_LORA_RANK.value,
-        lora_alpha=_LORA_ALPHA.value,
-        lora_dropout=_LORA_DROPOUT.value,
-        num_epochs=_NUM_EPOCHS.value,
-        batch_size=_BATCH_SIZE.value,
-        learning_rate=_LEARNING_RATE.value,
-    )
-  elif task == constants.CAUSAL_LANGUAGE_MODELING_LORA:
-    causal_language_modeling_lora.finetune_causal_language_modeling(
-        pretrained_model_id=pretrained_model_id,
-        dataset_name=_DATASET_NAME.value,
-        output_dir=_OUTPUT_DIR.value,
-        precision_mode=_PRECISION_MODE.value,
-        lora_rank=_LORA_RANK.value,
-        lora_alpha=_LORA_ALPHA.value,
-        lora_dropout=_LORA_DROPOUT.value,
-        target_modules=_TARGET_MODULES.value,
-        warmup_steps=_WARMUP_STEPS.value,
-        max_steps=_MAX_STEPS.value,
-        learning_rate=_LEARNING_RATE.value,
-        local_pretrained_model_id=local_pretrained_model_id,
-    )
-  elif task == constants.INSTRUCT_LORA:
-    instruct_lora.finetune_instruct(
-        pretrained_model_id=pretrained_model_id,
-        dataset_name=_DATASET_NAME.value,
-        output_dir=_OUTPUT_DIR.value,
-        lora_rank=_LORA_RANK.value,
-        lora_alpha=_LORA_ALPHA.value,
-        lora_dropout=_LORA_DROPOUT.value,
-        target_modules=_TARGET_MODULES.value,
-        warmup_ratio=_WARMUP_RATIO.value,
-        max_steps=_MAX_STEPS.value,
-        max_seq_length=_MAX_SEQ_LENGTH.value,
-        learning_rate=_LEARNING_RATE.value,
-    )
-  else:
-    raise ValueError('The task {} is not supported.'.format(task))
-
-
-if __name__ == '__main__':
-  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train.sh

@@ -1,6 +0,0 @@
-#!/bin/bash
-
-# Setup accelerate config before running trainer.
-python -c "from accelerate.utils import write_basic_config; write_basic_config(mixed_precision='fp16')"
-
-accelerate launch "$@"

community-content/vertex_model_garden/model_oss/peft/train/axolotl/dockerfile/train.Dockerfile

@@ -0,0 +1,16 @@
+# Dockerfile for axolotl training.
+#
+# To build:
+# docker build -f model_oss/peft/train/axolotol/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+FROM winglian/axolotl:main-latest
+
+RUN mkdir -p ./vertex_vision_model_garden/
+
+COPY model_oss/peft/train/axolotl/*.py ./vertex_vision_model_garden/
+
+ENTRYPOINT ["python3", "./vertex_vision_model_garden/train_entrypoint.py"]

community-content/vertex_model_garden/model_oss/peft/train/axolotl/scripts/build_train_docker.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+# Run copybara first:
+# cloud/ml/applications/vision/model_garden/copybara/run_copybara_local.sh
+# Run docker build:
+# cloud/ml/applications/vision/model_garden/model_oss/peft/train/axolotl/scripts/build_train_docker.sh
+
+set -x
+
+COPYBARA_DIR="/tmp/train_docker/"
+
+pushd "${COPYBARA_DIR}"
+
+PROJECT="cloud-nas-260507"
+IMAGE_TAG="gcr.io/${PROJECT}/axolotl-train:${USER}-test"
+
+docker build -f model_oss/peft/train/axolotl/dockerfile/train.Dockerfile . -t "${IMAGE_TAG}"
+docker push "${IMAGE_TAG}"
+
+popd

community-content/vertex_model_garden/model_oss/peft/train/axolotl/train_entrypoint.py

@@ -0,0 +1,88 @@
+"""Entrypoint for axolotl train docker."""
+
+import argparse
+import json
+import os
+import subprocess
+
+
+def _get_multi_node_flags(cluster_spec: str) -> list[str]:
+  """Returns the multi-node flags."""
+  print(f'CLUSTER_SPEC: {cluster_spec}')
+
+  cluster_data = json.loads(cluster_spec)
+
+  # Get primary node info
+  primary_node = cluster_data['cluster']['workerpool0'][0]
+  print(f'primary node: {primary_node}')
+  primary_node_addr, primary_node_port = primary_node.split(':')
+  print(f'primary node address: {primary_node_addr}')
+  print(f'primary node port: {primary_node_port}')
+
+  # Determine node rank of this machine
+  workerpool = cluster_data['task']['type']
+  if workerpool == 'workerpool0':
+    node_rank = 0
+  else:
+    node_rank = cluster_data['task']['index'] + 1
+  print(f'node rank: {node_rank}')
+
+  # Calculate total nodes
+  num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
+  num_nodes = num_worker_nodes + 1  # Add 1 for the primary node
+  print(f'num nodes: {num_nodes}')
+
+  return [
+      f'--machine_rank={node_rank}',
+      f'--num_machines={num_nodes}',
+      f'--main_process_ip={primary_node_addr}',
+      f'--main_process_port={primary_node_port}',
+      '--max_restarts=0',
+      '--monitor_interval=120',
+      '--dynamo_backend=no',
+  ]
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--config_file')
+  parser.add_argument('--huggingface_access_token')
+  args, unknown = parser.parse_known_args()
+
+  accelerate_flags = []
+
+  if args.config_file:
+    accelerate_flags.append(f'--config_file={args.config_file}')
+
+  if cluster_spec := os.getenv('CLUSTER_SPEC', default=None):
+    print('========== Launch on cloud multi nodes ==========')
+    accelerate_flags.extend(_get_multi_node_flags(cluster_spec))
+
+  cmd = (
+      [
+          'accelerate',
+          'launch',
+      ]
+      + accelerate_flags
+      + [
+          '-m',
+          'axolotl.cli.train',
+      ]
+      + unknown
+  )
+  print(f'{cmd=}', flush=True)
+
+  env = os.environ.copy()
+
+  if args.huggingface_access_token:
+    env['HF_TOKEN'] = args.huggingface_access_token
+
+  subprocess.run(
+      cmd,
+      check=True,
+      env=env,
+  )
+
+
+if __name__ == '__main__':
+  main()

community-content/vertex_model_garden/model_oss/peft/train/tests/docker_command_builder.py

@@ -0,0 +1,45 @@
+"""Class that bundles docker related flags."""
+
+import getpass
+import os
+import pwd
+
+
+class DockerCommandBuilder:
+  """Bundle docker related flags."""
+
+  def __init__(self, docker_uri, shm_size='128gb'):
+    self._docker_uri = [docker_uri]
+
+    self._defaults = [
+        'docker',
+        'run',
+        '--gpus=all',
+        '--net=host',
+        '--rm',
+        f'--shm-size={shm_size}',
+    ]
+
+    user = getpass.getuser()
+    # username ends with `_google_com` is managed by ldap and does not have a
+    # corresponding entry in /etc/passwd or /etc/group file. We cannot enable
+    # non-root docker user with below method.
+    if not user.endswith('_google_com'):
+      uid = os.getuid()
+      gid = pwd.getpwuid(uid).pw_gid
+      self._defaults += [
+          f'--user={uid}:{gid}',
+          '--volume=/etc/group:/etc/group:ro',
+          '--volume=/etc/passwd:/etc/passwd:ro',
+      ]
+    self._env_vars = []
+    self._mount_maps = []
+
+  def add_env_var(self, var, val):
+    self._env_vars.append(f'--env={var}={val}')
+
+  def add_mount_map(self, host_path, docker_path):
+    self._mount_maps.append(f'--volume={host_path}:{docker_path}')
+
+  def build_cmd(self) -> str:
+    return self._defaults + self._env_vars + self._mount_maps + self._docker_uri

community-content/vertex_model_garden/model_oss/peft/train/tests/instruct_lora_command_builder.py

@@ -0,0 +1,429 @@
+# pylint: disable=W,C,R
+
+# DO NOT MODIFY: this file is auto-generated
+# See go/vmg-oss-peft-tests#command-builder-genpy
+
+class InstructLoraCommandBuilder:
+
+  def __init__(self):
+    self._config_file = None
+    self._task = None
+    self._pretrained_model_id = None
+    self._dataset_name = None
+    self._train_split_name = None
+    self._template = None
+    self._instruct_column_in_dataset = None
+    self._output_dir = None
+    self._merge_base_and_lora_output_dir = None
+    self._logging_output_dir = None
+    self._per_device_train_batch_size = None
+    self._gradient_accumulation_steps = None
+    self._lora_rank = None
+    self._lora_alpha = None
+    self._lora_dropout = None
+    self._max_steps = None
+    self._num_epochs = None
+    self._max_seq_length = None
+    self._learning_rate = None
+    self._lr_scheduler_type = None
+    self._precision_mode = None
+    self._train_precision = None
+    self._enable_gradient_checkpointing = None
+    self._use_example_packing = None
+    self._attn_implementation = None
+    self._optimizer = None
+    self._warmup_ratio = None
+    self._report_to = None
+    self._save_steps = None
+    self._logging_steps = None
+    self._huggingface_access_token = None
+    self._eval_dataset_path = None
+    self._eval_column = None
+    self._eval_template = None
+    self._eval_split = None
+    self._eval_steps = None
+    self._eval_tasks = None
+    self._eval_metric_name = None
+    self._completion_only = None
+    self._max_grad_norm = None
+    self._logger_level = None
+    self._benchmark_out_file = None
+    self._tuning_data_stats_file = None
+    self._enable_peft = None
+    self._merge_model_precision_mode = None
+    self._target_modules = None
+
+  @property
+  def config_file(self):
+    return self._config_file
+
+  @config_file.setter
+  def config_file(self, val: str):
+    self._config_file = val
+
+  @property
+  def task(self):
+    return self._task
+
+  @task.setter
+  def task(self, val: str):
+    self._task = val
+
+  @property
+  def pretrained_model_id(self):
+    return self._pretrained_model_id
+
+  @pretrained_model_id.setter
+  def pretrained_model_id(self, val: str):
+    self._pretrained_model_id = val
+
+  @property
+  def train_dataset(self):
+    return self._dataset_name
+
+  @train_dataset.setter
+  def train_dataset(self, val: str):
+    self._dataset_name = val
+
+  @property
+  def train_split_name(self):
+    return self._train_split_name
+
+  @train_split_name.setter
+  def train_split_name(self, val: str):
+    self._train_split_name = val
+
+  @property
+  def template(self):
+    return self._template
+
+  @template.setter
+  def template(self, val: str):
+    self._template = val
+
+  @property
+  def instruct_column(self):
+    return self._instruct_column_in_dataset
+
+  @instruct_column.setter
+  def instruct_column(self, val: str):
+    self._instruct_column_in_dataset = val
+
+  @property
+  def ckpt_dir(self):
+    return self._output_dir
+
+  @ckpt_dir.setter
+  def ckpt_dir(self, val: str):
+    self._output_dir = val
+
+  @property
+  def merged_model_dir(self):
+    return self._merge_base_and_lora_output_dir
+
+  @merged_model_dir.setter
+  def merged_model_dir(self, val: str):
+    self._merge_base_and_lora_output_dir = val
+
+  @property
+  def logging_dir(self):
+    return self._logging_output_dir
+
+  @logging_dir.setter
+  def logging_dir(self, val: str):
+    self._logging_output_dir = val
+
+  @property
+  def per_device_batch_size(self):
+    return self._per_device_train_batch_size
+
+  @per_device_batch_size.setter
+  def per_device_batch_size(self, val: int):
+    self._per_device_train_batch_size = val
+
+  @property
+  def gradient_accumulation_steps(self):
+    return self._gradient_accumulation_steps
+
+  @gradient_accumulation_steps.setter
+  def gradient_accumulation_steps(self, val: int):
+    self._gradient_accumulation_steps = val
+
+  @property
+  def lora_rank(self):
+    return self._lora_rank
+
+  @lora_rank.setter
+  def lora_rank(self, val: int):
+    self._lora_rank = val
+
+  @property
+  def lora_alpha(self):
+    return self._lora_alpha
+
+  @lora_alpha.setter
+  def lora_alpha(self, val: int):
+    self._lora_alpha = val
+
+  @property
+  def lora_dropout(self):
+    return self._lora_dropout
+
+  @lora_dropout.setter
+  def lora_dropout(self, val: float):
+    self._lora_dropout = val
+
+  @property
+  def max_steps(self):
+    return self._max_steps
+
+  @max_steps.setter
+  def max_steps(self, val: int):
+    self._max_steps = val
+
+  @property
+  def num_epochs(self):
+    return self._num_epochs
+
+  @num_epochs.setter
+  def num_epochs(self, val: float):
+    self._num_epochs = val
+
+  @property
+  def max_seq_length(self):
+    return self._max_seq_length
+
+  @max_seq_length.setter
+  def max_seq_length(self, val: int):
+    self._max_seq_length = val
+
+  @property
+  def learning_rate(self):
+    return self._learning_rate
+
+  @learning_rate.setter
+  def learning_rate(self, val: float):
+    self._learning_rate = val
+
+  @property
+  def lr_scheduler_type(self):
+    return self._lr_scheduler_type
+
+  @lr_scheduler_type.setter
+  def lr_scheduler_type(self, val: str):
+    self._lr_scheduler_type = val
+
+  @property
+  def load_precision(self):
+    return self._precision_mode
+
+  @load_precision.setter
+  def load_precision(self, val: str):
+    self._precision_mode = val
+
+  @property
+  def train_precision(self):
+    return self._train_precision
+
+  @train_precision.setter
+  def train_precision(self, val: str):
+    self._train_precision = val
+
+  @property
+  def gradient_checkpointing(self):
+    return self._enable_gradient_checkpointing
+
+  @gradient_checkpointing.setter
+  def gradient_checkpointing(self, val: bool):
+    self._enable_gradient_checkpointing = val
+
+  @property
+  def example_packing(self):
+    return self._use_example_packing
+
+  @example_packing.setter
+  def example_packing(self, val: bool):
+    self._use_example_packing = val
+
+  @property
+  def attn_implementation(self):
+    return self._attn_implementation
+
+  @attn_implementation.setter
+  def attn_implementation(self, val: str):
+    self._attn_implementation = val
+
+  @property
+  def optimizer(self):
+    return self._optimizer
+
+  @optimizer.setter
+  def optimizer(self, val: str):
+    self._optimizer = val
+
+  @property
+  def warmup_ratio(self):
+    return self._warmup_ratio
+
+  @warmup_ratio.setter
+  def warmup_ratio(self, val: float):
+    self._warmup_ratio = val
+
+  @property
+  def report_to(self):
+    return self._report_to
+
+  @report_to.setter
+  def report_to(self, val: str):
+    self._report_to = val
+
+  @property
+  def save_steps(self):
+    return self._save_steps
+
+  @save_steps.setter
+  def save_steps(self, val: int):
+    self._save_steps = val
+
+  @property
+  def logging_steps(self):
+    return self._logging_steps
+
+  @logging_steps.setter
+  def logging_steps(self, val: int):
+    self._logging_steps = val
+
+  @property
+  def huggingface_access_token(self):
+    return self._huggingface_access_token
+
+  @huggingface_access_token.setter
+  def huggingface_access_token(self, val: str):
+    self._huggingface_access_token = val
+
+  @property
+  def eval_dataset(self):
+    return self._eval_dataset_path
+
+  @eval_dataset.setter
+  def eval_dataset(self, val: str):
+    self._eval_dataset_path = val
+
+  @property
+  def eval_instruct_column(self):
+    return self._eval_column
+
+  @eval_instruct_column.setter
+  def eval_instruct_column(self, val: str):
+    self._eval_column = val
+
+  @property
+  def eval_template(self):
+    return self._eval_template
+
+  @eval_template.setter
+  def eval_template(self, val: str):
+    self._eval_template = val
+
+  @property
+  def eval_split_name(self):
+    return self._eval_split
+
+  @eval_split_name.setter
+  def eval_split_name(self, val: str):
+    self._eval_split = val
+
+  @property
+  def eval_steps(self):
+    return self._eval_steps
+
+  @eval_steps.setter
+  def eval_steps(self, val: int):
+    self._eval_steps = val
+
+  @property
+  def eval_tasks(self):
+    return self._eval_tasks
+
+  @eval_tasks.setter
+  def eval_tasks(self, val: str):
+    self._eval_tasks = val
+
+  @property
+  def eval_metric_name(self):
+    return self._eval_metric_name
+
+  @eval_metric_name.setter
+  def eval_metric_name(self, val: str):
+    self._eval_metric_name = val
+
+  @property
+  def completion_only(self):
+    return self._completion_only
+
+  @completion_only.setter
+  def completion_only(self, val: bool):
+    self._completion_only = val
+
+  @property
+  def max_grad_norm(self):
+    return self._max_grad_norm
+
+  @max_grad_norm.setter
+  def max_grad_norm(self, val: float):
+    self._max_grad_norm = val
+
+  @property
+  def logger_level(self):
+    return self._logger_level
+
+  @logger_level.setter
+  def logger_level(self, val: str):
+    self._logger_level = val
+
+  @property
+  def benchmark_out_file(self):
+    return self._benchmark_out_file
+
+  @benchmark_out_file.setter
+  def benchmark_out_file(self, val: str):
+    self._benchmark_out_file = val
+
+  @property
+  def tuning_data_stats_file(self):
+    return self._tuning_data_stats_file
+
+  @tuning_data_stats_file.setter
+  def tuning_data_stats_file(self, val: str):
+    self._tuning_data_stats_file = val
+
+  @property
+  def enable_peft(self):
+    return self._enable_peft
+
+  @enable_peft.setter
+  def enable_peft(self, val: bool):
+    self._enable_peft = val
+
+  @property
+  def merge_model_precision_mode(self):
+    return self._merge_model_precision_mode
+
+  @merge_model_precision_mode.setter
+  def merge_model_precision_mode(self, val: str):
+    self._merge_model_precision_mode = val
+
+  @property
+  def target_modules(self):
+    return self._target_modules
+
+  @target_modules.setter
+  def target_modules(self, val: str):
+    self._target_modules = val
+
+  def build_cmd(self) -> str:
+    cmd = []
+    for k, v in self.__dict__.items():
+      if v is not None:
+        cmd.append(f'--{k[1:]}={v}')
+    return cmd

community-content/vertex_model_garden/model_oss/peft/train/tests/quantize_model_command_builder.py

@@ -0,0 +1,142 @@
+# pylint: disable=W,C,R
+
+# DO NOT MODIFY: this file is auto-generated
+# See go/vmg-oss-peft-tests#command-builder-genpy
+
+
+class QuantizeModelCommandBuilder:
+
+  def __init__(self):
+    self._task = None
+    self._pretrained_model_id = None
+    self._quantization_method = None
+    self._quantization_precision_mode = None
+    self._quantization_dataset_name = None
+    self._text_column_in_quantization_dataset = None
+    self._quantization_output_dir = None
+    self._device_map = None
+    self._max_memory = None
+    self._group_size = None
+    self._desc_act = None
+    self._damp_percent = None
+    self._cache_examples_on_gpu = None
+    self._awq_version = None
+
+  @property
+  def task(self):
+    return self._task
+
+  @task.setter
+  def task(self, val: str):
+    self._task = val
+
+  @property
+  def pretrained_model_id(self):
+    return self._pretrained_model_id
+
+  @pretrained_model_id.setter
+  def pretrained_model_id(self, val: str):
+    self._pretrained_model_id = val
+
+  @property
+  def quantization_method(self):
+    return self._quantization_method
+
+  @quantization_method.setter
+  def quantization_method(self, val: str):
+    self._quantization_method = val
+
+  @property
+  def quantization_precision_mode(self):
+    return self._quantization_precision_mode
+
+  @quantization_precision_mode.setter
+  def quantization_precision_mode(self, val: str):
+    self._quantization_precision_mode = val
+
+  @property
+  def quantization_dataset_name(self):
+    return self._quantization_dataset_name
+
+  @quantization_dataset_name.setter
+  def quantization_dataset_name(self, val: str):
+    self._quantization_dataset_name = val
+
+  @property
+  def text_column_in_quantization_dataset(self):
+    return self._text_column_in_quantization_dataset
+
+  @text_column_in_quantization_dataset.setter
+  def text_column_in_quantization_dataset(self, val: str):
+    self._text_column_in_quantization_dataset = val
+
+  @property
+  def quantization_output_dir(self):
+    return self._quantization_output_dir
+
+  @quantization_output_dir.setter
+  def quantization_output_dir(self, val: str):
+    self._quantization_output_dir = val
+
+  @property
+  def device_map(self):
+    return self._device_map
+
+  @device_map.setter
+  def device_map(self, val: str):
+    self._device_map = val
+
+  @property
+  def max_memory(self):
+    return self._max_memory
+
+  @max_memory.setter
+  def max_memory(self, val: str):
+    self._max_memory = val
+
+  @property
+  def group_size(self):
+    return self._group_size
+
+  @group_size.setter
+  def group_size(self, val: int):
+    self._group_size = val
+
+  @property
+  def desc_act(self):
+    return self._desc_act
+
+  @desc_act.setter
+  def desc_act(self, val: bool):
+    self._desc_act = val
+
+  @property
+  def damp_percent(self):
+    return self._damp_percent
+
+  @damp_percent.setter
+  def damp_percent(self, val: float):
+    self._damp_percent = val
+
+  @property
+  def cache_examples_on_gpu(self):
+    return self._cache_examples_on_gpu
+
+  @cache_examples_on_gpu.setter
+  def cache_examples_on_gpu(self, val: bool):
+    self._cache_examples_on_gpu = val
+
+  @property
+  def awq_version(self):
+    return self._awq_version
+
+  @awq_version.setter
+  def awq_version(self, val: str):
+    self._awq_version = val
+
+  def build_cmd(self) -> str:
+    cmd = []
+    for k, v in self.__dict__.items():
+      if v is not None:
+        cmd.append(f'--{k[1:]}={v}')
+    return cmd

community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_adapters.py

@@ -0,0 +1,106 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests adapters of PEFT train docker."""
+
+import inspect
+import os
+import time
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+from safetensors import safe_open
+import test_util
+
+
+class AdapterTest(test_util.TestBase):
+
+  # Needs to be accessible outside docker to check artifacts.
+  _TEST_OUTPUT_DIR = os.path.expanduser('~/output')
+  _MODULES_NEED_TO_BE_EXCLUDED_IN_ADAPTER = ['lm_head', 'embed_tokens']
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.test_suite_output_dir = os.path.join(
+        cls._TEST_OUTPUT_DIR,
+        os.path.splitext(os.path.basename(__file__))[0],
+        cls.__class__.__name__,
+    )
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'input_text'
+    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.max_steps = 1
+    self.task_cmd_builder.max_seq_length = 256
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.save_steps = 10
+    self.task_cmd_builder.max_steps = 3
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+
+  def setup_output_dir(self, testcase_name: str):
+    testcase_output_dir = os.path.join(
+        self.test_suite_output_dir, testcase_name, test_util.get_timestamp()
+    )
+    self.task_cmd_builder.ckpt_dir = os.path.join(
+        testcase_output_dir, 'adapter'
+    )
+    self.task_cmd_builder.logging_dir = os.path.join(
+        testcase_output_dir, 'logs'
+    )
+
+  def check_adapter_for_bad_modules(self, adapter_path):
+    unwanted_modules = set()
+    with safe_open(adapter_path, framework='pt', device='cpu') as f:
+      for key in f.keys():
+        for module in self._MODULES_NEED_TO_BE_EXCLUDED_IN_ADAPTER:
+          if module in key:
+            unwanted_modules.add(key)
+    assert (
+        not unwanted_modules
+    ), f'Adapter includes unwanted modules: {unwanted_modules}'
+
+  @parameterized.named_parameters(
+      ('llama3.1-8b', 'llama3.1-8b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+      ('llama2-7b', 'llama2-7b-hf'),
+  )
+  def test_llama_adapters(self, model_name):
+    test_function_name = inspect.stack()[0][3]
+    self.setup_output_dir(f'{test_function_name}-{model_name}')
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 1000)
+
+    adapter = os.path.join(
+        self.task_cmd_builder.ckpt_dir,
+        'checkpoint-final/adapter_model.safetensors',
+    )
+    self.check_adapter_for_bad_modules(adapter)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_features.py

@@ -0,0 +1,290 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests various features of PEFT train docker."""
+
+import os
+import time
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+import test_util
+
+
+class GcsUploadDownloadTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'input_text'
+    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.max_steps = 1
+    self.task_cmd_builder.max_seq_length = 256
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.ckpt_dir = '/tmp'
+
+  @parameterized.named_parameters(
+      (
+          'llama3_8b_gcs',
+          'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf',
+      ),
+      ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
+  )
+  def test_model_download_single_process(self, pretrained_model_id):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(pretrained_model_id)
+    )
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 5 * 60.0)
+
+  @parameterized.named_parameters(
+      (
+          'llama3_8b_gcs',
+          'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf',
+      ),
+      ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
+  )
+  def test_model_download_multi_process(self, pretrained_model_id):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(pretrained_model_id)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 5 * 60.0)
+
+  def test_70b_model_download(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        'gs://vertex-model-garden-public-us/llama3/llama3-70b-hf'
+    )
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 10 * 60.0)
+
+  @parameterized.named_parameters(
+      ('merged-without-upload', '/tmp/merged'),
+      ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
+  )
+  def test_model_merge(self, output_dir):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+
+    ckpt_dir = os.path.join(
+        output_dir,
+        f'output-{test_util.get_timestamp()}',
+    )
+    self.task_cmd_builder.ckpt_dir = ckpt_dir
+    self.task_cmd_builder.merged_model_dir = os.path.join(ckpt_dir, 'merged')
+    self.task_cmd_builder.logging_dir = '/tmp/logging'
+
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 5 * 60.0)
+
+  def test_model_fp8_conversion(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+
+    ckpt_dir = f'/tmp/output/output-{test_util.get_timestamp()}'
+    self.task_cmd_builder.ckpt_dir = ckpt_dir
+    self.task_cmd_builder.merged_model_dir = os.path.join(ckpt_dir, 'merged')
+    self.task_cmd_builder.logging_dir = os.path.join(ckpt_dir, 'logging')
+    self.task_cmd_builder.merge_model_precision_mode = 'float8'
+
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 5 * 60.0)
+
+  @parameterized.named_parameters(
+      ('merged-without-upload', '/tmp/merged'),
+      ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
+  )
+  def test_model_merge_and_upload_deepspeed(self, merged_model_dir):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero3_8gpu.yaml'
+    )
+    self.task_cmd_builder.merged_model_dir = os.path.join(
+        merged_model_dir, f'merged-{test_util.get_timestamp()}'
+    )
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 5 * 60.0)
+
+  @parameterized.named_parameters(
+      ('save-only-last', 10),
+      ('save-multiple-times', 1),
+  )
+  def test_llama3_8b_save_and_merge_8_gpus_fsdp(self, save_steps):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.save_steps = save_steps
+    self.task_cmd_builder.max_steps = 3
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.merged_model_dir = '/tmp/merged'
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 9 * 60.0)
+
+
+class TemplateAndDataStatsTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.max_steps = 1
+    self.task_cmd_builder.max_seq_length = 256
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.ckpt_dir = '/tmp'
+
+  def test_openai_chat_template(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'openai-multi-chat-example-data.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'messages'
+    self.task_cmd_builder.template = 'llama3'
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  def test_openai_completion_template(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'openai-completion-example-data.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'prompt'
+    self.task_cmd_builder.template = 'openai-completion'
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  def test_data_stats_chat_template(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'openai-multi-chat-example-data.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'messages'
+    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  def test_data_stats_completion_template(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'openai-completion-example-data.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'prompt'
+    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.assertEqual(self.run_cmd(), 0)
+
+
+class TargetModulesTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'input_text'
+    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.max_steps = 1
+    self.task_cmd_builder.max_seq_length = 256
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.ckpt_dir = '/tmp'
+
+  def test_target_modules(self):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.target_modules = 'q_proj, v_proj, k_proj'
+
+    self.assertEqual(self.run_cmd(), 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_throughput.py

@@ -0,0 +1,180 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests to check training throughput and GPU memory consumption."""
+
+import os
+import pathlib
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+import test_util
+
+
+class TrainerThroughputTest(test_util.TestBase):
+
+  _TEST_OUTPUT_DIR = os.path.expanduser('~/throughput_tests')
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.test_suite_output_dir = os.path.join(
+        cls._TEST_OUTPUT_DIR, os.path.splitext(os.path.basename(__file__))[0]
+    )
+    if not os.path.isdir(cls.test_suite_output_dir):
+      pathlib.Path(cls.test_suite_output_dir).mkdir(parents=True)
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = 10
+    self.task_cmd_builder.save_steps = 1000
+    self.task_cmd_builder.logging_steps = 1
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.example_packing = True
+    self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2-1k'
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'text'
+    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
+    self.task_cmd_builder.logging_dir = '/tmp/logs'
+
+  def run_cmd_and_handle_failure(self):
+    ret = self.run_cmd()
+    if ret != 0:
+      with open(self.task_cmd_builder.benchmark_out_file, 'a') as f:
+        max_seq_length = self.task_cmd_builder.max_seq_length
+        f.write(f'{max_seq_length/1024.0:.1f}k | failed | n/a\n')
+    return ret
+
+  @parameterized.product(
+      model_name=[
+          'llama3-70b-hf',
+          'llama3.1-70b-hf',
+          'Mistral-7B-v0.1',
+          'Mixtral-8x7B-v0.1',
+          'Gemma2-9b-it',
+      ],
+      precision=['4bit', '8bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+  )
+  def test_model_single_gpu(self, model_name, precision, max_seq_length):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    self.task_cmd_builder.benchmark_out_file = os.path.join(
+        self.test_suite_output_dir, f'bm_{model_name}_{precision}.txt'
+    )
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.assertEqual(self.run_cmd_and_handle_failure(), 0)
+
+  @parameterized.product(
+      model_name=[
+          'llama3-70b-hf',
+          'llama3.1-70b-hf',
+          'Mistral-7B-v0.1',
+          'Mixtral-8x7B-v0.1',
+          'Gemma2-9b-it',
+      ],
+      precision=['4bit', '8bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      num_gpus=[8],
+      config=['deepspeed_zero2', 'deepspeed_zero3'],
+  )
+  def test_model_multi_gpu_deepspeed(
+      self, model_name, precision, max_seq_length, num_gpus, config
+  ):
+    self.assertTrue(num_gpus == 4 or num_gpus == 8)
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    self.task_cmd_builder.benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+
+    self.task_cmd_builder.config_file = (
+        f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
+    )
+    self.docker_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+
+    self.assertEqual(self.run_cmd_and_handle_failure(), 0)
+
+  @parameterized.product(
+      model_name=['llama3.1-70b-hf'],
+      precision=['4bit', '8bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      num_gpus=[8],
+  )
+  def test_model_multi_gpu_fsdp_lora(
+      self, model_name, precision, max_seq_length, num_gpus
+  ):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    self.task_cmd_builder.benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_fsdp_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+
+    self.docker_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+
+    self.assertEqual(self.run_cmd_and_handle_failure(), 0)
+
+  @parameterized.product(
+      model_name=['llama3.1-70b-hf'],
+      precision=['bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      num_gpus=[8],
+  )
+  def test_model_multi_gpu_fsdp_full_finetuning(
+      self, model_name, precision, max_seq_length, num_gpus
+  ):
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    self.task_cmd_builder.benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_fsdp_full_finetuning_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.enable_peft = False
+
+    self.docker_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+
+    self.assertEqual(self.run_cmd_and_handle_failure(), 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_instruct_lora_trained_model_quality.py

@@ -0,0 +1,152 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests to make sure trained model achieves decent quality.
+
+Right now, the metric is loss decreasing and we'll eyeball the TB graphs.
+"""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+import test_util
+
+
+class TrainedModelQualityTest(test_util.TestBase):
+
+  _TEST_OUTPUT_DIR = os.path.expanduser('~/output')
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.test_suite_output_dir = os.path.join(
+        cls._TEST_OUTPUT_DIR,
+        os.path.splitext(os.path.basename(__file__))[0],
+        cls.__class__.__name__,
+    )
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.eval_tasks = 'builtin_eval'
+    self.task_cmd_builder.eval_metric_name = 'loss'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 8
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.num_epochs = 2.0
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = -1
+    self.task_cmd_builder.save_steps = 10
+    self.task_cmd_builder.eval_steps = 10
+    self.task_cmd_builder.max_seq_length = 4096
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.completion_only = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.report_to = 'tensorboard'
+
+  def setup_output_dir(self, testcase_name: str):
+    testcase_output_dir = os.path.join(
+        self.test_suite_output_dir, testcase_name
+    )
+    self.task_cmd_builder.ckpt_dir = os.path.join(
+        testcase_output_dir, 'adapter'
+    )
+    self.task_cmd_builder.logging_dir = os.path.join(
+        testcase_output_dir, 'logs'
+    )
+    self.task_cmd_builder.merged_model_dir = os.path.join(
+        testcase_output_dir, 'merged'
+    )
+
+  @parameterized.named_parameters(
+      ('llama3-8b', 'llama3-8b-hf'),
+      ('llama3.1-8b', 'llama3.1-8b-hf'),
+  )
+  def test_8b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'input_text'
+    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'peft_eval_sample.jsonl'
+    )
+    self.task_cmd_builder.eval_split_name = 'train'
+    self.task_cmd_builder.eval_instruct_column = (
+        self.task_cmd_builder.instruct_column
+    )
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'text'
+    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split_name = 'test'
+    self.task_cmd_builder.eval_instruct_column = (
+        self.task_cmd_builder.instruct_column
+    )
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_fsdp(self, model_name):
+    self.setup_output_dir(f'test_fsdp_{model_name}')
+    self.task_cmd_builder.pretrained_model_id = (
+        test_util.get_pretrained_model_id(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split_name = 'train'
+    self.task_cmd_builder.instruct_column = 'text'
+    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split_name = 'test'
+    self.task_cmd_builder.eval_instruct_column = (
+        self.task_cmd_builder.instruct_column
+    )
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_quantize_model.py

@@ -0,0 +1,49 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests quantize model task in PEFT docker."""
+
+import os
+import time
+
+from absl.testing import absltest
+import quantize_model_command_builder as task_cmd_builder
+import test_util
+
+
+class QuantizeModelTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '')
+    self.docker_builder.add_mount_map(
+        os.path.expanduser('~'), os.path.expanduser('~')
+    )
+
+    self.task_cmd_builder = task_cmd_builder.QuantizeModelCommandBuilder()
+    self.task_cmd_builder.task = 'quantize-model'
+    self.task_cmd_builder.pretrained_model_id = (
+        'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf'
+    )
+    self.task_cmd_builder.quantization_method = 'awq'
+    self.task_cmd_builder.quantization_precision_mode = '4bit'
+    self.task_cmd_builder.quantization_dataset_name = 'pileval'
+    self.task_cmd_builder.text_column_in_quantization_dataset = 'text'
+    self.task_cmd_builder.quantization_output_dir = '~/llama3-8b-hf-quantized'
+    self.task_cmd_builder.device_map = None
+    self.task_cmd_builder.max_memory = None
+    self.task_cmd_builder.group_size = 128
+    self.task_cmd_builder.desc_act = False
+    self.task_cmd_builder.damp_percent = 0.1
+    self.task_cmd_builder.cache_examples_on_gpu = False
+    self.task_cmd_builder.awq_version = 'GEMM'
+
+  def test_llama3_8b_model_awq_quantization(self):
+    start_time = time.time()
+    self.assertEqual(self.run_cmd(), 0)
+    end_time = time.time()
+    self.assertLess(end_time - start_time, 1.5 * 60 * 60)
+
+
+if __name__ == '__main__':
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/test_util.py

@@ -0,0 +1,134 @@
+"""Test util class."""
+
+import datetime
+import os
+import signal
+import subprocess
+import sys
+
+from absl import flags
+from absl import logging
+from absl.testing import parameterized
+import docker_command_builder as docker_cmd_builder
+
+_DOCKER_URI = flags.DEFINE_string(
+    'docker_uri', None, 'docker image uri', required=True
+)
+
+_DRY_RUN = flags.DEFINE_bool('dry_run', False, 'dry-run the commands')
+
+_LOCAL_INPUT_DIR = flags.DEFINE_string(
+    'local_input_dir',
+    os.path.expanduser('~/test_input'),
+    'local directory for storing input data.',
+)
+
+_LOCAL_OUTPUT_DIR = flags.DEFINE_string(
+    'local_output_dir',
+    '/tmp',
+    'local directory for storing test output.',
+)
+
+
+_GCS_INPUT_DIR = flags.DEFINE_string(
+    'gcs_input_dir',
+    'gs://peft-docker-test',
+    'GCS directory that stores model checkpoint, dataset and etc.',
+)
+
+_GCS_OUTPUT_DIR = flags.DEFINE_string(
+    'gcs_output_dir',
+    'gs://peft-docker-test/output',
+    'GCS directory that stores test output.',
+)
+
+
+class TestBase(parameterized.TestCase):
+  """Test base class that defines how to run commands."""
+
+  def setUp(self):
+    super().setUp()
+
+    self.docker_builder = docker_cmd_builder.DockerCommandBuilder(
+        _DOCKER_URI.value
+    )
+    self.docker_builder.add_mount_map(
+        os.path.expanduser('~'), os.path.expanduser('~')
+    )
+    self.docker_builder.add_mount_map(
+        self.local_input_dir(), self.local_input_dir()
+    )
+
+    self.task_cmd_builder = None
+
+  def cmd(self):
+    return self.docker_builder.build_cmd() + self.task_cmd_builder.build_cmd()
+
+  def run_cmd(self) -> int:
+    logging.info('running command: \n%s', ' \\\n'.join(self.cmd()))
+    if _DRY_RUN.value:
+      return 0
+
+    p = subprocess.Popen(self.cmd(), stdout=sys.stdout, stderr=sys.stderr)
+    try:
+      unused_output, unused_error = p.communicate()
+      return p.returncode
+    except KeyboardInterrupt:
+      p.send_signal(signal.SIGINT)
+      return 0
+
+  def gcs_output_dir(self):
+    return _GCS_OUTPUT_DIR.value
+
+  def local_output_dir(self):
+    return _LOCAL_OUTPUT_DIR.value
+
+  def local_input_dir(self):
+    return _LOCAL_INPUT_DIR.value
+
+
+def get_timestamp():
+  return datetime.datetime.now(datetime.timezone.utc).strftime(
+      '%Y%m%d_%H%M%S%Z'
+  )
+
+
+def get_test_data_path(name: str, download: bool = True) -> str:
+  """Gets test data path.
+
+  Args:
+    name: name of the test data
+    download: if True, then download data from GCS and returns its local path.
+
+  Returns:
+    test data path.
+  """
+
+  def _download_from_gcs(name):
+    if not os.path.exists(_LOCAL_INPUT_DIR.value):
+      os.mkdir(_LOCAL_INPUT_DIR.value)
+    subprocess.check_output([
+        'gsutil',
+        '-m',
+        'cp',
+        '-r',
+        os.path.join(_GCS_INPUT_DIR.value, name),
+        _LOCAL_INPUT_DIR.value,
+    ])
+
+  if not download:
+    return os.path.join(_GCS_INPUT_DIR.value, name)
+
+  local_data = os.path.join(_LOCAL_INPUT_DIR.value, name)
+  if not os.path.exists(local_data):
+    _download_from_gcs(name)
+
+  return local_data
+
+
+def get_pretrained_model_id(model_id: str) -> str:
+  # If `model_id` contains `/`, it is assumed to be HF model or model from GCS.
+  if '/' in model_id:
+    return model_id
+
+  return get_test_data_path(model_id, download=True)

community-content/vertex_model_garden/model_oss/peft/train/tests/test_validate_dataset_with_template.py

@@ -0,0 +1,327 @@
+"""Tests validate the dataset with template task in PEFT docker."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import test_util
+import validate_dataset_with_template_command_builder as task_cmd_builder
+
+
+class ValidateDatasetWithTemplateTest(test_util.TestBase):
+  """Test the validate dataset with template task in PEFT docker."""
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = (
+        task_cmd_builder.ValidateDatasetWithTemplateCommandBuilder()
+    )
+    self.task_cmd_builder.task = "validate-dataset-with-template"
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="valid_rows",
+          validate_top_k_rows=100,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="negative_rows",
+          validate_top_k_rows=-10,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="out_of_range_rows",
+          validate_top_k_rows=100000,
+          expected_result=1,
+      ),
+  )
+  def test_validate_dataset_with_template_top_k_rows(
+      self,
+      validate_top_k_rows,
+      expected_result,
+  ):
+    self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
+    self.task_cmd_builder.train_split_name = "train"
+    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.template = (
+        "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+    )
+    self.task_cmd_builder.validate_percentage_of_dataset = None
+    self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
+    self.task_cmd_builder.use_multiprocessing = True
+    result = self.run_cmd()
+    self.assertEqual(result, expected_result)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="valid_positive_x_percent",
+          validate_percentage_of_dataset=10,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="valid_negative_x_percent",
+          validate_percentage_of_dataset=-10,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="invalid_positive_x_percent",
+          validate_percentage_of_dataset=110,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="invalid_negative_x_percent",
+          validate_percentage_of_dataset=-110,
+          expected_result=1,
+      ),
+  )
+  def test_validate_dataset_with_template_x_percent(
+      self,
+      validate_percentage_of_dataset,
+      expected_result,
+  ):
+    self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
+    self.task_cmd_builder.train_split_name = "train"
+    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.template = (
+        "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+    )
+    self.task_cmd_builder.validate_percentage_of_dataset = (
+        validate_percentage_of_dataset
+    )
+    self.task_cmd_builder.validate_k_rows_of_dataset = None
+    self.task_cmd_builder.use_multiprocessing = True
+    result = self.run_cmd()
+    self.assertEqual(result, expected_result)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="invalid_default_input_column",
+          dataset_name="timdettmers/openassistant-guanaco",
+          split="train",
+          input_column="",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="invalid_percentage",
+          dataset_name="timdettmers/openassistant-guanaco",
+          split="train",
+          input_column="text",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=110,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="negative_percentage",
+          dataset_name="timdettmers/openassistant-guanaco",
+          split="train",
+          input_column="text",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=-110,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="empty_dataset",
+          dataset_name="",
+          split="train",
+          input_column="text",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="empty_split",
+          dataset_name="timdettmers/openassistant-guanaco",
+          split="",
+          input_column="text",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="empty_template",
+          dataset_name="timdettmers/openassistant-guanaco",
+          split="train",
+          input_column="text",
+          template="",
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="wrong_gcs_template",
+          dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
+          split="train",
+          input_column="text",
+          template="gs://cloud-nas-260507-tmp-20240724/sample-template.json",
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+      dict(
+          testcase_name="wrong_gcs_dataset_name",
+          dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft-train_sample.jsonl",
+          split="train",
+          input_column="text",
+          template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=1,
+      ),
+  )
+  def test_validate_dataset_with_template_invalid_input(
+      self,
+      dataset_name,
+      split,
+      input_column,
+      template,
+      validate_percentage_of_dataset,
+      validate_top_k_rows,
+      use_multiprocessing,
+      expected_result,
+  ):
+    self.task_cmd_builder.dataset_name = dataset_name
+    self.task_cmd_builder.train_split_name = split
+    self.task_cmd_builder.instruct_column_in_dataset = input_column
+    self.task_cmd_builder.template = template
+    self.task_cmd_builder.validate_percentage_of_dataset = (
+        validate_percentage_of_dataset
+    )
+    self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
+    self.task_cmd_builder.use_multiprocessing = use_multiprocessing
+    result = self.run_cmd()
+    self.assertEqual(result, expected_result)
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="full_hf_dataset_with_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="full_gcs_dataset_with_multiprocessing",
+          dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
+          template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="half_dataset_with_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=50,
+          validate_top_k_rows=None,
+          use_multiprocessing=True,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="top_100_rows_with_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=100,
+          use_multiprocessing=True,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="full_gcs_dataset_without_multiprocessing",
+          dataset_name="gs://cloud-nas-260507-tmp-20240724/model-evaluation/peft_train_sample.jsonl",
+          template="gs://cloud-nas-260507-tmp-20240724/sample_template.json",
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=False,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="full_hf_dataset_without_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=None,
+          use_multiprocessing=False,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="half_dataset_without_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=50,
+          validate_top_k_rows=None,
+          use_multiprocessing=False,
+          expected_result=0,
+      ),
+      dict(
+          testcase_name="top_100_rows_without_multiprocessing",
+          dataset_name="timdettmers/openassistant-guanaco",
+          template=(
+              "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+          ),
+          validate_percentage_of_dataset=None,
+          validate_top_k_rows=100,
+          use_multiprocessing=True,
+          expected_result=0,
+      ),
+  )
+  def test_validate_dataset_with_template_multiprocessing_option(
+      self,
+      dataset_name,
+      template,
+      validate_percentage_of_dataset,
+      validate_top_k_rows,
+      use_multiprocessing,
+      expected_result,
+  ):
+    self.task_cmd_builder.dataset_name = dataset_name
+    self.task_cmd_builder.train_split_name = "train"
+    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.template = template
+    self.task_cmd_builder.validate_percentage_of_dataset = (
+        validate_percentage_of_dataset
+    )
+    self.task_cmd_builder.validate_k_rows_of_dataset = validate_top_k_rows
+    self.task_cmd_builder.use_multiprocessing = use_multiprocessing
+    result = self.run_cmd()
+    self.assertEqual(result, expected_result)
+
+
+if __name__ == "__main__":
+  absltest.main()

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/command_builder_gen.py

@@ -0,0 +1,109 @@
+"""Tools to generate CommandBuilder class.
+
+See go/vmg-oss-peft-tests#commandbuilder-class-generation for details.
+"""
+
+import argparse
+import dataclasses
+from typing import List
+
+_DO_NOT_MODIFY_WARNING = """
+# DO NOT MODIFY: this file is auto-generated
+# See go/vmg-oss-peft-tests#command-builder-genpy
+"""
+
+_GETTER_TMPL = """
+    @property
+    def {}(self):
+        return self._{}
+"""
+
+_SETTER_TMPL = """
+    @{}.setter
+    def {}(self, val: {}):
+        self._{} = val
+"""
+
+_INIT_NAME = """
+    def __init__(self):"""
+
+_INIT_FIELDS = """
+        self._{} = None"""
+
+_BUILD_CMD = r"""
+    def build_cmd(self) -> str:
+        cmd = []
+        for k, v in self.__dict__.items():
+            if v is not None:
+                cmd.append(f'--{k[1:]}={v}')
+        return cmd
+"""
+
+
+@dataclasses.dataclass
+class FlagInfo:
+  api_name: str
+  impl_name: str
+  arg_type: str
+
+
+def get_flag_info(line: str) -> FlagInfo:
+  api_name, impl_name, arg_type = [x.strip() for x in line.split(',')]
+  return FlagInfo(api_name, impl_name, arg_type)
+
+
+def gen_getter(info: FlagInfo) -> str:
+  return _GETTER_TMPL.format(info.api_name, info.impl_name)
+
+
+def gen_setter(info: FlagInfo) -> str:
+  return _SETTER_TMPL.format(
+      info.api_name, info.api_name, info.arg_type, info.impl_name
+  )
+
+
+def gen_init(infos: List[FlagInfo]) -> str:
+  fields = [_INIT_FIELDS.format(i.impl_name) for i in infos]
+  return ''.join([_INIT_NAME] + fields)
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--flags_def', required=True, help='file path contain flags definition.'
+  )
+  parser.add_argument(
+      '--generated_file',
+      required=True,
+      help='file path to the generated command builder.',
+  )
+  parser.add_argument(
+      '--class_name',
+      required=True,
+      help='class name for command build',
+  )
+  args = parser.parse_args()
+
+  flags_info = []
+  with open(args.flags_def, 'r') as flags_f:
+    for line in flags_f:
+      if not line.startswith('#'):
+        flags_info.append(get_flag_info(line))
+
+  with open(args.generated_file, 'w') as gen_f:
+    # Disables pylint messages.
+    # See https://stackoverflow.com/a/43510297
+    print('# pylint: disable=W,C,R', file=gen_f)
+    print(_DO_NOT_MODIFY_WARNING, file=gen_f)
+    print(f'class {args.class_name}:', file=gen_f)
+    print(gen_init(flags_info), file=gen_f)
+    for info in flags_info:
+      print(gen_getter(info), file=gen_f)
+      print(gen_setter(info), file=gen_f)
+    print(_BUILD_CMD, file=gen_f)
+
+  print(f'file generated at {args.generated_file}')
+
+
+if __name__ == '__main__':
+  main()

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/instruct_lora_flags.txt

@@ -0,0 +1,50 @@
+# api_name, impl_name, value_type
+
+
+#                    eval related and etc.
+config_file, config_file, str
+task, task, str
+pretrained_model_id, pretrained_model_id, str
+train_dataset, dataset_name, str
+train_split_name, train_split_name, str
+template, template, str
+instruct_column, instruct_column_in_dataset, str
+ckpt_dir, output_dir, str
+merged_model_dir, merge_base_and_lora_output_dir, str
+logging_dir, logging_output_dir, str
+per_device_batch_size, per_device_train_batch_size, int
+gradient_accumulation_steps, gradient_accumulation_steps, int
+lora_rank, lora_rank, int
+lora_alpha, lora_alpha, int
+lora_dropout, lora_dropout, float
+max_steps, max_steps, int
+num_epochs, num_epochs, float
+max_seq_length, max_seq_length, int
+learning_rate, learning_rate, float
+lr_scheduler_type, lr_scheduler_type, str
+load_precision, precision_mode, str
+train_precision, train_precision, str
+gradient_checkpointing, enable_gradient_checkpointing, bool
+example_packing, use_example_packing, bool
+attn_implementation, attn_implementation, str
+optimizer, optimizer, str
+warmup_ratio, warmup_ratio, float
+report_to, report_to, str
+save_steps, save_steps, int
+logging_steps, logging_steps, int
+huggingface_access_token, huggingface_access_token, str
+eval_dataset, eval_dataset_path, str
+eval_instruct_column, eval_column, str
+eval_template, eval_template, str
+eval_split_name, eval_split, str
+eval_steps, eval_steps, int
+eval_tasks, eval_tasks, str
+eval_metric_name, eval_metric_name, str
+completion_only, completion_only, bool
+max_grad_norm, max_grad_norm, float
+logger_level, logger_level, str
+benchmark_out_file, benchmark_out_file, str
+tuning_data_stats_file, tuning_data_stats_file, str
+enable_peft, enable_peft, bool
+merge_model_precision_mode, merge_model_precision_mode, str
+target_modules, target_modules, str

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/quantize_model_flags.txt

@@ -0,0 +1,15 @@
+# api_name, impl_name, value_type
+task, task, str
+pretrained_model_id, pretrained_model_id, str
+quantization_method, quantization_method, str
+quantization_precision_mode, quantization_precision_mode, str
+quantization_dataset_name, quantization_dataset_name, str
+text_column_in_quantization_dataset, text_column_in_quantization_dataset, str
+quantization_output_dir, quantization_output_dir, str
+device_map, device_map, str
+max_memory, max_memory, str
+group_size, group_size, int
+desc_act, desc_act, bool
+damp_percent, damp_percent, float
+cache_examples_on_gpu, cache_examples_on_gpu, bool
+awq_version, awq_version, str

community-content/vertex_model_garden/model_oss/peft/train/tests/tools/validate_dataset_with_template_flags.txt

@@ -0,0 +1,9 @@
+# api_name, impl_name, value_type
+task, task, str
+template, template, str
+dataset_name, dataset_name, str
+train_split_name, train_split_name, str
+instruct_column_in_dataset, instruct_column_in_dataset, str
+use_multiprocessing, use_multiprocessing, bool
+validate_k_rows_of_dataset, validate_k_rows_of_dataset, int
+validate_percentage_of_dataset, validate_percentage_of_dataset, int

community-content/vertex_model_garden/model_oss/peft/train/tests/validate_dataset_with_template_command_builder.py

@@ -0,0 +1,88 @@
+# pylint: disable=W,C,R
+
+# DO NOT MODIFY: this file is auto-generated
+# See go/vmg-oss-peft-tests#command-builder-genpy
+
+
+class ValidateDatasetWithTemplateCommandBuilder:
+
+  def __init__(self):
+    self._task = None
+    self._template = None
+    self._dataset_name = None
+    self._train_split_name = None
+    self._instruct_column_in_dataset = None
+    self._use_multiprocessing = None
+    self._validate_k_rows_of_dataset = None
+    self._validate_percentage_of_dataset = None
+
+  @property
+  def task(self):
+    return self._task
+
+  @task.setter
+  def task(self, val: str):
+    self._task = val
+
+  @property
+  def template(self):
+    return self._template
+
+  @template.setter
+  def template(self, val: str):
+    self._template = val
+
+  @property
+  def dataset_name(self):
+    return self._dataset_name
+
+  @dataset_name.setter
+  def dataset_name(self, val: str):
+    self._dataset_name = val
+
+  @property
+  def train_split_name(self):
+    return self._train_split_name
+
+  @train_split_name.setter
+  def train_split_name(self, val: str):
+    self._train_split_name = val
+
+  @property
+  def instruct_column_in_dataset(self):
+    return self._instruct_column_in_dataset
+
+  @instruct_column_in_dataset.setter
+  def instruct_column_in_dataset(self, val: str):
+    self._instruct_column_in_dataset = val
+
+  @property
+  def use_multiprocessing(self):
+    return self._use_multiprocessing
+
+  @use_multiprocessing.setter
+  def use_multiprocessing(self, val: bool):
+    self._use_multiprocessing = val
+
+  @property
+  def validate_k_rows_of_dataset(self):
+    return self._validate_k_rows_of_dataset
+
+  @validate_k_rows_of_dataset.setter
+  def validate_k_rows_of_dataset(self, val: int):
+    self._validate_k_rows_of_dataset = val
+
+  @property
+  def validate_percentage_of_dataset(self):
+    return self._validate_percentage_of_dataset
+
+  @validate_percentage_of_dataset.setter
+  def validate_percentage_of_dataset(self, val: int):
+    self._validate_percentage_of_dataset = val
+
+  def build_cmd(self) -> str:
+    cmd = []
+    for k, v in self.__dict__.items():
+      if v is not None:
+        cmd.append(f'--{k[1:]}={v}')
+    return cmd

community-content/vertex_model_garden/model_oss/peft/train/vmg/callbacks.py

@@ -0,0 +1,86 @@
+"""Different trainer callbacks for PEFT Trainer."""
+
+import time
+
+from absl import logging
+import accelerate
+from transformers import TrainingArguments
+from transformers.trainer_callback import TrainerCallback
+from transformers.trainer_callback import TrainerControl
+from transformers.trainer_callback import TrainerState
+
+from vertex_vision_model_garden_peft.train.vmg import utils
+
+
+class TrainerStatsCallback(TrainerCallback):
+  """Trainer callback to report trainer stats."""
+
+  def __init__(self, max_seq_length, filename=None):
+    self._max_seq_length = max_seq_length
+    self._filename = filename
+
+    self._partial_state = accelerate.PartialState()
+    self._start_time = float('nan')
+    self._prev_time = float('nan')
+    self._peak_mem = 0.0
+    self._avg_throughput = 0.0
+
+  def on_step_end(
+      self,
+      args: TrainingArguments,
+      state: TrainerState,
+      control: TrainerControl,
+      **kwargs,
+  ):
+    if self._partial_state.is_main_process:
+      if state.global_step == 1:
+        self._prev_time = time.time()
+        delta_t = float('nan')
+      else:
+        cur_time = time.time()
+        delta_t = cur_time - self._prev_time
+        self._prev_time = cur_time
+        self._avg_throughput += (delta_t - self._avg_throughput) / (
+            state.global_step - 1
+        )
+
+      gpu_stats = utils.gpu_stats()
+      self._peak_mem = max(gpu_stats.total_mem, self._peak_mem)
+      logging.info(
+          'on_step_end: %s, throughput: %.2f s/it',
+          utils.gpu_stats_str(gpu_stats),
+          delta_t,
+      )
+
+  def on_train_begin(
+      self,
+      args: TrainingArguments,
+      state: TrainerState,
+      control: TrainerControl,
+      **kwargs,
+  ):
+    if self._partial_state.is_main_process:
+      self._start_time = time.time()
+      logging.info('on_train_begin: %s', utils.gpu_stats_str())
+
+  def on_train_end(
+      self,
+      args: TrainingArguments,
+      state: TrainerState,
+      control: TrainerControl,
+      **kwargs,
+  ):
+    if self._partial_state.is_main_process:
+      train_time = time.time() - self._start_time
+      logging.info(
+          'training time %.2f s, throughput: %.2f s/it, peak_mem: %.2f GB',
+          train_time,
+          self._avg_throughput,
+          self._peak_mem,
+      )
+      if self._filename:
+        with open(self._filename, 'a') as out_f:
+          out_f.write(
+              f'{self._max_seq_length/1024.0:.1f}k | {self._peak_mem:.2f} |'
+              f' {self._avg_throughput:.2f}\n'
+          )

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/custom_loglikelihood.yaml

@@ -0,0 +1,18 @@
+group:
+- vertex
+task: custom_loglikelihood
+dataset_path: json
+dataset_name: null
+output_type: loglikelihood
+training_split: null
+validation_split: null
+test_split: test
+doc_to_text: "Request: {{prompt}}\nResponse:"
+doc_to_target: " {{ground_truth}}"
+metric_list:
+- metric: perplexity
+  aggregation: perplexity
+  higher_is_better: false
+- metric: acc
+  aggregation: mean
+  higher_is_better: true

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/ddp_4gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/ddp_8gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+gpu_ids: all
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero2_4gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero2.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero2_8gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero2.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero3_4gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero3.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/deepspeed_zero3_8gpu.yaml

@@ -0,0 +1,17 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_config_file: /diffusers/examples/vertex_vision_model_garden_peft/zero3.json
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_128gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 16
+num_processes: 128
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_16gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 2
+num_processes: 16
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_24gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 3
+num_processes: 24
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_32gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 4
+num_processes: 32
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_fsdp_8gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_16gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: HYBRID_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 2
+num_processes: 16
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_24gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: HYBRID_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 3
+num_processes: 24
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/llama_hsdp_32gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+  fsdp_backward_prefetch: NO_PREFETCH
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: HYBRID_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_use_orig_params: false
+  fsdp_activation_checkpointing: false
+main_training_function: main
+mixed_precision: bf16
+machine_rank: 0
+num_machines: 4
+num_processes: 32
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/zero2.json

@@ -0,0 +1,24 @@
+{
+  "zero_optimization": {
+    "stage": 2,
+    "contiguous_gradients": false,
+    "overlap_comm": false
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+  "fp16": {
+    "enabled": "auto",
+    "auto_cast": false,
+    "loss_scale": 0,
+    "initial_scale_power": 32,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/configs/zero3.json

@@ -0,0 +1,31 @@
+{
+  "zero_optimization": {
+    "stage": 3,
+    "overlap_comm": false,
+    "contiguous_gradients": false,
+    "sub_group_size": 0,
+    "reduce_bucket_size": "auto",
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "stage3_max_live_parameters": 0,
+    "stage3_max_reuse_distance": 0,
+    "stage3_gather_16bit_weights_on_model_save": true
+  },
+  "bf16": {
+    "enabled": "auto"
+  },
+  "fp16": {
+    "enabled": "auto",
+    "auto_cast": false,
+    "loss_scale": 0,
+    "initial_scale_power": 32,
+    "loss_scale_window": 1000,
+    "hysteresis": 2,
+    "min_loss_scale": 1
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/merge_env.yaml

@@ -0,0 +1,43 @@
+# Doc about format of conda environment file
+# https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#create-env-file-manually
+name: merge
+channels:
+- nodefaults
+- conda-forge
+dependencies:
+- _libgcc_mutex=0.1=conda_forge
+- _openmp_mutex=4.5=2_gnu
+- bzip2=1.0.8=h4bc722e_7
+- ca-certificates=2024.7.4=hbcca054_0
+- ld_impl_linux-64=2.40=hf3520f5_7
+- libffi=3.4.2=h7f98852_5
+- libgcc-ng=14.1.0=h77fa898_0
+- libgomp=14.1.0=h77fa898_0
+- libnsl=2.0.1=hd590300_0
+- libsqlite=3.46.0=hde9e2c9_0
+- libuuid=2.38.1=h0b41bf4_0
+- libxcrypt=4.4.36=hd590300_1
+- libzlib=1.3.1=h4ab18f5_1
+- ncurses=6.5=h59595ed_0
+- openssl=3.3.1=h4bc722e_2
+- pip=24.2=pyhd8ed1ab_0
+- python=3.10.14=hd12c33a_0_cpython
+- readline=8.2=h8228510_1
+- setuptools=72.1.0=pyhd8ed1ab_0
+- tk=8.6.13=noxft_h4845f30_101
+- tzdata=2024a=h0c530f3_0
+- wheel=0.44.0=pyhd8ed1ab_0
+- xz=5.2.6=h166bdaf_0
+- pip:
+  - --extra-index-url https://download.pytorch.org/whl/cu121
+  - absl-py==2.1.0
+  - accelerate==0.33.0  # Needed for fp8
+  - datasets==2.19.2
+  - fbgemm-gpu==0.8.0+cu121  # Needed for fp8
+  - kfp==2.5.0
+  - peft==0.12.0
+  - protobuf==3.20.3
+  - pynvml==11.5.3
+  - torch==2.4.0+cu121  # Needed for fp8
+  - transformers==4.43.1
+  - trl==0.9.6

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/requirements.txt

@@ -0,0 +1,27 @@
+# Doc about format of requirement file
+# https://pip.pypa.io/en/stable/reference/requirements-file-format
+
+--extra-index-url https://download.pytorch.org/whl/cu118
+--extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
+
+# keep sorted
+accelerate==0.31.0
+auto_gptq==0.7.1+cu118
+autoawq==0.2.5
+bitsandbytes==0.43.2
+cloudml-hypertune==0.1.0.dev6
+datasets==2.19.2
+deepspeed==0.14.4
+diffusers==0.25.1
+fsspec==2024.3.1
+gcsfs==2024.3.1
+lm_eval==0.4.3
+ninja==1.11.1 # Needed to avoid `ninja 1.11.1.1 is not supported on this platform` error
+optimum==1.17.1
+peft==0.12.0
+pynvml==11.5.3
+torch==2.2.2+cu118
+torchvision==0.17.2+cu118
+transformers==4.43.1
+trl==0.9.6
+wandb==0.17.1

community-content/vertex_model_garden/model_oss/peft/train/vmg/dockerfile/train.Dockerfile

@@ -0,0 +1,76 @@
+# Dockerfile for PEFT Training.
+#
+# To build:
+# docker build -f model_oss/peft/dockerfile/train.Dockerfile . -t ${YOUR_IMAGE_TAG}
+#
+# To push to gcr:
+# docker tag ${YOUR_IMAGE_TAG} gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+# docker push gcr.io/${YOUR_PROJECT}/${YOUR_IMAGE_TAG}
+
+# Picked from https://cloud.google.com/deep-learning-containers/docs/choosing-container#pytorch
+FROM us-docker.pkg.dev/deeplearning-platform-release/gcr.io/pytorch-cu121.2-2.py310:m123
+RUN apt-get update && \
+    apt-get upgrade -y && \
+    apt-get install -y curl git wget software-properties-common vim libaio-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists*
+
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+
+# Install libraries.
+ENV PIP_ROOT_USER_ACTION=ignore
+RUN pip install --upgrade pip
+
+# Prefer to install with requirement file as much as possible for reasons
+# described in b/355034754.
+COPY model_oss/peft/train/vmg/dockerfile/requirements.txt /tmp/requirements.txt
+RUN pip install -r /tmp/requirements.txt
+
+# flash-attn cannot be installed with the requirement file approach above
+# because of the `no-build-isolation` requirement.
+#
+# It is OK to install it after other packages FOR NOW because it only has
+# limited dependencies. And there's no concern about it overwriting previously
+# installed packages.
+# https://github.com/Dao-AILab/flash-attention/blob/v2.6.3/setup.py#L523
+RUN pip install flash-attn==2.6.3 --no-build-isolation
+
+# Install `diffusers` library as editable and in root folder (/) on purpose.
+RUN git clone --depth 1 --branch v0.25.1 https://github.com/huggingface/diffusers.git
+# Remove `diffusers` (NOTE that the dependency libraries are kept).
+RUN pip uninstall -y diffusers
+# Using `--no-deps` option to make sure previously installed packages are not
+# overwritten.
+RUN pip install --no-deps -e /diffusers
+
+# Make sure there's no inconsistent pip libraries.
+RUN pip check
+
+# Install merge related packages in a separate env.
+COPY model_oss/peft/train/vmg/dockerfile/merge_env.yaml /tmp/merge_env.yaml
+RUN conda env create -n merge --yes --file /tmp/merge_env.yaml
+RUN conda init
+
+# Switch to diffusers examples folder.
+WORKDIR /diffusers/examples
+
+RUN mkdir -p ./vertex_vision_model_garden_peft/
+COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
+# custom `lm_eval` task.
+ARG LM_EVAL_DIR=$(python -c 'import site; print(site.getsitepackages()[0])')/lm_eval
+RUN mkdir -p $LM_EVAL_DIR/tasks/vertex && \
+  mv ./vertex_vision_model_garden_peft/custom_loglikelihood.yaml $LM_EVAL_DIR/tasks/vertex/
+COPY model_oss/peft/train/vmg/*.py ./vertex_vision_model_garden_peft/train/vmg/
+COPY model_oss/peft/train/vmg/templates /diffusers/examples/util/templates
+COPY model_oss/util /diffusers/examples/util
+COPY model_oss/notebook_util/dataset_validation_util.py /diffusers/examples/util
+COPY model_oss/peft/train/tests/*.py ./vertex_vision_model_garden_peft/tests/
+
+RUN chmod a+rwX -R /diffusers/examples/
+ENV PYTHONPATH /diffusers/examples/
+# Must disable torch XLA, otherwise runtime uses CPU even if GPU exists.
+ENV USE_TORCH_XLA 0
+
+ENTRYPOINT ["python3", "./vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py"]

community-content/vertex_model_garden/model_oss/peft/train/vmg/eval_lib.py

@@ -0,0 +1,183 @@
+"""Library for running evaluations during training."""
+
+import dataclasses
+from typing import Any, Optional, Type
+
+from absl import logging
+import datasets
+from lm_eval import evaluator
+from lm_eval import tasks
+from lm_eval import utils
+from lm_eval.api import model as lm_model
+from lm_eval.api import registry
+from lm_eval.models import huggingface
+from peft import peft_model
+import transformers
+from transformers import trainer
+
+from util import dataset_validation_util
+from util import constants
+
+
+_DESCRIPTION_EVALUATION = "evaluation"
+_BUILTIN_EVAL_TASK = "builtin_eval"
+
+
+@dataclasses.dataclass(frozen=True)
+class EvalConfig:
+  steps: int
+  tasks: list[str]
+  per_device_batch_size: int
+  num_fewshot: Optional[int]
+  limit: Optional[float]
+  metric_name: str
+  tokenize_dataset: bool
+  dataset_path: str = ""
+  split: str = "test"
+  template: str = ""
+  column: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET
+
+
+class PeftCausalLMModel(huggingface.HFLM):
+  """PeftCausalLMModel that supports loading an in-memory model."""
+
+  AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
+
+  def __init__(
+      self,
+      model: peft_model.PeftModelForCausalLM,
+      tokenizer: transformers.PreTrainedTokenizerBase,
+      batch_size_per_gpu: int,
+  ):
+    lm_model.LM.__init__(self)
+    self._model = model
+    self.tokenizer = tokenizer
+    self.vocab_size = tokenizer.vocab_size
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+    self._config = model.config
+    self.batch_size_per_gpu = batch_size_per_gpu
+    self._device = model.device
+    self._max_length = None  # Will be automatically determined from config.
+    self._add_special_tokens = (
+        None  # Will be automatically determined from AUTO_MODEL_CLASS.
+    )
+
+
+def create_trainer(
+    cls: Type[transformers.Trainer],
+    eval_config: Optional[EvalConfig],
+    tokenizer: Optional[transformers.PreTrainedTokenizerBase],
+    args: trainer.TrainingArguments,
+    **kwargs,
+) -> transformers.Trainer:
+  """Creates a trainer. If eval config is provided, injects evaluation loop."""
+  if not eval_config:
+    return cls(args=args, **kwargs)
+
+  args.eval_strategy = "steps"
+  args.eval_steps = eval_config.steps
+  args.per_device_eval_batch_size = eval_config.per_device_batch_size
+  kwargs["tokenizer"] = tokenizer
+
+  if eval_config.tasks == [_BUILTIN_EVAL_TASK]:
+    try:
+      eval_dataset = dataset_validation_util.load_dataset_with_template(
+          dataset_name=eval_config.dataset_path,
+          split=eval_config.split,
+          input_column=eval_config.column,
+          template=eval_config.template,
+          tokenizer=tokenizer,
+      )
+      if eval_config.limit is not None:
+        if eval_config.limit >= 1:
+          limit = int(eval_config.limit)
+        else:
+          limit = int(eval_config.limit * len(eval_dataset))
+        eval_dataset = eval_dataset.select(range(limit))
+      if eval_config.tokenize_dataset:
+        eval_dataset = eval_dataset.map(
+            lambda samples: tokenizer(samples[eval_config.column])
+        )
+      kwargs["eval_dataset"] = eval_dataset
+    except (OSError, ValueError, IndexError) as e:
+      logging.warning(
+          "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
+          eval_config.dataset_path,
+          e,
+      )
+      del args.evaluation_strategy
+      del args.eval_steps
+      del args.per_device_eval_batch_size
+    return cls(args=args, **kwargs)
+
+  class LMEvalTrainer(cls):
+    """Trainer with lm_eval injected as the eval library."""
+
+    def __init__(self, **kwargs):
+      super().__init__(**kwargs)
+      task_names = utils.pattern_match(eval_config.tasks, registry.ALL_TASKS)
+      logging.info("Selected Eval Tasks: %s", task_names)
+      task_args = {}
+      if eval_config.num_fewshot is not None:
+        task_args["num_fewshot"] = eval_config.num_fewshot
+      if eval_config.dataset_path:
+        task_args["dataset_path"] = "json"
+        task_args["dataset_kwargs"] = {
+            "data_files": {"test": eval_config.dataset_path},
+        }
+      self._eval_task_dict = tasks.get_task_dict(task_names, **task_args)
+
+    def evaluation_loop(
+        self,
+        dataloader: trainer.DataLoader,
+        description: str,
+        prediction_loss_only: Optional[bool] = None,
+        ignore_keys: Optional[list[str]] = None,
+        metric_key_prefix: str = "eval",
+    ) -> trainer.EvalLoopOutput:
+      """Custom evaluation loop that invokes lm_eval."""
+      if description.lower() != _DESCRIPTION_EVALUATION:
+        return super().evaluation_loop(
+            dataloader,
+            description,
+            prediction_loss_only,
+            ignore_keys,
+            metric_key_prefix,
+        )
+
+      model = self._wrap_model(self.model, training=False)
+      lm = PeftCausalLMModel(
+          model,
+          self.tokenizer or self.data_collator.tokenizer,
+          eval_config.per_device_batch_size,
+      )
+      results: dict[str, Any] = evaluator.evaluate(
+          lm=lm,
+          task_dict=self._eval_task_dict,
+          limit=eval_config.limit,
+      )["results"]
+      metric_name = eval_config.metric_name
+      # Compute average value if there are multiple tasks.
+      metric_values: list[float] = []
+      for result in results.values():
+        for key, value in result.items():
+          if key.split(",")[0] == metric_name:
+            metric_values.append(value)
+      if not metric_values:
+        raise ValueError(
+            f"Metric {metric_name} not found in eval response: {results}"
+        )
+      metric_average = sum(metric_values) / len(metric_values)
+      logging.info("%s value: %f\n%s", metric_name, metric_average, results)
+      return trainer.EvalLoopOutput(
+          # Only metrics field is set. Other fields are dummy values.
+          predictions=None,
+          label_ids=None,
+          metrics={f"{metric_key_prefix}_{metric_name}": metric_average},
+          num_samples=0,
+      )
+
+  # Use empty eval dataset as a placeholder.
+  return LMEvalTrainer(
+      args=args, eval_dataset=datasets.Dataset.from_dict({"test": []}), **kwargs
+  )

community-content/vertex_model_garden/model_oss/peft/train/vmg/instruct_lora.py

@@ -0,0 +1,794 @@
+"""Instruct/Chat with LoRA models."""
+
+import dataclasses
+import datetime
+import json
+import os
+from typing import Any, Dict, Optional, Sequence
+import warnings
+
+from absl import app
+from absl import flags
+from absl import logging
+from accelerate import DistributedType
+from accelerate import PartialState
+import bitsandbytes as bnb
+import hypertune
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import TrainingArguments
+from trl import DataCollatorForCompletionOnlyLM
+from trl import SFTTrainer
+import wandb
+
+from util import dataset_validation_util
+from vertex_vision_model_garden_peft.train.vmg import callbacks
+from vertex_vision_model_garden_peft.train.vmg import eval_lib
+from vertex_vision_model_garden_peft.train.vmg import utils
+from util import constants
+from util import fileutils
+
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    'pretrained_model_id',
+    None,
+    'The pretrained model id. Supported models can be causal language modeling'
+    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
+    ' might be different paddings for different models. This tool assumes the'
+    ' pretrained_model_id contains model name, and then choose proper padding'
+    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    required=True,
+)
+
+_HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
+    'huggingface_access_token',
+    None,
+    'The access token for loading huggingface gated models.',
+)
+
+_DATASET_NAME = flags.DEFINE_string(
+    'dataset_name',
+    None,
+    'The dataset name in huggingface.',
+)
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'The output directory.',
+)
+
+_LOGGING_OUTPUT_DIR = flags.DEFINE_string(
+    'logging_output_dir',
+    '',
+    'The logging output directory, which defaults to same as output_dir.',
+)
+
+_PRECISION_MODE = flags.DEFINE_enum(
+    'precision_mode',
+    constants.PRECISION_MODE_16,
+    [
+        constants.PRECISION_MODE_4,
+        constants.PRECISION_MODE_8,
+        constants.PRECISION_MODE_16,
+        constants.PRECISION_MODE_16B,
+        constants.PRECISION_MODE_32,
+    ],
+    'Precision to load model weights for finetuning.',
+)
+
+_LORA_RANK = flags.DEFINE_integer(
+    'lora_rank',
+    16,
+    'The rank of the update matrices, expressed in int. Lower rank results in'
+    ' smaller update matrices with fewer trainable parameters, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_ALPHA = flags.DEFINE_integer(
+    'lora_alpha',
+    32,
+    'LoRA scaling factor, referring to'
+    ' https://huggingface.co/docs/peft/conceptual_guides/lora.',
+)
+
+_LORA_DROPOUT = flags.DEFINE_float(
+    'lora_dropout',
+    0.05,
+    'dropout probability of the LoRA layers, referring to'
+    ' https://huggingface.co/docs/peft/task_guides/token-classification-lora.',
+)
+
+_WARMUP_STEPS = flags.DEFINE_integer(
+    'warmup_steps',
+    10,
+    'Number of steps for the warmup in the learning rate scheduler.',
+)
+
+_WARMUP_RATIO = flags.DEFINE_float(
+    'warmup_ratio',
+    0.03,
+    'The warmup ratio in the learning rate scheduler.',
+)
+
+_WEIGHT_DECAY = flags.DEFINE_float(
+    'weight_decay',
+    0.001,
+    'The weight decay in the learning rate scheduler.',
+)
+
+_NUM_EPOCHS = flags.DEFINE_float(
+    'num_epochs',
+    None,
+    'The number of training epochs. Only used for'
+    ' "sequence-classification-lora" with an integer value and for'
+    ' "instruct-lora" with a float value allowed.',
+)
+
+_MAX_STEPS = flags.DEFINE_integer(
+    'max_steps',
+    None,
+    'Total number of training steps. Overrides num_epochs if set. Only used for'
+    ' "instruct-lora."',
+)
+
+_MAX_SEQ_LENGTH = flags.DEFINE_integer(
+    'max_seq_length',
+    512,
+    'The maximum sequence length.',
+)
+
+_LEARNING_RATE = flags.DEFINE_float(
+    'learning_rate',
+    2e-4,
+    'The learning rate after the potential warmup period.',
+)
+
+_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
+    'instruct_column_in_dataset',
+    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    'The instruct column in dataset.',
+)
+
+_REPORT_TO = flags.DEFINE_string(
+    'report_to',
+    constants.REPORT_TO_NONE,
+    'Where logging is reported to, which can be tensorboard or none.',
+)
+
+_PER_DEVICE_TRAIN_BATCH_SIZE = flags.DEFINE_integer(
+    'per_device_train_batch_size',
+    4,
+    'The per device train batch size.',
+)
+
+_GRADIENT_ACCUMULATION_STEPS = flags.DEFINE_integer(
+    'gradient_accumulation_steps',
+    4,
+    'The gradient accumulation steps.',
+)
+
+_ENABLE_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
+    'enable_gradient_checkpointing',
+    False,
+    'Whether to enable gradient checkpointing.',
+)
+
+_ENABLE_PEFT = flags.DEFINE_boolean(
+    'enable_peft',
+    True,
+    'Whether to enable peft.',
+)
+_TEMPLATE = flags.DEFINE_string(
+    'template',
+    None,
+    'Template for formatting language model training data. Must be a filename'
+    ' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
+    ' Cloud Storage URI to a JSON file.',
+)
+
+_OPTIMIZER = flags.DEFINE_string(
+    'optimizer',
+    'adamw_torch',
+    'The optimizer.',
+)
+
+_LR_SCHEDULER_TYPE = flags.DEFINE_string(
+    'lr_scheduler_type',
+    'cosine',
+    'The learning rate scheduler type.',
+)
+
+_SAVE_STEPS = flags.DEFINE_integer(
+    'save_steps',
+    10,
+    'The save steps.',
+)
+
+_LOGGING_STEPS = flags.DEFINE_integer(
+    'logging_steps',
+    10,
+    'The logging steps.',
+)
+
+_EVAL_STEPS = flags.DEFINE_integer(
+    'eval_steps',
+    10,
+    'The number of training steps between evaluations.',
+)
+
+_TRAIN_SPLIT_NAME = flags.DEFINE_string(
+    'train_split_name',
+    'train',
+    'The train split name.',
+)
+
+_EVAL_TASKS = flags.DEFINE_list(
+    'eval_tasks',
+    None,
+    'List of eval task names (can have wildcards) as in'
+    ' https://github.com/EleutherAI/lm-evaluation-harness. Will not run'
+    ' evaluation if not set. Runs the built-in trainer evaluation loop if set'
+    ' to `builtin_eval`.',
+)
+
+_EVAL_PER_DEVICE_BATCH_SIZE = flags.DEFINE_integer(
+    'eval_per_device_batch_size',
+    1,
+    'The per device batch size for model evaluation.',
+)
+
+_EVAL_NUM_FEWSHOT = flags.DEFINE_integer(
+    'eval_num_fewshot',
+    None,
+    'Run N-shot language model evaluation. Not implemented in `builtin_eval`.',
+)
+
+_EVAL_LIMIT = flags.DEFINE_float(
+    'eval_limit',
+    None,
+    'Limit the number of examples per task. If <1, limit is a percentage of the'
+    ' total number of examples.',
+)
+
+_EVAL_METRIC_NAME = flags.DEFINE_string(
+    'eval_metric_name',
+    'acc',
+    'The metric name to aggregate during model evaluation.',
+)
+
+_EVAL_DATASET_PATH = flags.DEFINE_string(
+    'eval_dataset_path',
+    None,
+    'Overrides the default evaluation dataset path. In `builtin_eval` mode,'
+    ' this can be any Hugging Face dataset name.',
+)
+
+# We set the default eval split as `test`, based on observation from
+# https://huggingface.co/datasets/timdettmers/openassistant-guanaco/viewer/default/test.
+_EVAL_SPLIT = flags.DEFINE_string(
+    'eval_split',
+    'test',
+    'Eval split name in the eval dataset for `builtin_eval`.',
+)
+
+_EVAL_TEMPLATE = flags.DEFINE_string(
+    'eval_template',
+    None,
+    'Template for formatting language model evaluation data for `builtin_eval`.'
+    ' Must be a filename under `templates` folder, without `.json` extension,'
+    ' e.g. `alpaca`, or a Cloud Storage URI to a JSON file.',
+)
+
+_EVAL_COLUMN = flags.DEFINE_string(
+    'eval_column',
+    None,
+    'Eval column name in the eval dataset for `builtin_eval`.',
+)
+
+_TRAIN_PRECISION = flags.DEFINE_enum(
+    'train_precision',
+    constants.PRECISION_MODE_16B,
+    [
+        constants.PRECISION_MODE_16,
+        constants.PRECISION_MODE_16B,
+        constants.PRECISION_MODE_32,
+    ],
+    'Precision to train the model.',
+)
+
+_USE_EXAMPLE_PACKING = flags.DEFINE_boolean(
+    'use_example_packing',
+    False,
+    'Enables example packing during training, which uses '
+    '`ConstantLengthDataset` under the hood.',
+)
+
+_COMPLETION_ONLY = flags.DEFINE_boolean(
+    'completion_only',
+    False,
+    'If set, it uses DataCollatorForCompletionOnlyLM to train the model on the'
+    ' generated prompts only, i.e., masking out the input',
+)
+
+_ATTN_IMPLEMENTATION = flags.DEFINE_string(
+    'attn_implementation',
+    None,
+    'Attention implementation, can be `eager`, `sdpa` or `flash_attention_2`',
+)
+
+_MAX_GRAD_NORM = flags.DEFINE_float(
+    'max_grad_norm',
+    0.3,
+    'Maximum gradient norm used for gradient clipping',
+)
+
+_WARNINGS_FILTER = flags.DEFINE_string(
+    'warnings_filter',
+    'ignore',
+    'Warning filter as defined in '
+    'https://docs.python.org/3/library/warnings.html#the-warnings-filter',
+)
+
+_LOGGER_LEVEL = flags.DEFINE_string(
+    'logger_level',
+    'passive',
+    'logging level passed to TrainingArguments. Note that this is for python'
+    ' logging module, NOT the one from absl',
+)
+
+_BENCHMARK_OUT_FILE = flags.DEFINE_string(
+    'benchmark_out_file', None, 'file path for writing benchmark result'
+)
+
+
+_NCCL_TIMEOUT = flags.DEFINE_integer(
+    'nccl_timeout', 6000, 'nccl timeout in seconds'
+)
+
+_TUNING_DATA_STATS_FILE = flags.DEFINE_string(
+    'tuning_data_stats_file', None, 'file path for writing tuning data stats.'
+)
+
+_TARGET_MODULES = flags.DEFINE_list(
+    'target_modules', None, 'The names of the modules to apply LoRA adapter to.'
+)
+
+
+@flags.multi_flags_validator(
+    [
+        _COMPLETION_ONLY.name,
+        _USE_EXAMPLE_PACKING.name,
+    ],
+    message=(
+        '`use_example_packing=True` does not work with `completion_only=True`'
+    ),
+)
+def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
+  """Check to make sure example packing is enabled properly.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If `use_example_packing` is set properly.
+  """
+  if (
+      flags_dict[_COMPLETION_ONLY.name]
+      and flags_dict[_USE_EXAMPLE_PACKING.name]
+  ):
+    return False
+  return True
+
+
+@flags.multi_flags_validator(
+    [
+        _COMPLETION_ONLY.name,
+        _TEMPLATE.name,
+    ],
+    message='`template` should be provided if using `completion_only=True`',
+)
+def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
+  """Check to make sure completion_only is enabled properly.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If `completion_only` is set properly
+  """
+  if flags_dict[_COMPLETION_ONLY.name] and flags_dict[_TEMPLATE.name] is None:
+    return False
+  return True
+
+
+# References:
+# Huggingface SFT trainer example:
+# https://github.com/huggingface/trl/blob/main/examples/scripts/sft_trainer.py.
+# Huggingface sagemaker example:
+# https://github.com/huggingface/notebooks/blob/main/sagemaker/28_train_llms_with_qlora/scripts/run_clm.py.
+
+
+# Copied from https://github.com/artidoro/qlora/blob/main/qlora.py.
+def find_all_linear_names(
+    model: AutoModelForCausalLM, precision_mode: str
+) -> list[str]:
+  """Finds all linear module names."""
+  if precision_mode == constants.PRECISION_MODE_4:
+    cls = bnb.nn.Linear4bit
+  elif precision_mode == constants.PRECISION_MODE_8:
+    cls = bnb.nn.Linear8bitLt
+  else:
+    cls = torch.nn.Linear
+  lora_module_names = set()
+  for name, module in model.named_modules():
+    if isinstance(module, cls):
+      names = name.split('.')
+      lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+  if 'lm_head' in lora_module_names:  # needed for 16-bit
+    lora_module_names.remove('lm_head')
+  return list(lora_module_names)
+
+
+def finetune_instruct(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    logging_output_dir: str,
+    lora_rank: int = 64,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    warmup_ratio: int = 0.03,
+    num_epochs: Optional[float] = None,
+    max_steps: Optional[int] = None,
+    warmup_steps: int = 10,
+    max_seq_length: int = 512,
+    learning_rate: float = 2e-4,
+    precision_mode: str = None,
+    instruct_column_in_dataset: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    per_device_train_batch_size: int = 4,
+    gradient_accumulation_steps: int = 4,
+    optim: str = 'paged_adamw_32bit',
+    weight_decay: float = 0.001,
+    enable_gradient_checkpointing: bool = False,
+    enable_peft: bool = True,
+    template: str = None,
+    lr_scheduler_type: str = 'constant',
+    save_steps: int = 10,
+    logging_steps: int = 10,
+    train_split_name: str = 'train',
+    eval_config: Optional[eval_lib.EvalConfig] = None,
+    report_to: str = constants.REPORT_TO_NONE,
+    access_token: Optional[str] = None,
+    train_precision: str = constants.PRECISION_MODE_16B,
+    use_example_packing: bool = False,
+    attn_implementation: Optional[str] = None,
+    max_grad_norm: float = 0.3,
+    completion_only: bool = False,
+    logger_level: str = 'passive',
+    benchmark_out_file: Optional[str] = None,
+    tuning_data_stats_file: Optional[str] = None,
+    target_modules: Optional[str] = None,
+) -> None:
+  """Finetunes instruct."""
+  logging.info('on entering instruct_lora, %s', utils.gpu_stats_str())
+  gradient_checkpointing_kwargs = {}
+  # DDP provides limited support with the reentrant variant of gradient
+  # checkpoint [1]. Below is an indirect way of checking whether DDP will be
+  # used. It is "indirect" because there are complex logic under the hood of
+  # `SFTTrainer` and since those are not public API, they might change as we
+  # update the library.
+  if PartialState().distributed_type == DistributedType.MULTI_GPU:
+    gradient_checkpointing_kwargs['use_reentrant'] = False
+
+  tokenizer = utils.load_tokenizer(
+      pretrained_model_id,
+      'right',
+      access_token=access_token,
+  )
+
+  train_dataset = dataset_validation_util.load_dataset_with_template(
+      dataset_name,
+      split=train_split_name,
+      input_column=instruct_column_in_dataset,
+      template=template,
+      tokenizer=tokenizer,
+  )
+
+  if tuning_data_stats_file:
+    with PartialState().main_process_first():
+      effective_batch_size = (
+          per_device_train_batch_size
+          * gradient_accumulation_steps
+          * PartialState().num_processes
+      )
+      logging.info(
+          'getting tuning data stats with effective batch size %s',
+          effective_batch_size,
+      )
+      train_dataset_stats = utils.get_dataset_stats(
+          train_dataset,
+          tokenizer,
+          instruct_column_in_dataset,
+          effective_batch_size,
+      )
+      logging.info('stats: %s', train_dataset_stats)
+      tuning_data_stats_file = dataset_validation_util.force_gcs_fuse_path(
+          tuning_data_stats_file
+      )
+      with open(tuning_data_stats_file, 'w') as out_f:
+        json.dump(dataclasses.asdict(train_dataset_stats), out_f)
+
+  model = utils.load_model(
+      pretrained_model_id=pretrained_model_id,
+      tokenizer=tokenizer,
+      precision_mode=precision_mode,
+      enable_gradient_checkpointing=enable_gradient_checkpointing,
+      access_token=access_token,
+      attn_implementation=attn_implementation,
+      train_precision=train_precision,
+  )
+
+  if enable_peft:
+    if target_modules is None:
+      target_modules = find_all_linear_names(
+          model, precision_mode=precision_mode
+      )
+    logging.info('applying lora adapters to modules: %s', target_modules)
+    peft_config = LoraConfig(
+        lora_alpha=lora_alpha,
+        lora_dropout=lora_dropout,
+        r=lora_rank,
+        bias='none',
+        task_type='CAUSAL_LM',
+        target_modules=target_modules,
+    )
+    # If we pass in `peft_config` to SFTTrainer, it does a lot of magic under
+    # the hood, e.g., calling `prepare_model_for_kbit_training` before calling
+    # `get_peft_model`, which may revert other changes we did before. That's why
+    # we are calling `get_peft_model` explicitly here.
+    model = get_peft_model(model, peft_config)
+
+    # This is to work-around mix-precision training. This issue is not fixed as
+    # of transformers==4.41.2.
+    # See b/332760883#comment30 for more details.
+    if precision_mode in (
+        constants.PRECISION_MODE_16,
+        constants.PRECISION_MODE_16B,
+    ):
+      for param in filter(lambda p: p.requires_grad, model.parameters()):
+        param.data = param.data.to(torch.float32)
+
+  if not logging_output_dir:
+    logging_output_dir = output_dir
+
+  # To use singleton PartialState() without re-initializing it. See
+  # b/357970482#comment3
+  accelerator_config = {'use_configured_state': True}
+
+  training_arguments = TrainingArguments(
+      report_to=report_to,
+      output_dir=output_dir,
+      per_device_train_batch_size=per_device_train_batch_size,
+      gradient_accumulation_steps=gradient_accumulation_steps,
+      optim=optim,
+      save_steps=save_steps,
+      save_strategy='steps',
+      save_total_limit=3,
+      logging_dir=os.path.join(logging_output_dir, 'logs'),
+      logging_steps=logging_steps,
+      learning_rate=learning_rate,
+      fp16=(train_precision == constants.PRECISION_MODE_16),
+      bf16=(train_precision == constants.PRECISION_MODE_16B),
+      max_grad_norm=max_grad_norm,
+      num_train_epochs=num_epochs if num_epochs else -1,
+      max_steps=max_steps if max_steps else -1,
+      warmup_ratio=warmup_ratio,
+      warmup_steps=warmup_steps,
+      group_by_length=False,
+      lr_scheduler_type=lr_scheduler_type,
+      gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
+      weight_decay=weight_decay,
+      log_level=logger_level,
+      accelerator_config=accelerator_config,
+  )
+  trainer_kwargs = {}
+  if completion_only and template:
+    template_json = dataset_validation_util.get_template(template_path=template)
+    instruction_sep = dataset_validation_util.get_instruction_separator(
+        template_json
+    )
+    response_sep = dataset_validation_util.get_response_separator(template_json)
+    if not response_sep:
+      raise ValueError(
+          '`response_separator` must be provided to use'
+          ' `DataCollatorForCompletionOnlyLM`'
+      )
+
+    trainer_kwargs['data_collator'] = DataCollatorForCompletionOnlyLM(
+        instruction_template=instruction_sep,
+        response_template=response_sep,
+        tokenizer=tokenizer,
+    )
+    logging.info('using DataCollatorForCompletionOnlyLM')
+
+  trainer_stats_callback = callbacks.TrainerStatsCallback(
+      max_seq_length, benchmark_out_file
+  )
+  trainer = eval_lib.create_trainer(
+      cls=SFTTrainer,
+      eval_config=eval_config,
+      model=model,
+      train_dataset=train_dataset,
+      dataset_text_field=instruct_column_in_dataset,
+      max_seq_length=max_seq_length,
+      tokenizer=tokenizer,
+      args=training_arguments,
+      packing=use_example_packing,
+      callbacks=[trainer_stats_callback],
+      **trainer_kwargs,
+  )
+
+  # `eval_lib.create_trainer` might modify the training args. Printing here
+  # should capture what will be used by the trainer.
+  if PartialState().is_main_process:
+    logging.info('training args: %s', trainer.args)
+
+  if enable_peft:
+    trainer.model.print_trainable_parameters()
+
+  if trainer.is_fsdp_enabled:
+    logging.info('Trainer running with FSDP.')
+  elif trainer.is_deepspeed_enabled:
+    logging.info('Trainer running with DeepSpeed.')
+  else:
+    logging.info('Trainer running without parallelism.')
+
+  trainer.train()
+
+  # Always save the final checkpoint.
+  final_checkpoint = utils.get_final_checkpoint_path(output_dir)
+  logging.info('The final checkpoint is: %s.', final_checkpoint)
+
+  if trainer.is_fsdp_enabled:
+    trainer.accelerator.state.fsdp_plugin.set_state_dict_type('FULL_STATE_DICT')
+    # This method saves the sharded weights like `accelerator.save_state`, see
+    # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
+    trainer.save_model(output_dir)
+    model = trainer.model.cpu()  # Avoids GPU OOM
+    state_dict = trainer.accelerator.get_state_dict(model)
+    # To aggregate the weights from all the devices, we need to use
+    # `state_dict=state_dict`.
+    model.save_pretrained(
+        final_checkpoint,
+        state_dict=state_dict,
+        is_main_process=PartialState().is_main_process,
+        save_embedding_layers=False,  # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
+    )
+    model.cuda()  # Move back to GPU to do eval.
+  else:
+    trainer.model.save_pretrained(
+        final_checkpoint,
+        is_main_process=PartialState().is_main_process,
+        save_embedding_layers=False,  # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
+    )
+
+  if eval_config is not None and trainer.eval_dataset is not None:
+    metrics = trainer.evaluate(metric_key_prefix='eval')
+    # Both `log_metrics` and `save_metrics` are multiple process safe.
+    # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L911 #pylint: disable=line-too-long
+    # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L1001 #pylint: disable=line-too-long
+    trainer.log_metrics('eval', metrics)
+    trainer.save_metrics('eval', metrics)
+    if PartialState().is_main_process:
+      hp_metric = metrics[f'eval_{eval_config.metric_name}']
+      hpt = hypertune.HyperTune()
+      hpt.report_hyperparameter_tuning_metric(
+          hyperparameter_metric_tag=constants.HP_METRIC_TAG,
+          metric_value=hp_metric,
+      )
+      logging.info('Send HP metric: %f to hyperparameter tuning.', hp_metric)
+    PartialState().wait_for_everyone()
+
+  if not enable_peft:
+    tokenizer.save_pretrained(
+        final_checkpoint, is_main_process=PartialState().is_main_process
+    )
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  # This needs to be called before any other PartialState() calls.
+  utils.init_partial_state(
+      timeout=datetime.timedelta(seconds=_NCCL_TIMEOUT.value)
+  )
+
+  utils.print_library_versions()
+  warnings.simplefilter(_WARNINGS_FILTER.value)
+
+  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
+  if dataset_validation_util.is_gcs_path(pretrained_model_id):
+    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
+        pretrained_model_id
+    )
+
+  output_dir = utils.GcsOrLocalDirectory(
+      _OUTPUT_DIR.value, check_empty=True, upload_from_all_nodes=True
+  )
+
+  # GCS Fuse does not sync flushed files if not closed. See b/361771727.
+  logging_output_dir = fileutils.force_gcs_path(_LOGGING_OUTPUT_DIR.value)
+
+  # Creates evaluation config.
+  if _EVAL_TASKS.value:
+    eval_config = eval_lib.EvalConfig(
+        tasks=_EVAL_TASKS.value,
+        per_device_batch_size=_EVAL_PER_DEVICE_BATCH_SIZE.value,
+        num_fewshot=_EVAL_NUM_FEWSHOT.value,
+        limit=_EVAL_LIMIT.value,
+        metric_name=_EVAL_METRIC_NAME.value,
+        steps=_EVAL_STEPS.value,
+        dataset_path=dataset_validation_util.force_gcs_fuse_path(
+            _EVAL_DATASET_PATH.value
+        ),
+        split=_EVAL_SPLIT.value,
+        template=_EVAL_TEMPLATE.value,
+        column=_EVAL_COLUMN.value,
+        tokenize_dataset=False,
+    )
+  else:
+    eval_config = None
+
+  if _REPORT_TO.value == constants.REPORT_TO_WANDB:
+    wandb.login()
+
+  finetune_instruct(
+      pretrained_model_id=pretrained_model_id,
+      dataset_name=_DATASET_NAME.value,
+      output_dir=output_dir.local_dir,
+      logging_output_dir=logging_output_dir,
+      precision_mode=_PRECISION_MODE.value,
+      lora_rank=_LORA_RANK.value,
+      lora_alpha=_LORA_ALPHA.value,
+      lora_dropout=_LORA_DROPOUT.value,
+      warmup_ratio=_WARMUP_RATIO.value,
+      num_epochs=_NUM_EPOCHS.value,
+      warmup_steps=_WARMUP_STEPS.value,
+      max_steps=_MAX_STEPS.value,
+      max_seq_length=_MAX_SEQ_LENGTH.value,
+      learning_rate=_LEARNING_RATE.value,
+      instruct_column_in_dataset=_INSTRUCT_COLUMN_IN_DATASET.value,
+      per_device_train_batch_size=_PER_DEVICE_TRAIN_BATCH_SIZE.value,
+      optim=_OPTIMIZER.value,
+      weight_decay=_WEIGHT_DECAY.value,
+      gradient_accumulation_steps=_GRADIENT_ACCUMULATION_STEPS.value,
+      enable_gradient_checkpointing=_ENABLE_GRADIENT_CHECKPOINTING.value,
+      enable_peft=_ENABLE_PEFT.value,
+      template=_TEMPLATE.value,
+      lr_scheduler_type=_LR_SCHEDULER_TYPE.value,
+      save_steps=_SAVE_STEPS.value,
+      logging_steps=_LOGGING_STEPS.value,
+      train_split_name=_TRAIN_SPLIT_NAME.value,
+      eval_config=eval_config,
+      report_to=_REPORT_TO.value,
+      access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
+      train_precision=_TRAIN_PRECISION.value,
+      use_example_packing=_USE_EXAMPLE_PACKING.value,
+      attn_implementation=_ATTN_IMPLEMENTATION.value,
+      max_grad_norm=_MAX_GRAD_NORM.value,
+      completion_only=_COMPLETION_ONLY.value,
+      logger_level=_LOGGER_LEVEL.value,
+      benchmark_out_file=_BENCHMARK_OUT_FILE.value,
+      tuning_data_stats_file=_TUNING_DATA_STATS_FILE.value,
+      target_modules=_TARGET_MODULES.value,
+  )
+  # Frees the model from GPU.
+  utils.force_gc()
+
+  output_dir.upload_to_gcs(skip_if_exists=True)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/merge_causal_language_model_lora.py

@@ -0,0 +1,132 @@
+"""Script to merge PEFT adapter with base model."""
+
+from typing import Any, Dict, Sequence
+
+from absl import app
+from absl import flags
+
+from util import dataset_validation_util
+from vertex_vision_model_garden_peft.train.vmg import utils
+from util import constants
+from util import fileutils
+
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    'pretrained_model_id',
+    None,
+    'The pretrained model id. Supported models can be causal language modeling'
+    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
+    ' might be different paddings for different models. This tool assumes the'
+    ' pretrained_model_id contains model name, and then choose proper padding'
+    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    required=True,
+)
+
+_MERGE_BASE_AND_LORA_OUTPUT_DIR = flags.DEFINE_string(
+    'merge_base_and_lora_output_dir',
+    None,
+    'The directory to store the merged model with the base and lora adapter.',
+)
+
+_MERGE_MODEL_PRECISION_MODE = flags.DEFINE_enum(
+    'merge_model_precision_mode',
+    constants.PRECISION_MODE_16,
+    [
+        constants.PRECISION_MODE_4,
+        constants.PRECISION_MODE_8,
+        constants.PRECISION_MODE_FP8,
+        constants.PRECISION_MODE_16,
+        constants.PRECISION_MODE_16B,
+        constants.PRECISION_MODE_32,
+    ],
+    'Merging model precision mode.',
+)
+
+_FINETUNED_LORA_MODEL_DIR = flags.DEFINE_string(
+    'finetuned_lora_model_dir',
+    None,
+    'The directory storing finetuned LoRA model weights.',
+)
+
+_RESTRICT_MODEL_UPLOAD_DOCKER_URI = flags.DEFINE_string(
+    'restrict_model_upload_docker_uri',
+    '',
+    'If set, mark output model as only uploadable to Model Registry with the'
+    ' specified Docker URI.',
+)
+
+_EXECUTOR_INPUT = flags.DEFINE_string(
+    'executor_input',
+    '',
+    'For internal use. Kubeflow pipeline context when running trainer as part'
+    ' of an internal pipeline.',
+)
+
+_HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
+    'huggingface_access_token',
+    None,
+    'The access token for loading huggingface gated models.',
+)
+
+
+@flags.multi_flags_validator(
+    [
+        _PRETRAINED_MODEL_ID.name,
+        _FINETUNED_LORA_MODEL_DIR.name,
+        _MERGE_BASE_AND_LORA_OUTPUT_DIR.name,
+    ],
+)
+def check_merge_lora_model_flags(flags_dict: Dict[str, Any]) -> bool:
+  """Check if required flags are set on merge model LoRA task.
+
+  Args:
+    flags_dict: Dictionary containing task and flags to check.
+
+  Returns:
+    If required flags are not None.
+  """
+  return all(map(lambda x: x is not None, flags_dict.values()))
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
+  if dataset_validation_util.is_gcs_path(pretrained_model_id):
+    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
+        pretrained_model_id
+    )
+
+  finetuned_lora_model_dir = utils.GcsOrLocalDirectory(
+      _FINETUNED_LORA_MODEL_DIR.value
+  )
+
+  merge_base_and_lora_output_dir = utils.GcsOrLocalDirectory(
+      _MERGE_BASE_AND_LORA_OUTPUT_DIR.value
+  )
+
+  utils.merge_causal_language_model_with_lora(
+      pretrained_model_id=pretrained_model_id,
+      precision_mode=_MERGE_MODEL_PRECISION_MODE.value,
+      finetuned_lora_model_dir=finetuned_lora_model_dir.local_dir,
+      merged_model_output_dir=merge_base_and_lora_output_dir.local_dir,
+      access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
+  )
+
+  if _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value:
+    utils.write_first_party_model_metadata(
+        merge_base_and_lora_output_dir.local_dir,
+        _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value,
+    )
+
+  if _EXECUTOR_INPUT.value:
+    utils.write_kfp_outputs(
+        _EXECUTOR_INPUT.value,
+        {
+            'saved_model': _MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
+        },
+    )
+
+  merge_base_and_lora_output_dir.upload_to_gcs(skip_if_exists=True)
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/quantize_model.py

@@ -0,0 +1,349 @@
+"""Quantizes the model."""
+
+import json
+import os
+from typing import Any, Dict, List, Sequence, Union
+
+from absl import app
+from absl import flags
+from absl import logging
+from auto_gptq import AutoGPTQForCausalLM
+from auto_gptq import BaseQuantizeConfig
+from awq import AutoAWQForCausalLM
+from optimum.gptq.data import get_dataset
+from transformers import AutoTokenizer
+
+from util import dataset_validation_util
+from vertex_vision_model_garden_peft.train.vmg import utils
+from util import constants
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    'pretrained_model_id',
+    None,
+    'The pretrained model id. Supported models can be causal language modeling'
+    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
+    ' might be different paddings for different models. This tool assumes the'
+    ' pretrained_model_id contains model name, and then choose proper padding'
+    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+)
+
+_QUANTIZATION_METHOD = flags.DEFINE_enum(
+    'quantization_method',
+    None,
+    [constants.GPTQ, constants.AWQ],
+    'The quantization method. Choose from ["gtpq", "awq"].',
+)
+
+_QUANTIZATION_PRECISION_MODE = flags.DEFINE_enum(
+    'quantization_precision_mode',
+    constants.PRECISION_MODE_4,
+    [
+        constants.PRECISION_MODE_8,
+        constants.PRECISION_MODE_4,
+        constants.PRECISION_MODE_3,
+        constants.PRECISION_MODE_2,
+    ],
+    'Quantization precision mode.',
+)
+
+_QUANTIZATION_DATASET_NAME = flags.DEFINE_string(
+    'quantization_dataset_name',
+    None,
+    'The dataset used for quantization. You can provide your own dataset in a'
+    ' list of string or just use the original datasets used in GPTQ paper'
+    ' ["wikitext2","c4","c4-new","ptb","ptb-new"] for GPTQ quantization. Using'
+    " a dataset more appropriate to the model's training can improve"
+    ' quantisation accuracy. Note that the GPTQ dataset is not the same as the'
+    ' dataset used to train the model.',
+)
+
+_TEXT_COLUMN_IN_QUANTIZATION_DATASET = flags.DEFINE_string(
+    'text_column_in_quantization_dataset',
+    constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
+    'The text column in quantization dataset.',
+)
+
+_QUANTIZATION_OUTPUT_DIR = flags.DEFINE_string(
+    'quantization_output_dir',
+    None,
+    'The directory to store the quantized model.',
+)
+
+_QUANTIZATION_DEVICE_MAP = flags.DEFINE_string(
+    'device_map', None, 'The device map.'
+)
+
+_QUANTIZATION_MAX_MEMORY = flags.DEFINE_string(
+    'max_memory', None, 'The maximum memory.'
+)
+
+_GROUP_SIZE = flags.DEFINE_integer(
+    'group_size',
+    None,
+    'The group size to use for quantization. Recommended value is 128 and -1'
+    ' uses per-column quantization. Higher numbers use less VRAM, but have'
+    ' lower quantisation accuracy. "None" is the lowest possible value.',
+)
+
+_DESC_ACT = flags.DEFINE_boolean(
+    'desc_act',
+    False,
+    'Whether to quantize columns in order of decreasing activation size.'
+    ' Setting it to False can significantly speed up inference but the'
+    ' perplexity may become slightly worse. Also known as act-order.',
+)
+
+_DAMP_PERCENT = flags.DEFINE_float(
+    'damp_percent',
+    0.1,
+    'The percent of the average Hessian diagonal to use for dampening.',
+)
+
+_CACHE_EXAMPLES_ON_GPU = flags.DEFINE_boolean(
+    'cache_examples_on_gpu',
+    True,
+    'Whether to cache the examples on GPU. Disabling will reduce VRAM usage,'
+    ' but increase quantization time.',
+)
+
+_AWQ_VERSION = flags.DEFINE_enum(
+    'awq_version',
+    constants.GEMM,
+    [constants.GEMM, constants.GEMV],
+    'The version of the AWQ to use. It determines how matrix multiplication'
+    ' runs under the hood. GEMV is 20% faster than GEMM, only at batch size 1'
+    ' (not good for large contexts). GEMM is much faster than FP16 at batch'
+    ' sizes below 8 (good with large contexts).',
+)
+
+
+@flags.multi_flags_validator(
+    [
+        _PRETRAINED_MODEL_ID.name,
+        _QUANTIZATION_METHOD.name,
+        _QUANTIZATION_PRECISION_MODE.name,
+        _QUANTIZATION_DATASET_NAME.name,
+        _QUANTIZATION_OUTPUT_DIR.name,
+    ],
+)
+def check_quantization_flags(flags_dict: Dict[str, Any]) -> bool:
+  """Check if required flags are set on quantization task.
+
+  Args:
+    flags_dict: Dictionary containing task and flags to check.
+
+  Returns:
+    If required flags are not None.
+  """
+  required_flags = [
+      _QUANTIZATION_METHOD.name,
+      _PRETRAINED_MODEL_ID.name,
+      _QUANTIZATION_PRECISION_MODE.name,
+      _QUANTIZATION_DATASET_NAME.name,
+      _QUANTIZATION_OUTPUT_DIR.name,
+  ]
+
+  return all(map(lambda x: flags_dict[x] is not None, required_flags))
+
+
+def quantize_model(
+    quantization_method: str,
+    pretrained_model_id: str,
+    quantization_output_dir: str,
+    quantization_precision_mode: str = None,
+    quantization_dataset_name: Union[List[str]] = None,
+    text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
+    group_size: int = None,
+    desc_act: bool = True,
+    damp_percent: float = 0.1,
+    awq_version: str = 'GEMM',
+    device_map: str = None,
+    max_memory: Dict[Any, str] = None,
+    cache_examples_on_gpu: bool = True,
+) -> None:
+  """Quantizes the model using `quantization_method`."""
+  if quantization_method == constants.GPTQ:
+    gptq_quantize_model(
+        pretrained_model_id=pretrained_model_id,
+        gptq_output_dir=quantization_output_dir,
+        gptq_precision_mode=quantization_precision_mode,
+        gptq_dataset_name=quantization_dataset_name,
+        group_size=group_size,
+        desc_act=desc_act,
+        damp_percent=damp_percent,
+        cache_examples_on_gpu=cache_examples_on_gpu,
+    )
+  elif quantization_method == constants.AWQ:
+    awq_quantize_model(
+        pretrained_model_id=pretrained_model_id,
+        quantization_output_dir=quantization_output_dir,
+        quantization_precision_mode=quantization_precision_mode,
+        quantization_dataset_name=quantization_dataset_name,
+        text_column_in_quantization_dataset=text_column_in_quantization_dataset,
+        group_size=group_size,
+        awq_version=awq_version,
+        device_map=device_map,
+        max_memory=max_memory,
+    )
+
+
+def awq_quantize_model(
+    pretrained_model_id: str,
+    quantization_output_dir: str,
+    quantization_precision_mode: str = None,
+    quantization_dataset_name: Union[List[str]] = None,
+    text_column_in_quantization_dataset: str = constants.DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET,
+    group_size: int = None,
+    awq_version: str = 'GEMM',
+    device_map: str = None,
+    max_memory: Dict[Any, str] = None,
+) -> None:
+  """Quantizes the model using AWQ."""
+  if quantization_precision_mode != constants.PRECISION_MODE_4:
+    raise ValueError(
+        f'Invalid precision mode: {quantization_precision_mode} for AWQ. 4bit'
+        ' quantization must be used.'
+    )
+  else:
+    bits = 4
+  if not group_size:
+    group_size = 128
+  if not device_map:
+    device_map = 'cpu'
+  if dataset_validation_util.is_gcs_path(quantization_dataset_name):
+    logging.info('Using custom dataset: %s', quantization_dataset_name)
+    with open(
+        dataset_validation_util.force_gcs_fuse_path(quantization_dataset_name),
+        'r',
+    ) as f:
+      quantization_dataset = [line.rstrip('\n') for line in f]
+  else:
+    quantization_dataset = quantization_dataset_name
+  quant_config = {
+      'zero_point': True,
+      'q_group_size': group_size,
+      'w_bit': bits,
+      'version': awq_version,
+  }
+  logging.info('Quantization config: %s', quant_config)
+  model = AutoAWQForCausalLM.from_pretrained(
+      pretrained_model_id,
+      trust_remote_code=True,
+      device_map=device_map,
+      max_memory=max_memory,
+      low_cpu_mem_usage=True,
+  )
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, trust_remote_code=True
+  )
+  model.quantize(
+      tokenizer,
+      quant_config=quant_config,
+      calib_data=quantization_dataset,
+      text_column=text_column_in_quantization_dataset,
+  )
+  model.save_quantized(quantization_output_dir)
+  tokenizer.save_pretrained(quantization_output_dir)
+
+
+def gptq_quantize_model(
+    pretrained_model_id: str,
+    gptq_output_dir: str,
+    gptq_precision_mode: str = None,
+    gptq_dataset_name: Union[List[str]] = None,
+    group_size: int = -1,
+    desc_act: bool = False,
+    damp_percent: float = 0.1,
+    cache_examples_on_gpu: bool = True,
+) -> None:
+  """Quantizes the model using GPTQ."""
+  logging.info(
+      'PYTORCH_CUDA_ALLOC_CONF: %s',
+      os.environ.get('PYTORCH_CUDA_ALLOC_CONF', ''),
+  )
+  if dataset_validation_util.is_gcs_path(gptq_dataset_name):
+    logging.info('Using custom dataset: %s', gptq_dataset_name)
+    with open(
+        dataset_validation_util.force_gcs_fuse_path(gptq_dataset_name), 'r'
+    ) as f:
+      gptq_dataset = [line.rstrip('\n') for line in f]
+  else:
+    gptq_dataset = gptq_dataset_name
+  if gptq_precision_mode == constants.PRECISION_MODE_8:
+    bits = 8
+  elif gptq_precision_mode == constants.PRECISION_MODE_4:
+    bits = 4
+  elif gptq_precision_mode == constants.PRECISION_MODE_3:
+    bits = 3
+  elif gptq_precision_mode == constants.PRECISION_MODE_2:
+    bits = 2
+  else:
+    raise ValueError(f'Invalid precision mode: {gptq_precision_mode} for GPTQ.')
+  if not group_size:
+    group_size = -1
+
+  tokenizer = AutoTokenizer.from_pretrained(pretrained_model_id)
+  gptq_dataset = get_dataset(gptq_dataset, tokenizer)
+
+  quantization_config = BaseQuantizeConfig(
+      bits=bits,
+      group_size=group_size,
+      damp_percent=damp_percent,
+      desc_act=desc_act,
+  )
+
+  logging.info('Quantization config: %s', quantization_config.to_dict())
+
+  model = AutoGPTQForCausalLM.from_pretrained(
+      pretrained_model_id,
+      quantization_config,
+      low_cpu_mem_usage=True,
+      torch_dtype='auto',
+      trust_remote_code=True,
+  )
+  model.quantize(
+      examples=gptq_dataset,
+      cache_examples_on_gpu=cache_examples_on_gpu,
+  )
+
+  if utils.should_add_pad_token(pretrained_model_id):
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    model.resize_token_embeddings(len(tokenizer))
+  model.save_pretrained(gptq_output_dir)
+  tokenizer.save_pretrained(gptq_output_dir)
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  pretrained_model_id = _PRETRAINED_MODEL_ID.value
+  if dataset_validation_util.is_gcs_path(pretrained_model_id):
+    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
+        pretrained_model_id
+    )
+  pretrained_model_id = dataset_validation_util.force_gcs_fuse_path(
+      pretrained_model_id
+  )
+
+  if _QUANTIZATION_MAX_MEMORY.value:
+    max_memory = json.loads(_QUANTIZATION_MAX_MEMORY.value)
+  else:
+    max_memory = None
+
+  quantize_model(
+      quantization_method=_QUANTIZATION_METHOD.value,
+      pretrained_model_id=pretrained_model_id,
+      quantization_output_dir=_QUANTIZATION_OUTPUT_DIR.value,
+      quantization_precision_mode=_QUANTIZATION_PRECISION_MODE.value,
+      quantization_dataset_name=_QUANTIZATION_DATASET_NAME.value,
+      text_column_in_quantization_dataset=_TEXT_COLUMN_IN_QUANTIZATION_DATASET.value,
+      group_size=_GROUP_SIZE.value,
+      desc_act=_DESC_ACT.value,
+      damp_percent=_DAMP_PERCENT.value,
+      awq_version=_AWQ_VERSION.value,
+      device_map=_QUANTIZATION_DEVICE_MAP.value,
+      max_memory=max_memory,
+      cache_examples_on_gpu=_CACHE_EXAMPLES_ON_GPU.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/scripts/build_train_docker.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+# Run copybara first:
+# cloud/ml/applications/vision/model_garden/copybara/run_copybara_local.sh
+# Run docker build:
+# cloud/ml/applications/vision/model_garden/model_oss/peft/train/vmg/scripts/build_train_docker.sh
+
+set -x
+set -e
+
+COPYBARA_DIR="/tmp/train_docker/"
+
+pushd "${COPYBARA_DIR}"
+
+PROJECT="cloud-nas-260507"
+IMAGE_TAG="gcr.io/${PROJECT}/pytorch-peft-train:${USER}-test"
+
+docker build -f model_oss/peft/train/vmg/dockerfile/train.Dockerfile . -t "${IMAGE_TAG}"
+docker push "${IMAGE_TAG}"
+
+popd

community-content/vertex_model_garden/model_oss/peft/train/vmg/sequence_classification_lora.py

@@ -1,7 +1,9 @@
 """Sequence classification with LoRA models."""
 
-# pylint: disable=g-importing-member
+from typing import Sequence
 
+from absl import app
+from absl import flags
 from datasets import load_dataset
 import evaluate
 from peft import get_peft_model
@@ -14,6 +16,71 @@ from transformers import AutoModelForSequenceClassification
 from transformers import AutoTokenizer
 from transformers import get_linear_schedule_with_warmup
 
+from util import dataset_validation_util
+
+
+_PRETRAINED_MODEL_ID = flags.DEFINE_string(
+    "pretrained_model_id",
+    None,
+    "The pretrained model id. Supported models can be causal language modeling"
+    " models from https://github.com/huggingface/peft/tree/main. Note, there"
+    " might be different paddings for different models. This tool assumes the"
+    " pretrained_model_id contains model name, and then choose proper padding"
+    " methods. e.g. it must contain `llama` for `Llama2 models`.",
+)
+
+_OUTPUT_DIR = flags.DEFINE_string(
+    "output_dir",
+    None,
+    "The output directory.",
+)
+
+_DATASET_NAME = flags.DEFINE_string(
+    "dataset_name",
+    None,
+    "The dataset name in huggingface.",
+)
+
+_LORA_RANK = flags.DEFINE_integer(
+    "lora_rank",
+    16,
+    "The rank of the update matrices, expressed in int. Lower rank results in"
+    " smaller update matrices with fewer trainable parameters, referring to"
+    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
+)
+
+_LORA_ALPHA = flags.DEFINE_integer(
+    "lora_alpha",
+    32,
+    "LoRA scaling factor, referring to"
+    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
+)
+
+_LORA_DROPOUT = flags.DEFINE_float(
+    "lora_dropout",
+    0.05,
+    "dropout probability of the LoRA layers, referring to"
+    " https://huggingface.co/docs/peft/task_guides/token-classification-lora.",
+)
+
+_NUM_EPOCHS = flags.DEFINE_integer(
+    "num_epochs",
+    None,
+    "The number of training epochs.",
+)
+
+_BATCH_SIZE = flags.DEFINE_integer(
+    "batch_size",
+    32,
+    "The batch size.",
+)
+
+_LEARNING_RATE = flags.DEFINE_float(
+    "learning_rate",
+    2e-4,
+    "The learning rate after the potential warmup period.",
+)
+
 
 def finetune_sequence_classification(
     pretrained_model_id: str,
@@ -131,3 +198,32 @@ def finetune_sequence_classification(
     print(f"epoch {epoch}:", eval_metric)
 
   model.save_pretrained(output_dir)
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  if dataset_validation_util.is_gcs_path(_PRETRAINED_MODEL_ID.value):
+    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
+        _PRETRAINED_MODEL_ID.value
+    )
+  else:
+    pretrained_model_id = _PRETRAINED_MODEL_ID.value
+  pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
+      pretrained_model_id
+  )
+  output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
+
+  finetune_sequence_classification(
+      pretrained_model_id=pretrained_model_path,
+      dataset_name=_DATASET_NAME.value,
+      output_dir=output_dir,
+      lora_rank=_LORA_RANK.value,
+      lora_alpha=_LORA_ALPHA.value,
+      lora_dropout=_LORA_DROPOUT.value,
+      num_epochs=int(_NUM_EPOCHS.value),
+      batch_size=_BATCH_SIZE.value,
+      learning_rate=_LEARNING_RATE.value,
+  )
+
+
+if __name__ == "__main__":
+  app.run(main)

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/README.md

@@ -0,0 +1,98 @@
+# Vertex Model Garden Training Dataset Template
+
+## Overview
+
+Vertex Model Garden training provides templates for streamlined preprocessing of
+datasets. Although datasets often have intricate structures, the supported LLM
+models accept only flat strings. A template facilitates parsing a dataset and
+preprocessing it to be compatible with the model.
+
+When fine-tuning a pretrained model, it is advisable to maintain the same format
+as the original training data. A template helps replicate the format, ensuring
+consistency and potentially enhancing the fine-tuning process.
+
+Both multi-turn messages and single instruction-response pairs are supported.
+Multi-turn messages are accommodated using a more general `chat_template` field,
+whereas simple instruction-response pair datasets are supported through the
+`prompt_input` field.
+
+A template is a JSON file consisting of string key-value pairs. Refer to the
+following for the definitions of the supported fields.
+
+## Template field documentation
+
+**description**: An explanation of the template.
+
+**source**: Information about the origin of the template.
+
+**chat_template**: A
+[jinja template](https://jinja.palletsprojects.com/en/3.1.x/templates/) that can
+be used to parse a chat dataset. This is the same format as
+[HF chat templates](https://huggingface.co/docs/transformers/main/en/chat_templating).
+To create a chat_template, use the `messages` variable to be filled with the
+sample. The flag `--instruct_column_in_dataset` identifies which column will be
+passed to the `messages` variable in the chat_template. This field is mutually
+exclusive with `prompt_input` and `prompt_no_input`.
+
+**prompt_input**: A string template that is used when value for the input column
+exists in the sample. It should be able to be formatted with the
+[str.format](https://docs.python.org/3/library/stdtypes.html#str.format) method.
+The input column is specified with the flag `--instruct_column_in_dataset`. Used
+for instruction dataset. This field is mutually exclusive with `chat_template`.
+
+**prompt_no_input**: A string template that is used when value for the input
+column does not exist in the sample. It should be able to be formatted with the
+[str.format](https://docs.python.org/3/library/stdtypes.html#str.format) method.
+The input column is specified with the flag `--instruct_column_in_dataset`. Used
+for instruction dataset. This field is mutually exclusive with `chat_template`.
+
+**instruction_separator**: A unique string used to indicate the start of the
+instructions. If not specified, every token after response_separator will be
+treated as a response, and every token before the first response_separator will
+be treated as instruction.
+
+**response_separator**: A unique string used to indicate the start of the
+response. This field is required if `--completion_only` flag is set to `True`.
+
+## Example templates
+
+-   See the list of all supported templates [here](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content/vertex_model_garden/model_oss/peft/train/vmg/templates).
+-   For an example with `chat_template` see the JSON template below.
+
+```
+{
+  "description": "Chat template used by Llama 3.",
+  "source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
+  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
+}
+```
+
+-   For an example with `prompt_input` see the JSON template below. In this case
+    the flag `--instruct_column_in_dataset=text` should be set, and there must
+    be a column named `text` in the dataset.
+
+```
+{
+  "description": "Template for openassistant-guanaco dataset.",
+  "source": "https://huggingface.co/datasets/timdettmers/openassistant-guanaco",
+  "prompt_input": "{text}",
+  "instruction_separator": "### Human:",
+  "response_separator": "### Assistant:"
+}
+```
+
+-   For an example with `prompt_no_input` see the JSON template below. In this
+    case the flag `--instruct_column_in_dataset=input` should be set, and there
+    must be columns named `input` and `instruction` in the dataset.
+
+```
+{
+  "description": "Template used by Alpaca-LoRA.",
+  "source": "https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json",
+  "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+  "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
+  "response_separator": "### Response:"
+}
+```

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/alpaca.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template used by Alpaca-LoRA.",
+  "source": "https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json",
+  "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+  "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
+  "response_separator": "### Response:"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/alpaca_short.json

@@ -0,0 +1,7 @@
+{
+  "description": "A shorter template to experiment with.",
+  "source": "https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca_short.json",
+  "prompt_input": "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+  "prompt_no_input": "### Instruction:\n{instruction}\n\n### Response:\n",
+  "response_separator": "### Response:"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/gemma.json

@@ -0,0 +1,7 @@
+{
+  "description": "Chat template used by Gemma. 'assistant' role is replaced by 'model'",
+  "source": "https://huggingface.co/google/gemma-1.1-2b-it/blob/bf4924f313df5166dee1467161e886e55f2eb4d4/tokenizer_config.json#L1507",
+  "chat_template": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}",
+  "instruction_separator": "<start_of_turn>user\n",
+  "response_separator": "<start_of_turn>model\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/llama3-text-bison.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template used by Llama 3, accepting text-bison format.",
+  "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format",
+  "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|>",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/llama3.json

@@ -0,0 +1,7 @@
+{
+  "description": "Chat template used by Llama 3.",
+  "source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
+  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/mistral.json

@@ -0,0 +1,7 @@
+{
+  "description": "Chat template used by Mistral.",
+  "source": "https://github.com/OpenAccess-AI-Collective/axolotl/blob/main/src/axolotl/utils/chat_templates.py",
+  "chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}",
+  "instruction_separator": "[INST]",
+  "response_separator": "[/INST]"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/openai-completion.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template used by openai completion.",
+  "source": "https://platform.openai.com/docs/api-reference/fine-tuning/completions-input",
+  "prompt_input": "\n\n<|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|>\n\n<|start_header_id|>assistant<|end_header_id|>\n\n{completion}<|eot_id|>",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>assistant<|end_header_id|>\n\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/openassistant-guanaco.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template for openassistant-guanaco dataset.",
+  "source": "https://huggingface.co/datasets/timdettmers/openassistant-guanaco",
+  "prompt_input": "{text}",
+  "instruction_separator": "### Human:",
+  "response_separator": "### Assistant:"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/templates/vertex-chat.json

@@ -0,0 +1,7 @@
+{
+  "description": "Template used for chat based models.",
+  "source": "https://huggingface.co/meta-llama/Meta-Llama-3-70B-Instruct/blob/a5a71a7527eac1d651bb145436c72026887fb68e/tokenizer_config.json#L2053",
+  "chat_template": "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>model<|end_header_id|>\n\n' }}{% endif %}",
+  "instruction_separator": "<|start_header_id|>user<|end_header_id|>\n\n",
+  "response_separator": "<|start_header_id|>model<|end_header_id|>\n\n"
+}

community-content/vertex_model_garden/model_oss/peft/train/vmg/train_entrypoint.py

@@ -0,0 +1,242 @@
+"""Entrypoint for peft train docker.
+
+Dispatches to different scripts based on `task` type.
+
+For task type in `_TASK_TO_SCRIPT`, if `--config_file` is specified, the script
+will dispatch the call to `accelerate`, which is friendly for multi-GPU
+environment. Otherwise, `python3` is used.
+"""
+
+import argparse
+import json
+import os
+import subprocess
+from typing import List, Optional, Sequence
+from absl import app
+from absl import flags
+from absl import logging
+from util import dataset_validation_util
+from vertex_vision_model_garden_peft.train.vmg import utils
+from util import constants
+from util import hypertune_utils
+
+
+_TEXT_TO_IMAGE_TASKS_SCRIPTS = {
+    constants.TEXT_TO_IMAGE: 'text_to_image/train_text_to_image.py',
+    constants.TEXT_TO_IMAGE_LORA: 'text_to_image/train_text_to_image_lora.py',
+    constants.TEXT_TO_IMAGE_DREAMBOOTH: 'dreambooth/train_dreambooth.py',
+    constants.TEXT_TO_IMAGE_DREAMBOOTH_LORA: (
+        'dreambooth/train_dreambooth_lora.py'
+    ),
+    constants.TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL: (
+        'dreambooth/train_dreambooth_lora_sdxl.py'
+    ),
+}
+
+_TASK_TO_SCRIPT = {
+    constants.INSTRUCT_LORA: (
+        'vertex_vision_model_garden_peft/train/vmg/instruct_lora.py'
+    ),
+    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: 'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py',
+    constants.QUANTIZE_MODEL: (
+        'vertex_vision_model_garden_peft/train/vmg/quantize_model.py'
+    ),
+    constants.SEQUENCE_CLASSIFICATION_LORA: 'vertex_vision_model_garden_peft/train/vmg/sequence_classification_lora.py',
+    constants.VALIDATE_DATASET_WITH_TEMPLATE: 'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py',
+}
+
+
+def launch_script_cmd(
+    script: str,
+    config_file: Optional[str],
+    accelerate_args: argparse.Namespace = argparse.Namespace(),
+) -> List[str]:
+  """Returns the command to launch the script."""
+  if config_file:
+    cmd = [
+        'accelerate',
+        'launch',
+        '--config_file={}'.format(config_file),
+    ]
+  else:
+    cmd = ['python3']
+
+  _append_args_to_command_in_place(accelerate_args, cmd)
+  cmd.append(script)
+
+  return cmd
+
+
+def _get_accelerate_args() -> argparse.Namespace:
+  """Returns the accelerate args."""
+  # For the format of the cluster spec, see
+  # https://cloud.google.com/vertex-ai/docs/training/distributed-training#cluster-spec-format # pylint: disable=line-too-long
+  cluster_spec = os.getenv('CLUSTER_SPEC', default=None)
+  if not cluster_spec:
+    return argparse.Namespace()
+  logging.info('CLUSTER_SPEC: %s', cluster_spec)
+
+  cluster_data = json.loads(cluster_spec)
+  if (
+      'workerpool1' not in cluster_data['cluster']
+      or not cluster_data['cluster']['workerpool1']
+  ):
+    return argparse.Namespace()
+
+  # Get primary node info
+  primary_node = cluster_data['cluster']['workerpool0'][0]
+  logging.info('primary node: %s', primary_node)
+  primary_node_addr, primary_node_port = primary_node.split(':')
+  logging.info('primary node address: %s', primary_node_addr)
+  logging.info('primary node port: %s', primary_node_port)
+
+  # Determine node rank of this machine
+  workerpool = cluster_data['task']['type']
+  if workerpool == 'workerpool0':
+    node_rank = 0
+  elif workerpool == 'workerpool1':
+    # Add 1 for the primary node, since `index` is the index of workerpool1.
+    node_rank = cluster_data['task']['index'] + 1
+  else:
+    raise ValueError(
+        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
+        f' {workerpool}'
+    )
+  logging.info('node rank: %s', node_rank)
+
+  # Calculate total nodes
+  num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
+  num_nodes = num_worker_nodes + 1  # Add 1 for the primary node
+  logging.info('num nodes: %s', num_nodes)
+
+  accelerate_args = argparse.Namespace()
+  accelerate_args.machine_rank = node_rank
+  accelerate_args.num_machines = num_nodes
+  accelerate_args.main_process_ip = primary_node_addr
+  accelerate_args.main_process_port = primary_node_port
+  accelerate_args.max_restarts = 0
+  accelerate_args.monitor_interval = 120
+
+  return accelerate_args
+
+
+def _append_args_to_command_in_place(
+    args: argparse.Namespace, command: List[str]
+):
+  for key, value in vars(args).items():
+    # If not specified, skip.
+    if value is not None:
+      command.append(f'--{key}={value}')
+
+
+def _get_train_cmd_and_maybe_merge_cmd(
+    task: str, config_file: str, unknown: Sequence[str]
+) -> Sequence[Sequence[str]]:
+  """Returns the training command and maybe the merge command if applicable."""
+
+  # Only populated when multi-node is used.
+  accelerate_args = _get_accelerate_args()
+  training_cmd = launch_script_cmd(
+      _TASK_TO_SCRIPT[task],
+      config_file,
+      accelerate_args=accelerate_args,
+  )
+
+  # Training only flag.
+  train_parser = argparse.ArgumentParser()
+  train_parser.add_argument('--output_dir', required=True)
+  training_args, unknown = train_parser.parse_known_args(unknown)
+  # Checks for `hypertune_utils._ENVIRONMENT_VARIABLE_FOR_TRIAL_ID` env var and
+  # appends the trial id if it exists.
+  training_args.output_dir = hypertune_utils.maybe_append_trial_id(
+      dataset_validation_util.force_gcs_fuse_path(training_args.output_dir)
+  )
+
+  # Merge only flags.
+  merge_parser = argparse.ArgumentParser()
+  merge_parser.add_argument('--merge_model_precision_mode')
+  merge_parser.add_argument('--executor_input')
+  merge_parser.add_argument('--restrict_model_upload_docker_uri')
+  merge_parser.add_argument('--merge_base_and_lora_output_dir')
+  merge_args, unknown = merge_parser.parse_known_args(unknown)
+
+  # Common flags shared by merging and training.
+  common_parser = argparse.ArgumentParser()
+  common_parser.add_argument('--pretrained_model_id', required=True)
+  common_parser.add_argument('--huggingface_access_token')
+  common_args, remaining = common_parser.parse_known_args(unknown)
+
+  # Add flags for training.
+  _append_args_to_command_in_place(training_args, training_cmd)
+  _append_args_to_command_in_place(common_args, training_cmd)
+  training_cmd.extend(remaining)  # Remaining args are passed to training cmd.
+  commands = [training_cmd]
+
+  # Only the main node runs merging.
+  if (
+      merge_args.merge_base_and_lora_output_dir
+      and getattr(accelerate_args, 'machine_rank', 0) == 0
+  ):
+    lora_dir = utils.get_final_checkpoint_path(training_args.output_dir)
+
+    merge_cmd = [
+        'WORLD_SIZE=1',  # To ignore other nodes in multi-node setting.
+        'python3',
+        _TASK_TO_SCRIPT[constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA],
+        f'--finetuned_lora_model_dir={lora_dir}',
+    ]
+    _append_args_to_command_in_place(merge_args, merge_cmd)
+    _append_args_to_command_in_place(common_args, merge_cmd)
+
+    # Run in a conda environment.
+    conda_run_cmd = [
+        '/bin/bash',
+        '-c',
+        f'conda run -n merge {" ".join(merge_cmd)}',
+    ]
+    commands.append(conda_run_cmd)
+
+  return commands
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--config_file')
+  parser.add_argument('--task')
+  args, unknown = parser.parse_known_args()
+
+  task = args.task
+
+  if task in _TEXT_TO_IMAGE_TASKS_SCRIPTS:
+    # Setup accelerate config before running trainer.
+    config_gen_cmd = [
+        'python',
+        '-c',
+        (
+            'from accelerate.utils import write_basic_config;'
+            ' write_basic_config(mixed_precision="fp16")'
+        ),
+    ]
+    task_cmd = [
+        'accelerate',
+        'launch',
+        _TEXT_TO_IMAGE_TASKS_SCRIPTS[task],
+    ] + list(map(dataset_validation_util.force_gcs_fuse_path, unknown))
+    commands = [config_gen_cmd, task_cmd]
+  elif task in [constants.INSTRUCT_LORA]:
+    commands = _get_train_cmd_and_maybe_merge_cmd(
+        task=task, config_file=args.config_file, unknown=unknown
+    )
+  else:
+    assert task in _TASK_TO_SCRIPT
+    cmd = launch_script_cmd(_TASK_TO_SCRIPT[task], args.config_file)
+    cmd.extend(unknown)
+    commands = [cmd]
+
+  for cmd in commands:
+    logging.info('launching task=%s with cmd: \n%s', task, ' \\\n'.join(cmd))
+    subprocess.run(cmd, check=True)
+
+
+if __name__ == '__main__':
+  app.run(main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True))

community-content/vertex_model_garden/model_oss/peft/train/vmg/utils.py

@@ -0,0 +1,742 @@
+"""Common libraries for PEFT."""
+
+import dataclasses
+import datetime
+import gc
+import multiprocessing as mp
+import os
+import subprocess
+from typing import Any, Dict, Optional, Sequence
+
+from absl import logging
+import accelerate
+from accelerate import DistributedType
+from accelerate import PartialState
+from google.protobuf import json_format
+from kfp.pipeline_spec import pipeline_spec_pb2
+import numpy as np
+import peft
+from peft import PeftModel
+from peft import prepare_model_for_kbit_training
+import pynvml
+import torch
+import transformers
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+from transformers import BitsAndBytesConfig
+from transformers import FbgemmFp8Config
+from transformers.integrations import is_deepspeed_zero3_enabled
+import trl
+
+from util import dataset_validation_util
+from util import constants
+from util import fileutils
+
+
+_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
+_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
+_LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
+_LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
+
+
+
+class GcsOrLocalDirectory(os.PathLike):
+  """A class to represent a directory with upload support if GCS path is given.
+
+  This class is used to represent a directory. It can be used for a temporary
+  local directory and for uploading files to the GCS directory later if the
+  given path is a GCS directory. If the given path is a local directory, a call
+  to gcs_dir attribute will raise an error. This class has multi-node and
+  multi-process support with accelerate.
+
+  Attributes:
+    local_dir: The local directory to store the files.
+    gcs_dir: The path to the GCS directory.
+  """
+
+  def __init__(
+      self,
+      path: str,
+      check_empty: bool = False,
+      upload_from_all_nodes: bool = False,
+  ):
+    """Initializes the GcsOrLocalDirectory.
+
+    Args:
+      path: The path to the directory.
+      check_empty: If True, check if the GCS directory is empty. No-op for local
+        directory.
+      upload_from_all_nodes: If True, upload the local directory to GCS from all
+        nodes.
+    """
+    if len(path) > 1:
+      path = path.rstrip("/")
+
+    self._upload_from_all_nodes = upload_from_all_nodes
+
+    if path.startswith(constants.GCS_URI_PREFIX) or path.startswith(
+        constants.GCSFUSE_URI_PREFIX
+    ):
+      self._is_gcs_path = True
+      self._local_dir = _get_local_dir_from_gcs_dir(path)
+      self._gcs_dir = fileutils.force_gcs_path(path)
+      os.makedirs(self.local_dir, exist_ok=True)
+
+      with PartialState().main_process_first():
+        if (
+            check_empty
+            and PartialState().is_main_process
+            and not _is_gcs_dir_empty(self._gcs_dir)
+        ):
+          raise ValueError(f"{self._gcs_dir} needs to be empty.")
+    else:
+      self._is_gcs_path = False
+      self._local_dir = path
+      self._gcs_dir = path
+
+  def __fspath__(self) -> str:
+    return self.local_dir
+
+  @property
+  def local_dir(self) -> str:
+    return self._local_dir
+
+  @property
+  def gcs_dir(self) -> str:
+    """Returns the GCS directory path.
+
+    Returns:
+      The GCS directory path.
+
+    Raises:
+      ValueError: If the path is not a GCS path.
+    """
+    if not self._is_gcs_path:
+      raise ValueError(f"{self._gcs_dir} is not a GCS path.")
+    return self._gcs_dir
+
+  def upload_to_gcs(
+      self,
+      skip_if_exists: bool = True,
+      force_upload: bool = False,
+  ):
+    """Uploads the local directory to GCS."""
+    if not self._is_gcs_path:
+      logging.info(
+          "Not uploading to GCS since %s is not a GCS path.", self.local_dir
+      )
+      return
+
+    if not os.listdir(self.local_dir):
+      logging.info("Not uploading to GCS since %s is empty.", self.local_dir)
+      return
+
+    target = os.path.dirname(self.gcs_dir) + "/"
+    # Avoid race condition uploading the same file from multiple processes.
+    with PartialState().main_process_first():
+      if not PartialState().is_local_main_process:
+        # Non local main processes don't upload.
+        pass
+      elif self._upload_from_all_nodes or PartialState().is_main_process:
+        logging.info("Uploading %s to %s...", self.local_dir, target)
+        cmd = [
+            "gsutil",
+            "-m",
+            "cp",
+            "-r",
+        ]
+        if skip_if_exists:
+          cmd.append("-n")
+        if force_upload:
+          cmd.append("-f")
+        cmd.extend([self.local_dir, target])
+        subprocess.check_output(cmd)
+        logging.info("%s uploaded.", self.local_dir)
+
+
+def _get_local_dir_from_gcs_dir(path: str) -> str:
+  return os.path.join(
+      constants.LOCAL_OUTPUT_DIR,
+      dataset_validation_util.force_gcs_fuse_path(path)[1:],
+  )
+
+
+def _is_gcs_dir_empty(path: str) -> bool:
+  """Checks if a GCS directory is empty.
+
+  Args:
+    path: The GCS directory path.
+
+  Returns:
+    True if the directory is empty.
+
+  Raises:
+    subprocess.CalledProcessError: If the gsutil command failure reason is not
+      because the dir is empty.
+  """
+  path = path.rstrip("/") + "/"
+  try:
+    subprocess.check_output(["gsutil", "ls", path], stderr=subprocess.STDOUT)
+  except subprocess.CalledProcessError as e:
+    if (
+        str(e.output, encoding="utf-8")
+        == "CommandException: One or more URLs matched no objects.\n"
+    ):
+      return True
+    else:
+      logging.info(str(e.output, encoding="utf-8"))
+      raise
+  else:
+    return False
+
+
+def load_tokenizer(
+    pretrained_model_id: str,
+    padding_side: Optional[str] = None,
+    access_token: Optional[str] = None,
+) -> AutoTokenizer:
+  """Loads tokenizer based on `pretrained_model_id`."""
+  tokenizer_kwargs = {}
+  if should_add_eos_token(pretrained_model_id):
+    tokenizer_kwargs["add_eos_token"] = True
+  if padding_side:
+    tokenizer_kwargs["padding_side"] = padding_side
+
+  with PartialState().local_main_process_first():
+    tokenizer = AutoTokenizer.from_pretrained(
+        pretrained_model_id,
+        trust_remote_code=False,
+        use_fast=True,
+        token=access_token,
+        **tokenizer_kwargs,
+    )
+
+  if should_add_pad_token(pretrained_model_id):
+    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+  return tokenizer
+
+
+def load_model(
+    pretrained_model_id: str,
+    tokenizer: AutoTokenizer,
+    precision_mode: str = None,
+    enable_gradient_checkpointing: bool = False,
+    gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
+    access_token: Optional[str] = None,
+    attn_implementation: Optional[str] = None,
+    train_precision: Optional[str] = None,
+    device_map: Optional[str] = None,
+    is_training: bool = True,
+) -> AutoModelForCausalLM:
+  """Loads models from the local dir if specified or from huggingface."""
+  # The `distributed_type` we got through `PartialState` is incorrect for FSDP.
+  # And that's why `Accelerator` is used here.
+  # See b/357138252 for more details.
+  accelerator = accelerate.Accelerator()
+  logging.info("using distributed_type %s", accelerator.distributed_type)
+
+  if device_map is None:
+    if accelerator.distributed_type == DistributedType.MULTI_GPU:
+      # https://github.com/artidoro/qlora/issues/186#issuecomment-1943045599
+      # and b/342038175.
+      device_map = {"": accelerator.process_index}
+    elif accelerator.distributed_type == DistributedType.DEEPSPEED:
+      # Deepspeed Zero3 does not allow setting device_map.
+      # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/modeling_utils.py#L2941-L2943
+      device_map = None
+    elif accelerator.distributed_type == DistributedType.FSDP:
+      if precision_mode in [
+          constants.PRECISION_MODE_4,
+          constants.PRECISION_MODE_8,
+      ]:
+        device_map = trl.get_kbit_device_map()
+      else:
+        device_map = None
+    elif (
+        accelerator.distributed_type == DistributedType.NO
+        and torch.cuda.device_count() > 1
+    ):
+      # Setting device map to None to avoid using model parallelism (MP) when
+      # there are multiple GPUs, which can have very inefficient GPU utilization
+      # (b/342252819). This setting should trigger torch's nn.DataParallel
+      # instead, which has better GPU utilization.
+      device_map = None
+    else:
+      device_map = "auto"
+  logging.info("using device_map %s", device_map)
+
+  if train_precision == constants.PRECISION_MODE_32:
+    train_dtype = torch.float32
+  elif train_precision == constants.PRECISION_MODE_16:
+    train_dtype = torch.float16
+  elif train_precision == constants.PRECISION_MODE_16B:
+    train_dtype = torch.bfloat16
+  else:
+    train_dtype = "auto"
+
+  quantization_config = None
+  # Note: use_cache is False when enable gradient checkpointing.
+  if precision_mode == constants.PRECISION_MODE_32:
+    torch_dtype = torch.float32
+  elif precision_mode == constants.PRECISION_MODE_16:
+    torch_dtype = torch.float16
+  elif precision_mode == constants.PRECISION_MODE_16B:
+    torch_dtype = torch.bfloat16
+  elif precision_mode == constants.PRECISION_MODE_8:
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True, int8_threshold=0
+    )
+    torch_dtype = train_dtype
+  elif precision_mode == constants.PRECISION_MODE_4:
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=train_dtype,
+    )
+    # `bnb_4bit_quant_storage` must be set when using FSDP.
+    # https://huggingface.co/docs/bitsandbytes/main/en/fsdp_qlora
+    if accelerator.distributed_type == DistributedType.FSDP:
+      quantization_config.bnb_4bit_quant_storage = train_dtype
+    torch_dtype = train_dtype
+  else:
+    raise ValueError(f"Invalid precision mode: {precision_mode}")
+  logging.info("using torch_type=%s", torch_dtype)
+
+  model = AutoModelForCausalLM.from_pretrained(
+      pretrained_model_id,
+      use_cache=not enable_gradient_checkpointing,
+      device_map=device_map,
+      torch_dtype=torch_dtype,
+      quantization_config=quantization_config,
+      trust_remote_code=True,
+      token=access_token,
+      attn_implementation=attn_implementation,
+  )
+
+  if precision_mode in (constants.PRECISION_MODE_4, constants.PRECISION_MODE_8):
+    model = prepare_model_for_kbit_training(
+        model,
+        use_gradient_checkpointing=enable_gradient_checkpointing,
+        gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
+    )
+
+  if enable_gradient_checkpointing:
+    model.gradient_checkpointing_enable(
+        gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
+    )
+
+  # Flash attention only supports fp16 or bf16 [1].
+  # prepare_model_for_kbit_training will force cast some layers to float32 [2]
+  #
+  # [1]: https://github.com/Dao-AILab/flash-attention/issues/882
+  # [2]: https://github.com/huggingface/peft/blob/v0.10.0/src/peft/utils/other.py#L79-L81 # pylint: disable=line-too-long
+  if attn_implementation == "flash_attention_2" and precision_mode in (
+      constants.PRECISION_MODE_4,
+      constants.PRECISION_MODE_8,
+  ):
+    for _, param in model.named_parameters():
+      if param.dtype == torch.float32:
+        param.data = param.data.to(torch_dtype)
+
+  if is_training:
+    # KV cache is useless during training
+    # https://stackoverflow.com/a/77408076
+    model.config.use_cache = False
+
+  if should_add_pad_token(pretrained_model_id):
+    model.resize_token_embeddings(len(tokenizer))
+    if is_training:
+      # The following is needed since we added a new token that needs to be
+      # learned.
+      # https://github.com/QwenLM/Qwen/issues/405#issuecomment-1751680291
+      model.enable_input_require_grads()
+
+  return model
+
+
+def _merge_causal_language_model_with_lora_internal(
+    pretrained_model_id: str,
+    merge_precision_mode: str,
+    finetuned_lora_model_dir: str,
+    merged_model_output_dir: str,
+    access_token: Optional[str] = None,
+) -> None:
+  """Internal function to merges the base model with the lora adapter."""
+  logging.info("loading tokenizer...")
+  tokenizer = load_tokenizer(pretrained_model_id)
+
+  # Note: merging peft adapter requires loading model in 16 bits, so merging
+  # is done on CPU on purpose in case one GPU cannot hold the base model.
+  logging.info("loading model %s...", pretrained_model_id)
+  device_map = "cpu"
+  model = load_model(
+      pretrained_model_id=pretrained_model_id,
+      tokenizer=tokenizer,
+      precision_mode=merge_precision_mode,
+      access_token=access_token,
+      device_map=device_map,
+      is_training=False,
+  )
+
+  logging.info("loading LoRA model...")
+  model = PeftModel.from_pretrained(
+      model, finetuned_lora_model_dir, device_map=device_map
+  )
+
+  logging.info("merging base model with finetuned LoRA model...")
+  model = model.merge_and_unload()
+
+  logging.info("saving model to %s...", merged_model_output_dir)
+  model.save_pretrained(
+      merged_model_output_dir,
+      safe_serialization=False,
+      is_main_process=PartialState().is_main_process,
+  )
+
+  logging.info("saving tokenizer to %s...", merged_model_output_dir)
+  tokenizer.save_pretrained(
+      merged_model_output_dir,
+      is_main_process=PartialState().is_main_process,
+  )
+
+
+def merge_causal_language_model_with_lora_fsdp(
+    pretrained_model_id: str,
+    merge_precision_mode: str,
+    finetuned_lora_model_dir: str,
+    merged_model_output_dir: str,
+    access_token: Optional[str] = None,
+) -> None:
+  """Merges the base model with the lora adapter for FSDP.
+
+  Only the main process should call this function.
+
+  Args:
+    pretrained_model_id: Predefined base model name or path to directory
+      containing model checkpoints.
+    merge_precision_mode: Precision mode for saving model weights.
+    finetuned_lora_model_dir: Path to directory containing PEFT-finetuned model
+      weights.
+    merged_model_output_dir: Path to directory to save the merged model.
+    access_token: Access token for accessing the model.
+  """
+  assert PartialState().is_main_process
+  _merge_causal_language_model_with_lora_internal(
+      pretrained_model_id=pretrained_model_id,
+      merge_precision_mode=merge_precision_mode,
+      finetuned_lora_model_dir=finetuned_lora_model_dir,
+      merged_model_output_dir=merged_model_output_dir,
+      access_token=access_token,
+  )
+
+
+def merge_causal_language_model_with_lora(
+    pretrained_model_id: str,
+    precision_mode: str,
+    finetuned_lora_model_dir: str,
+    merged_model_output_dir: str,
+    access_token: Optional[str] = None,
+) -> None:
+  """Merges the base model with the lora adapter."""
+
+  # Set merge related variables.
+  if precision_mode == constants.PRECISION_MODE_FP8:
+    # Merge as FP16. FP8 requires conversion after merge.
+    merge_precision_mode = constants.PRECISION_MODE_16
+    local_merged_model_dir = _LOCAL_MERGED_MODEL_DIR
+  else:
+    merge_precision_mode = precision_mode
+    local_merged_model_dir = merged_model_output_dir
+
+  if PartialState().is_main_process:
+    logging.info("Starting merging job...")
+    # When deepspeed Zero3 is enabled, users are not allowed to specify
+    # `device_map` when loading the model (even on CPU).
+    #
+    # To work-around this, we kick off another process (from the
+    # is_main_process) and set up the environment to avoid using Deepspeed when
+    # doing the merging.
+    if is_deepspeed_zero3_enabled():
+      ctx = mp.get_context("spawn")
+      os.environ["ACCELERATE_USE_DEEPSPEED"] = "false"
+      merge_job = ctx.Process(
+          target=_merge_causal_language_model_with_lora_internal,
+          args=(
+              pretrained_model_id,
+              merge_precision_mode,
+              finetuned_lora_model_dir,
+              local_merged_model_dir,
+          ),
+          kwargs={
+              "access_token": access_token,
+          },
+      )
+      merge_job.start()
+      merge_job.join()
+      os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
+    else:
+      _merge_causal_language_model_with_lora_internal(
+          pretrained_model_id=pretrained_model_id,
+          merge_precision_mode=merge_precision_mode,
+          finetuned_lora_model_dir=finetuned_lora_model_dir,
+          merged_model_output_dir=local_merged_model_dir,
+          access_token=access_token,
+      )
+
+    logging.info("merging job is done")
+
+  # Wait for all processes to sync here.
+  PartialState().wait_for_everyone()
+
+  if precision_mode == constants.PRECISION_MODE_FP8:
+    convert_model_to_fp8(
+        pretrained_model_name_or_path=pretrained_model_id,
+        merged_model_output_dir=local_merged_model_dir,
+        quantized_model_output_dir=merged_model_output_dir,
+        access_token=access_token,
+    )
+
+
+def convert_model_to_fp8(
+    pretrained_model_name_or_path: str,
+    merged_model_output_dir: str,
+    quantized_model_output_dir: str,
+    access_token: Optional[str] = None,
+) -> None:
+  """Converts the model to fp8.
+
+  Args:
+    pretrained_model_name_or_path: Original base model name or path.
+    merged_model_output_dir: Path to directory containing the merged model.
+    quantized_model_output_dir: Path to directory to save the quantized model.
+    access_token: Access token for accessing the model.
+  """
+  if PartialState().is_main_process:
+    quantization_config = FbgemmFp8Config(
+        modules_to_not_convert=_maybe_get_modules_to_not_convert_by_model_id(
+            pretrained_model_name_or_path
+        )
+    )
+    quantized_model = AutoModelForCausalLM.from_pretrained(
+        merged_model_output_dir,
+        device_map="cpu",
+        quantization_config=quantization_config,
+        trust_remote_code=False,
+        token=access_token,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(merged_model_output_dir)
+
+    quantized_model.save_pretrained(quantized_model_output_dir)
+    tokenizer.save_pretrained(quantized_model_output_dir)
+  PartialState().wait_for_everyone()
+
+
+@dataclasses.dataclass
+class TuningDataStats:
+  tuning_dataset_example_count: int
+  total_billable_token_count: int
+  tuning_step_count: int
+
+
+def get_dataset_stats(
+    dataset: Any,
+    tokenizer: transformers.PreTrainedTokenizer,
+    column: str,
+    effective_batch_size: int,
+) -> TuningDataStats:
+  """Calculates dataset statistics, e.g., total number of tokens."""
+  tokenized_dataset = dataset.map(lambda x: tokenizer(x[column]))
+  inputs = tokenized_dataset["input_ids"]
+  tuning_dataset_example_count = int(len(inputs))
+  total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
+  tuning_step_count = (
+      tuning_dataset_example_count + effective_batch_size - 1
+  ) // effective_batch_size
+  return TuningDataStats(
+      tuning_dataset_example_count,
+      total_billable_token_count,
+      tuning_step_count,
+  )
+
+
+def force_gc():
+  """Collects garbage immediately to release unused CPU/GPU resources."""
+  gc.collect()
+  torch.cuda.empty_cache()
+
+
+def should_add_pad_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special pad token."""
+  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
+
+
+def should_add_eos_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special eos token."""
+  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
+
+
+def write_kfp_outputs(
+    executor_input: str, output_artifacts: Dict[str, str]
+) -> None:
+  """Writes KFP outputs given a dict of output artifact names and URIs."""
+  # Only the main process writes to avoid race condition.
+  if PartialState().is_main_process:
+    executor_input = json_format.Parse(
+        executor_input, pipeline_spec_pb2.ExecutorInput()
+    )
+    outputs = executor_input.outputs
+    # set all artifacts
+    for name, uri in output_artifacts.items():
+      artifact_list = outputs.artifacts.get(name)
+      if not artifact_list or not artifact_list.artifacts:
+        raise ValueError(f"Artifact name={name} does not exist.")
+      artifact_list.artifacts[0].uri = uri
+
+    # write output file
+    executor_output = pipeline_spec_pb2.ExecutorOutput(
+        artifacts=outputs.artifacts
+    )
+    os.makedirs(os.path.dirname(outputs.output_file), exist_ok=True)
+    with open(outputs.output_file, "w") as f:
+      f.write(json_format.MessageToJson(executor_output, indent=None))
+
+  # Wait for the main process to finish before moving on to the next task.
+  PartialState().wait_for_everyone()
+
+
+def upload_local_dir_to_gcs(local_dir: str, gcs_path: str):
+  """Uploads local dir to GCS."""
+
+  if PartialState().is_main_process:
+    logging.info("uploading %s to %s...", local_dir, gcs_path)
+    subprocess.check_output([
+        "gsutil",
+        "-m",
+        "cp",
+        "-r",
+        local_dir,
+        gcs_path,
+    ])
+    logging.info("%s uploaded.", local_dir)
+
+  PartialState().wait_for_everyone()
+
+
+def write_first_party_model_metadata(output_dir: str, docker_uri: str) -> None:
+  """Multi-process friendly version of fileutils.write_first_party_model_metadata."""
+  if PartialState().is_main_process:
+    fileutils.write_first_party_model_metadata(output_dir, docker_uri)
+  PartialState().wait_for_everyone()
+
+
+@dataclasses.dataclass
+class GpuStats:
+  """Holds information about GPU usage stats.
+
+  For memory related, see
+  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
+  """
+
+  # total memory
+  total_mem: float
+  # memory occupied.
+  occupied: float
+  # memory reserved, but not used.
+  unused: float
+  # nvidia-smi usually reports more memory usages than pytorch (for driver,
+  # kernel and etc). `smi_diff` tracks this difference.
+  smi_diff: float
+  # Gpu utilization.
+  util: float
+
+  # Allows unpacking operation like
+  # total_mem, occupied, unused, smi_diff, util = GpuStats(...)
+  # See https://stackoverflow.com/a/70753113
+  def __iter__(self):
+    return iter(dataclasses.astuple(self))
+
+
+def gpu_stats() -> GpuStats:
+  """Reports GPU memory usage and utilization."""
+  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
+  bytes_per_gb = 1024.0**3
+  device = torch.cuda.current_device()
+  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
+  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
+  unused = reserved - occupied
+
+  def smi_mem(device):
+    try:
+      pynvml.nvmlInit()
+      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
+      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+      return info.used / bytes_per_gb
+    except pynvml.NVMLError:
+      return 0.0
+
+  mem_used_smi = smi_mem(device)
+  smi_diff = mem_used_smi - reserved
+
+  util = torch.cuda.utilization(device)
+  return GpuStats(mem_used_smi, occupied, unused, smi_diff, util)
+
+
+def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
+  if stats is None:
+    stats = gpu_stats()
+  total, occupied, unused, smi_diff, util = stats
+  return (
+      f"GPU memory: {total:.2f}({occupied=:.2f}, {unused=:.2f},"
+      f" {smi_diff=:.2f}) GB. Utilization: {util:.2f}%"
+  )
+
+
+def init_partial_state(
+    timeout: datetime.timedelta = datetime.timedelta(seconds=600),
+) -> None:
+  """Initializes the partial state with timeout."""
+  # This needs to be called before any other PartialState() calls, and
+  # TrainingArguments needs `use_configured_state`. See b/357970482#comment3
+  # for more details.
+  PartialState(timeout=timeout)
+
+
+def print_library_versions():
+  if PartialState().is_main_process:
+    logging.info("======================")
+    logging.info("library versions")
+    logging.info("======================")
+    logging.info("accelerate: %s", accelerate.__version__)
+    logging.info("peft: %s", peft.__version__)
+    logging.info("transformers: %s", transformers.__version__)
+    logging.info("trl: %s", trl.__version__)
+  PartialState().wait_for_everyone()
+
+
+def get_final_checkpoint_path(output_dir: str) -> str:
+  """Returns the final checkpoint path."""
+  return os.path.join(output_dir, constants.FINAL_CHECKPOINT_DIRNAME)
+
+
+def _maybe_get_modules_to_not_convert_by_model_id(
+    pretrained_model_name_or_path: str,
+) -> Optional[Sequence[str]]:
+  """Returns the modules to not convert for the model."""
+  if _LLAMA_3_1_405B_MODEL_ID in pretrained_model_name_or_path:
+    return _get_llama_3_1_405b_modules_to_not_convert()
+  else:
+    return None
+
+
+def _get_llama_3_1_405b_modules_to_not_convert() -> Sequence[str]:
+  """Returns the modules to not convert for Llama 3.1 405B model."""
+  modules_to_not_convert = ["lm_head"]
+  for idx in range(126):
+    for proj_name in ["k_proj", "o_proj", "q_proj", "v_proj"]:
+      modules_to_not_convert.append(f"model.layers.{idx}.self_attn.{proj_name}")
+  for proj_name in ["down_proj", "gate_proj", "up_proj"]:
+    modules_to_not_convert.append(f"model.layers.0.mlp.{proj_name}")
+    modules_to_not_convert.append(f"model.layers.125.mlp.{proj_name}")
+  return tuple(modules_to_not_convert)

community-content/vertex_model_garden/model_oss/peft/train/vmg/validate_dataset_with_template.py

@@ -0,0 +1,77 @@
+"""Validate the dataset with the template."""
+
+from typing import Sequence
+
+from absl import app
+from absl import flags
+
+from util import dataset_validation_util
+from vertex_vision_model_garden_peft.train.vmg import utils
+from util import constants
+
+
+_DATASET_NAME = flags.DEFINE_string(
+    'dataset_name',
+    None,
+    'The dataset name in huggingface.',
+    required=True,
+)
+
+_TRAIN_SPLIT_NAME = flags.DEFINE_string(
+    'train_split_name',
+    'train',
+    'The train split name.',
+)
+
+_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
+    'instruct_column_in_dataset',
+    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    'The instruct column in dataset.',
+)
+
+_TEMPLATE = flags.DEFINE_string(
+    'template',
+    None,
+    'Template for formatting language model training data. Must be a filename'
+    ' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
+    ' Cloud Storage URI to a JSON file.',
+    required=True,
+)
+
+_VALIDATE_PERCENTAGE_OF_DATASET = flags.DEFINE_integer(
+    'validate_percentage_of_dataset',
+    None,
+    'The percentage of the dataset to validate with the template. If set to'
+    ' -1, it loads the full dataset.',
+)
+
+_VALIDATE_K_ROWS_OF_DATASET = flags.DEFINE_integer(
+    'validate_k_rows_of_dataset',
+    None,
+    'The top k rows of the dataset to validate with the template. If set to -1,'
+    ' it loads the full dataset.',
+)
+
+_USE_MULTIPROCESSING = flags.DEFINE_boolean(
+    'use_multiprocessing',
+    False,
+    'Whether to use multiprocessing for loading the dataset.',
+)
+
+
+def main(unused_argv: Sequence[str]) -> None:
+  utils.print_library_versions()
+
+  dataset_validation_util.validate_dataset_with_template(
+      dataset_name=_DATASET_NAME.value,
+      split=_TRAIN_SPLIT_NAME.value,
+      input_column=_INSTRUCT_COLUMN_IN_DATASET.value,
+      template=_TEMPLATE.value,
+      use_multiprocessing=_USE_MULTIPROCESSING.value,
+      validate_percentage_of_dataset=_VALIDATE_PERCENTAGE_OF_DATASET.value,
+      validate_k_rows_of_dataset=_VALIDATE_K_ROWS_OF_DATASET.value,
+  )
+
+
+if __name__ == '__main__':
+  app.run(main)

community-content/vertex_model_garden/model_oss/util/commons.py

@@ -1,6 +1,6 @@
 """Common utility lib for prediction on images."""
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
 
 import numpy as np
 from PIL import Image
@@ -10,6 +10,28 @@ import yaml
 from util import image_format_converter
 
 
+def convert_list_to_label_map(
+    input_list: List[str],
+) -> Tuple[Dict[str, Dict[int, str]], List[int]]:
+  """Converts a list of labels to a dictionary and numerical encoding.
+
+  Args:
+    input_list: A list of strings representing class labels.
+
+  Returns:
+    A tuple containing:
+      label_map: A dictionary mapping unique labels to integer indices.
+      encoded_list: A list of integers corresponding to the labels in the input
+      list.
+  """
+  unique_labels = set(input_list)
+  label_map_reverse = {label: idx for idx, label in enumerate(unique_labels)}
+  label_map = {idx: label for idx, label in enumerate(unique_labels)}
+  encoded_list = [label_map_reverse[label] for label in input_list]
+
+  return {"label_map": label_map}, encoded_list
+
+
 def get_prediction_instances(image: Image.Image) -> List[Dict[str, Any]]:
   """Gets prediction instances.
 
@@ -40,14 +62,14 @@ def get_label_map(label_map_yaml_filepath: str) -> Dict[str, Any]:
 
 
 def get_object_detection_endpoint_predictions(
-    detection_endpoint: ...,
+    detector_endpoint: ...,
     input_image: np.ndarray,
     detection_thresh: float = 0.2,
 ) -> np.ndarray:
   """Gets endpoint predictions.
 
   Args:
-      detection_endpoint: image object detection endpoint.
+      detector_endpoint: image object detection endpoint.
       input_image: Input image.
       detection_thresh: Detection threshold.
 
@@ -55,9 +77,10 @@ def get_object_detection_endpoint_predictions(
       Object detection predictions from endpoints.
   """
   height, width, _ = input_image.shape
-  predictions = detection_endpoint.predict(
+  predictions = detector_endpoint.predict(
       get_prediction_instances(Image.fromarray(input_image))
   ).predictions
+
   detection_scores = np.array(predictions[0]["detection_scores"])
   detection_classes = np.array(predictions[0]["detection_classes"])
   detection_boxes = np.array(
@@ -66,6 +89,29 @@ def get_object_detection_endpoint_predictions(
           for b in predictions[0]["detection_boxes"]
       ]
   )
+  return merge_boxes_and_classes(
+      detection_scores, detection_boxes, detection_classes, detection_thresh
+  )
+
+
+def merge_boxes_and_classes(
+    detection_scores: np.ndarray,
+    detection_boxes: np.ndarray,
+    detection_classes: np.ndarray,
+    detection_thresh: float = 0.2,
+) -> np.ndarray:
+  """Merges prediction boxes and classes.
+
+  Args:
+      detection_scores: array of detection scores.
+      detection_boxes: array of detection boxes.
+      detection_classes: array of detection classes.
+      detection_thresh: float indicating the detection threshold.
+
+  Returns:
+      preds_merge_cls: a numpy array containing the detection boxes, scores and
+      classes.
+  """
   thresh_indices = [
       x for x, val in enumerate(detection_scores) if val > detection_thresh
   ]
@@ -76,4 +122,5 @@ def get_object_detection_endpoint_predictions(
   preds_merge_cls = np.column_stack(
       (preds_merge_conf, detection_classes[thresh_indices])
   )
+
   return preds_merge_cls

community-content/vertex_model_garden/model_oss/util/constants.py

@@ -36,6 +36,12 @@ BEST_CKPT_METRIC_COMP = 'higher'
 
 # Reported hyperparameter tuning metric tag.
 HP_METRIC_TAG = 'model_performance'
+HP_LOSS_TAG = 'model_loss'
+
+# Reported places.
+REPORT_TO_NONE = 'none'
+REPORT_TO_WANDB = 'wandb'
+REPORT_TO_TENSORBOARD = 'tensorboard'
 
 # HPT trial prefix.
 TRIAL_PREFIX = 'trial_'
@@ -45,7 +51,7 @@ ML_USE_TRAINING = 'training'
 ML_USE_VALIDATION = 'validation'
 ML_USE_TEST = 'test'
 
-# COCO json keys
+# COCO json keys.
 COCO_JSON_ANNOTATIONS = 'annotations'
 COCO_JSON_ANNOTATION_IMAGE_ID = 'image_id'
 COCO_JSON_ANNOTATION_CATEGORY_ID = 'category_id'
@@ -60,36 +66,88 @@ COCO_JSON_IMAGE_HEIGHT = 'height'
 COCO_JSON_IMAGE_COCO_URL = 'coco_url'
 COCO_ANNOTATION_BBOX = 'bbox'
 
-# GCS prefixes
+# GCS prefixes.
 GCS_URI_PREFIX = 'gs://'
 GCSFUSE_URI_PREFIX = '/gcs/'
 
 LOCAL_EVALUATION_RESULT_DIR = '/tmp/evaluation_result_dir'
 LOCAL_MODEL_DIR = '/tmp/model_dir'
+LOCAL_LORA_DIR = '/tmp/lora_dir'
 LOCAL_BASE_MODEL_DIR = '/tmp/base_model_dir'
 LOCAL_DATA_DIR = '/tmp/data'
+LOCAL_OUTPUT_DIR = '/tmp/output_dir'
+LOCAL_PREDICTION_RESULT_DIR = '/tmp/prediction_result_dir'
+SHARED_MEM_DIR = '/dev/shm'
 
 # Huggingface files.
 HF_MODEL_WEIGHTS_SUFFIX = '.bin'
 
 # PEFT finetuning constants.
+TEXT_TO_IMAGE = 'text-to-image'
 TEXT_TO_IMAGE_LORA = 'text-to-image-lora'
+TEXT_TO_IMAGE_DREAMBOOTH = 'text-to-image-dreambooth'
+TEXT_TO_IMAGE_DREAMBOOTH_LORA = 'text-to-image-dreambooth-lora'
+TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL = 'text-to-image-dreambooth-lora-sdxl'
 SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
-CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora'
+MERGE_CAUSAL_LANGUAGE_MODEL_LORA = 'merge-causal-language-model-lora'
+QUANTIZE_MODEL = 'quantize-model'
 INSTRUCT_LORA = 'instruct-lora'
-CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES = [
-  "q_proj",
-  "v_proj",
-]
-INSTRUCT_LORA_TARGET_MODULES = [
-    "query_key_value",
-    "dense",
-    "dense_h_to_4h",
-    "dense_4h_to_h",
-]
+VALIDATE_DATASET_WITH_TEMPLATE = 'validate-dataset-with-template'
+DEFAULT_TEXT_COLUMN_IN_DATASET = 'quote'
+DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET = 'text'
+DEFAULT_INSTRUCT_COLUMN_IN_DATASET = 'text'
+
+FINAL_CHECKPOINT_DIRNAME = 'checkpoint-final'
+
+# ImageBind inference constants.
+FEATURE_EMBEDDING_GENERATION = 'feature-embedding-generation'
+ZERO_SHOT_CLASSIFICATION = 'zero-shot-classification'
 
 # Precision modes for loading model weights.
+PRECISION_MODE_2 = '2bit'
+PRECISION_MODE_3 = '3bit'
 PRECISION_MODE_4 = '4bit'
 PRECISION_MODE_8 = '8bit'
+PRECISION_MODE_FP8 = 'float8'  # to use fbgemm_fp8 quantization
 PRECISION_MODE_16 = 'float16'
+PRECISION_MODE_16B = 'bfloat16'
 PRECISION_MODE_32 = 'float32'
+
+# Quantization modes.
+GPTQ = 'gptq'
+AWQ = 'awq'
+
+# AWQ versions.
+GEMM = 'GEMM'
+GEMV = 'GEMV'
+
+# Environment variable keys.
+PRIVATE_BUCKET_ENV_KEY = 'AIP_PRIVATE_BUCKET_NAME'
+
+# Kfp pipeline constants.
+TFVISION_TRAIN_OUTPUT_ARTIFACT_NAME = 'checkpoint_dir'
+
+# Vertex IOD type.
+AUTOML = 'AUTOML'
+MODEL_GARDEN = 'MODEL_GARDEN'
+
+# LRU Disk Cache constants.
+MD5_HASHMAP_FILENAME = 'md5_hashmap.json'
+
+# Prediction request keys.
+PREDICT_INSTANCE_KEY = 'instances'
+PREDICT_INSTANCE_IMAGE_KEY = 'image'
+PREDICT_INSTANCE_POSE_IMAGE_KEY = 'pose_image'
+PREDICT_INSTANCE_TEXT_KEY = 'text'
+PREDICT_INSTANCE_PROMPT_KEY = 'prompt'
+
+PREDICT_PARAMETERS_KEY = 'parameters'
+PREDICT_PARAMETERS_NUM_INFERENCE_STEPS_KEY = 'num_inference_steps'
+PREDICT_PARAMETERS_HEIGHT_KEY = 'height'
+PREDICT_PARAMETERS_WIDTH_KEY = 'width'
+PREDICT_PARAMETERS_GUIDANCE_SCALE_KEY = 'guidance_scale'
+PREDICT_PARAMETERS_NEGATIVE_PROMPT_KEY = 'negative_prompt'
+PREDICT_PARAMETERS_LORA_ID_KEY = 'lora_id'
+PREDICT_PARAMETERS_IGNORE_LORA_CACHE_KEY = 'ignore_lora_cache'
+
+PREDICT_OUTPUT_KEY = 'output'

community-content/vertex_model_garden/model_oss/util/fileutils.py

@@ -1,10 +1,10 @@
 """Fileutil lib to copy files between gcs and local."""
 
-import glob
+import fnmatch
 import os
 import pathlib
 import shutil
-from typing import Tuple
+from typing import List, Optional, Tuple
 import uuid
 
 from absl import logging
@@ -13,6 +13,17 @@ from google.cloud import storage
 from util import constants
 
 
+_GCS_CLIENT = None
+
+
+def _get_gcs_client() -> storage.Client:
+  """Gets the default GCS client."""
+  global _GCS_CLIENT
+  if _GCS_CLIENT is None:
+    _GCS_CLIENT = storage.Client()
+  return _GCS_CLIENT
+
+
 def generate_tmp_path(extension: str = '') -> str:
   """Generates a temporary file path with UUID.
 
@@ -36,6 +47,16 @@ def force_gcs_fuse_path(gcs_uri: str) -> str:
     return gcs_uri
 
 
+def force_gcs_path(uri: str) -> str:
+  """Converts /gcs/ uris to their gs:// equivalents. No-op for other uris."""
+  if uri.startswith(constants.GCSFUSE_URI_PREFIX):
+    return uri.replace(
+        constants.GCSFUSE_URI_PREFIX, constants.GCS_URI_PREFIX, 1
+    )
+  else:
+    return uri
+
+
 def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
   """Download a gcs file to a local dir.
 
@@ -62,15 +83,47 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str):
     raise ValueError(
         f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
     )
-  client = storage.Client()
+  client = _get_gcs_client()
   os.makedirs(os.path.dirname(local_path), exist_ok=True)
   with open(local_path, 'wb') as f:
     client.download_blob_to_file(gcs_uri, f)
 
 
+def download_gcs_file_list_to_local(
+    gcs_uri_list: List[str], local_dir: str
+) -> List[str]:
+  """Downloads a list of GCS files to a local directory.
+
+  Args:
+    gcs_uri_list: A list of GCS file paths.
+    local_dir: Local directory in which the GCS files are saved.
+
+  Returns:
+    The local file paths corresponding to the input GCS file paths.
+
+  Raises:
+    ValueError: An input file path is not a GCS path.
+  """
+  local_paths = []
+  for gcs_uri in gcs_uri_list:
+    if not is_gcs_path(gcs_uri):
+      raise ValueError(
+          f'{gcs_uri} is not a GCS path starting with'
+          f' {constants.GCS_URI_PREFIX}.'
+      )
+    local_path = os.path.join(local_dir, gcs_uri.replace('gs://', ''))
+    download_gcs_file_to_local(gcs_uri, local_path)
+    local_paths.append(local_path)
+  return local_paths
+
+
 def download_gcs_dir_to_local(
-    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
-):
+    gcs_dir: str,
+    local_dir: str,
+    skip_hf_model_bin: bool = False,
+    allow_patterns: Optional[List[str]] = None,
+    log: bool = True,
+) -> None:
   """Downloads files in a GCS directory to a local directory.
 
   For example:
@@ -78,16 +131,21 @@ def download_gcs_dir_to_local(
     gs://bucket/foo/a -> /tmp/bar/a
     gs://bucket/foo/b/c -> /tmp/bar/b/c
 
-  Arguments:
+  Args:
     gcs_dir: A string of directory path on GCS.
     local_dir: A string of local directory path.
     skip_hf_model_bin: True to skip downloading HF model bin files.
+    allow_patterns: A list of allowed patterns. If provided, only files matching
+      one or more patterns are downloaded.
+    log: True to log each downloaded file.
   """
   if not is_gcs_path(gcs_dir):
     raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.')
   bucket_name = gcs_dir.split('/')[2]
-  prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
-  client = storage.Client()
+  prefix = (
+      gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/') + '/'
+  )
+  client = _get_gcs_client()
   blobs = client.list_blobs(bucket_name, prefix=prefix)
   for blob in blobs:
     if blob.name[-1] == '/':
@@ -95,43 +153,63 @@ def download_gcs_dir_to_local(
     file_path = blob.name[len(prefix) :].strip('/')
     local_file_path = os.path.join(local_dir, file_path)
     os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+    if allow_patterns and all(
+        [not fnmatch.fnmatch(file_path, p) for p in allow_patterns]
+    ):
+      continue
     if (
         file_path.endswith(constants.HF_MODEL_WEIGHTS_SUFFIX)
         and skip_hf_model_bin
     ):
-      logging.info('Skip downloading model bin %s', file_path)
+      if log:
+        logging.info('Skip downloading model bin %s', file_path)
       with open(local_file_path, 'w') as f:
-        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}/{file_path}')
+        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}{file_path}')
     else:
-      logging.info('Downloading %s to %s', file_path, local_file_path)
+      if log:
+        logging.info('Downloading %s to %s', file_path, local_file_path)
       blob.download_to_filename(local_file_path)
 
 
+def _get_relative_paths(base_dir: str) -> List[str]:
+  """Gets relative paths of all files in a local base directory."""
+  path = pathlib.Path(base_dir)
+  relative_paths = []
+  for local_file in path.rglob('*'):
+    if os.path.isfile(local_file):
+      relative_path = os.path.relpath(local_file, base_dir)
+      relative_paths.append(relative_path)
+  return relative_paths
+
+
+def _upload_local_files_to_gcs(
+    relative_paths: List[str], local_dir: str, gcs_dir: str
+):
+  """Uploads local files to gcs."""
+  bucket_name = gcs_dir.split('/')[2]
+  blob_dir = '/'.join(gcs_dir.split('/')[3:])
+  client = _get_gcs_client()
+  bucket = client.bucket(bucket_name)
+  for relative_path in relative_paths:
+    blob = bucket.blob(os.path.join(blob_dir, relative_path))
+    blob.upload_from_filename(os.path.join(local_dir, relative_path))
+
+
 def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
   """Uploads local dir to gcs.
 
   For example:
     upload_local_dir_to_gcs(/tmp/bar, gs://bucket/foo)
-    gs://bucket/foo/a -> /tmp/bar/a
-    gs://bucket/foo/b/c -> /tmp/bar/b/c
+    /tmp/bar/a -> gs://bucket/foo/a
+    /tmp/bar/b/c -> gs://bucket/foo/b/c
 
   Arguments:
     local_dir: A string of local directory path.
     gcs_dir: A string of directory path on GCS.
   """
-  bucket_name = gcs_dir.split('/')[2]
-  blob_dir = '/'.join(gcs_dir.split('/')[3:])
-  client = storage.Client()
-  bucket = client.bucket(bucket_name)
-  for local_file in glob.glob(local_dir + '/**'):
-    if os.path.isfile(local_file):
-      logging.info(
-          'Uploading %s to %s',
-          local_file,
-          os.path.join(constants.GCS_URI_PREFIX, bucket_name, blob_dir),
-      )
-      blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file)))
-      blob.upload_from_filename(local_file)
+  # Relative paths of all files in local_dir.
+  relative_paths = _get_relative_paths(local_dir)
+  _upload_local_files_to_gcs(relative_paths, local_dir, gcs_dir)
 
 
 def upload_file_to_gcs_path(
@@ -155,7 +233,7 @@ def upload_file_to_gcs_path(
   if not source_path_obj.exists():
     raise RuntimeError(f'Source path does not exist: {source_path}')
 
-  storage_client = storage.Client()
+  storage_client = _get_gcs_client()
   source_file_path = source_path
   destination_file_uri = destination_uri
   logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
@@ -174,7 +252,9 @@ def is_gcs_path(input_path: str) -> bool:
   Returns:
       True if the input path is a GCS path, False otherwise.
   """
-  return input_path.startswith(constants.GCS_URI_PREFIX)
+  return input_path is not None and input_path.startswith(
+      constants.GCS_URI_PREFIX
+  )
 
 
 def release_text_assets(
@@ -232,13 +312,10 @@ def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
   """
   _, local_video_file_name = os.path.split(video_file_path)
   file_extension = os.path.splitext(video_file_path)[1]
-  if file_extension:
-    remote_video_file_name = local_video_file_name.replace(
-        file_extension, '_overlay.mp4'
-    )
-  else:
-    remote_video_file_name = local_video_file_name + '_overlay.mp4'
-  local_file_path = generate_tmp_path(file_extension)
+  remote_video_file_name = local_video_file_name.replace(
+      file_extension, '_overlay.mp4'
+  )
+  local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1])
   logging.info('Downloading %s to %s...', video_file_path, local_file_path)
   download_gcs_file_to_local(video_file_path, local_file_path)
   return local_file_path, remote_video_file_name
@@ -254,10 +331,28 @@ def get_output_video_file(video_output_file_path: str) -> str:
     str: Local video output file path.
   """
   file_extension = os.path.splitext(video_output_file_path)[1]
-  if file_extension:
-    out_local_video_file_name = video_output_file_path.replace(
-        file_extension, '_overlay' + file_extension
-    )
-  else:
-    out_local_video_file_name = video_output_file_path + '_overlay'
+  out_local_video_file_name = video_output_file_path.replace(
+      file_extension, '_overlay' + file_extension
+  )
   return out_local_video_file_name
+
+
+def write_first_party_model_metadata(
+    output_path: str, required_container_uri: str
+) -> None:
+  """Write Vertex internal model metadata for first party artifacts."""
+  model_metadata_fname = 'model_metadata.jsonl'
+  if len(required_container_uri) > 126:
+    raise ValueError(f'Docker URI exceeds 126 chars: {required_container_uri}')
+  payload = '\n{}{}'.format(  # serialized proto
+      chr(len(required_container_uri)),
+      required_container_uri,
+  )
+  os.makedirs(output_path, exist_ok=True)
+  output_dirs = [output_path]
+  if output_path.startswith('/gcs'):
+    # include all parent dirs, except "/", "/gcs"
+    output_dirs.extend([str(p) for p in pathlib.Path(output_path).parents][:-2])
+  for output_dir in output_dirs:
+    with open(os.path.join(output_dir, model_metadata_fname), 'w') as f:
+      f.write(payload)

community-content/vertex_model_garden/model_oss/util/hypertune_utils.py

@@ -20,3 +20,11 @@ def get_trial_id_from_environment() -> str:
         _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID,
     )
   return os.environ.get(_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID, '0')
+
+
+def maybe_append_trial_id(path: str) -> str:
+  """Appends trial_N to path if running in a Hyperparameter Tuning Job."""
+  trial_id = os.environ.get(_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID)
+  if trial_id is None:
+    return path
+  return os.path.join(path, f'trial_{trial_id}')

community-content/vertex_pipeline_kfpv2/customjob.py

@@ -0,0 +1,41 @@
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def customjob(
+    project_id: str,
+    location: str,
+    staging_bucket: str,
+    experiment: str,
+    job_name: str,
+    script_path: str,
+    container_uri: str,
+    machine_type: str,
+):
+    import os
+    from google.cloud import aiplatform
+
+    aiplatform.init(
+        project=project_id,
+        location=location,
+        staging_bucket=staging_bucket,
+        experiment=experiment,
+    )
+    job = aiplatform.CustomJob.from_local_script(
+        display_name=job_name,
+        script_path=os.path.join(os.getcwd(), script_path),
+        container_uri=container_uri,
+        machine_type=machine_type,
+    )
+    job.run()
+
+@dsl.pipeline(name='run-customjob')
+def pipeline_customjob():
+    customjob("990000000009", "us-west1", "gs://staging-bucket/customjob",
+        "run-experiment", "custom-job", "customjob.py",
+        "gcr.io/path/to/model_name:latest", "n1-standard-4")
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_customjob,
+        package_path='customjob.json')

community-content/vertex_pipeline_kfpv2/pipelinejob.py

@@ -0,0 +1,43 @@
+from kfp.v2 import dsl
+
+@dsl.component(base_image='python:3.8',packages_to_install=['google-cloud-aiplatform==1.36.0'])
+def pipelineJob(
+    project_id: str,
+    location: str,
+    display_name: str,
+    json_file: str,
+    pipeline_root: str,
+):
+    import os
+    from google.cloud import aiplatform
+
+    aiplatform.init(
+        project=project_id,
+        location=location,
+    )
+
+    job = aiplatform.PipelineJob(
+        display_name=display_name,
+        template_path=json_file,
+        pipeline_root=pipeline_root,
+        enable_caching=False,
+    ).run()
+
+    job.delete()
+
+
+@dsl.pipeline(name='pipelineJobs')
+def pipeline_run_jobs():
+    # 1. create endpoint
+    pipelineJob("990000000009", "us-west1", "Pipeline-create endpoint",
+        "create_endpoint.json", "gs://pipeline-root-bucket/pipelines")
+        
+    # 2. deploy model to endpoint
+    pipelineJob("990000000009", "us-west1", "Pipeline-deploy model",
+        "deploy_model.json", "gs://pipeline-root-bucket/pipelines")
+
+if __name__ == "__main__":
+    from kfp.v2 import compiler
+    compiler.Compiler().compile(
+        pipeline_func=pipeline_run_jobs,
+        package_path='pipelineJobs.json')

notebooks/community/CODEOWNERS

@@ -98,6 +98,7 @@
 /notebooks/community/model_garden/model_garden_pytorch_text_to_video.ipynb @KCFindstr
 /notebooks/community/generative_ai/text_embedding_api_cloud_next_new_models.ipynb @xqr-g
 /notebooks/community/generative_ai/text_embedding_api_semantic_search_with_scann.ipynb @henrytansetiawan
+/notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb @pemujo
 /notebooks/community/bigquery_ml_inference/bq_ml_with_vision_translation_nlp.ipynb @deaconsmith
 /notebooks/community/model_garden/model_garden_keras_stable_diffusion.ipynb @genquan9
 /notebooks/community/model_garden/model_garden_keras_yolov8.ipynb @@dstnluong-google

notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb

@@ -0,0 +1,663 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "63c7b05c4717"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "id": "670bbc2007a2"
+      },
+      "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": "890ac0f4e121"
+      },
+      "source": [
+        "# Backoff and retry for LLM\n",
+        "\n",
+        "<table align=\"left\">\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
+        "    </a>\n",
+        "  </td>\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%2Fgenerative_ai%2Fbackoff_and_retry_for_LLMs.ipynb\">\n",
+        "      <img width=\"32px\" src=\"https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>    \n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb\">\n",
+        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/generative_ai/backoff_and_retry_for_LLMs.ipynb\">\n",
+        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "81ce710a5836"
+      },
+      "source": [
+        "NOTE: This notebook has been tested in the following environment:\n",
+        "\n",
+        "Python version = 3.10"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ff76ac47eb9b"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how sending large amounts of traffic to Gemini-1.5-Pro can cause \"429 Quota Exceeded Errors\" and how implementing a backoff-and-retry strategy can help complete jobs without interrupting operations.\n",
+        "\n",
+        "This notebook provides examples for the blog post: [Don't let resource exhaustion leave your users hanging: A guide to handling 429 errors](https://cloud.google.com/blog/products/ai-machine-learning/learn-how-to-handle-429-resource-exhaustion-errors-in-your-llms?e=48754805)\n",
+        "\n",
+        "This tutorial uses the following Google Cloud ML service:\n",
+        "\n",
+        "- Vertex LLM SDK\n",
+        "\n",
+        "The steps performed include:\n",
+        "\n",
+        "- Installation and imports\n",
+        "- Asynchronously calling the Gemini model\n",
+        "- Using the Tenacity retry decorator to implement backoff and retry"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4e3e949c0bdd"
+      },
+      "source": [
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing),\n",
+        "and use the [Pricing Calculator](https://cloud.google.com/products/calculator/)\n",
+        "to generate a cost estimate based on your projected usage.\n",
+        "\n",
+        "**This notebook sends large amount of tokens to Gemini for inference, reduce the number of attempts or use smaller video to reduce costs.**"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0316df526f8"
+      },
+      "source": [
+        "## Get started"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "FyyMdUeAJIVv"
+      },
+      "source": [
+        "## Install Vertex AI SDK for Python and other required packages\n",
+        "\n",
+        "Install the following packages required to execute this notebook.\n",
+        "\n",
+        "**Remember to restart the runtime after installation.**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 16,
+      "metadata": {
+        "id": "snBUuUamoJPz"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install --upgrade --quiet google-cloud-aiplatform tenacity google-cloud-storage"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WX3CHZitmSJM"
+      },
+      "source": [
+        "### Restart runtime (Colab only)\n",
+        "To use the newly installed packages, you must restart the runtime on Google Colab."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f09b4dff629a"
+      },
+      "outputs": [],
+      "source": [
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    import IPython\n",
+        "\n",
+        "    app = IPython.Application.instance()\n",
+        "    app.kernel.do_shutdown(True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "SbmM4z7FOBpM"
+      },
+      "source": [
+        "<div class=\"alert alert-block alert-warning\">\n",
+        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step.  ⚠️</b>\n",
+        "</div>\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "sBCra4QMA2wR"
+      },
+      "source": [
+        "### Authenticate your Google Cloud account\n",
+        "\n",
+        "Depending on your Jupyter environment, you may have to manually authenticate. Follow the relevant instructions below."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "74ccc9e52986"
+      },
+      "source": [
+        "**1. Vertex AI Workbench**\n",
+        "* Do nothing as you are already authenticated."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "de775a3773ba"
+      },
+      "source": [
+        "**2. Local JupyterLab instance, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "254614fa0c46"
+      },
+      "outputs": [],
+      "source": [
+        "# ! gcloud auth login"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ef21552ccea8"
+      },
+      "source": [
+        "**3. Colab, uncomment and run:**"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "603adbbf0532"
+      },
+      "outputs": [],
+      "source": [
+        "# import sys\n",
+        "\n",
+        "# if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "#     from google.colab import auth\n",
+        "\n",
+        "#     auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dae340cb-0583-4e7e-a562-6817ee4d7f6d"
+      },
+      "source": [
+        "### Import libraries"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "412d00f1-08db-4880-8ced-52a9583757b8"
+      },
+      "outputs": [],
+      "source": [
+        "import asyncio\n",
+        "import time\n",
+        "\n",
+        "import nest_asyncio\n",
+        "import vertexai\n",
+        "\n",
+        "nest_asyncio.apply()\n",
+        "from google.cloud import storage\n",
+        "from tenacity import retry, wait_random_exponential\n",
+        "from vertexai.generative_models import GenerationConfig, GenerativeModel, Part"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "DF4l8DTdWgPY"
+      },
+      "source": [
+        "### Set Google Cloud project information and initialize Vertex AI SDK for Python\n",
+        "\n",
+        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com). Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 2,
+      "metadata": {
+        "id": "3EdtdqnoldX4"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Updated property [core/project].\n"
+          ]
+        }
+      ],
+      "source": [
+        "PROJECT_ID = \"[your-project-id]\"  # @param {type:\"string\"}\n",
+        "DEFAUL_MODEL_NAME = \"gemini-1.5-pro-001\"  # @param {type:\"string\"}\n",
+        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
+        "\n",
+        "\n",
+        "# Initiate Vertex AI\n",
+        "vertexai.init(project=PROJECT_ID, location=REGION)\n",
+        "config = GenerationConfig(temperature=0.5, max_output_tokens=512)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f50f22f3-ec85-463e-b6fe-5c8e6b80b07b"
+      },
+      "source": [
+        "### Helper functions"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 9,
+      "metadata": {
+        "id": "b18a366df00b"
+      },
+      "outputs": [],
+      "source": [
+        "def get_images_uri_from_bucket(bucket_name, prefix, delimiter=None):\n",
+        "    \"\"\"Lists all the images with extension '.jpg', 'jpeg' or 'png' in the bucket that begin with the prefix (folder).\"\"\"\n",
+        "    storage_client = storage.Client()\n",
+        "    blobs = storage_client.list_blobs(bucket_name, prefix=prefix, delimiter=delimiter)\n",
+        "    images = [\n",
+        "        f\"gs://{bucket_name}/{blob.name}\"\n",
+        "        for blob in blobs\n",
+        "        if blob.name.endswith(tuple([\".jpg\", \"jpeg\", \"png\"]))\n",
+        "    ]\n",
+        "    return images\n",
+        "\n",
+        "\n",
+        "async def async_ask_gemini(contents, model_name=DEFAUL_MODEL_NAME):\n",
+        "    # This basic function calls Gemini asynchronously without a retry logic\n",
+        "    multimodal_model = GenerativeModel(model_name)\n",
+        "    response = await multimodal_model.generate_content_async(\n",
+        "        contents=contents, generation_config=config\n",
+        "    )\n",
+        "    return response.text\n",
+        "\n",
+        "\n",
+        "@retry(wait=wait_random_exponential(multiplier=1, max=60))\n",
+        "async def retry_async_ask_gemini(contents, model_name=DEFAUL_MODEL_NAME):\n",
+        "    \"\"\"This is the same code as the async_ask_gemini function but implements a retry logic using tenacity decorator.\n",
+        "    wait_random_exponential(multiplier=1, max=60) means that it will\n",
+        "    Retry “Randomly wait up to 2^x * 1 seconds between each retry until the range reaches 60 seconds, then randomly up to 60 seconds afterwards.\n",
+        "    \"\"\"\n",
+        "\n",
+        "    multimodal_model = GenerativeModel(model_name)\n",
+        "    response = await multimodal_model.generate_content_async(\n",
+        "        contents=contents, generation_config=config\n",
+        "    )\n",
+        "    return response.text\n",
+        "\n",
+        "\n",
+        "async def load_test_gemini(function, model_name, attempts=5):\n",
+        "    failed_attempts = 0\n",
+        "    print(f\"Testing with model: {model_name} and function: {function.__name__}\")\n",
+        "    for i in range(attempts):\n",
+        "        try:\n",
+        "            time_start = time.time()\n",
+        "            get_gemini_responses = [\n",
+        "                function(\n",
+        "                    [\n",
+        "                        prompt,\n",
+        "                        video_part,\n",
+        "                        Part.from_uri(image_uri, mime_type=\"image/jpeg\"),\n",
+        "                    ],\n",
+        "                    model_name=MODEL_NAME,\n",
+        "                )\n",
+        "                for image_uri in images_list\n",
+        "            ]\n",
+        "            async_poems = await asyncio.gather(*get_gemini_responses)\n",
+        "            time_taken = time.time() - time_start\n",
+        "            print(f\"{len(async_poems)} Poems written in {time_taken:.0f} seconds\")\n",
+        "        except Exception as error:\n",
+        "            failed_attempts += 1\n",
+        "            print(\"An error occurred:\", error)\n",
+        "\n",
+        "    print(\n",
+        "        f\"{failed_attempts} out of {attempts} failed\"\n",
+        "    ) if failed_attempts > 0 else print(f\"All {attempts} attempts succeded\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c01755042c4b"
+      },
+      "source": [
+        "### Getting images and videos used for testing"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "id": "204228ea941e"
+      },
+      "outputs": [],
+      "source": [
+        "# The images and video used for this test are stored in a public GCS bucket: \"cloud-samples-data\"\n",
+        "bucket_name = \"cloud-samples-data\"\n",
+        "image_prefix = \"generative-ai/image/\"\n",
+        "images_list = get_images_uri_from_bucket(bucket_name, image_prefix, delimiter=\"/\")\n",
+        "\n",
+        "prompt = \"Get the elements from the image, get all the animals from the video, print all the animals and elements found on a numbered list, and then write a poem about them\\n\"\n",
+        "small_video_uri = \"gs://cloud-samples-data/generative-ai/video/animals.mp4\"\n",
+        "large_video_uri = (\n",
+        "    \"gs://cloud-samples-data/generative-ai/video/behind_the_scenes_pixel.mp4\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4f08049c98e5"
+      },
+      "source": [
+        "## Load testing Gemini "
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "91ffee55f2d7"
+      },
+      "source": [
+        "### Test without retry and default quota for Gemini-1.5-pro-001 of 60 QPM\n",
+        "\n",
+        "4 out of 5 tests fail due to 429 Quota exceeded"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 8,
+      "metadata": {
+        "id": "7add4399f0a9"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Testing with model: gemini-1.5-pro-001 and function: async_ask_gemini\n",
+            "72 Poems written in 23 seconds\n",
+            "An error occurred: 429 Quota exceeded for aiplatform.googleapis.com/generate_content_requests_per_minute_per_project_per_base_model with base model: gemini-1.5-pro. Please submit a quota increase request. https://cloud.google.com/vertex-ai/docs/generative-ai/quotas-genai.\n",
+            "An error occurred: 429 Quota exceeded for aiplatform.googleapis.com/generate_content_requests_per_minute_per_project_per_base_model with base model: gemini-1.5-pro. Please submit a quota increase request. https://cloud.google.com/vertex-ai/docs/generative-ai/quotas-genai.\n",
+            "An error occurred: 429 Quota exceeded for aiplatform.googleapis.com/generate_content_requests_per_minute_per_project_per_base_model with base model: gemini-1.5-pro. Please submit a quota increase request. https://cloud.google.com/vertex-ai/docs/generative-ai/quotas-genai.\n",
+            "An error occurred: 429 Quota exceeded for aiplatform.googleapis.com/generate_content_input_tokens_per_minute_per_base_model with base model: gemini-1.5-pro. Please submit a quota increase request. https://cloud.google.com/vertex-ai/docs/generative-ai/quotas-genai.\n",
+            "4 out of 5 failed\n"
+          ]
+        }
+      ],
+      "source": [
+        "video_part = Part.from_uri(small_video_uri, mime_type=\"video/mp4\")\n",
+        "MODEL_NAME = \"gemini-1.5-pro-001\"\n",
+        "# Uncomment line below to re-run the test. Beware of costs since it will make multiple calls to Gemini\n",
+        "# await (load_test_gemini(async_ask_gemini, MODEL_NAME, attempts=5))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4121c249d591"
+      },
+      "source": [
+        "### Re-testing with backoff and retry mechanism enabled \n",
+        "\n",
+        "All tests finallize correctly"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 6,
+      "metadata": {
+        "id": "2338b49fd72d"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Testing with model: gemini-1.5-pro-001 and function: retry_async_ask_gemini\n",
+            "72 Poems written in 21 seconds\n",
+            "72 Poems written in 167 seconds\n",
+            "72 Poems written in 18 seconds\n",
+            "72 Poems written in 149 seconds\n",
+            "72 Poems written in 22 seconds\n",
+            "All 5 attempts succeded\n"
+          ]
+        }
+      ],
+      "source": [
+        "MODEL_NAME = \"gemini-1.5-pro-001\"\n",
+        "# Uncomment line below to re-run the test. Beware of costs since it will make multiple calls to Gemini\n",
+        "# await (load_test_gemini(retry_async_ask_gemini, MODEL_NAME, attempts=5))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d26a20615065"
+      },
+      "source": [
+        "### Testing without retry but with Dynamic Shared Quota using Gemini-1.5-pro-002 \n",
+        "\n",
+        "All 5 attempts succeded with a small video as input"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 5,
+      "metadata": {
+        "id": "37cd3facf381"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Testing with model: gemini-1.5-pro-002 and function: async_ask_gemini\n",
+            "72 Poems written in 23 seconds\n",
+            "72 Poems written in 21 seconds\n",
+            "72 Poems written in 19 seconds\n",
+            "72 Poems written in 17 seconds\n",
+            "72 Poems written in 22 seconds\n",
+            "All 5 attempts succeded\n"
+          ]
+        }
+      ],
+      "source": [
+        "video_part = Part.from_uri(small_video_uri, mime_type=\"video/mp4\")\n",
+        "MODEL_NAME = \"gemini-1.5-pro-002\"\n",
+        "# Uncomment line below to re-run the test. Beware of costs since it will make multiple calls to Gemini\n",
+        "# await (load_test_gemini(async_ask_gemini, MODEL_NAME, attempts=5))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a5fa686b7c52"
+      },
+      "source": [
+        "### Re-testing Dynamic Shared quota with larger video\n",
+        "\n",
+        "Without backoff and retry, testing Gemini-1.5-pro-002 with larger context window caused all tests to fail with 429 reason code."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 7,
+      "metadata": {
+        "id": "d241524f5071"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Testing with model: gemini-1.5-pro-002 and function: async_ask_gemini\n",
+            "An error occurred: 429 Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#error-code-429 for more details.\n",
+            "An error occurred: 429 Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#error-code-429 for more details.\n",
+            "An error occurred: 429 Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#error-code-429 for more details.\n",
+            "An error occurred: 429 Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#error-code-429 for more details.\n",
+            "An error occurred: 429 Resource exhausted. Please try again later. Please refer to https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#error-code-429 for more details.\n",
+            "5 out of 5 failed\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Larger video  used to increase token input size\n",
+        "video_part = Part.from_uri(large_video_uri, mime_type=\"video/mp4\")\n",
+        "MODEL_NAME = \"gemini-1.5-pro-002\"\n",
+        "# Uncomment line below to re-run the test. Beware of costs since it will make multiple calls to Gemini\n",
+        "# await (load_test_gemini(async_ask_gemini, MODEL_NAME, attempts=5))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "999a6cd90e5d"
+      },
+      "source": [
+        "### Adding Backoff and Retry to Dynamic Shared Quota Testing\n",
+        "\n",
+        "Adding backoff and retry mechanisms significantly increased inference time, but all tests completed successfully even with much larger context window.\n",
+        "\n",
+        "Provisioned Throughput should be used to guarantee the capacity and therefore reduce latency.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 13,
+      "metadata": {
+        "id": "91c4d04ab9b4"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Testing with model: gemini-1.5-pro-002 and function: retry_async_ask_gemini\n",
+            "72 Poems written in 188 seconds\n",
+            "72 Poems written in 205 seconds\n",
+            "72 Poems written in 216 seconds\n",
+            "All 3 attempts succeded\n"
+          ]
+        }
+      ],
+      "source": [
+        "video_part = Part.from_uri(large_video_uri, mime_type=\"video/mp4\")\n",
+        "MODEL_NAME = \"gemini-1.5-pro-002\"\n",
+        "# Uncomment line below to re-run the test. Beware of costs since it will make multiple calls to Gemini\n",
+        "# await (load_test_gemini(retry_async_ask_gemini, MODEL_NAME, attempts=3))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9cabeb585f56"
+      },
+      "source": [
+        "## Summary\n",
+        "\n",
+        "These basic tests demonstrate how Dynamic Shared Quota reduces the frequency of \"429 Resource Exhausted\" errors.  The results highlight the importance of always using backoff and retry mechanisms when calling LLMs, regardless of the model version.  Combining this with Provisioned Throughput further enhances reliability by guaranteeing capacity."
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "backoff_and_retry_for_LLMs.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/migration/UJ7 AutoML for natural language with Vertex AI Text Entity Extraction.ipynb

@@ -109,27 +109,27 @@
         "id": "3Sq3sGfdt89E"
       },
       "source": [
-        "## Before you begin\r\n",
-        "\r\n",
-        "### GPU run-time\r\n",
-        "\r\n",
-        "*Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select* **Runtime > Change Runtime Type > GPU**\r\n",
-        "\r\n",
-        "### Set up your GCP project\r\n",
-        "\r\n",
-        "**The following steps are required, regardless of your notebook environment.**\r\n",
-        "\r\n",
-        "1. [Select or create a GCP project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\r\n",
-        "\r\n",
-        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\r\n",
-        "\r\n",
-        "3. [Enable the Vertex APIs and Compute Engine APIs.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,compute_component)\r\n",
-        "\r\n",
-        "4. [Google Cloud SDK](https://cloud.google.com/sdk) is already installed in Google Cloud Notebooks.\r\n",
-        "\r\n",
-        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\r\n",
-        "Cloud SDK uses the right project for all the commands in this notebook.\r\n",
-        "\r\n",
+        "## Before you begin\n",
+        "\n",
+        "### GPU run-time\n",
+        "\n",
+        "*Make sure you're running this notebook in a GPU runtime if you have that option. In Colab, select* **Runtime > Change Runtime Type > GPU**\n",
+        "\n",
+        "### Set up your GCP project\n",
+        "\n",
+        "**The following steps are required, regardless of your notebook environment.**\n",
+        "\n",
+        "1. [Select or create a GCP project](https://console.cloud.google.com/cloud-resource-manager). When you first create an account, you get a $300 free credit towards your compute/storage costs.\n",
+        "\n",
+        "2. [Make sure that billing is enabled for your project.](https://cloud.google.com/billing/docs/how-to/modify-project)\n",
+        "\n",
+        "3. [Enable the Vertex APIs and Compute Engine APIs.](https://console.cloud.google.com/flows/enableapi?apiid=ml.googleapis.com,compute_component)\n",
+        "\n",
+        "4. [Google Cloud SDK](https://cloud.google.com/sdk) is already installed in Google Cloud Notebooks.\n",
+        "\n",
+        "5. Enter your project ID in the cell below. Then run the  cell to make sure the\n",
+        "Cloud SDK uses the right project for all the commands in this notebook.\n",
+        "\n",
         "**Note**: Jupyter runs lines prefixed with `!` as shell commands, and it interpolates Python variables prefixed with `$` into these commands."
       ]
     },
@@ -240,11 +240,11 @@
         "id": "9zpjPUOhvRQz"
       },
       "source": [
-        "### Authenticate your GCP account\r\n",
-        "\r\n",
-        "**If you are using Google Cloud Notebooks**, your environment is already\r\n",
-        "authenticated. Skip this step.\r\n",
-        "\r\n",
+        "### Authenticate your GCP account\n",
+        "\n",
+        "**If you are using Google Cloud Notebooks**, your environment is already\n",
+        "authenticated. Skip this step.\n",
+        "\n",
         "*Note: If you are on an Vertex notebook and run the cell, the cell knows to skip executing the authentication steps.*"
       ]
     },
@@ -1459,8 +1459,7 @@
         "id": "gM-YixlLmDy9"
       },
       "source": [
-        "### Make a batch prediction file\r\n",
-        "\r\n"
+        "### Make a batch prediction file\n"
       ]
     },
     {
@@ -1489,19 +1488,6 @@
         "! gsutil cat $gcs_test_item"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sPupiwqN_jAB"
-      },
-      "source": [
-        "*Example output*:\n",
-        "```\n",
-        "{\"content\": \"gs://migration-ucaip-trainingaip-20210301154552/test.txt\", \"mime_type\": \"text/plain\"}\n",
-        "Molecular basis of hexosaminidase A deficiency and pseudodeficiency in the Berks County Pennsylvania Dutch.\\tFollowing the birth of two infants with Tay-Sachs disease ( TSD ) , a non-Jewish , Pennsylvania Dutch kindred was screened for TSD carriers using the biochemical assay . A high frequency of individuals who appeared to be TSD heterozygotes was detected ( Kelly et al . , 1975 ) . Clinical and biochemical evidence suggested that the increased carrier frequency was due to at least two altered alleles for the hexosaminidase A alpha-subunit . We now report two mutant alleles in this Pennsylvania Dutch kindred , and one polymorphism . One allele , reported originally in a French TSD patient ( Akli et al . , 1991 ) , is a GT-- > AT transition at the donor splice-site of intron 9 . The second , a C-- > T transition at nucleotide 739 ( Arg247Trp ) , has been shown by Triggs-Raine et al . ( 1992 ) to be a clinically benign \" pseudodeficient \" allele associated with reduced enzyme activity against artificial substrate . Finally , a polymorphism [ G-- > A ( 759 ) ] , which leaves valine at codon 253 unchanged , is described\n",
-        "```\n"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1557,45 +1543,6 @@
         ")"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sPupiwqN_jAB"
-      },
-      "source": [
-        "*Example output*:\n",
-        "```\n",
-        "{\n",
-        "  \"parent\": \"projects/migration-ucaip-training/locations/us-central1\",\n",
-        "  \"batchPredictionJob\": {\n",
-        "    \"displayName\": \"ten_20210301154552\",\n",
-        "    \"model\": \"projects/116273516712/locations/us-central1/models/4400738115568795648\",\n",
-        "    \"inputConfig\": {\n",
-        "      \"instancesFormat\": \"jsonl\",\n",
-        "      \"gcsSource\": {\n",
-        "        \"uris\": [\n",
-        "          \"gs://migration-ucaip-trainingaip-20210301154552/test.jsonl\"\n",
-        "        ]\n",
-        "      }\n",
-        "    },\n",
-        "    \"outputConfig\": {\n",
-        "      \"predictionsFormat\": \"jsonl\",\n",
-        "      \"gcsDestination\": {\n",
-        "        \"outputUriPrefix\": \"gs://migration-ucaip-trainingaip-20210301154552/batch_output/\"\n",
-        "      }\n",
-        "    },\n",
-        "    \"dedicatedResources\": {\n",
-        "      \"machineSpec\": {\n",
-        "        \"machineType\": \"n1-standard-2\"\n",
-        "      },\n",
-        "      \"startingReplicaCount\": 1,\n",
-        "      \"maxReplicaCount\": 1\n",
-        "    }\n",
-        "  }\n",
-        "}\n",
-        "```\n"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -1638,42 +1585,6 @@
         "print(MessageToJson(request.__dict__[\"_pb\"]))"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sPupiwqN_jAB"
-      },
-      "source": [
-        "*Example output*:\n",
-        "```\n",
-        "{\n",
-        "  \"name\": \"projects/116273516712/locations/us-central1/batchPredictionJobs/3588251799200464896\",\n",
-        "  \"displayName\": \"ten_20210301154552\",\n",
-        "  \"model\": \"projects/116273516712/locations/us-central1/models/4400738115568795648\",\n",
-        "  \"inputConfig\": {\n",
-        "    \"instancesFormat\": \"jsonl\",\n",
-        "    \"gcsSource\": {\n",
-        "      \"uris\": [\n",
-        "        \"gs://migration-ucaip-trainingaip-20210301154552/test.jsonl\"\n",
-        "      ]\n",
-        "    }\n",
-        "  },\n",
-        "  \"outputConfig\": {\n",
-        "    \"predictionsFormat\": \"jsonl\",\n",
-        "    \"gcsDestination\": {\n",
-        "      \"outputUriPrefix\": \"gs://migration-ucaip-trainingaip-20210301154552/batch_output/\"\n",
-        "    }\n",
-        "  },\n",
-        "  \"state\": \"JOB_STATE_PENDING\",\n",
-        "  \"completionStats\": {\n",
-        "    \"incompleteCount\": \"-1\"\n",
-        "  },\n",
-        "  \"createTime\": \"2021-03-01T17:59:42.777083Z\",\n",
-        "  \"updateTime\": \"2021-03-01T17:59:42.777083Z\"\n",
-        "}\n",
-        "```\n"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1739,42 +1650,6 @@
         "print(MessageToJson(request.__dict__[\"_pb\"]))"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "sPupiwqN_jAB"
-      },
-      "source": [
-        "*Example output*:\n",
-        "```\n",
-        "{\n",
-        "  \"name\": \"projects/116273516712/locations/us-central1/batchPredictionJobs/3588251799200464896\",\n",
-        "  \"displayName\": \"ten_20210301154552\",\n",
-        "  \"model\": \"projects/116273516712/locations/us-central1/models/4400738115568795648\",\n",
-        "  \"inputConfig\": {\n",
-        "    \"instancesFormat\": \"jsonl\",\n",
-        "    \"gcsSource\": {\n",
-        "      \"uris\": [\n",
-        "        \"gs://migration-ucaip-trainingaip-20210301154552/test.jsonl\"\n",
-        "      ]\n",
-        "    }\n",
-        "  },\n",
-        "  \"outputConfig\": {\n",
-        "    \"predictionsFormat\": \"jsonl\",\n",
-        "    \"gcsDestination\": {\n",
-        "      \"outputUriPrefix\": \"gs://migration-ucaip-trainingaip-20210301154552/batch_output/\"\n",
-        "    }\n",
-        "  },\n",
-        "  \"state\": \"JOB_STATE_PENDING\",\n",
-        "  \"completionStats\": {\n",
-        "    \"incompleteCount\": \"-1\"\n",
-        "  },\n",
-        "  \"createTime\": \"2021-03-01T17:59:42.777083Z\",\n",
-        "  \"updateTime\": \"2021-03-01T17:59:42.777083Z\"\n",
-        "}\n",
-        "```\n"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -1798,19 +1673,6 @@
         "    time.sleep(60)"
       ]
     },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "trainingpipelines_create:migration,new,response,icn"
-      },
-      "source": [
-        "*Example output*:\n",
-        "```\n",
-        "gs://migration-ucaip-trainingaip-20210301154552/batch_output/prediction-ten_20210301154552-2021-03-01T17:59:42.638222Z/predictions_00001.jsonl\n",
-        "{\"instance\":{\"content\":\"gs://migration-ucaip-trainingaip-20210301154552/test.txt\",\"mimeType\":\"text/plain\"},\"prediction\":{\"ids\":[\"7806436899697983488\",\"7806436899697983488\",\"7806436899697983488\",\"4347672385877442560\",\"4347672385877442560\",\"4347672385877442560\"],\"displayNames\":[\"SpecificDisease\",\"SpecificDisease\",\"SpecificDisease\",\"Modifier\",\"Modifier\",\"Modifier\"],\"textSegmentStartOffsets\":[\"149\",\"19\",\"169\",\"236\",\"688\",\"330\"],\"textSegmentEndOffsets\":[\"165\",\"45\",\"171\",\"238\",\"690\",\"332\"],\"confidences\":[0.99957836,0.9995628,0.9995044,0.9993287,0.9993144,0.99927235]}}\n",
-        "```\n"
-      ]
-    },
     {
       "cell_type": "markdown",
       "metadata": {
@@ -2338,11 +2200,11 @@
         "id": "bQ-VVaSxJjkd"
       },
       "source": [
-        "# Cleaning up\r\n",
-        "\r\n",
-        "To clean up all GCP resources used in this project, you can [delete the GCP\r\n",
-        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\r\n",
-        "\r\n",
+        "# Cleaning up\n",
+        "\n",
+        "To clean up all GCP resources used in this project, you can [delete the GCP\n",
+        "project](https://cloud.google.com/resource-manager/docs/creating-managing-projects#shutting_down_projects) you used for the tutorial.\n",
+        "\n",
         "Otherwise, you can delete the individual resources you created in this tutorial."
       ]
     },
@@ -2404,7 +2266,7 @@
   ],
   "metadata": {
     "colab": {
-      "name": "UJ7 unified AutoML for natural language with Vertex AI Text Entity Extraction.ipynb",
+      "name": "UJ7 AutoML for natural language with Vertex AI Text Entity Extraction.ipynb",
       "toc_visible": true
     },
     "kernelspec": {

notebooks/community/model_garden/benchmarking_reports/jax_vit_benchmarking_report.md

@@ -0,0 +1,253 @@
+# ViT PyTorch vs JAX training benchmarks on Vertex AI Training Platform
+
+Lav Rai, Software Engineer, Google Cloud
+
+Xiang Xu, Software Engineer, Google Cloud
+
+Andreas Steiner, Software Engineer, Google DeepMind
+
+Tao Wang, Software Engineer, Google DeepMind
+
+Alexander Kolesnikov, Research Engineer, Google DeepMind
+
+## Introduction
+
+Many repositories now offer both PyTorch and JAX versions of a model. For
+example, [Hugging Face offers many models such as GPT2, BERT][1]
+etc. Other examples are [OpenLLaMa][2] and [ViT][3]
+models which were first developed in JAX and then their corresponding PyTorch
+versions were made available. **Given both the PyTorch and JAX options for a
+model, it may not be obvious as to which option to choose**. To make such a
+decision, it is important for one to know about the training cost, effectiveness
+and efficiency for each choice.
+
+Apart from the framework choice, the other choice that one faces on Vertex AI
+training platform is the type and count of the accelerators. Although the
+[Vertex AI pricing table][4] lists the price per hour for each
+machine, **one may not know beforehand about the training speed of JAX and
+PyTorch frameworks for different types and count of the accelerators**.
+
+If one has access to some training benchmark numbers for the same model
+under (a) PyTorch and JAX frameworks and (b) for different types and count of
+the accelerators, then it will be easier for them to make a cost effective
+decision. Such a benchmark will also aid the developers in identifying strength
+and weakness of different choices and then figure out recipes to remove those
+weaknesses if possible.
+
+This blog uses the ViT [classification models][5] of varying sizes
+to benchmark the training performance of PyTorch and JAX versions on the Vertex
+AI Platform under different machine configurations. The goal is to:
+
+- Benchmark OSS ViT training for both PyTorch and JAX frameworks.
+- Benchmark OSS ViT L16, H14, g14, and G14 models.
+- Benchmark OSS ViT PyTorch training with A100 GPUs.
+- Benchmark OSS ViT JAX training with A100 GPUs and TPU V3 accelerators.
+
+## Benchmarking setup
+
+This section lays out the benchmarking set up for the [PyTorch][6] and [JAX][7]
+frameworks and provides a reasoning for choosing those settings.
+
+### PyTorch GPU
+
+#### Machine configuration
+
+We run training jobs on [Vertex AI Custom Training][8] using 1
+single node with 8 A100-40GB GPUs.
+
+- Machine type: [a2-highgpu-8g][9]
+- Machine count: 1
+- Accelerator type: [NVIDIA_TESLA_A100 (40GB)][10]
+- Accelerator count: 8
+
+#### Modeling
+
+We benchmark 4 variants of ViT model in different sizes:
+
+- [ViT-L16, 300M params][11]
+- [ViT-H14, 630M params][12]
+- [ViT-g14, 1B params][13]
+- [ViT-G14, 1.8B params][14]
+
+We use the Huggingface [transformers library][15] for ViT L16 and
+H14 variants, and the [TIMM library][16] for ViT g14 and G14
+variants.
+
+#### Dataset
+
+We run training against the [cifar10][17] dataset with 50K training
+images and 10K test images. To factor out network communication overhead for
+data loading, we copy the whole dataset to the local disk then load data from
+the local disk during training.
+
+#### Training parameters
+
+- Trainer
+  - We use [PyTorch Lightning][18] as the trainer for the
+    boilerplate data loading and train loop coding.
+- Precision
+  - Float16
+- Input resolution
+  - 224 x 224
+- Strategy
+  - We use [DDP][19] for models which can be entirely loaded to one
+    GPU, use [Deepspeed-ZeRO][20] otherwise:
+    - ViT-L16: DDP
+    - ViT-H14: DDP
+    - ViT-g14: DDP
+    - ViT-G14: Deepspeed-ZeRO stage-3
+- Batch size
+  - We use the max batch size as power of 2 without CUDA OOM for each model:
+    - ViT-L16: 64 per GPU
+    - ViT-H14: 16 per GPU
+    - ViT-g14: 16 per GPU
+    - ViT-G14: 32 per GPU
+- Compilation
+  - We apply [torch.compile][21] to model whenever it's applicable:
+    - ViT-L16: torch.compile
+    - ViT-H14: torch.compile
+    - ViT-g14: torch.compile
+    - ViT-G14: N/A
+
+### JAX TPU and GPU
+
+#### Machine configuration
+
+All the TPU and GPU training jobs are run on [Vertex AI Custom
+Training][8]. The following machine configurations were used for the
+TPU and GPU experiments:
+
+**Note**: TPU V3 POD requires multi-host supporting training code. For example,
+a 32 core POD runs on 4 hosts with each host using 8 cores.
+
+**Note**: 8 A100 are similar to TPU V3 32 cores in terms of [Vertex AI
+pricing][4].
+
+**Note**: [Each TPU v3 chip has 2 cores which can use 32 GB high-bandwidth
+memory][22] (16 GB per core) so total memory for 32 cores is 16x32 =
+512 GB. Therefore for the same price, TPUs offer more memory than 8 A100-40GB
+GPUs.
+
+#### Modeling
+
+We decided to use an OSS code repository for model implementation. Using an OSS
+repository helps anyone to independently verify the benchmarking results and
+also relate to the results well. For JAX, we selected the
+[Big Vision][23] code repository.
+
+Same as the PyTorch modeling, we benchmark 4 variants of ViT model in different
+sizes:
+
+- [ViT-L16, 300M params][24]
+- [ViT-H14, 630M params][24]
+- [ViT-g14, 1B params][24]
+- [ViT-G14, 1.8B params][24]
+
+**Note**: The [Big Vision code repo][23] has not made the
+checkpoints publicly available for the models larger than the ViT-L16. Therefore
+for the rest of the three variants, the experiments only used random
+initialization for benchmarking the training speed.
+
+#### Dataset
+
+We use training against the [cifar10 TensorFlow dataset][25] with
+50K training images and 10K test images. This dataset is the same as the one
+used for PyTorch experiments except that it is loaded as a TensorFlow dataset.
+Similar to the PyTorch experiments, we copy the whole dataset to the docker
+image to factor out network communication overhead for data loading.
+
+#### Training parameters
+
+- Precision
+  - "bfloat16" setting was used.
+- Input resolution
+  - 224 x 224 after resize (to 448x448) and random crop (to 224x224) before
+    training.
+    - This resolution for training was the same as the PyTorch settings.
+- Strategy
+  - Used DDP for all models except ViT-G14. ViT-G14 used the FSDP strategy.
+- Batch size
+  - We use the max batch size as power of 2 without OOM for each model. The
+    [Benchmarking results][26] section shows the final
+    batch size for each experiment.
+  - Once a maximum batch-size for TPU V3 8 cores was determined, we just scaled
+    it linearly for 32 cores.
+  - Once a maximum batch-size for 1 A100 GPU was determined, we just scaled it
+    linearly for 8 A100 GPUs.
+- Compilation
+  - [jax.jit() compilation][27] is used in JAX codes for efficient
+    execution in XLA.
+- GPU related flags
+  - The following flags are set in the dockerfile for the GPU runs.
+    - Note: _xla_gpu_enable_pipelined_collectives_ is set to false for the
+      ViT-G14 FSDP run.
+
+### Evaluation metric
+
+For both the PyTorch and JAX experiments, the following evaluation metrics are
+collected:
+
+- Throughput: Images-per-second observed for training.
+- Cost: The training-cost-per-epoch (USD).
+
+**Note**: The above metrics are not biased against any framework or machine
+configurations. In addition, these metrics will help one decide the most
+efficient training configurations on Vertex AI.
+
+## Benchmarking results
+
+The lowest cost experiment for each model is marked in **bold** in the last
+column.
+
+![vit_benchmarking_table](images/vit_benchmarking_table.png)
+
+The following bar charts summarize the performance visually:
+
+![vit_training_time](images/vit_training_time.png)
+
+![vit_training_cost](images/vit_training_cost.png)
+
+The following section provides observations and conclusions for these results.
+
+## Observation and Conclusions
+
+- Training with JAX TPU V3 POD with 32 cores costs 33% less than the PyTorch GPU
+  8 A100-40GBs runs.
+- Training with JAX GPU 8 A100-40GBs costs 23% less than the PyTorch GPU 8
+  A100-40GBs runs.
+- JAX TPU V3 POD with 32 cores was 4x faster and slightly more cost-effective
+  than the JAX TPU V3 8 core run for the ViT-large model. This indicates that it
+  might be better to use more cores. The JAX TPU V3 speed scales very well with
+  the number of cores.
+- Cloud TPU VM training speed numbers were the same as the Vertex AI for
+  TPU V3 8 cores. The dataset was copied to the docker in both the cases.
+- The training-cost-per-epoch increases with the model size irrespective of the
+  framework.
+
+[1]: https://github.com/huggingface/transformers/blob/main/examples/research_projects/jax-projects/README.md#quickstart-flax-and-jax-in-transformers
+[2]: https://github.com/openlm-research/open_llama
+[3]: https://github.com/google-research/vision_transformer
+[4]: https://cloud.google.com/vertex-ai/pricing#custom-trained_models
+[5]: https://arxiv.org/abs/2010.11929
+[6]: #pytorch-gpu
+[7]: #jax-tpu-and-gpu
+[8]: https://cloud.google.com/vertex-ai/docs/training/overview
+[9]: https://cloud.google.com/vertex-ai/docs/training/configure-compute#machine-types
+[10]: https://cloud.google.com/vertex-ai/docs/training/configure-compute#specifying_gpus
+[11]: https://huggingface.co/google/vit-large-patch16-224-in21k
+[12]: https://huggingface.co/google/vit-huge-patch14-224-in21k
+[13]: https://github.com/huggingface/pytorch-image-models/blob/v0.9.2/timm/models/vision_transformer.py#L1308
+[14]: https://github.com/huggingface/pytorch-image-models/blob/v0.9.2/timm/models/vision_transformer.py#L1312
+[15]: https://huggingface.co/docs/transformers/main/model_doc/vit#transformers.ViTModel
+[16]: https://github.com/huggingface/pytorch-image-models
+[17]: https://huggingface.co/datasets/cifar10
+[18]: https://lightning.ai/docs/pytorch/stable/
+[19]: https://pytorch.org/docs/stable/notes/ddp.html
+[20]: https://www.deepspeed.ai/tutorials/zero/
+[21]: https://pytorch.org/tutorials/intermediate/torch_compile_tutorial.html
+[22]: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#tpu_v3
+[23]: https://github.com/google-research/big_vision
+[24]: https://screenshot.googleplex.com/BximJgxsgvBVu38
+[25]: https://www.tensorflow.org/datasets/catalog/cifar10
+[26]: #benchmarking-results
+[27]: https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html

notebooks/community/model_garden/benchmarking_reports/pytorch_openllama_hyperparameter_tuning_benchmark_report.md

@@ -0,0 +1,188 @@
+# Benchmark report on hyperparameter tuning the OpenLLaMA models on Google Cloud Vertex Model Garden
+
+Changyu Zhu, Software Engineer, Google Cloud
+
+Dustin Luong, Software Engineer, Google Cloud
+
+Gary Wei, Software Engineer, Google Cloud
+
+Genquan Duan, Software Engineer, Google Cloud
+
+## Introduction
+
+Fine-tuning of LLMs can be non-trivial to find an optimal configuration of
+machine types, training parameters, and other hyperparameters that achieves a
+good balance between cost efficiency and model performance. To facilitate users
+in conducting tuning experiments, this report benchmarks fine-tuning OpenLLaMA
+models with [Vertex AI Hyperparameter Tuning Service](https://cloud.google.com/vertex-ai/docs/training/hyperparameter-tuning-overview), demonstrating both efficiency
+and effectiveness. Similar hyperparameter tuning techniques can apply to other models as well.
+
+## Key takeaways
+
+- **The hyperparameter tuning service finds good parameters**: The best model found by the hyperparameter tuning service has an average improvement of around 4% in accuracy in *ARC*, *HellaSwag*, and *TruthfulQA* datasets, while only tuning the learning rate.
+
+- **Hyperparameter tuning works with QLoRA on limited resources**: 4bit QLoRA is sufficient for hyperparameter tuning to find a set of good parameters. In this way, all OpenLLaMA models can run on 1 single `NVIDIA_L4` GPU. It is also possible to train for more steps on the good parameters discovered by hyperparameter tuning, avoiding the waste of computing resources on fine-tuning with suboptimal hyperparameters.
+
+- **Hyperparameter tuning is cost-effective**: While `NVIDIA_L4` is slower than `NVIDIA_TESLA_V100`, it costs less and avoids the overhead of multi-GPU training since it has more GPU memory. Finding a good 3B/7B/13B OpenLLaMA model costs $28.5671, $47.8016, and $87.9208, respectively.
+
+## Benchmarking setup
+
+This section describes the experiment setup of the hyperparameter tuning experiments. The default tuning parameters are:
+
+### Machine configuration
+
+- Machine type: g2-standard-8
+- Machine count: 1
+- Accelerator type: NVIDIA_L4
+- Accelerator count: 1
+
+### Modeling
+
+We benchmark all 3 OpenLLaMA models:
+
+- [open_llama_3b](https://huggingface.co/openlm-research/open_llama_3b)
+- [open_llama_7b](https://huggingface.co/openlm-research/open_llama_7b)
+- [open_llama_13b](https://huggingface.co/openlm-research/open_llama_13b)
+
+We use the Huggingface [PEFT](https://github.com/huggingface/peft) library for fine-tuning.
+
+### Training dataset
+
+We use the dataset [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) loaded directly via Huggingface.
+
+### Training parameters
+
+The set of training parameters used during benchmarking:
+
+- Batch size: 4
+- Precision mode: 4bit QLoRA
+- LoRA rank: 32
+- LoRA alpha: 64
+- Max sequence length: 512
+- Max train steps: 1000
+
+### Evaluation dataset
+
+We use the [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) library injected into the training loop for evaluation. The hyperparameter tuning job will pick the model according to the evaluation metrics.
+
+- Eval task: [ARC Challenge](https://huggingface.co/datasets/ai2_arc)
+- Eval metric: acc_norm
+- Max eval examples: 10000
+
+### Standalone evaluation dataset
+
+After finding the best model with Vertex hyperparameter tuning service, we run standalone evaluations with the model on the following datasets:
+
+- [ARC Challenge](https://huggingface.co/datasets/ai2_arc)
+- [HellaSwag](https://huggingface.co/datasets/Rowan/hellaswag)
+- [TruthfulQA](https://huggingface.co/datasets/EleutherAI/truthful_qa_mc)
+
+### Hyperparameter tuning
+
+We only tune the learning rate hyperparameter. It is considered a floating point value in the continuous range [1e-5, 1e-4]. We run 8 trials in total, with a parallelism of 1 or 2.
+
+### Code example
+
+The following code example launches an example hyperparameter tuning job of OpenLLaMA 7B model.
+
+```py
+from google.cloud import aiplatform
+from google.cloud.aiplatform import hyperparameter_tuning as hpt
+
+
+TRAIN_DOCKER_URI = 'us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20231130_0936_RC00'
+output_dir = "gs://path/to/output/dir"
+base_model_id = "openlm-research/open_llama_7b"
+dataset_name = "timdettmers/openassistant-guanaco"
+hpt_precision_mode = "4bit"
+machine_type = "g2-standard-8"
+accelerator_type = "NVIDIA_L4"
+accelerator_count = 1
+eval_task = "arc_challenge"
+eval_metric_name = "acc_norm"
+max_steps = 1000
+eval_limit = 10000
+
+flags = {
+    "learning_rate": 1e-5,
+    "precision_mode": hpt_precision_mode,
+    "task": "instruct-lora",
+    "pretrained_model_id": base_model_id,
+    "output_dir": output_dir,
+    "warmup_steps": 10,
+    "max_steps": max_steps,
+    "lora_rank": 32,
+    "lora_alpha": 64,
+    "lora_dropout": 0.05,
+    "dataset_name": dataset_name,
+    "eval_steps": max_steps + 1,  # Only evaluates at the end.
+    "eval_tasks": eval_task,
+    "eval_limit": eval_limit,
+    "eval_metric_name": eval_metric_name,
+}
+worker_pool_specs = [
+    {
+        "machine_spec": {
+            "machine_type": machine_type,
+            "accelerator_type": accelerator_type,
+            "accelerator_count": accelerator_count,
+        },
+        "replica_count": 1,
+        "container_spec": {
+            "image_uri": TRAIN_DOCKER_URI,
+            "args": ["--{}={}".format(k, v) for k, v in flags.items()],
+        },
+    }
+]
+metric_spec = {"model_performance": "maximize"}
+parameter_spec = {
+    "learning_rate": hpt.DoubleParameterSpec(
+        min=1e-5, max=1e-4, scale="linear"
+    ),
+}
+
+train_job = aiplatform.CustomJob(
+    display_name=job_name,
+    worker_pool_specs=worker_pool_specs,
+    staging_bucket=STAGING_BUCKET,
+)
+
+train_hpt_job = aiplatform.HyperparameterTuningJob(
+    display_name=f"{job_name}_hpt",
+    custom_job=train_job,
+    metric_spec=metric_spec,
+    parameter_spec=parameter_spec,
+    max_trial_count=8,
+    parallel_trial_count=2,
+)
+
+train_hpt_job.run()
+```
+
+## Benchmark results
+
+### Fine-tuning cost
+
+The fine-tuning cost is calculated from `us-central1` pricing and may be subject to changes.
+
+| Model         | Train time | Trials | Parallel Trials | Hourly cost | Cost     | Eval acc_norm (ARC-Challenge) |
+|---------------|------------|--------|-----------------|-------------|----------|-------------------------------|
+| OpenLLaMA 3B  | 16 hrs     | 8      | 2               | $1.7072     | $28.5671 | 39.9%                         |
+| OpenLLaMA 7B  | 28 hrs     | 8      | 2               | $1.7072     | $47.8016 | 45.8%                         |
+| OpenLLaMA 13B | 103 hrs    | 8      | 1               | $0.8536     | $87.9208 | 47.6%                         |
+
+### Fine-tuning performance
+
+Here are the evaluation results of the best model found by hyperparameter tuning, compared with the baseline model. The column `Eval acc_norm` is calculated during training, which is always lower than that during standalone evaluation, because the model is loaded and evaluated at a lower precision (4bit during training / float16 during standalone evaluation).
+
+| Model         | Eval acc_norm (ARC-Challenge) | ARC    | hellaswag | Truthfulqa_mc | ∆ARC   | ∆Hellaswag | ∆Truthfulqa_mc | ∆Average |
+|---------------|-------------------------------|--------|-----------|---------------|--------|------------|----------------|----------|
+| OpenLLaMA 3B  | 39.9%                         | 41.47% | 69.97%    | 38.31%        | +1.62% | +7.32%     | +3.34%         | +4.09%   |
+| OpenLLaMA 7B  | 45.8%                         | 49.83% | 75.53%    | 41.53%        | +2.82% | +3.55%     | +6.68%         | +4.35%   |
+| OpenLLaMA 13B | 47.6%                         | 52.20% | 78.90%    | 44.27%        | +1.01% | +3.67%     | +6.19%         | +3.62%   |
+
+## Related documents
+
+1. [Benchmark report on fine tuning the OpenLLaMA 7B model on Google Cloud Vertex Model Garden
+](
+https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/community-content/vertex_model_garden/benchmarking_reports/pytorch_openllama_7b_finetune_benchmark_report.md)

notebooks/community/model_garden/benchmarking_reports/stable_diffusion_v1-5_benchmarking_report.md

@@ -0,0 +1,218 @@
+# Benchmark Stable Diffusion v1-5 Fine Tuning and Serving With Google Cloud Vertex Model Garden
+
+Dustin Luong, Software Engineer, Google Cloud
+Gary Wei, Software Engineer, Google Cloud
+Changyu Zhu, Software Engineer, Google Cloud
+Genquan Duan, Software Engineer, Google Cloud
+
+## Introduction
+[The public notebook][1] shows the full examples of fine tuning and serving of Stable diffusion v1-5. [The github repo][2] contains examples of building training and serving dockers for Google Cloud Vertex Model Garden. This report benchmarks Stable diffusion v1-5 fine tuning and serving in Google Cloud Vertex AI, showing both efficiencies and effectiveness.
+
+### Benchmark Highlights
+- Fine tuning
+  - Stable diffusion v1-5 with LoRA and Gradient checkpointing only requires ~10G GPU memory. Larger batch sizes, or larger resolutions require more GPU memories, but not does not change much for different LoRA ranks.
+  - The fine tuning speed is fast in ~11 minutes for 1k steps, and costs less than $1 in 1 A100. The fine tuning speed increases with batch sizes, decreases with resolution, but is not affected much by LoRA ranks.
+  - LoRA tunes a few percent (only 0.1% with LoRA rank=8) of all parameters, and the tuned models are very small (only 3.1MB with LoRA rank=8).
+  - Dreambooth+LoRA and Dreambooth can achieve similar performances, but Dreambooth LoRA can require much less GPU.
+  - Increasing batch size, reducing training steps, and increasing learning rate can result in models with the same performance for less cost.
+- Inference
+  - The optimized serving docker pytorch-peft-serve can speed up inference by 2x than current pytorch-diffuser-serve, and support both base models and fine tuned lora models.
+  - The optimized serving docker pytorch-peft-serve can generate 4 512*512 images in 4.1 seconds on 1 V100 and 1.7 seconds on 1 A100.
+
+Benchmark details are below.
+
+## Fine Tuning Benchmarks
+
+### Experiment Setup
+We mainly compare two tuning algorithms:
+- parameter efficient finetuning based on [dreambooth][3] and [LoRA][4] (shorten as Dreambooth+LoRA below)
+- full parameter fine tuning based on [dreambooth][3] (shorten as Dreambooth below)
+
+And then  report benchmark results on GPU memories, tuning parameters, tuning speeds, costs and accuracy, using the public oxford flowers dataset: [train][5] and [test][6], where the column blip_caption as texts, and column image as images. We also benchmark subject and prompt fidelity using the [dataset][7] from the Dreambooth paper.
+
+The default tuning parameters during benchmark are:
+- Hardware: 1 A100 40G
+- batch size: 4
+- lora_rank: 8
+- resolution: 512
+- max_train_steps: 10
+- use_lora: False
+- gradient_checkpointing: False
+
+```
+# Examples to start finetuning dockers.
+MODEL_NAME="runwayml/stable-diffusion-v1-5"
+OUTPUT_DIR=<OUTPUT_DIR>
+INSTANCE_DATA_DIR=<INSTANCE_DATA_DIR>
+INSTANCE_PROMPT=<INSTANCE_PROMPT>
+IMAGE="us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train"
+docker run \
+  --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0 \
+  --rm --name "test_gpu" \
+  -it ${IMAGE} \
+  --task=text-to-image-dreambooth-lora-peft \
+  --pretrained_model_name_or_path=$MODEL_NAME \
+  --resolution=512 \
+  --instance_data_dir=$INSTANCE_DATA_DIR \
+  --instance_prompt=$INSTANCE_PROMPT \
+  --train_batch_size=4 \
+  --max_train_steps=10 \
+  --output_dir=${OUTPUT_DIR} \
+  --use_lora \
+  --lora_r=8 \
+  --gradient_checkpointing
+```
+
+### GPU Memories
+Many various factors will impact GPU memory usages. In this benchmark, we mainly benchmark with different finetuning algorithms, batch sizes, lora rank, resolution, and then recommended max batch size on different GPUs.
+
+
+![sd_v1-5_peak_gpu_algorithm](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_algorithm.png)
+
+![sd_v1-5_peak_gpu_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_batch_size.png)
+
+![sd_v1-5_peak_gpu_lora_rank](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_lora_rank.png)
+
+![sd_v1-5_peak_gpu_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_peak_gpu_resolution.png)
+
+- LoRA tuning reduced about 47% peak RAM and 42% peak VRAM for GPU memory, compared to full parameter fine tuning.
+- Gradient checkpointing decreases about 1% peak RAM and 31% peak VRAM for GPU memory further, compared without gradient checkpointing.
+- The GPU memory does not change much for different LoRA ranks.
+- Larger batch sizes require more GPU memories.
+- Larger resolutions require more GPU memories.
+- Dreambooth+LoRA+Gradient_Checkpointing can support max batch size as 32, or max resolution as 2048, but Dreambooth can only support max batch size as 8, or max resolution as 1024.
+
+### Fine Tuning Parameters
+This section shows the percentage of trainable parameters, and tuned model sizes.
+
+- LoRA tunes quite a few percent (only 0.1% with LoRA rank=8) of all parameters, and the tuned models are very small (only 3.1MB with LoRA rank=8).
+
+| LoRA Rank | Trainable parameters | Total parameters | Trainable Parameter Percentage | Fine tuned model size (MB) |
+|---|---|---|---|---|
+| 4 | 398592 | 859919556 | 0.05% | 1.57 |
+|8 | 797184 | 860318148 | 0.09% | 3.09 |
+| 16 | 1594368| 861115332| 0.19%| 6.13|
+| 32| 3188736| 862709700| 0.37%| 12.21|
+### Fine Tuning Speed And Costs
+Fine tuning speeds and costs are affected by many different factors, such as batch size, tuning parameters, image resolutions, GPUs, and datasets. In order to make the report easy to understand, we set the following values in this section:
+- Hardware: 1 A100 40G
+- use_lora: True
+- gradient_checkpointing: True
+
+![sd_v1-5_training_speed_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_batch_size.png)
+
+![sd_v1-5_training_speed_lora_rank](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_lora_rank.png)
+
+![sd_v1-5_training_speed_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_speed_resolution.png)
+
+![sd_v1-5_training_cost_max_steps](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_training_cost_max_steps.png)
+
+- The fine tuning speed increases with batch sizes, decreases with resolution, but is not affected much by LoRA ranks.
+- The fine tuning speed is about 11 minutes for 1k steps, and costs less than $1 in 1 A100.
+
+### Fine Tuning Quality
+In this benchmark, we mainly benchmark Dreambooth and Dreambooth+LoRA to compare fine tuning quality. We compare [subject fidelity scored (DINO)][8], how well the subject is represented in the generated images, and [prompt fidelity scores (CoCa)][9], how well the generated images match the given prompt, for a single subject, a [dog][10] from the dataset released with the original Dreambooth paper. In practice, we recommend saving checkpoints periodically and inspecting validation prompts visually. We fine tuned the unet without fine tuning the text encoder and used the following hyperparameters:
+
+Dreambooth
+- Learning rate: 5e-6
+- Batch size: 1
+
+Dreambooth+LoRA
+- Learning rate: 1e-4
+- Batch size: 1
+
+![sd_v1-5_finetuning_quality_subject_fidelity](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_finetuning_quality_subject_fidelity.png)
+
+![sd_v1-5_finetuning_quality_prompt_fidelity](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_finetuning_quality_prompt_fidelity.png)
+
+- Fine tuning with Dreambooth or Dreambooth+LoRA can result in models with comparable performance. The base model produced images of the class rather than the instance.
+- Dreambooth+LoRA is able to achieve the same subject fidelity score as Dreambooth if trained for more epochs.
+- Increasing the number of training steps results in better subject fidelity but at the cost of prompt fidelity.
+
+### Suggested Max Batch Sizes By Resolutions
+We benchmarked and suggested max batch sizes by resolutions on 1 A100 and 1 V100 as below. This is with LoRA and gradient checkpointing enabled.
+
+![sd_v1-5_batch_size_by_resolution](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_batch_size_by_resolution.png)
+
+### Fine Tuning Cost Optimization
+Increasing batch size allows for more images to be considered at each training step for fine tuning. This allows models to be trained in fewer training steps. In this benchmark, we aim to show how batch size can be increased to reduce training costs while still preserving subject and prompt fidelity.
+
+Since the training dataset consists of 5 images, we train with a batch size of 5 and reduce the number of training steps from 400 to 80. Doing so results in a model that has not learned the subject since we’ve decreased the number of training steps. Conceptually, the model is taking a more precise step at each iteration, but it is taking fewer steps. To compensate for this, we increased the learning rate from 5e-6 and observed the best results at 1e-5 for full parameter finetuning.
+
+![sd_v1-5_subject_fidelity_batch_size_5_learning_rate](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_subject_fidelity_batch_size_5_learning_rate.png)
+
+![sd_v1-5_prompt_fidelity_batch_size_5_learning_rate](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_prompt_fidelity_batch_size_5_learning_rate.png)
+
+Comparing cost of training the “best” model for batch size 1 vs. batch size 5
+
+![sd_v1-5_cost_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_cost_batch_size.png)
+
+
+| Train method| Training parameters| Sample image| CoCa (prompt fidelity)| DINO (subject fidelity) | Cost of training on A100 |
+|---|---|---|---|---|---|
+| dreambooth| dreambooth, num_train_steps=400, batch_size=1, lr=5e-6| ![dog1](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog1.png) | 0.12215| 0.76531| $0.26 |
+| dreambooth | dreambooth, num_train_steps=80, batch_size=5,lr=1e-5| ![dog2](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog2.png)| 0.12644| 0.74697 | $0.15 |
+| dreambooth-lora| num_train_steps=500, batch_size=1, lr=1e-4, gc|![dog3](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog3.png)| 0.12856| 0.78148 | $0.26|
+| dreambooth-lora | num_train_steps=50, batch_size=5, lr=1e-3, gc | ![dog4](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_dog4.png) | 0.12566 | 0.75479 | $0.09 |
+
+
+
+A followup question is that since finetuning can be run on a single GPU, should finetuning be run on 1 V100 or A100?
+
+Setup:
+- num_train_steps=800 / batch_size
+- Resolution=512
+
+![sd_v1-5_cost_training_method_batch_size](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_cost_training_method_batch_size.png)
+- Although V100 has a lower $/hr cost than an A100, the same training setup takes longer. Even given the longer training time, the cost on V100 is still lower.
+- Dreambooth+LoRA enables training with larger batch sizes, however, larger batch sizes will not necessarily mean faster training time.
+- It is possible to fine tune with 1 V100 on 512 resolution with Dreambooth+LoRA.
+- Dreambooth fine tuning must be run on 1 A100 at 512 resolution.
+
+## Inference Benchmarks
+We provide two serving dockers in vertex model garden for stable diffusion:
+- pytorch-diffuser-serve:
+  - us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve
+  - This serving docker only serves base stable diffusion models and does not contain any optimizations yet.
+- pytorch-peft-serve:
+  - us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve
+  - This serving docker can serve base stable diffusion models, and base stable diffusion models with fine tuned lora models, and contains optimization for serving.
+
+We run the two serving dockers on T4/V100/A100 to generate 4 512*512 images, and compare the inference speed without network considerations as:
+
+![sd_v1-5_inference_speed_gpu](images/stable_diffusion_v1-5_benchmarking_report/sd_v1-5_inference_speed_gpu.png)
+The speed up of optimized pytorch-peft-serve is about 2x than current pytorch-diffuser-serve.
+
+### Serving cost comparison
+
+Pytorch-diffuser-serve (without any optimizations)
+
+| GPU type|  Time required to generate 4 512x512 images | Machine unit price ($ / hour) | Cost per image ($) |
+|---|---|---|---|
+| T4 | 28.6 | 0.4025| 0.00080 |
+| V100 | 8.8 | 2.852| 0.00174|
+| A100 | 4.2 | 4.2245 | 0.00123 |
+
+Pytorch-peft-serve (with optimizations)
+
+| GPU type | Time required to generate 4 512x512 images | Machine unit price ($ / hour) | Cost per image ($) |
+|--- |---|---|---|
+| T4 | 12.6 | 0.4025 | 0.00035 |
+| V100 | 4.1 | 2.852 | 0.00081 |
+|  A100 | 1.7 | 4.2245 | 0.00050 |
+
+- The optimized pytorch-peft-serve has approximately half the price per image, compared with the un-optimized pytorch-diffuser-serve.
+- Serving the model with a T4 is most cost effective, however, serving with an A100 still has the best throughput and fastest predictions.
+
+
+
+[1]: https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_stable_diffusion.ipynb
+[2]: https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/main/community-content/vertex_model_garden/model_oss
+[3]: https://arxiv.org/abs/2208.12242
+[4]: https://arxiv.org/abs/2106.09685
+[5]: https://huggingface.co/datasets/Multimodal-Fatima/OxfordFlowers_train
+[6]: https://huggingface.co/datasets/Multimodal-Fatima/OxfordFlowers_test_facebook_opt_6.7b_Attributes_ns_6149
+[7]: https://github.com/google/dreambooth
+[8]: https://arxiv.org/abs/2104.14294
+[9]: https://arxiv.org/abs/2205.01917
+[10]: https://github.com/google/dreambooth/tree/main/dataset/dog6

notebooks/community/model_garden/model_garden_camp_zipnerf.ipynb

@@ -127,6 +127,7 @@
         "# @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 uuid\n",
         "from datetime import datetime\n",
         "\n",
         "from google.cloud import aiplatform\n",
@@ -142,14 +143,16 @@
         "# 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",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
         "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
         "\n",
         "# Create a unique GCS bucket for this notebook, if not specified by the user.\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",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_URI} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\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",
@@ -160,7 +163,6 @@
         "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
         "\n",
         "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
         "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
@@ -176,7 +178,6 @@
         "RENDERING_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/jax-cloudnerf-render:latest\"\n",
         "\n",
         "import subprocess\n",
-        "from datetime import datetime\n",
         "from typing import Any, List\n",
         "\n",
         "IMAGE_EXTENSIONS = (\".png\", \".jpg\", \".jpeg\", \".gif\", \".bmp\")\n",
@@ -658,7 +659,7 @@
         "\n",
         "delete_bucket = False  # @param {type:\"boolean\"}\n",
         "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_URI\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME\n",
         "\n",
         "# Delete pose estimation, training and rendering custom jobs.\n",
         "if data_calibration_custom_job.list(\n",

notebooks/community/model_garden/model_garden_camp_zipnerf_gradio.ipynb

@@ -401,7 +401,9 @@
         "                    else (\n",
         "                        \"BOOL\"\n",
         "                        if key == \"Use_Hierarchical_Mapper\"\n",
-        "                        else \"FLOAT64\" if key == \"Render_Focal\" else \"TIMESTAMP\"\n",
+        "                        else \"FLOAT64\"\n",
+        "                        if key == \"Render_Focal\"\n",
+        "                        else \"TIMESTAMP\"\n",
         "                    )\n",
         "                )\n",
         "            ),\n",
@@ -421,7 +423,6 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "cellView": "form",
         "id": "j1evctm2h34g"
       },
       "outputs": [],
@@ -484,7 +485,7 @@
         "    level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\"\n",
         ")\n",
         "\n",
-        "# track unique experiments\n",
+        "# Track unique experiments\n",
         "unique_experiments = set()\n",
         "\n",
         "# Define dataset and table IDs\n",
@@ -959,7 +960,7 @@
         "        ]\n",
         "        return pd.DataFrame(data).sort_values(by=\"Experiment ID\").reset_index(drop=True)\n",
         "    except Exception as e:\n",
-        "        logging.info(f\"Please upload a dataset to obtain an {e}.\", exc_info=True)\n",
+        "        logging.info(f\"Exception encountered in {e}.\", exc_info=True)\n",
         "        return pd.DataFrame()\n",
         "\n",
         "\n",
@@ -985,7 +986,7 @@
         "        ]\n",
         "        return pd.DataFrame(data).sort_values(by=\"Experiment ID\").reset_index(drop=True)\n",
         "    except Exception as e:\n",
-        "        logging.info(f\"Please upload a dataset to obtain an {e}.\", exc_info=True)\n",
+        "        logging.info(f\"Exception encountered in {e}.\", exc_info=True)\n",
         "        return pd.DataFrame()\n",
         "\n",
         "\n",
@@ -1020,7 +1021,7 @@
         "        ]\n",
         "        return pd.DataFrame(data).sort_values(by=\"Experiment ID\").reset_index(drop=True)\n",
         "    except Exception as e:\n",
-        "        logging.info(f\"Please upload a dataset to obtain an {e}.\", exc_info=True)\n",
+        "        logging.info(f\"Exception encountered in {e}.\", exc_info=True)\n",
         "        return pd.DataFrame()\n",
         "\n",
         "\n",
@@ -1313,7 +1314,6 @@
         "\n",
         "\n",
         "def create_pose_estimation_workshop():\n",
-        "\n",
         "    def get_worker_pool_specs(docker_uri, args, machine_type, accelerator_type):\n",
         "        return [\n",
         "            {\n",
@@ -1444,10 +1444,6 @@
         "        frame_rate=1,\n",
         "        progress=gr.Progress(),\n",
         "    ):\n",
-        "        if not scene_name:\n",
-        "            gr.Warning(\"Please provide a name for the scene. This field is required.\")\n",
-        "            return get_bq_folders_dataframe_colmap(colmap_table_id)\n",
-        "\n",
         "        if gcs_folder:\n",
         "            prepare_instance_images_from_gcs(\n",
         "                scene_name, experiment_name, gcs_folder, progress\n",
@@ -1839,7 +1835,6 @@
         "\n",
         "\n",
         "def create_training_workshop():\n",
-        "\n",
         "    def get_worker_pool_specs(docker_uri, args, machine_type, accelerator_type):\n",
         "        return [\n",
         "            {\n",
@@ -2194,7 +2189,6 @@
         "\n",
         "\n",
         "def create_rendering_workshop():\n",
-        "\n",
         "    def get_worker_pool_specs(docker_uri, args, machine_type, accelerator_type):\n",
         "        return [\n",
         "            {\n",
@@ -2601,7 +2595,6 @@
         "\n",
         "\n",
         "def create_pipeline_workshop():\n",
-        "\n",
         "    def get_worker_pool_specs(docker_uri, args, machine_type, accelerator_type):\n",
         "        return [\n",
         "            {\n",
@@ -3127,9 +3120,10 @@
         "\n",
         "    with gr.Tabs():\n",
         "        with gr.TabItem(\"Pose Estimation\") as tab1:\n",
-        "            pose_estimation_workshop, folders_dataframe_colmap = (\n",
-        "                create_pose_estimation_workshop()\n",
-        "            )\n",
+        "            (\n",
+        "                pose_estimation_workshop,\n",
+        "                folders_dataframe_colmap,\n",
+        "            ) = create_pose_estimation_workshop()\n",
         "            gr.on(\n",
         "                [tab1.select],\n",
         "                get_bq_folders_dataframe_colmap,\n",
@@ -3145,9 +3139,10 @@
         "                outputs=[folders_dataframe_training],\n",
         "            )\n",
         "        with gr.TabItem(\"Rendering\") as tab3:\n",
-        "            rendering_workshop, folders_dataframe_rendering = (\n",
-        "                create_rendering_workshop()\n",
-        "            )\n",
+        "            (\n",
+        "                rendering_workshop,\n",
+        "                folders_dataframe_rendering,\n",
+        "            ) = create_rendering_workshop()\n",
         "            gr.on(\n",
         "                [tab3.select],\n",
         "                get_bq_folders_dataframe_rendering,\n",

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -4,11 +4,12 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# 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",
@@ -94,17 +95,25 @@
         "\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",
+        "# @markdown 2. By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
-        "# @markdown 3. By default, the quota for TPU deployment `Custom model serving TPU v5e cores per region` is 4. TPU quota is only available in `us-west1`. You can request for higher TPU quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "# @markdown 3. **[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 4. **[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",
         "# Import the necessary packages\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
+        "import datetime\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",
@@ -114,12 +123,15 @@
         ")\n",
         "\n",
         "models, endpoints = {}, {}\n",
+        "# Dedicated endpoint not supported yet\n",
+        "use_dedicated_endpoint = False\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",
+        "if not REGION:\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",
@@ -128,8 +140,7 @@
         "# 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",
+        "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",
@@ -165,13 +176,14 @@
         "# 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"
+        "! 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": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "45c8c5438737"
@@ -200,7 +212,7 @@
         "\n",
         "# @markdown *--- Or ---*\n",
         "\n",
-        "# @markdown ##Enable the Vertex AI API## Access CodeGemma models on HuggingFace\n",
+        "# @markdown #### Access CodeGemma models on HuggingFace\n",
         "# @markdown You must provide a Hugging Face User Access Token (read) to access the CodeGemma models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
         "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "if LOAD_MODEL_FROM == \"Hugging Face\":\n",
@@ -233,7 +245,7 @@
       "source": [
         "## Deploy CodeGemma models with Hex-LLM on TPU\n",
         "\n",
-        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud."
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud. Learn more about Hex-LLM [here](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm)."
       ]
     },
     {
@@ -249,38 +261,54 @@
         "\n",
         "# @markdown Set the model to deploy.\n",
         "\n",
-        "MODEL_ID = \"codegemma-7b-it\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
+        "MODEL_ID = \"codegemma-2b\"  # @param [\"codegemma-2b\", \"codegemma-7b\", \"codegemma-7b-it\"]\n",
+        "IMPROVE_2B_LATENCY_AND_THROUGHPUT_WITH_4CHIPS = True  # @param {type:\"boolean\"}\n",
         "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
         "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
+        "hf_model_id = \"google/\" + MODEL_ID\n",
         "\n",
         "# The pre-built serving docker image for Hex-LLM.\n",
-        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:deploy\"\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
         "\n",
         "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
         "# @markdown https://cloud.google.com/vertex-ai/docs/predictions/use-tpu#deploy_a_model.\n",
         "if \"2b\" in model_id:\n",
-        "    # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy CodeGemma 2B models.\n",
-        "    machine_type = \"ct5lp-hightpu-1t\"\n",
-        "    accelerator_type = \"TPU_V5e\"\n",
-        "    # Note: 1 TPU V5 chip has only one core.\n",
-        "    accelerator_count = 1\n",
+        "    if IMPROVE_2B_LATENCY_AND_THROUGHPUT_WITH_4CHIPS:\n",
+        "        # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy CodeGemma 2B models.\n",
+        "        # Using a tensor parallel size of 2 improves latency.\n",
+        "        # Using a data parallel size of 2 approximately increases throughput by a factor of 2.\n",
+        "        machine_type = \"ct5lp-hightpu-4t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU V5 chip has only one core.\n",
+        "        data_parallel_size = 2\n",
+        "        tensor_parallel_size = 2\n",
+        "    else:\n",
+        "        # Sets ct5lp-hightpu-1t (1 TPU chip) to deploy CodeGemma 2B models.\n",
+        "        machine_type = \"ct5lp-hightpu-1t\"\n",
+        "        accelerator_type = \"TPU_V5e\"\n",
+        "        # Note: 1 TPU V5 chip has only one core.\n",
+        "        data_parallel_size = 1\n",
+        "        tensor_parallel_size = 1\n",
         "else:\n",
         "    # Sets ct5lp-hightpu-4t (4 TPU chips) to deploy CodeGemma 7B models.\n",
         "    machine_type = \"ct5lp-hightpu-4t\"\n",
         "    accelerator_type = \"TPU_V5e\"\n",
         "    # Note: 1 TPU V5 chip has only one core.\n",
-        "    accelerator_count = 4\n",
+        "    data_parallel_size = 1\n",
+        "    tensor_parallel_size = 4\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
         "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    accelerator_count=int(data_parallel_size * tensor_parallel_size),\n",
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Set enable_prefix_cache_hbm to False if you don't want to use [prefix caching](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm#prefix-caching).\n",
+        "enable_prefix_cache_hbm = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "# Server parameters.\n",
-        "tensor_parallel_size = accelerator_count\n",
         "hbm_utilization_factor = 0.6  # A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
         "max_running_seqs = 256\n",
         "\n",
@@ -292,15 +320,23 @@
         "def deploy_model_hexllm(\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",
+        "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
         "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tpu_topology: str = \"1x1\",\n",
+        "    disagg_topology: str = None,\n",
         "    hbm_utilization_factor: float = 0.6,\n",
         "    max_running_seqs: int = 256,\n",
+        "    max_model_len: int = 4096,\n",
+        "    enable_prefix_cache_hbm: bool = False,\n",
         "    endpoint_id: str = \"\",\n",
         "    min_replica_count: int = 1,\n",
         "    max_replica_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
         "    if endpoint_id:\n",
@@ -312,6 +348,7 @@
         "        endpoint = aiplatform.Endpoint.create(\n",
         "            display_name=f\"{model_name}-endpoint\",\n",
         "            location=TPU_DEPLOYMENT_REGION,\n",
+        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
         "        )\n",
         "\n",
         "    if not base_model_id:\n",
@@ -320,26 +357,32 @@
         "    if not tensor_parallel_size:\n",
         "        tensor_parallel_size = int(machine_type[-2])\n",
         "\n",
+        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
+        "\n",
+        "    # Learn more about the supported arguments and environment variables at https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm#config-server.\n",
         "    hexllm_args = [\n",
         "        \"--host=0.0.0.0\",\n",
         "        \"--port=7080\",\n",
-        "        \"--log_level=INFO\",\n",
         "        f\"--model={model_id}\",\n",
+        "        f\"--data_parallel_size={data_parallel_size}\",\n",
         "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
-        "        \"--enable_jit\",\n",
-        "        \"--load_format=auto\",\n",
+        "        f\"--num_hosts={num_hosts}\",\n",
         "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
         "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "        f\"--max_model_len={max_model_len}\",\n",
         "    ]\n",
+        "    if disagg_topology:\n",
+        "        hexllm_args.append(f\"--disagg_topo={disagg_topology}\")\n",
+        "    if enable_prefix_cache_hbm and not disagg_topology:\n",
+        "        hexllm_args.append(\"--enable_prefix_cache_hbm\")\n",
         "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
-        "        \"PJRT_DEVICE\": \"TPU\",\n",
-        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
-        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
+        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\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.update({\"HF_TOKEN\": HF_TOKEN})\n",
@@ -358,28 +401,42 @@
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
         "        location=TPU_DEPLOYMENT_REGION,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
         "    )\n",
         "\n",
         "    model.deploy(\n",
         "        endpoint=endpoint,\n",
         "        machine_type=machine_type,\n",
+        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
         "        min_replica_count=min_replica_count,\n",
         "        max_replica_count=max_replica_count,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\n",
+        "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
+        "\n",
         "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
         "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"codegemma\",\n",
+        "    base_model_id=hf_model_id,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
+        "    data_parallel_size=data_parallel_size,\n",
         "    tensor_parallel_size=tensor_parallel_size,\n",
         "    hbm_utilization_factor=hbm_utilization_factor,\n",
         "    max_running_seqs=max_running_seqs,\n",
+        "    enable_prefix_cache_hbm=enable_prefix_cache_hbm,\n",
         "    min_replica_count=min_replica_count,\n",
         "    max_replica_count=max_replica_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
         "# @markdown Click \"Show code\" to see more details."
@@ -428,7 +485,9 @@
         "        \"top_k\": top_k,\n",
         "    },\n",
         "]\n",
-        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
+        "response = endpoints[\"hexllm_tpu\"].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",
@@ -455,7 +514,6 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "tQIEisUajS6t"
@@ -507,6 +565,8 @@
         "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",
@@ -515,28 +575,69 @@
         "    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",
         ") -> 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",
+        "    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=7080\",\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",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -553,12 +654,15 @@
         "        display_name=model_name,\n",
         "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
         "        serving_container_args=vllm_args,\n",
-        "        serving_container_ports=[7080],\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",
@@ -570,20 +674,27 @@
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\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=\"codegemma-serve-vllm\"),\n",
         "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"codegemma\",\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",
         "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
         "# @markdown Click \"Show code\" to see more details."
@@ -630,6 +741,7 @@
         "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",
@@ -642,7 +754,9 @@
         "        \"raw_response\": raw_response,\n",
         "    },\n",
         "]\n",
-        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\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",

notebooks/community/model_garden/model_garden_e5.ipynb

@@ -211,6 +211,7 @@
         "        serving_container_args=tei_args,\n",
         "        serving_container_ports=[7080],\n",
         "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=\"publishers/intfloat/models/e5\"\n",
         "    )\n",
         "\n",
         "    model.deploy(\n",
@@ -220,6 +221,9 @@
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_e5.ipynb\"\n",
+        "        },\n",
         "    )\n",
         "    return model, endpoint"
       ]

notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

@@ -5,11 +5,12 @@
       "execution_count": null,
       "language": "python",
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# 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",
@@ -110,16 +111,24 @@
         "\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",
+        "# @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",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. **[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",
         "# Import the necessary packages\n",
         "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
+        "import datetime\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",
@@ -134,7 +143,8 @@
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
         "# Get the default region for launching jobs.\n",
-        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "if not REGION:\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",
@@ -143,8 +153,7 @@
         "# 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",
+        "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",
@@ -181,6 +190,8 @@
         "! 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\"\n",
         "\n",
         "# @markdown ## Access Gemma 2 Models\n",
         "\n",
@@ -191,141 +202,7 @@
         "    HF_TOKEN\n",
         "), \"Provide a read HF_TOKEN to load models from Hugging Face, or select a different model source.\"\n",
         "\n",
-        "model_path_prefix = \"google/\"\n",
-        "\n",
-        "# The pre-built serving docker images.\n",
-        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:gemma2\"\n",
-        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
-        "\n",
-        "def deploy_model_hexllm(\n",
-        "    model_name: str,\n",
-        "    model_id: str,\n",
-        "    service_account: str,\n",
-        "    base_model_id: str = None,\n",
-        "    tensor_parallel_size: int = 1,\n",
-        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
-        "    hbm_utilization_factor: float = 0.6,\n",
-        "    max_running_seqs: int = 256,\n",
-        "    endpoint_id: str = \"\",\n",
-        "    min_replica_count: int = 1,\n",
-        "    max_replica_count: int = 1,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
-        "    if endpoint_id:\n",
-        "        aip_endpoint_name = (\n",
-        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
-        "        )\n",
-        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
-        "    else:\n",
-        "        endpoint = aiplatform.Endpoint.create(\n",
-        "            display_name=f\"{model_name}-endpoint\",\n",
-        "            location=TPU_DEPLOYMENT_REGION,\n",
-        "        )\n",
-        "\n",
-        "    if not base_model_id:\n",
-        "        base_model_id = model_id\n",
-        "\n",
-        "    if not tensor_parallel_size:\n",
-        "        tensor_parallel_size = int(machine_type[-2])\n",
-        "\n",
-        "    hexllm_args = [\n",
-        "        \"--host=0.0.0.0\",\n",
-        "        \"--port=7080\",\n",
-        "        \"--log_level=INFO\",\n",
-        "        f\"--model={model_id}\",\n",
-        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
-        "        \"--enable_jit\",\n",
-        "        \"--load_format=auto\",\n",
-        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
-        "        f\"--max_running_seqs={max_running_seqs}\",\n",
-        "    ]\n",
-        "\n",
-        "    env_vars = {\n",
-        "        \"MODEL_ID\": base_model_id,\n",
-        "        \"PJRT_DEVICE\": \"TPU\",\n",
-        "        \"RAY_DEDUP_LOGS\": \"0\",\n",
-        "        \"RAY_USAGE_STATS_ENABLED\": \"0\",\n",
-        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
-        "    }\n",
-        "\n",
-        "    try:\n",
-        "        if HF_TOKEN:\n",
-        "            env_vars.update({\"HF_TOKEN\": HF_TOKEN})\n",
-        "    except:\n",
-        "        pass\n",
-        "\n",
-        "    model = aiplatform.Model.upload(\n",
-        "        display_name=model_name,\n",
-        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
-        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
-        "        serving_container_args=hexllm_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",
-        "        location=TPU_DEPLOYMENT_REGION,\n",
-        "    )\n",
-        "\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "        min_replica_count=min_replica_count,\n",
-        "        max_replica_count=max_replica_count,\n",
-        "    )\n",
-        "    return model, endpoint\n",
-        "\n",
-        "\n",
-        "def deploy_model_tgi(\n",
-        "    model_name: str,\n",
-        "    model_id: str,\n",
-        "    service_account: str,\n",
-        "    machine_type: str = \"g2-standard-8\",\n",
-        "    accelerator_type: str = \"NVIDIA_L4\",\n",
-        "    accelerator_count: int = 1,\n",
-        "    max_input_length: int = 2047,\n",
-        "    max_total_tokens: int = 2048,\n",
-        "    max_batch_prefill_tokens: int = 2048,\n",
-        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
-        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
-        "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
-        "\n",
-        "    env_vars = {\n",
-        "        \"MODEL_ID\": model_id,\n",
-        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
-        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
-        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
-        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\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=TGI_DOCKER_URI,\n",
-        "        serving_container_ports=[8080],\n",
-        "        serving_container_environment_variables=env_vars,\n",
-        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\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"
+        "model_path_prefix = \"google/\""
       ]
     },
     {
@@ -336,7 +213,7 @@
       "source": [
         "## Deploy Gemma 2 models with Hex-LLM on TPU\n",
         "\n",
-        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud.\n",
+        "**Hex-LLM** is a **H**igh-**E**fficiency **L**arge **L**anguage **M**odel (LLM) TPU serving solution built with **XLA**, which is being developed by Google Cloud. Learn more about Hex-LLM [here](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm).\n",
         "\n",
         "Refer to the \"Request for TPU quota\" section for TPU quota."
       ]
@@ -353,6 +230,9 @@
         "# @title Deploy\n",
         "# @markdown Set the model ID. Model weights can be loaded from HuggingFace or from a GCS bucket.\n",
         "\n",
+        "# The pre-built serving docker images.\n",
+        "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
+        "\n",
         "# @markdown Select one of the four model variations.\n",
         "MODEL_ID = \"gemma-2-2b-it\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
         "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
@@ -392,13 +272,123 @@
         "hbm_utilization_factor = 0.6  # Fraction of HBM memory allocated for KV cache after model loading. A larger value improves throughput but gives higher risk of TPU out-of-memory errors with long prompts.\n",
         "max_running_seqs = 256  # Maximum number of running sequences in a continuous batch.\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "# Endpoint configurations.\n",
         "min_replica_count = 1\n",
         "max_replica_count = 1\n",
         "\n",
+        "\n",
+        "def deploy_model_hexllm(\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",
+        "    data_parallel_size: int = 1,\n",
+        "    tensor_parallel_size: int = 1,\n",
+        "    machine_type: str = \"ct5lp-hightpu-1t\",\n",
+        "    tpu_topology: str = \"1x1\",\n",
+        "    disagg_topology: str = None,\n",
+        "    hbm_utilization_factor: float = 0.6,\n",
+        "    max_running_seqs: int = 256,\n",
+        "    max_model_len: int = 4096,\n",
+        "    enable_prefix_cache_hbm: bool = False,\n",
+        "    endpoint_id: str = \"\",\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Hex-LLM on TPU in Vertex AI.\"\"\"\n",
+        "    if endpoint_id:\n",
+        "        aip_endpoint_name = (\n",
+        "            f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_id}\"\n",
+        "        )\n",
+        "        endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "    else:\n",
+        "        endpoint = aiplatform.Endpoint.create(\n",
+        "            display_name=f\"{model_name}-endpoint\",\n",
+        "            location=TPU_DEPLOYMENT_REGION,\n",
+        "            dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "        )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    if not tensor_parallel_size:\n",
+        "        tensor_parallel_size = int(machine_type[-2])\n",
+        "\n",
+        "    num_hosts = int(tpu_topology.split(\"x\")[0])\n",
+        "\n",
+        "    # Learn more about the supported arguments and environment variables at https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm#config-server.\n",
+        "    hexllm_args = [\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=7080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--data_parallel_size={data_parallel_size}\",\n",
+        "        f\"--tensor_parallel_size={tensor_parallel_size}\",\n",
+        "        f\"--num_hosts={num_hosts}\",\n",
+        "        f\"--hbm_utilization_factor={hbm_utilization_factor}\",\n",
+        "        f\"--max_running_seqs={max_running_seqs}\",\n",
+        "        f\"--max_model_len={max_model_len}\",\n",
+        "    ]\n",
+        "    if disagg_topology:\n",
+        "        hexllm_args.append(f\"--disagg_topo={disagg_topology}\")\n",
+        "    if enable_prefix_cache_hbm and not disagg_topology:\n",
+        "        hexllm_args.append(\"--enable_prefix_cache_hbm\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"HEX_LLM_LOG_LEVEL\": \"info\",\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.update({\"HF_TOKEN\": HF_TOKEN})\n",
+        "    except:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=HEXLLM_DOCKER_URI,\n",
+        "        serving_container_command=[\"python\", \"-m\", \"hex_llm.server.api_server\"],\n",
+        "        serving_container_args=hexllm_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",
+        "        location=TPU_DEPLOYMENT_REGION,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        min_replica_count=min_replica_count,\n",
+        "        max_replica_count=max_replica_count,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "        },\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
         "models[\"hexllm_tpu\"], endpoints[\"hexllm_tpu\"] = deploy_model_hexllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
         "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma2\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    tensor_parallel_size=tensor_parallel_size,\n",
@@ -406,6 +396,7 @@
         "    max_running_seqs=max_running_seqs,\n",
         "    min_replica_count=min_replica_count,\n",
         "    max_replica_count=max_replica_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")"
       ]
     },
@@ -463,7 +454,9 @@
         "        \"top_k\": top_k,\n",
         "    },\n",
         "]\n",
-        "response = endpoints[\"hexllm_tpu\"].predict(instances=instances)\n",
+        "response = endpoints[\"hexllm_tpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "for prediction in response.predictions:\n",
         "    print(prediction)"
@@ -494,6 +487,10 @@
       "outputs": [],
       "source": [
         "# @title Deploy\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu121.2-1.ubuntu2204.py310\"\n",
+        "\n",
         "MODEL_ID = \"gemma-2-2b\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
         "model_id = os.path.join(model_path_prefix, MODEL_ID)\n",
         "\n",
@@ -550,9 +547,73 @@
         "max_total_tokens = 2048\n",
         "max_batch_prefill_tokens = 2048\n",
         "\n",
+        "\n",
+        "def deploy_model_tgi(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    service_account: str,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    max_input_length: int = 2047,\n",
+        "    max_total_tokens: int = 2048,\n",
+        "    max_batch_prefill_tokens: int = 2048,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with TGI on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"NUM_SHARD\": f\"{accelerator_count}\",\n",
+        "        \"MAX_INPUT_LENGTH\": f\"{max_input_length}\",\n",
+        "        \"MAX_TOTAL_TOKENS\": f\"{max_total_tokens}\",\n",
+        "        \"MAX_BATCH_PREFILL_TOKENS\": f\"{max_batch_prefill_tokens}\",\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=TGI_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\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",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "        },\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
         "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
         "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma2\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
@@ -560,6 +621,7 @@
         "    max_input_length=max_input_length,\n",
         "    max_total_tokens=max_total_tokens,\n",
         "    max_batch_prefill_tokens=max_batch_prefill_tokens,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")"
       ]
     },
@@ -617,7 +679,9 @@
         "        },\n",
         "    },\n",
         "]\n",
-        "response = endpoints[\"tgi\"].predict(instances=instances)\n",
+        "response = endpoints[\"tgi\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
         "\n",
         "for prediction in response.predictions:\n",
         "    print(prediction)"

notebooks/community/model_garden/model_garden_gemma2_finetuning_on_vertex.ipynb

@@ -5,6 +5,7 @@
       "execution_count": null,
       "language": "python",
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
@@ -66,6 +67,9 @@
         "- Finetune and deploy Gemma 2 models with Vertex AI Custom Training Jobs.\n",
         "- Send prediction requests to your finetuned Gemma 2 model.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
         "\n",
         "### Costs\n",
         "\n",
@@ -86,6 +90,30 @@
         "## Before you begin"
       ]
     },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "aS52SK74gDoB"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Install Python Packages for Finetuning\n",
+        "\n",
+        "# @markdown 1. Install google-cloud-aiplatform package and restart the session if instructed.\n",
+        "! pip install --upgrade --quiet 'google-cloud-aiplatform>=1.66.0'\n",
+        "\n",
+        "# @markdown 2. Install packages to validate dataset with template.\n",
+        "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
+        "! pip install --upgrade --quiet accelerate==0.31.0\n",
+        "! pip install --upgrade --quiet transformers==4.43.1\n",
+        "! pip install --upgrade --quiet datasets==2.19.2\n",
+        "\n",
+        "# Load local tensorboard.\n",
+        "%load_ext tensorboard"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": null,
@@ -97,19 +125,39 @@
       "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",
+        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). 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. 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",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\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-8g | 8 NVIDIA_H100_80GB | us-central1, us-west1, europe-west4, 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",
+        "# Import the necessary packages\n",
+        "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! cd vertex-ai-samples && git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
+        "\n",
+        "import datetime\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",
+        "from google.cloud.aiplatform.compat.types import \\\n",
+        "    custom_job as gca_custom_job_compat\n",
         "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
@@ -121,7 +169,8 @@
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
         "# Get the default region for launching jobs.\n",
-        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "if not REGION:\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",
@@ -130,8 +179,7 @@
         "# 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",
+        "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",
@@ -168,6 +216,8 @@
         "! 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\"\n",
         "\n",
         "# @markdown ## Access Gemma 2 Models\n",
         "\n",
@@ -265,6 +315,81 @@
     {
       "cell_type": "code",
       "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "NoueJWi72OSo"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set model\n",
+        "\n",
+        "# @markdown Select a model variant of Gemma 2.\n",
+        "base_model_id = \"gemma-2-2b-it\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {isTemplate: true}\n",
+        "pretrained_model_id = os.path.join(model_path_prefix, base_model_id)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-NlLSiCOvru1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Validate Dataset with Template\n",
+        "\n",
+        "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
+        "\n",
+        "import transformers\n",
+        "\n",
+        "dataset_validation_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
+        ")\n",
+        "\n",
+        "if dataset_validation_util.is_gcs_path(pretrained_model_id):\n",
+        "    # Download tokenizer.\n",
+        "    ! mkdir tokenizer\n",
+        "    ! gsutil cp {pretrained_model_id}/tokenizer.json ./tokenizer\n",
+        "    ! gsutil cp {pretrained_model_id}/config.json ./tokenizer\n",
+        "    tokenizer_path = \"./tokenizer\"\n",
+        "    access_token = \"\"\n",
+        "else:\n",
+        "    tokenizer_path = pretrained_model_id\n",
+        "    access_token = HF_TOKEN\n",
+        "\n",
+        "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
+        "    tokenizer_path,\n",
+        "    trust_remote_code=False,\n",
+        "    use_fast=True,\n",
+        "    token=access_token,\n",
+        ")\n",
+        "\n",
+        "# Validate the train dataset.\n",
+        "dataset_validation_util.validate_dataset_with_template(\n",
+        "    dataset_name=train_dataset_name,\n",
+        "    split=train_split_name,\n",
+        "    input_column=instruct_column_in_dataset,\n",
+        "    template=template,\n",
+        "    use_multiprocessing=False,\n",
+        "    tokenizer=tokenizer,\n",
+        ")\n",
+        "\n",
+        "# Validate the eval dataset.\n",
+        "dataset_validation_util.validate_dataset_with_template(\n",
+        "    dataset_name=eval_dataset_name,\n",
+        "    split=eval_split_name,\n",
+        "    input_column=instruct_column_in_dataset,\n",
+        "    template=template,\n",
+        "    use_multiprocessing=False,\n",
+        "    tokenizer=tokenizer,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "ivVGS9dHXPOz"
@@ -272,21 +397,48 @@
       "outputs": [],
       "source": [
         "# @title Finetune\n",
-        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown This section demonstrates how to finetune the Gemma 2 model and merge the finetuned LoRA adapter with the base model on Vertex AI. It uses the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown The training job takes approximately between 10 to 20 mins to set-up. Once done, the training job is expected to take around 20 mins with the default configuration. To find the training time, throughput, and memory usage of your training job, you can go to the training logs and check the log line of the last training epoch.\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 takes 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 ~ 1.5 hours.\n",
         "\n",
-        "# @markdown This section demonstrates how to finetune the Gemma 2 model and merge the finetuned LoRA adapter with the base model on Vertex AI.\n",
-        "\n",
-        "# @markdown Select one of the four model variations.\n",
-        "base_model_id = \"gemma-2-2b-it\"  # @param [\"gemma-2-2b\", \"gemma-2-2b-it\", \"gemma-2-9b\", \"gemma-2-9b-it\", \"gemma-2-27b\", \"gemma-2-27b-it\"] {allow-input: true, isTemplate: true}\n",
-        "pretrained_model_id = os.path.join(model_path_prefix, base_model_id)\n",
+        "# @markdown Accelerator type to use for training.\n",
+        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
         "\n",
         "# The pre-built training docker image.\n",
-        "TRAIN_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-train:20240724_0936_RC00\"\n",
+        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
+        "    is_restricted_image = True\n",
+        "    is_dynamic_workload_scheduler = False\n",
+        "    dws_kwargs = {}\n",
+        "else:\n",
+        "    repo = \"us-docker.pkg.dev/vertex-ai\"\n",
+        "    is_restricted_image = False\n",
+        "    is_dynamic_workload_scheduler = True\n",
+        "    dws_kwargs = {\n",
+        "        \"max_wait_duration\": 1800,  # 30 minutes\n",
+        "        \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
+        "    }\n",
+        "\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20240909\"\n",
+        ")\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    per_node_accelerator_count = 8\n",
+        "    machine_type = \"a2-ultragpu-8g\"\n",
+        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    per_node_accelerator_count = 8\n",
+        "    machine_type = \"a3-highgpu-8g\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `per_node_accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
         "\n",
         "# @markdown Batch size for finetuning.\n",
         "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
@@ -324,23 +476,16 @@
         "# Train precision of the model.\n",
         "train_precision = \"bfloat16\"\n",
         "\n",
-        "# Worker pool spec for 4bit finetuning.\n",
-        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param[\"NVIDIA_A100_80GB\"]\n",
-        "\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
-        "    accelerator_count = 8\n",
-        "    machine_type = \"a2-ultragpu-8g\"\n",
-        "else:\n",
-        "    raise ValueError(f\"Unsupported accelerator type: {accelerator_type}\")\n",
-        "\n",
         "replica_count = 1\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
         "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    accelerator_count=per_node_accelerator_count * replica_count,\n",
         "    is_for_training=True,\n",
+        "    is_restricted_image=is_restricted_image,\n",
+        "    is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
         ")\n",
         "\n",
         "job_name = common_util.get_job_name_with_datetime(\"gemma2-lora-train\")\n",
@@ -352,6 +497,16 @@
         "# finetuned LORA adapter.\n",
         "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
         "\n",
+        "# Add labels for the finetuning job.\n",
+        "labels = {\n",
+        "    \"mg-source\": \"notebook\",\n",
+        "    \"mg-notebook-name\": \"model_garden_gemma2_finetuning_on_vertex.ipynb\".split(\".\")[0],\n",
+        "}\n",
+        "\n",
+        "labels[\"mg-tune\"] = \"publishers-google-models-gemma-2\"\n",
+        "versioned_model_id = base_model_id.lower().replace(\".\", \"-\")\n",
+        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
+        "\n",
         "eval_args = [\n",
         "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
         "    f\"--eval_column={instruct_column_in_dataset}\",\n",
@@ -363,7 +518,7 @@
         "]\n",
         "\n",
         "train_job_args = [\n",
-        "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_4gpu.yaml\",\n",
+        "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml\",\n",
         "    \"--task=instruct-lora\",\n",
         "    \"--completion_only=True\",\n",
         "    f\"--pretrained_model_id={pretrained_model_id}\",\n",
@@ -396,33 +551,33 @@
         "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
         "] + eval_args\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",
         "# 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",
-        "    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",
+        "    accelerator_count=per_node_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",
+        "    sync=False,  # Non-blocking call to run.\n",
+        "    **dws_kwargs,\n",
         ")\n",
         "\n",
-        "print(\"LoRA adapter was saved in:\", lora_output_dir)\n",
-        "print(\"Trained and merged models were saved in:\", merged_model_output_dir)\n",
+        "# Wait until resource has been created.\n",
+        "train_job.wait_for_resource_creation()\n",
+        "\n",
+        "print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
+        "print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -430,6 +585,27 @@
     {
       "cell_type": "code",
       "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lu-uxrFBmZ0s"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run TensorBoard\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": "code",
+      "execution_count": null,
+      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "qmHW6m8xG_4U"
@@ -439,16 +615,20 @@
         "# @title Deploy\n",
         "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
         "\n",
+        "if train_job.end_time is None:\n",
+        "    print(\"Waiting for the training job to finish...\")\n",
+        "    train_job.wait()\n",
+        "    print(\"The training job has finished.\")\n",
+        "\n",
         "print(\"Deploying models in:\", merged_model_output_dir)\n",
         "\n",
         "# The pre-built serving docker image for vLLM.\n",
         "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240815_1634_RC00\"\n",
         "\n",
+        "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
+        "# @markdown Accelerator type to use for serving.\n",
         "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
         "\n",
-        "# Find Vertex AI prediction supported accelerators and regions in\n",
-        "# https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
         "if \"2b\" in base_model_id:\n",
         "    if accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 1 L4 (24G) to deploy Gemma 2 2B models.\n",
@@ -492,6 +672,9 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to True if the endpoint is [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) enabled.\n",
+        "use_dedicated_endpoint = False  # @param {type:\"boolean\"}\n",
+        "\n",
         "gpu_memory_utilization = 0.85\n",
         "max_model_len = 4096  # Maximum context length.\n",
         "\n",
@@ -507,28 +690,55 @@
         "    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",
+        "    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",
         ") -> 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",
+        "    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=7080\",\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 model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -545,7 +755,7 @@
         "        display_name=model_name,\n",
         "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
         "        serving_container_args=vllm_args,\n",
-        "        serving_container_ports=[7080],\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",
@@ -577,6 +787,7 @@
         "    accelerator_count=accelerator_count,\n",
         "    gpu_memory_utilization=gpu_memory_utilization,\n",
         "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         ")\n",
         "\n",
         "# @markdown Click \"Show code\" to see more details."
@@ -618,11 +829,12 @@
         "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
         "\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",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\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",
+        "# @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 = False  # @param {type:\"boolean\"}\n",
         "\n",
         "# Overrides parameters for inferences.\n",
@@ -636,7 +848,9 @@
         "        \"raw_response\": raw_response,\n",
         "    },\n",
         "]\n",
-        "response = endpoints[\"vllm_gpu\"].predict(instances=instances)\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",