Update run_linter.sh

Fix spelling error `notebooked` → `notebooks`

.cloud-build/execute_changed_notebooks_helper.py

@@ -238,7 +238,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
 
             # Look for the python version specification pattern
             re_match = re.search(
-                r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
+                "python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
             )
             if re_match:
                 # get the version number

.cloud-build/test_notebook_vm.txt

@@ -1,3 +1,5 @@
+notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
+notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
 notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
 notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
 .cloud-build/tests/python_version_test.ipynb

.github/dependabot.yml

@@ -1,10 +0,0 @@
-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.13
+FROM python:3.12
 
 WORKDIR setup
 

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==25.1.0
-pyupgrade==3.19.1
-isort==6.0.1
-flake8==7.2.0
-nbqa==1.9.1
+black==24.4.2
+pyupgrade==3.16.0
+isort==5.13.2
+flake8==7.1.0
+nbqa==1.8.5
 

community-content/CODEOWNERS

@@ -20,7 +20,6 @@
 /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,40 +1,16 @@
-# 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 - && \
-    apt-get update -y && \
-    apt-get install -y google-cloud-sdk
-
-# Copy application code
-COPY . /trainer
-
-# Set working directory
-WORKDIR /trainer
-
-# Install Python dependencies
-RUN pip install --no-cache-dir -r requirements.txt
-
-# 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 - && \
+    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/*
+    apt-get install google-cloud-sdk -y
 
-# Copy from the builder stage
-COPY --from=builder /trainer /trainer
+COPY . /trainer
 
-# Set working directory
 WORKDIR /trainer
 
-# Set the entry point
-ENTRYPOINT ["python", "-m", "task"]
+RUN pip install -r requirements.txt
+
+ENTRYPOINT ["python", "-m", "task"]

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt

@@ -1,3 +1,3 @@
-torch==2.2.0
+torch==1.13.1
 torchvision==0.9.1
 tensorboard==2.5.0

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

@@ -1,4 +1,4 @@
 google-cloud-bigquery==2.20.0
-tensorflow==2.12.1
+tensorflow==2.7.2
 pillow==10.3.0
 tf-agents==0.8.0

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

@@ -1,4 +1,4 @@
 google-cloud-pubsub==2.5.0
 pillow==10.3.0
 tf-agents==0.8.0
-tensorflow==2.12.1
+tensorflow==2.7.2

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.12.1
+tensorflow==2.7.2
 pillow==10.3.0
 tf-agents==0.8.0

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/requirements.txt

@@ -1 +1 @@
-tensorflow==2.12.1
+tensorflow==2.7.2

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/trainer/requirements.txt

@@ -1 +1 @@
-tensorflow==2.12.1
+tensorflow==2.7.2

community-content/vertex_cpr_samples/sklearn/predictor_Classifier.py

@@ -1,33 +0,0 @@
-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_LinearRegression.py

@@ -1,33 +0,0 @@
-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 make_blobs
-from sklearn.linear_model import LinearRegression 
-
-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 = LogisticRegression()
-        X, y = make_blobs(n_samples=100, centers=2, n_features=2, random_state=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_proba(instances)
-
-    def postprocess(self, prediction_results: np.ndarray) -> dict:
-        return {"predictions": prediction_results.tolist()}

community-content/vertex_cpr_samples/sklearn/predictor_SGDClassifier.py

@@ -1,33 +0,0 @@
-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

@@ -1,34 +0,0 @@
-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

@@ -1,37 +0,0 @@
-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

@@ -1,41 +0,0 @@
-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/cloudnerf/requirements.txt

@@ -1,6 +1,6 @@
 --find-links https://download.pytorch.org/whl/torch_stable.html
 
-torch==2.2.0
+torch==2.0.1+cu118
 numpy==1.26.1
 absl_py==2.0.0
 accelerate==0.24.0

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

@@ -1,16 +1,13 @@
 """Common util functions for notebook."""
 
 import base64
-from collections.abc import Sequence
 import datetime
 import io
 import json
 import os
 import subprocess
-import time
-from typing import Any
+from typing import Any, Dict, Sequence
 
-from google import auth
 from google.cloud import storage
 import matplotlib.pyplot as plt
 import numpy as np
@@ -88,9 +85,7 @@ def get_job_name_with_datetime(prefix: str) -> str:
   Returns:
     A job name.
   """
-  now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-  job_name = f"{prefix}-{now}".replace("_", "-")
-  return job_name
+  return prefix + datetime.datetime.now().strftime("_%Y%m%d_%H%M%S")
 
 
 def create_job_name(prefix: str) -> str:
@@ -104,7 +99,7 @@ def create_job_name(prefix: str) -> str:
   """
   user = os.environ.get("USER")
   now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-  job_name = f"{prefix}-{user}-{now}".replace("_", "-")
+  job_name = f"{prefix}-{user}-{now}"
   return job_name
 
 
@@ -233,23 +228,23 @@ def download_image(url: str) -> str:
     base64 encoded image.
   """
   response = requests.get(url)
-  return Image.open(io.BytesIO(response.content))  # pytype: disable=bad-return-type  # pillow-102-upgrade
+  return Image.open(io.BytesIO(response.content))
 
 
 def resize_image(image: Any, new_width: int = 1000) -> Any:
-  """Resizes an image to a certain width.
+    """Resizes an image to a certain width.
 
-  Args:
-    image: The image which has to be resized.
-    new_width: New width of the image.
+    Args:
+      image: The image which has to be resized.
+      new_width: New width of the image.
 
-  Returns:
-    New resized image.
-  """
-  width, height = image.size
-  new_height = int(height * new_width / width)
-  new_img = image.resize((new_width, new_height))
-  return new_img
+    Returns:
+      New resized image.
+    """
+    width, height = image.size
+    new_height = int(height * new_width / width)
+    new_img = image.resize((new_width, new_height))
+    return new_img
 
 
 def load_img(path: str) -> Any:
@@ -284,7 +279,7 @@ def decode_image(
   return image
 
 
-def get_label_map(label_map_yaml_filepath: str) -> dict[int, str]:
+def get_label_map(label_map_yaml_filepath: str) -> Dict[int, str]:
   """Returns class id to label mapping given a filepath to the label map.
 
   Args:
@@ -328,150 +323,6 @@ def get_prediction_instances(test_filepath: str, new_width: int = -1) -> Any:
   return instances
 
 
-def vqa_predict(
-    endpoint: Any,
-    question_prompts: Sequence[str],
-    image: Any,
-    language_code: str = "en",
-    new_width: int = 1000,
-    use_dedicated_endpoint: bool = False,
-) -> Sequence[str]:
-  """Predicts the answer to a question about an image using an Endpoint."""
-  # Resize and convert image to base64 string.
-  resized_image = resize_image(image, new_width)
-  resized_image_base64 = image_to_base64(resized_image)
-
-  instances = []
-  if question_prompts:
-    # Format question prompt
-    question_prompt_format = "answer {} {}\n"
-    for question_prompt in question_prompts:
-      if question_prompt:
-        instances.append({
-            "prompt": question_prompt_format.format(
-                language_code, question_prompt
-            ),
-            "image": resized_image_base64,
-        })
-  else:
-    instances.append({
-        "image": resized_image_base64,
-    })
-
-  response = endpoint.predict(
-      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
-  )
-  return [pred.get("response") for pred in response.predictions]
-
-
-def caption_predict(
-    endpoint: Any,
-    language_code: str,
-    image: Any,
-    caption_prompt: bool = False,
-    new_width: int = 1000,
-    use_dedicated_endpoint: bool = False,
-) -> str:
-  """Predicts a caption for a given image using an Endpoint."""
-  # Resize and convert image to base64 string.
-  resized_image = resize_image(image, new_width)
-  resized_image_base64 = image_to_base64(resized_image)
-
-  instance = {"image": resized_image_base64}
-
-  if caption_prompt:
-    # Format caption prompt
-    caption_prompt_format = "caption {}\n"
-    instance["prompt"] = caption_prompt_format.format(language_code)
-
-  instances = [instance]
-  response = endpoint.predict(
-      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
-  )
-  return response.predictions[0].get("response")
-
-
-def ocr_predict(
-    endpoint: Any,
-    ocr_prompt: str,
-    image: Any,
-    new_width: int = 1000,
-    use_dedicated_endpoint: bool = False,
-) -> str:
-  """Extracts text from a given image using an Endpoint."""
-  # Resize and convert image to base64 string.
-  resized_image = resize_image(image, new_width)
-  resized_image_base64 = image_to_base64(resized_image)
-
-  instance = {"image": resized_image_base64}
-  if ocr_prompt:
-    instance["prompt"] = ocr_prompt
-  instances = [instance]
-
-  response = endpoint.predict(
-      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
-  )
-  return response.predictions[0].get("response")
-
-
-def detect_predict(
-    endpoint: Any,
-    detect_prompt: str,
-    image: Any,
-    new_width: int = 1000,
-    use_dedicated_endpoint: bool = False,
-) -> str:
-  """Predicts the answer to a question about an image using an Endpoint."""
-  # Resize and convert image to base64 string.
-  resized_image = resize_image(image, new_width)
-  resized_image_base64 = image_to_base64(resized_image)
-
-  instance = {"image": resized_image_base64}
-  if detect_prompt:
-    instance["prompt"] = detect_prompt
-  instances = [instance]
-
-  response = endpoint.predict(
-      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
-  )
-  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.
 
@@ -510,17 +361,6 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
   ):
     return -1
   all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
-
-  # If the quota data does not have dimensions, it is global quota. However,
-  # global quota may be overridden by regional quota. So we need to check the
-  # global quota first.
-  global_quota = -1
-  if (
-      all_regions_data
-      and "dimensions" not in all_regions_data[0]
-      and "effectiveLimit" in all_regions_data[0]
-  ):
-    global_quota = int(all_regions_data[0]["effectiveLimit"])
   for region_data in all_regions_data:
     if (
         region_data.get("dimensions")
@@ -530,92 +370,47 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
         return int(region_data["effectiveLimit"])
       else:
         return 0
-  return global_quota
+  return -1
 
 
-def get_resource_id(
-    accelerator_type: str,
-    is_for_training: bool,
-    is_spot: bool = False,
-    is_restricted_image: bool = False,
-    is_dynamic_workload_scheduler: bool = False,
-) -> str:
+def get_resource_id(accelerator_type: str, is_for_training: bool) -> str:
   """Returns the resource id for a given accelerator type and the use case.
 
   Args:
     accelerator_type: The accelerator type.
     is_for_training: Whether the resource is used for training. Set false for
       serving use case.
-    is_spot: Whether the resource is used with Spot.
-    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_H100_MEGA_80GB": "nvidia_h100_mega_gpus",
-      "NVIDIA_H200_141GB": "nvidia_h200_gpus",
-      "NVIDIA_TESLA_T4": "nvidia_t4_gpus",
-      "TPU_V6e": "tpu_v6e",
-      "TPU_V5e": "tpu_v5e",
-      "TPU_V3": "tpu_v3",
-  }
-  default_training_accelerator_map = {
-      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",
-  }
-  spot_serving_accelerator_map = {
-      key: f"custom_model_serving_preemptible_{accelerator_suffix_map[key]}"
-      for key in accelerator_suffix_map
+  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_TESLA_T4": "custom_model_training_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_training_tpu_v5e",
+      "TPU_V3": "custom_model_training_tpu_v3",
   }
   serving_accelerator_map = {
-      key: f"custom_model_serving_{accelerator_suffix_map[key]}"
-      for key in accelerator_suffix_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_TESLA_T4": "custom_model_serving_nvidia_t4_gpus",
+      "TPU_V5e": "custom_model_serving_tpu_v5e",
   }
-
   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:
-      if is_dynamic_workload_scheduler:
-        return dws_training_accelerator_map[accelerator_type]
-      else:
-        return training_accelerator_map[accelerator_type]
+      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.")
-    accelerator_map = (
-        spot_serving_accelerator_map if is_spot else serving_accelerator_map
-    )
-    if accelerator_type in accelerator_map:
-      return accelerator_map[accelerator_type]
+    if accelerator_type in serving_accelerator_map:
+      return serving_accelerator_map[accelerator_type]
     else:
       raise ValueError(
           f"Could not find accelerator type: {accelerator_type} for serving."
@@ -628,31 +423,9 @@ def check_quota(
     accelerator_type: str,
     accelerator_count: int,
     is_for_training: bool,
-    is_spot: bool = False,
-    is_restricted_image: bool = False,
-    is_dynamic_workload_scheduler: bool = False,
-) -> None:
-  """Checks if the project and the region has the required quota.
-
-  Args:
-    project_id: The project id.
-    region: The region.
-    accelerator_type: The accelerator type.
-    accelerator_count: The number of accelerators to check quota for.
-    is_for_training: Whether the resource is used for training. Set false for
-      serving use case.
-    is_spot: Whether the resource is used with Spot.
-    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.
-  """
-  resource_id = get_resource_id(
-      accelerator_type,
-      is_for_training=is_for_training,
-      is_spot=is_spot,
-      is_restricted_image=is_restricted_image,
-      is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,
-  )
+):
+  """Checks if the project and the region has the required quota."""
+  resource_id = get_resource_id(accelerator_type, is_for_training)
   quota = get_quota(project_id, region, resource_id)
   quota_request_instruction = (
       "Either use "
@@ -672,76 +445,3 @@ def check_quota(
         f"Quota not enough for {resource_id} in {region}: {quota} <"
         f" {accelerator_count}. {quota_request_instruction}"
     )
-
-
-def get_deploy_source() -> str:
-  """Gets deploy_source string based on running environment."""
-  vertex_product = os.environ.get("VERTEX_PRODUCT", "")
-  match vertex_product:
-    case "COLAB_ENTERPRISE":
-      return "notebook_colab_enterprise"
-    case "WORKBENCH_INSTANCE":
-      return "notebook_workbench"
-    case _:
-      # Legacy workbench, legacy colab, or other custom environments.
-      return "notebook_environment_unspecified"
-
-
-def _is_operation_done(op_name: str, region: str) -> bool:
-  """Checks if the operation is done.
-
-  Args:
-    op_name: The name of the operation to poll.
-    region: The region of the operation.
-
-  Returns:
-    True if the operation is done, False otherwise.
-
-  Raises:
-    ValueError: If the operation failed.
-  """
-  creds, _ = auth.default()
-  auth_req = auth.transport.requests.Request()
-  creds.refresh(auth_req)
-  headers = {
-      "Authorization": f"Bearer {creds.token}",
-  }
-  url = f"https://{region}-aiplatform.googleapis.com/ui/{op_name}"
-  response = requests.get(url, headers=headers)
-  operation_data = response.json()
-  if "error" in operation_data:
-    raise ValueError(f"Operation failed: {operation_data['error']}")
-  return operation_data.get("done", False)
-
-
-def poll_and_wait(
-    op_name: str, region: str, total_wait: int, interval: int = 60
-) -> None:
-  """Polls the operation and waits for it to complete.
-
-  Args:
-    op_name: The name of the operation to poll.
-    region: The region of the operation.
-    total_wait: The total wait time in seconds.
-    interval: The interval between each poll in seconds.
-
-  Raises:
-    TimeoutError: If the operation times out.
-  """
-  start_time = time.time()
-  while True:
-    if _is_operation_done(op_name, region):
-      break
-    time_elapsed = time.time() - start_time
-    if time_elapsed > total_wait:
-      raise TimeoutError(
-          f"Operation timed out after {int(time_elapsed)} seconds."
-      )
-    print(
-        "\rStill waiting for operation... Elapsed time in seconds:"
-        f" {int(time_elapsed):<6}",
-        end="",
-        flush=True,
-    )
-    time.sleep(interval)
-

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

@@ -1,570 +0,0 @@
-"""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, Tuple, 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,
-    check_path_exists: bool = True,
-) -> 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.
-    check_path_exists: Whether to check if the path exists.
-
-  Returns:
-    Local path to target folder/file.
-  """
-  target = os.path.join(
-      destination_dir,
-      os.path.basename(os.path.normpath(gcs_uri)),
-  )
-  if check_path_exists and 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 format_dataset(
-    dataset: datasets.Dataset,
-    input_column: str,
-    template: str = None,
-    tokenizer: transformers.PreTrainedTokenizer | None = None,
-) -> datasets.Dataset:
-  """Takes a raw dataset and formats it using a template and tokenizer.
-
-  Args:
-    dataset: The raw (unprocessed) dataset to format.
-    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.
-  """
-  return dataset.map(
-      _format_template_fn(
-          template,
-          input_column=input_column,
-          tokenizer=tokenizer,
-      )
-  )
-
-
-def load_dataset_with_template(
-    dataset_name: str,
-    split: str,
-    input_column: str,
-    template: str = None,
-    tokenizer: transformers.PreTrainedTokenizer | None = None,
-) -> Tuple[Any, 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:
-    The raw dataset and the dataset compatible with the template.
-  """
-  raw = _get_dataset(dataset_name, split=split)
-  if template:
-    templated = format_dataset(raw, input_column, template, tokenizer)
-  else:
-    templated = None
-
-  return raw, templated
-
-
-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 = format_dataset(
-      _get_dataset(dataset_name, split, num_proc),
-      input_column,
-      template_path,
-      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

@@ -0,0 +1,142 @@
+"""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 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,
+    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=["q_proj", "v_proj"],
+      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

@@ -1,28 +0,0 @@
-# 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

@@ -1,22 +0,0 @@
-# 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,40 +7,39 @@
 # 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.11.0-gpu
+FROM pytorch/torchserve:0.7.0-gpu
 
 USER root
 
-ENV INFER_PORT=7080
-ENV MNG_PORT=7081
-ENV MODEL="peft_serving"
+ENV infer_port=7080
+ENV mng_port=7081
+ENV model_name="peft_serving"
 ENV PATH="/home/model-server/:${PATH}"
 
-RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
+RUN apt-get update && 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 --index-url https://download.pytorch.org/whl/cu118
+RUN pip install --upgrade torch==2.0.1
 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 pip install peft==0.5.0
+RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft.git
 RUN pip install datasets==2.14.4
-RUN pip install triton==3.0.0
+RUN pip install triton==2.0.0.dev20221120
 RUN pip install xformers==0.0.20
-RUN pip install google-cloud-storage
-RUN pip install absl-py
+RUN pip install google-cloud-storage==2.7.0
+RUN pip install absl-py==1.4.0
 RUN pip install scipy==1.10.1
 RUN pip install evaluate==0.4.0
 RUN pip install scikit-learn==1.2.2
@@ -48,43 +47,52 @@ 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 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
+
+# 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
 
 # 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 BASE_MODEL_ID ""
-ENV MODEL_ID ""
+ENV MODEL_ID "openlm-research/open_llama_7b"
 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} \
+--model-name=${model_name} \
 --version=1.0 \
 --handler=/home/model-server/handler.py \
 --runtime=python3 \
@@ -95,5 +103,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}=${MODEL}.mar", \
+     "--models", "${model_name}=${model_name}.mar", \
      "--model-store", "/home/model-server/model-store"]

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

@@ -0,0 +1,111 @@
+# 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,39 +72,15 @@ class PeftHandler(BaseHandler):
         "PRECISION_LOADING_MODE", constants.PRECISION_MODE_16
     )
     self.task = os.environ.get("TASK", CAUSAL_LANGUAGE_MODELING_LORA)
-    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.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", "")
     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):
@@ -125,7 +101,8 @@ class PeftHandler(BaseHandler):
 
     logging.info(
         f"Using task:{self.task}, base model:{self.model_id}, lora model:"
-        f" {self.finetuned_lora_model_path}, precision {self.precision_mode}."
+        f" {self.finetuned_lora_model_path}, and precision"
+        f" {self.precision_mode}."
     )
 
     self.pipeline = None
@@ -168,18 +145,11 @@ class PeftHandler(BaseHandler):
     elif (
         self.task == CAUSAL_LANGUAGE_MODELING_LORA or self.task == INSTRUCT_LORA
     ):
-      tokenizer = AutoTokenizer.from_pretrained(
-          self.model_id,
-          trust_remote_code=self.trust_remote_code,
-      )
-
+      tokenizer = AutoTokenizer.from_pretrained(self.model_id)
       logging.debug("Initialized the tokenizer.")
       if self.task == CAUSAL_LANGUAGE_MODELING_LORA:
         if self.quantization == constants.AWQ:
-          model = AutoAWQForCausalLM.from_quantized(
-              self.model_id,
-              trust_remote_code=self.trust_remote_code,
-          )
+          model = AutoAWQForCausalLM.from_quantized(self.model_id)
         elif self.quantization == constants.GPTQ or not self.quantization:
           if self.precision_mode == constants.PRECISION_MODE_32:
             model = AutoModelForCausalLM.from_pretrained(
@@ -187,7 +157,6 @@ 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(
@@ -195,7 +164,6 @@ 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(
@@ -203,7 +171,6 @@ 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(
@@ -215,7 +182,6 @@ 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(
@@ -229,7 +195,6 @@ 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}")
@@ -238,14 +203,14 @@ class PeftHandler(BaseHandler):
           model = AutoModelForCausalLM.from_pretrained(
               self.model_id,
               torch_dtype=torch.bfloat16,
-              trust_remote_code=self.trust_remote_code,
+              trust_remote_code=True,
               device_map="auto",
           )
         except:  # pylint: disable=bare-except
           model = AutoModelForCausalLM.from_pretrained(
               self.model_id,
               torch_dtype=torch.bfloat16,
-              trust_remote_code=self.trust_remote_code,
+              trust_remote_code=True,
               device_map="auto",
           )
       logging.debug("Initialized the base model.")
@@ -364,4 +329,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

@@ -0,0 +1,131 @@
+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

@@ -0,0 +1,97 @@
+"""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
+
+
+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,
+    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=[
+          "query_key_value",
+          "dense",
+          "dense_h_to_4h",
+          "dense_4h_to_h",
+      ],
+  )
+
+  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

@@ -0,0 +1,177 @@
+"""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.',
+)
+
+_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,
+        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,
+        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/sequence_classification_lora.py

@@ -0,0 +1,133 @@
+"""Sequence classification with LoRA models."""
+
+# pylint: disable=g-importing-member
+
+from datasets import load_dataset
+import evaluate
+from peft import get_peft_model
+from peft import LoraConfig
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoTokenizer
+from transformers import get_linear_schedule_with_warmup
+
+
+def finetune_sequence_classification(
+    pretrained_model_id: str,
+    dataset_name: str,
+    output_dir: str,
+    lora_rank: int = 8,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.1,
+    num_epochs: int = 20,
+    batch_size: int = 32,
+    learning_rate: float = 3e-4,
+) -> None:
+  """Finetunes sequence classification."""
+  task = "mrpc"
+  device = "cuda"
+
+  peft_config = LoraConfig(
+      task_type="SEQ_CLS",
+      inference_mode=False,
+      r=lora_rank,
+      lora_alpha=lora_alpha,
+      lora_dropout=lora_dropout,
+  )
+  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
+    padding_side = "left"
+  else:
+    padding_side = "right"
+
+  tokenizer = AutoTokenizer.from_pretrained(
+      pretrained_model_id, padding_side=padding_side
+  )
+  if getattr(tokenizer, "pad_token_id") is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+  datasets = load_dataset(dataset_name, task)
+  metric = evaluate.load(dataset_name, task)
+
+  def tokenize_function(examples):
+    # max_length=None => use the model max length (it's actually the default)
+    outputs = tokenizer(
+        examples["sentence1"],
+        examples["sentence2"],
+        truncation=True,
+        max_length=None,
+    )
+    return outputs
+
+  tokenized_datasets = datasets.map(
+      tokenize_function,
+      batched=True,
+      remove_columns=["idx", "sentence1", "sentence2"],
+  )
+
+  # We also rename the 'label' column to 'labels' which is the expected name for
+  # labels by the models of the transformers library.
+  tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+  def collate_fn(examples):
+    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+  # Instantiate dataloaders.
+  train_dataloader = DataLoader(
+      tokenized_datasets["train"],
+      shuffle=True,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+  eval_dataloader = DataLoader(
+      tokenized_datasets["validation"],
+      shuffle=False,
+      collate_fn=collate_fn,
+      batch_size=batch_size,
+  )
+
+  model = AutoModelForSequenceClassification.from_pretrained(
+      pretrained_model_id, return_dict=True
+  )
+  model = get_peft_model(model, peft_config)
+  model.print_trainable_parameters()
+
+  optimizer = AdamW(params=model.parameters(), lr=learning_rate)
+
+  # Instantiate scheduler
+  lr_scheduler = get_linear_schedule_with_warmup(
+      optimizer=optimizer,
+      num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
+      num_training_steps=(len(train_dataloader) * num_epochs),
+  )
+
+  model.to(device)
+  for epoch in range(num_epochs):
+    model.train()
+    for _, batch in enumerate(tqdm(train_dataloader)):
+      batch.to(device)
+      outputs = model(**batch)
+      loss = outputs.loss
+      loss.backward()
+      optimizer.step()
+      lr_scheduler.step()
+      optimizer.zero_grad()
+
+    model.eval()
+    for _, batch in enumerate(tqdm(eval_dataloader)):
+      batch.to(device)
+      with torch.no_grad():
+        outputs = model(**batch)
+      predictions = outputs.logits.argmax(dim=-1)
+      references = batch["labels"]
+      metric.add_batch(
+          predictions=predictions,
+          references=references,
+      )
+
+    eval_metric = metric.compute()
+    print(f"epoch {epoch}:", eval_metric)
+
+  model.save_pretrained(output_dir)

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

@@ -0,0 +1,6 @@
+#!/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

@@ -1,16 +0,0 @@
-# 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

@@ -1,20 +0,0 @@
-#!/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

@@ -1,88 +0,0 @@
-"""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/test_utils/command_builder.py

@@ -1,95 +0,0 @@
-"""Class that bundles docker related flags."""
-
-import getpass
-import os
-import pwd
-
-
-class CommandBuilder:
-  """Base class for building commands."""
-
-  def __init__(self):
-    self._defaults = []
-    self._env_vars = {}
-
-  def add_env_var(self, var: str, val: str) -> None:
-    """Add environment variable to the command.
-
-    Args:
-      var: environment variable name.
-      val: environment variable value.
-    """
-    self._env_vars[var] = val
-
-  def add_mount_map(self, host_path, docker_path):
-    pass
-
-
-class DockerCommandBuilder(CommandBuilder):
-  """Bundle docker related flags."""
-
-  def __init__(self, docker_uri: str, shm_size: str = '128gb'):
-    super().__init__()
-    self._docker_uri = [docker_uri]
-    self.privilege_mode = []
-    self.entrypoint = []
-
-    self._defaults = [
-        'docker',
-        'run',
-        '--gpus=all',
-        '--net=host',
-        '--rm',
-        f'--shm-size={shm_size}',
-    ]
-
-    self._mount_maps = []
-    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',
-      ]
-
-  def add_mount_map(self, host_path, docker_path):
-    self._mount_maps.append(f'--volume={host_path}:{docker_path}')
-
-  def add_privilege_mode(self):
-    self.privilege_mode = ['--privileged']
-
-  def add_entrypoint(self, entrypoint: list[str]):
-    self.entrypoint = entrypoint
-
-  def build_cmd(self) -> str:
-    return (
-        self._defaults
-        + [f'--env={var}={val}' for var, val in self._env_vars.items()]
-        + self._mount_maps
-        + self.privilege_mode
-        + self._docker_uri
-        + self.entrypoint
-    )
-
-
-class PythonCommandBuilder(CommandBuilder):
-  """Bundle Python test command related flags."""
-
-  def __init__(self):
-    super().__init__()
-    self._defaults = [
-        'python3',
-        './vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py',
-    ]
-
-  def build_cmd(self) -> str:
-    os.environ.update(self._env_vars)
-    return self._defaults
-
-  def add_entrypoint(self, entrypoint: list[str]):
-    self._defaults = entrypoint

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

@@ -1,471 +0,0 @@
-"""Test util class."""
-
-import copy
-import dataclasses
-import datetime
-import inspect
-import os
-import signal
-import subprocess
-import sys
-from absl import flags
-from absl import logging
-from absl.testing import parameterized
-import command_builder
-import immutabledict
-import torch
-
-_DOCKER_URI = flags.DEFINE_string('docker_uri', None, 'docker image uri')
-
-_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://vmg-tuning-docker-test',
-    'GCS directory that stores model checkpoint, dataset and etc.',
-)
-
-_GCS_OUTPUT_DIR = flags.DEFINE_string(
-    'gcs_output_dir',
-    'gs://vmg-tuning-docker-test/output',
-    'GCS directory that stores test output.',
-)
-
-_GCS_TESTDATA_DIR = 'peft-train-image-test'
-
-_THROUGHPUT_TEST_EXCEPTIONS = immutabledict.immutabledict({
-    ('bm_deepspeed_zero3_8gpu_gemma-2-9b-it_4bit.txt', '12.0'): float('inf'),
-    ('bm_fsdp_8gpu_llama3.1-70b-hf_4bit.txt', '20.0'): float('inf'),
-    ('bm_deepspeed_zero2_8gpu_gemma-2-2b-it_bfloat16.txt', '12.0'): 20.0,
-    ('bm_deepspeed_zero3_8gpu_gemma-2-2b-it_4bit.txt', '4.0'): 20.0,
-    ('bm_deepspeed_zero3_8gpu_gemma-2-27b-it_4bit.txt', '4.0'): 20.0,
-})
-
-
-@dataclasses.dataclass
-class BenchmarkStats:
-  """Class to store the benchmark result.
-
-  Attributes:
-    peak_mem: peak memory in GB.
-    throughput: throughput in tokens/sec.
-  """
-
-  peak_mem: float
-  throughput: float
-
-
-class TestBase(parameterized.TestCase):
-  """Test base class that defines how to run commands."""
-
-  def setUp(self):
-    super().setUp()
-
-    # Create a copy of the environment variables
-    self.old_env_var = copy.deepcopy(os.environ)
-    if _DOCKER_URI.value:
-      self.command_builder = command_builder.DockerCommandBuilder(
-          _DOCKER_URI.value
-      )
-    else:
-      self.command_builder = command_builder.PythonCommandBuilder()
-
-    self.command_builder.add_mount_map(
-        os.path.expanduser('~'), os.path.expanduser('~')
-    )
-    self.command_builder.add_mount_map(
-        self.local_input_dir(), self.local_input_dir()
-    )
-
-    self.task_cmd_builder = None
-
-  def tearDown(self):
-    super().tearDown()
-    # Restore the original environment variables
-    os.environ.clear()
-    os.environ.update(self.old_env_var)
-
-  def cmd(self):
-    return self.command_builder.build_cmd() + self.task_cmd_builder.build_cmd()
-
-  def run_cmd(self) -> int:
-    return run_cmd(self.cmd(), output_file=None)
-
-  def gcs_output_dir(self):
-    return _GCS_OUTPUT_DIR.value
-
-  def local_input_dir(self):
-    """Returns local input dir in host/docker."""
-    return _LOCAL_INPUT_DIR.value
-
-  def local_output_dir(self):
-    """Returns local output dir in host/docker."""
-    return _LOCAL_OUTPUT_DIR.value
-
-  def get_testcase_name(self):
-    """Returns the function name at the calling site."""
-    # https://docs.python.org/3/library/inspect.html#inspect.FrameInfo
-    cur_frame = inspect.currentframe()
-    # https://stackoverflow.com/a/17366561
-    return cur_frame.f_back.f_code.co_name
-
-
-def get_timestamp():
-  return datetime.datetime.now(datetime.timezone.utc).strftime(
-      '%Y%m%d_%H%M%S%Z'
-  )
-
-
-def download_from_gcs(gcs_uri: str, local_dir: str):
-  if not os.path.exists(local_dir):
-    os.mkdir(local_dir)
-  subprocess.check_output([
-      'gcloud',
-      'storage',
-      'cp',
-      '-r',
-      gcs_uri,
-      local_dir,
-  ])
-
-
-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.
-  """
-  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):
-    # If `name` is a file in sub-folders, then create the sub-folders under
-    # `_LOCAL_INPUT_DIR`.
-    local_data_dir = os.path.dirname(local_data)
-    if not os.path.exists(local_data_dir):
-      os.makedirs(local_data_dir)
-
-    download_from_gcs(os.path.join(_GCS_INPUT_DIR.value, name), local_data_dir)
-
-  return local_data
-
-
-def run_cmd(cmd: list[str], output_file: str = None) -> int:
-  """Runs the command and returns the return code.
-
-  Args:
-    cmd: The command to run.
-    output_file: The file to write the output to.
-
-  Returns:
-    The return code of the command.
-  """
-  logging.info('running command: \n%s', ' \\\n'.join(cmd))
-  if _DRY_RUN.value:
-    return 0
-  stdout = sys.stdout if output_file is None else open(output_file, 'w')
-  p = subprocess.Popen(cmd, stdout=stdout, stderr=sys.stderr)
-  try:
-    unused_output, unused_error = p.communicate()
-    return_code = p.returncode
-  except KeyboardInterrupt:
-    p.send_signal(signal.SIGINT)
-    return_code = 0
-  finally:
-    if output_file is not None:
-      stdout.close()
-  return return_code
-
-
-def get_pretrained_model_name_or_path(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)
-
-
-def is_gpu_h100():
-  """Checks if the GPU is H100."""
-  return 'H100' in torch.cuda.get_device_name()
-
-
-def is_gpu_a100():
-  """Checks if the GPU is A100."""
-  return 'A100' in torch.cuda.get_device_name()
-
-
-def _get_formatted_string(max_seq_length: int) -> str:
-  """Returns the formatted string for max_seq_length.
-
-  Args:
-    max_seq_length: max sequence length to get the formatted string.
-
-  Returns:
-    formatted string for max_seq_length.
-  """
-  return f'{max_seq_length/1024.0:.1f}'
-
-
-def get_benchmark_results(
-    benchmark_file_path: str, max_seq_length: int
-) -> BenchmarkStats:
-  """Gets benchmark results from the benchmark file.
-
-  Args:
-    benchmark_file_path: path to the benchmark file.
-    max_seq_length: max sequence length to get the benchmark results.
-
-  Returns:
-    peak_mem: peak memory in GB.
-    throughput: throughput in tokens/sec.
-  """
-  formatted_max_seq_length = _get_formatted_string(max_seq_length)
-  peak_mem, throughput = None, None
-  with open(benchmark_file_path, 'r') as f:
-    for line in f:
-      if line.startswith(formatted_max_seq_length):
-        metrics = line.split('|')
-        try:
-          peak_mem = float(metrics[1].strip())
-        except ValueError:
-          pass
-        try:
-          throughput = float(metrics[2].strip())
-        except ValueError:
-          pass
-        break
-    else:
-      logging.error(
-          'No metrics found for max_seq_length %s in %s',
-          formatted_max_seq_length,
-          benchmark_file_path,
-      )
-  return BenchmarkStats(peak_mem, throughput)
-
-
-def print_benchmark_file(file_path: str) -> None:
-  """Prints the contents of the file.
-
-  Args:
-    file_path: path to the file.
-  """
-  with open(file_path, 'r') as f:
-    for line in f:
-      logging.info(line.strip())
-
-
-def print_benchmark_results(
-    benchmark_file_path: str, benchmark_type: str
-) -> None:
-  """Prints the benchmark results.
-
-  Args:
-    benchmark_file_path: path to the benchmark file.
-    benchmark_type: type of the benchmark.
-  """
-  benchmark_filename = os.path.basename(benchmark_file_path)
-  logging.info('--------------------------------------------------------------')
-  logging.info('%s benchmark for %s', benchmark_type, benchmark_filename)
-  logging.info('--------------------------------------------------------------')
-  print_benchmark_file(benchmark_file_path)
-
-
-def _calculate_percent_change(
-    actual_value: float, expected_value: float
-) -> float:
-  """Calculates the percent change between the actual and expected values.
-
-  Args:
-    actual_value: actual value to compare.
-    expected_value: expected value to compare.
-
-  Returns:
-    percent change between the actual and expected values.
-  """
-  return ((actual_value - expected_value) / expected_value) * 100.0
-
-
-def compare_benchmark_results(
-    expected_benchmark_file_path: str,
-    actual_benchmark_file_path: str,
-    allowed_threshold: float,
-    max_seq_length: int,
-) -> bool:
-  """Compares if the benchmark results are the similar.
-
-  Args:
-    expected_benchmark_file_path: path to the expected benchmark file.
-    actual_benchmark_file_path: path to the actual benchmark file.
-    allowed_threshold: allowed percent range of the benchmark results.
-    max_seq_length: max sequence length to get the benchmark results.
-
-  Returns:
-    True if the benchmark results are the similar, False otherwise.
-  """
-  benchmark_filename = os.path.basename(expected_benchmark_file_path)
-  expected_results = get_benchmark_results(
-      expected_benchmark_file_path, max_seq_length
-  )
-  expected_peak_mem, expected_throughput = (
-      expected_results.peak_mem,
-      expected_results.throughput,
-  )
-  actual_results = get_benchmark_results(
-      actual_benchmark_file_path, max_seq_length
-  )
-  actual_peak_mem, actual_throughput = (
-      actual_results.peak_mem,
-      actual_results.throughput,
-  )
-  formatted_max_seq_length = _get_formatted_string(max_seq_length)
-
-  # Case 1: both peak mem and throughput are None(ideally due to OOM)
-  if expected_peak_mem is None and actual_peak_mem is None:
-    logging.info(
-        'Both peak mem and throughput are None for max_seq_length %d.',
-        max_seq_length,
-    )
-    return True
-
-  check_oom_exception = _THROUGHPUT_TEST_EXCEPTIONS.get(
-      (benchmark_filename, formatted_max_seq_length), 0.0
-  ) == float('inf')
-  # Case 2: When something strated to fail recently, or something which failed
-  # before but is working now.
-  if expected_peak_mem is None and actual_peak_mem is not None:
-    if check_oom_exception:
-      return True
-    logging.error(
-        'One of the failing benchmarks in %s is passing now for max_seq_length'
-        ' %d. The expected peak mem and throughput are None, but the actual'
-        ' peak mem is %f and actual throughput is %f',
-        benchmark_filename,
-        max_seq_length,
-        actual_peak_mem,
-        actual_throughput,
-    )
-    return False
-  if actual_peak_mem is None and expected_peak_mem is not None:
-    if check_oom_exception:
-      return True
-    logging.error(
-        'One of the passing benchmarks in %s is failing now for max_seq_length'
-        ' %d. The actual peak mem and throughput are None, but the expected'
-        ' peak mem is %f and expected throughput is %f',
-        benchmark_filename,
-        max_seq_length,
-        expected_peak_mem,
-        expected_throughput,
-    )
-    return False
-  # Case 3: When both actual peak mem and throughput lies within the range
-  # of their respective expected values.
-  mem_percent_change = _calculate_percent_change(
-      actual_peak_mem, expected_peak_mem
-  )
-  throughput_percent_change = _calculate_percent_change(
-      actual_throughput, expected_throughput
-  )
-  allowed_threshold = _THROUGHPUT_TEST_EXCEPTIONS.get(
-      (benchmark_filename, formatted_max_seq_length), allowed_threshold
-  )
-
-  if abs(mem_percent_change) > allowed_threshold:
-    logging.error(
-        'The peak memory is changing by more than %f%% for max_seq_length %d.'
-        ' Expected: %f, Actual: %f',
-        allowed_threshold,
-        max_seq_length,
-        expected_peak_mem,
-        actual_peak_mem,
-    )
-    return False
-  if abs(throughput_percent_change) > allowed_threshold:
-    logging.error(
-        'The throughput is changing by more than %f%% for max_seq_length %d.'
-        ' Expected throughput: %f, Actual throughput: %f',
-        allowed_threshold,
-        max_seq_length,
-        expected_throughput,
-        actual_throughput,
-    )
-    return False
-
-  return True
-
-
-def check_benchmark_results(
-    actual_benchmark_file_path: str,
-    model_family: str,
-    allowed_threshold: float,
-    max_seq_length: int,
-) -> bool:
-  """Checks the benchmark result between the actual and expected benchmark files.
-
-  Args:
-    actual_benchmark_file_path: path to the actual benchmark file.
-    model_family: family of the model.
-    allowed_threshold: allowed range of the benchmark results in percent.
-    max_seq_length: max sequence length to get the benchmark results.
-
-  Returns:
-    True if the benchmark results are the similar, False otherwise.
-  """
-  benchmark_filename = os.path.basename(actual_benchmark_file_path)
-  get_test_data_path(_GCS_TESTDATA_DIR)
-  expected_benchmark_file_path = os.path.join(
-      _LOCAL_INPUT_DIR.value,
-      _GCS_TESTDATA_DIR,
-      model_family,
-      benchmark_filename,
-  )
-  print_benchmark_results(expected_benchmark_file_path, 'Expected')
-  print_benchmark_results(actual_benchmark_file_path, 'Actual')
-
-  return compare_benchmark_results(
-      expected_benchmark_file_path,
-      actual_benchmark_file_path,
-      allowed_threshold,
-      max_seq_length,
-  )
-
-
-def list_gcs_directories(bucket: str, directory: str) -> list[str]:
-  """Lists GCS files."""
-  output = subprocess.check_output([
-      'gcloud',
-      'storage',
-      'ls',
-      f'gs://{bucket}/{directory}',
-  ])
-  return output.decode('utf-8').splitlines()
-
-
-def delete_gcs_object(gcs_directory: str):
-  """Deletes GCS object."""
-  subprocess.check_output([
-      'gcloud',
-      'storage',
-      'rm',
-      '-r',
-      f'{gcs_directory}',
-  ])

community-content/vertex_model_garden/model_oss/peft/train/util/cluster_spec.py

@@ -1,79 +0,0 @@
-"""Get cluster info from environment variables."""
-
-import dataclasses
-import json
-import os
-
-from absl import logging
-
-
-@dataclasses.dataclass
-class ClusterInfo:
-  """Contains information about the cluster.
-
-  Attributes:
-    primary_node_addr: The address of the primary node.
-    primary_node_port: The port of the primary node.
-    node_rank: The rank of the node.
-    num_nodes: The number of nodes in the cluster.
-  """
-
-  primary_node_addr: str | None = None
-  primary_node_port: str | None = None
-  node_rank: int = 0
-  num_nodes: int = 1
-
-  # Allows unpacking operation like
-  # primary_node_addr, primary_node_port, _, _ = ClusterInfo()
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def get_cluster_spec() -> ClusterInfo:
-  """Parses CLUSTER_SPEC environment variable and returns the cluster info.
-
-  Returns:
-    A ClusterInfo object.
-  """
-  cluster_spec = os.getenv('CLUSTER_SPEC', None)
-
-  # If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
-  if not cluster_spec:
-    cluster_info = ClusterInfo(
-        primary_node_addr=os.getenv('MASTER_ADDR', None),
-        primary_node_port=os.getenv('MASTER_PORT', None),
-        node_rank=int(os.getenv('RANK', '0')),
-        num_nodes=int(os.getenv('NNODES', '1')),
-    )
-    return cluster_info
-
-  cluster_data = json.loads(cluster_spec)
-  # 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_nodes = 1  # For the primary node.
-  if 'workerpool1' in cluster_data['cluster']:
-    num_nodes += len(cluster_data['cluster']['workerpool1'])
-  logging.info('num nodes: %s', num_nodes)
-
-  return ClusterInfo(primary_node_addr, primary_node_port, node_rank, num_nodes)

community-content/vertex_model_garden/model_oss/peft/train/util/common_util.py

@@ -1,24 +0,0 @@
-"""Utility functions."""
-
-import logging
-import subprocess
-import sys
-import time
-
-
-def run_cmd(cmd: list[str]) -> float:
-  """Runs the command and logs the output.
-
-  Args:
-    cmd: The command to run.
-
-  Returns:
-    The time it took to run the command.
-  """
-  cmd_str = ' \\\n'.join(cmd)
-  logging.info('launching cmd: \n%s', cmd_str)
-  start_time = time.time()
-  subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
-  elapsed_time = round(time.time() - start_time, 2)
-  logging.info('Command %s finished in %0.2f seconds.', cmd_str, elapsed_time)
-  return elapsed_time

community-content/vertex_model_garden/model_oss/peft/train/util/dataset_stats.py

@@ -1,197 +0,0 @@
-"""Calculate dataset statistics like token, example and character counts."""
-
-from collections.abc import Mapping, Sequence
-import dataclasses
-import json
-from typing import Any
-import datasets
-import numpy as np
-import transformers
-from util import dataset_validation_util
-
-_MAX_NUM_DATASET_SAMPLES = 6
-
-
-@dataclasses.dataclass
-class SupervisedTuningDatasetBucket:
-  """Represents a histogram bucket for tuning dataset distribution stats."""
-
-  count: float = 0
-  left: float = 0
-  right: float = 0
-
-
-@dataclasses.dataclass
-class SupervisedTuningDatasetDistribution:
-  """Represents a histogram with summary statistics for tuning dataset distribution stats."""
-
-  sum: int = 0
-  billable_sum: int = 0
-  min: float = 0
-  max: float = 0
-  mean: float = 0
-  median: float = 0
-  p5: float = 0
-  p95: float = 0
-  buckets: list[SupervisedTuningDatasetBucket] = dataclasses.field(
-      default_factory=list
-  )
-
-
-# Represents detailed tuning dataset statistics.
-@dataclasses.dataclass
-class SupervisedTuningDataStats:
-  """Represents detailed tuning dataset stats."""
-
-  tuning_dataset_example_count: int = 0
-  total_tuning_character_count: int = 0
-  total_billable_token_count: int = 0
-  tuning_step_count: int = 0
-  # Represents a histogram and some summary statistics of the number of input
-  # tokens across examples.
-  user_input_token_distribution: SupervisedTuningDatasetDistribution | None = (
-      None
-  )
-  # Represents a histogram and some summary statistics for the number of output
-  # tokens across examples.
-  user_output_token_distribution: SupervisedTuningDatasetDistribution | None = (
-      None
-  )
-  # Represents the number of "messages" (a single-turn conversation will have a
-  # single message) across examples.
-  user_message_per_example_distribution: (
-      SupervisedTuningDatasetDistribution | None
-  ) = None
-  user_dataset_examples: list[str] = dataclasses.field(default_factory=list)
-
-
-def get_dataset_stats(
-    *,
-    raw: Any,
-    templated: Any,
-    template: str,
-    tokenizer: transformers.PreTrainedTokenizer,
-    column: str,
-    effective_batch_size: int,
-) -> Mapping[str, Any]:
-  """Calculates dataset statistics for managed fine-tuning, e.g., total number of tokens."""
-  tokenized_dataset = templated.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]))
-  total_tuning_character_count = int(
-      np.sum([len(ex[column]) for ex in templated])
-  )
-  tuning_step_count = (
-      tuning_dataset_example_count + effective_batch_size - 1
-  ) // effective_batch_size
-
-  # Assume that data is represented as ChatCompletions or Vertex Text-Bison
-  # formats to extract per-example input/output tokens.
-  user_inputs = []
-  user_outputs = []
-  user_input_messages_counts = []
-
-  for ex in raw:
-    if "messages" in ex:
-      messages = ex["messages"]
-      if messages:
-        # For ChatCompletions assume the last turn (i.e. the instruction
-        # response) is the expected output.
-        user_inputs.append({**ex, "messages": messages[:-1]})
-        user_outputs.append({**ex, "messages": messages[-1:]})
-        # Exclude everything but the last message for the number of input
-        # messages.
-        user_input_messages_counts.append(len(messages[:-1]))
-    elif "input_text" in ex:
-      # For Vertex Text-Bison, the `output_text` field is the expected output.
-      user_inputs.append({**ex, "output_text": ""})
-      user_outputs.append(
-          {**ex, "input_text": ex["output_text"], "output_text": ""}
-      )
-      # Vertex Text-Bison goes from input -> output; i.e. there is only a single
-      # input "message".
-      user_input_messages_counts.append(1)
-
-  def calc_histogram(
-      counts: Sequence[int],
-  ) -> SupervisedTuningDatasetDistribution:
-    mean = np.mean(counts)
-    median = np.median(counts).item()
-    max_count = np.max(counts).item()
-    min_count = np.min(counts).item()
-    count_sum = np.sum(counts).item()
-    p5 = np.percentile(counts, 0.05).item()
-    p95 = np.percentile(counts, 0.95).item()
-    hist, bin_edges = np.histogram(counts, bins=10)
-
-    return SupervisedTuningDatasetDistribution(
-        sum=count_sum,
-        billable_sum=count_sum,
-        min=min_count,
-        max=max_count,
-        mean=mean,
-        median=median,
-        p5=p5,
-        p95=p95,
-        buckets=[
-            SupervisedTuningDatasetBucket(
-                count=hist[i].item(),
-                left=bin_edges[i].item(),
-                right=bin_edges[i + 1].item(),
-            )
-            for i in range(len(hist))
-        ],
-    )
-
-  # Tokenize input and output messages separately to generate separate summary
-  # statistics about them.
-  user_input_token_distribution = None
-  if user_inputs:
-    user_input_dataset = dataset_validation_util.format_dataset(
-        datasets.Dataset.from_list(user_inputs), column, template, tokenizer
-    )
-    user_input_tokenized_dataset = user_input_dataset.map(
-        lambda x: tokenizer(x[column])
-    )
-    user_input_tokens = user_input_tokenized_dataset["input_ids"]
-    user_input_token_counts = np.array([len(ex) for ex in user_input_tokens])
-    user_input_token_distribution = calc_histogram(user_input_token_counts)
-
-  user_output_token_distribution = None
-  if user_outputs:
-    user_output_dataset = dataset_validation_util.format_dataset(
-        datasets.Dataset.from_list(user_outputs), column, template, tokenizer
-    )
-    user_output_tokenized_dataset = user_output_dataset.map(
-        lambda x: tokenizer(x[column])
-    )
-    user_output_tokens = user_output_tokenized_dataset["input_ids"]
-    user_output_token_counts = np.array([len(ex) for ex in user_output_tokens])
-    user_output_token_distribution = calc_histogram(user_output_token_counts)
-
-  user_messages_per_example_distribution = None
-  if user_input_messages_counts:
-    user_input_messages_counts = np.array(user_input_messages_counts)
-    user_messages_per_example_distribution = calc_histogram(
-        user_input_messages_counts
-    )
-
-  user_dataset_examples = [
-      json.dumps(ex)
-      for ex in raw.shuffle().select(
-          range(min(len(raw), _MAX_NUM_DATASET_SAMPLES))
-      )
-  ]
-
-  dataset_stats = SupervisedTuningDataStats(
-      tuning_dataset_example_count=tuning_dataset_example_count,
-      total_tuning_character_count=total_tuning_character_count,
-      total_billable_token_count=total_billable_token_count,
-      tuning_step_count=tuning_step_count,
-      user_input_token_distribution=user_input_token_distribution,
-      user_output_token_distribution=user_output_token_distribution,
-      user_message_per_example_distribution=user_messages_per_example_distribution,
-      user_dataset_examples=user_dataset_examples,
-  )
-  return dataclasses.asdict(dataset_stats)

community-content/vertex_model_garden/model_oss/peft/train/util/device_stats.py

@@ -1,140 +0,0 @@
-"""Util functions for reporting device (GPU, CPU) stats."""
-
-import dataclasses
-
-import psutil
-import pynvml
-import torch
-
-
-@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
-  """
-
-  # device id
-  device_id: int
-  # memory reserved.
-  reserved: 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
-  # device_id, reserved, 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(device, reserved, occupied, unused, smi_diff, util)
-
-
-def gpu_stats_str(stats: GpuStats | None = None) -> str:
-  if stats is None:
-    stats = gpu_stats()
-  device, reserved, occupied, unused, smi_diff, util = stats
-  return (
-      f"GPU ({device=}) memory: {reserved:.2f}({occupied=:.2f}, {unused=:.2f}),"
-      f" {smi_diff=:.2f} GB. Utilization: {util:.2f}%"
-  )
-
-
-@dataclasses.dataclass
-class CpuStats:
-  """Holds information about CPU usage stats."""
-
-  # Total CPU virtual memory i.e. virtual memory allocated + unallocated.
-  total_virtual_mem: float
-  # CPU virtual memory available for use.
-  unallocated_virtual_mem: float
-  # CPU virtual memory already used.
-  allocated_virtual_mem: float
-  # Total CPU swap memory i.e. swap memory allocated + unallocated.
-  total_swap_mem: float
-  # CPU swap memory available for use.
-  unallocated_swap_mem: float
-  # CPU swap memory already used.
-  allocated_swap_mem: float
-  # CPU utilization percentage.
-  utilization: float
-
-
-def cpu_stats() -> CpuStats:
-  """Reports CPU memory usage and utilization."""
-
-  # https://psutil.readthedocs.io/en/latest/#memory
-  gb = 1024.0**3
-  vmem = psutil.virtual_memory()
-  vmem_total = vmem.total / gb
-  vmem_available = vmem.available / gb
-  vmem_used = vmem_total - vmem_available
-  smem = psutil.swap_memory()
-  swap_total = smem.total / gb
-  swap_free = smem.free / gb
-  swap_used = smem.used / gb
-  # https://psutil.readthedocs.io/en/latest/#psutil.cpu_percent
-  cpu_util = psutil.cpu_percent(interval=1e-6)
-  return CpuStats(
-      total_virtual_mem=vmem_total,
-      unallocated_virtual_mem=vmem_available,
-      allocated_virtual_mem=vmem_used,
-      total_swap_mem=swap_total,
-      unallocated_swap_mem=swap_free,
-      allocated_swap_mem=swap_used,
-      utilization=cpu_util,
-  )
-
-
-def cpu_stats_str(stats: CpuStats | None = None) -> str:
-  """Returns a string representation of the CPU stats."""
-
-  if stats is None:
-    stats = cpu_stats()
-  total, occupied, unused = (
-      stats.total_virtual_mem,
-      stats.allocated_virtual_mem,
-      stats.unallocated_virtual_mem,
-  )
-  virtual_mem = (
-      f"CPU virtual memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
-  )
-  total, occupied, unused = (
-      stats.total_swap_mem,
-      stats.allocated_swap_mem,
-      stats.unallocated_swap_mem,
-  )
-  swap_mem = f"CPU swap memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
-  percent = stats.utilization
-  return f"{virtual_mem} {swap_mem} CPU Utilization: {percent:.2f}%"

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

@@ -1,126 +0,0 @@
-"""Different trainer callbacks for PEFT Trainer."""
-
-from collections.abc import MutableMapping
-import math
-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 util import device_stats
-
-
-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_log(
-      self,
-      args: TrainingArguments,
-      state: TrainerState,
-      control: TrainerControl,
-      logs: MutableMapping[str, float] | None = None,
-      **kwargs,
-  ) -> None:
-    """Calculates perplexity from train loss.
-
-    Args:
-      args: Arguments passed to the trainer.
-      state: State of the trainer.
-      control: Control of the trainer.
-      logs: A dict of logs from the training loop.
-      **kwargs: Additional keyword arguments, not used in this callback.
-    """
-    del kwargs  # Unused.
-    if self._partial_state.is_main_process:
-      train_loss = logs.get('loss') if logs is not None else None
-      if train_loss is not None:
-        perplexity = round(float(math.exp(train_loss)), 4)
-        logs['perplexity'] = perplexity
-
-  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()
-        self._prev_num_token = state.num_input_tokens_seen
-        throughput = 0.0
-      else:
-        cur_time = time.time()
-        cur_num_token = state.num_input_tokens_seen
-        throughput = (cur_num_token - self._prev_num_token) / (
-            cur_time - self._prev_time
-        )
-        self._prev_time = cur_time
-        self._prev_num_token = cur_num_token
-        self._avg_throughput += (throughput - self._avg_throughput) / (
-            state.global_step - 1
-        )
-
-      gpu_stats = device_stats.gpu_stats()
-      self._peak_mem = max(
-          gpu_stats.reserved + gpu_stats.smi_diff, self._peak_mem
-      )
-      logging.info(
-          'on_step_end: Throughput: %.2f token/s. %s, %s',
-          throughput,
-          device_stats.gpu_stats_str(gpu_stats),
-          device_stats.cpu_stats_str(),
-      )
-
-  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, %s',
-          device_stats.gpu_stats_str(),
-          device_stats.cpu_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
-      throughput = state.num_input_tokens_seen / train_time
-      logging.info(
-          'training time %.2f s, throughput (including overhead, e.g., ckpt'
-          ' saving): %.2f token/s, peak_mem: %.2f GB',
-          train_time,
-          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} | {self._peak_mem:.2f} |'
-              f' {self._avg_throughput:.2f}\n'
-          )

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

@@ -1,18 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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

@@ -1,17 +0,0 @@
-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/gemma2_fsdp_8gpu.yaml

@@ -1,28 +0,0 @@
-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: Gemma2DecoderLayer
-  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_fsdp_128gpu.yaml

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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/qwen2_fsdp_8gpu.yaml

@@ -1,28 +0,0 @@
-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: Qwen2DecoderLayer
-  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/zero2.json

@@ -1,24 +0,0 @@
-{
-  "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

@@ -1,31 +0,0 @@
-{
-  "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

@@ -1,43 +0,0 @@
-# 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.34.2  # 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.47.1
-  - trl==0.11.2

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

@@ -1,32 +0,0 @@
-# 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.34.2
-auto_gptq==0.7.1+cu118
-autoawq==0.2.8
-bitsandbytes==0.43.2
-cloudml-hypertune==0.1.0.dev6
-datasets==2.20.0
-deepspeed==0.15.2
-diffusers==0.25.1
-evaluate==0.4.3
-fsspec==2024.3.1
-gcsfs==2024.3.1
-immutabledict==4.2.1
-ninja==1.11.1 # Needed to avoid `ninja 1.11.1.1 is not supported on this platform` error
-nltk==3.9.1
-optimum==1.17.1
-peft==0.12.0
-pynvml==11.5.3
-rouge_score==0.1.2
-torch==2.2.2+cu118
-torchvision==0.17.2+cu118
-transformers==4.47.1
-trl==0.11.2
-wandb==0.17.1
-ydata-profiling==4.7.0 # Upgrade the version from 4.6.0 to 4.7.0 to fix the old `pydantic` package error.
-psutil==6.0.0

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

@@ -1,83 +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}
-
-# 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
-
-# Remove packages that are not needed and are causing conflicts.
-# dataproc_jupyter_plugin was installed as a part of pytorch-cu121.2-2.py310
-# container which we don't need. It depends on ibis-framework and bigframes.
-# The package and its dependencies request lower versions of pyarrow/pydantic
-# than deepspeed/datasets. So, dataproc_jupyter_plugin conflicts with
-# deepspeed/datasets.
-RUN pip uninstall -y dataproc_jupyter_plugin ibis-framework bigframes
-
-# 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/
-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/peft/train/util/*.py /diffusers/examples/util/
-COPY model_oss/util/* /diffusers/examples/util/
-COPY model_oss/notebook_util/dataset_validation_util.py /diffusers/examples/util
-COPY model_oss/peft/train/vmg/tests/*.py ./vertex_vision_model_garden_peft/tests/
-COPY model_oss/peft/train/test_utils/test_util.py ./vertex_vision_model_garden_peft/tests/
-COPY model_oss/peft/train/test_utils/command_builder.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

@@ -1,284 +0,0 @@
-"""Library for running evaluations during training."""
-
-from collections.abc import Callable, Mapping, MutableMapping, Sequence
-import dataclasses
-import string
-from typing import Type
-
-from absl import logging
-import evaluate
-import numpy as np
-import torch
-import transformers
-
-from util import dataset_validation_util
-from util import constants
-
-
-_STRING_TRANSLATOR = str.maketrans("", "", string.punctuation)
-
-_GREATER_IS_BETTER_MAP = {
-    "loss": False,
-    "perplexity": False,
-    "bleu": True,
-    "google_bleu": True,
-    "rouge1": True,
-    "rouge2": True,
-    "rougeL": True,
-    "rougeLsum": True,
-}
-
-
-@dataclasses.dataclass(frozen=True)
-class EvalConfig:
-  """Configuration for running evaluations during training.
-
-  Attributes:
-    steps: The number of steps to run evaluation.
-    tasks: The list of tasks to run evaluation on.
-    per_device_batch_size: The per device batch size for evaluation.
-    limit: The maximum number of examples to evaluate.
-    metric_name: The name of the metric to compute.
-    tokenize_dataset: Whether to tokenize the dataset.
-    dataset_path: The path to the dataset.
-    split: The split of the dataset to evaluate.
-    template: The template to use for the dataset.
-    column: The column name of the dataset.
-    metric_for_best_model: The metric to use for loading the best model.
-  """
-
-  steps: int
-  per_device_batch_size: int
-  limit: float | None
-  metric_name: Sequence[str]
-  tokenize_dataset: bool
-  dataset_path: str = ""
-  split: str = "test"
-  template: str = ""
-  column: str = constants.DEFAULT_TRAIN_COLUMN
-  metric_for_best_model: str | None = None
-
-
-def create_trainer(
-    cls: Type[transformers.Trainer],
-    eval_config: EvalConfig | None,
-    tokenizer: transformers.PreTrainedTokenizerBase | None,
-    args: transformers.TrainingArguments,
-    **kwargs,
-) -> transformers.Trainer:
-  """Creates a trainer. If eval config is provided, injects evaluation loop.
-
-  Args:
-    cls: The trainer class.
-    eval_config: The evaluation config.
-    tokenizer: The tokenizer.
-    args: The training arguments.
-    **kwargs: The keyword arguments.
-
-  Returns:
-    A trainer.
-  """
-  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
-  args.metric_for_best_model = eval_config.metric_for_best_model
-  args.greater_is_better = _GREATER_IS_BETTER_MAP.get(
-      eval_config.metric_for_best_model, None
-  )
-  args.save_strategy = (
-      transformers.trainer_utils.SaveStrategy.STEPS
-      if eval_config.metric_for_best_model is None
-      else transformers.trainer_utils.SaveStrategy.BEST
-  )
-
-  kwargs["tokenizer"] = tokenizer
-
-  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 tokenizer is not None:
-      eval_dataset = dataset_validation_util.get_filtered_dataset(
-          dataset=eval_dataset,
-          input_column=eval_config.column,
-          max_seq_length=kwargs["max_seq_length"],
-          tokenizer=tokenizer,
-      )
-    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)
-
-
-def _cleanup_text(text: str) -> str:
-  """Cleans up the prediction and references text.
-
-  Args:
-    text: The text to clean up.
-
-  Returns:
-    Cleaned up text.
-  """
-  text = text.translate(_STRING_TRANSLATOR)
-  text = text.strip()
-  text = " ".join(text.split())
-  return text.lower()
-
-
-def create_compute_metrics(
-    tokenizer: transformers.PreTrainedTokenizerBase,
-    eval_metrics: Mapping[str, evaluate.EvaluationModule],
-) -> Callable[[transformers.EvalPrediction], MutableMapping[str, float]]:
-  """Creates a compute_metrics function using Hugging Face evaluate library.
-
-  Args:
-    tokenizer: The tokenizer for decoding predictions.
-    eval_metrics: The eval metrics to compute.
-
-  Returns:
-    Function that computes comprehensive metrics.
-  """
-
-  def _preprocess_data(
-      predictions: np.ndarray, labels: np.ndarray
-  ) -> tuple[Sequence[str], Sequence[str]]:
-    """Preprocesses predictions and lavels before evaluation.
-
-    Args:
-      predictions: The predictions to preprocess.
-      labels: The labels to preprocess.
-
-    Returns:
-      A tuple (preprocessed predictions, labels).
-    """
-    # Handle padding and special tokens.
-    predictions = np.where(
-        predictions != -100, predictions, tokenizer.pad_token_id
-    )
-    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
-
-    # Decode to text.
-    pred_texts = tokenizer.batch_decode(predictions, skip_special_tokens=True)
-    label_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
-
-    # Clean up text.
-    cleaned_pred_texts = [_cleanup_text(text) for text in pred_texts]
-    cleaned_label_texts = [_cleanup_text(text) for text in label_texts]
-
-    return cleaned_pred_texts, cleaned_label_texts
-
-  def _compute_metrics_with_tokenizer(
-      eval_pred: transformers.EvalPrediction,
-  ) -> MutableMapping[str, float]:
-    """Computes metrics using Hugging Face evaluate library.
-
-    Args:
-      eval_pred: The evaluation prediction.
-
-    Returns:
-      A dictionary of metrics.
-    """
-    predictions, perplexities = eval_pred.predictions
-    labels = eval_pred.label_ids
-
-    pred_texts, label_texts = _preprocess_data(predictions, labels)
-
-    metrics = {}
-
-    for eval_metric, computed_eval_metric in eval_metrics.items():
-      match eval_metric:
-        case "perplexity":
-          # We don't use the perplexity from HF Evaluate since it loads the
-          # model again. This causes an increase in the GPU utilization and
-          # hence an OOM. Due to this, we compute the perplexity ourselves
-          # using the eval_loss over the unmasked tokens in
-          # preprocess_logits_for_metrics fn.
-          metrics[eval_metric] = np.mean(perplexities)
-        case "bleu" | "google_bleu":
-          num_valid_labels = len(list(filter(None, label_texts)))
-          if num_valid_labels:
-            eval_score = computed_eval_metric.compute(
-                predictions=pred_texts,
-                references=[[text] for text in label_texts],
-            )
-            metrics[eval_metric] = eval_score[eval_metric]
-          else:
-            metrics[eval_metric] = 0.0
-        case "rouge1" | "rouge2" | "rougeL" | "rougeLsum":
-          rouge_scores = computed_eval_metric.compute(
-              predictions=pred_texts,
-              references=label_texts,
-              use_stemmer=True,
-          )
-          metrics[eval_metric] = rouge_scores[eval_metric]
-
-    pred_lengths = [len(pred.split()) for pred in pred_texts]
-    label_lengths = [len(label.split()) for label in label_texts]
-
-    metrics["gen_len"] = np.mean(pred_lengths)
-    metrics["ref_len"] = np.mean(label_lengths)
-    metrics["length_ratio"] = np.mean(
-        [len(p) / len(r) if r else 0 for p, r in zip(pred_texts, label_texts)]
-    )
-
-    # Round all metrics to 4 decimal places.
-    return {k: round(float(v), 4) for k, v in metrics.items()}
-
-  return _compute_metrics_with_tokenizer
-
-
-def preprocess_logits_for_metrics(
-    logits: torch.Tensor, labels: torch.Tensor
-) -> tuple[torch.Tensor, torch.Tensor]:
-  """Preprocesses the logits before caching them for eval metric calculation.
-
-  Args:
-    logits: Logits predicted by the model.
-    labels: Ground truth labels.
-
-  Returns:
-    A tuple (pred_ids, perplexities).
-  """
-  # Calculate prediction IDs.
-  pred_ids = logits.argmax(dim=-1)
-
-  # This step shifts the logits and labels to align them correctly, where we are
-  # predicting the next token in a sequence. The last logit doesn't have a
-  # corresponding label, and the first label doesn't have a preceding logit to
-  # predict it. This calculation of perplexity is inspired from
-  # https://github.com/huggingface/evaluate/blob/main/metrics/perplexity/perplexity.py.
-  shift_logits = logits[..., :-1, :].contiguous()
-  shift_labels = labels[..., 1:].contiguous()
-  attn_mask = shift_labels != -100
-  loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
-
-  perplexities = torch.exp(
-      (loss_fct(shift_logits.transpose(1, 2), shift_labels) * attn_mask).sum(1)
-      / attn_mask.sum(1)
-  )
-
-  return (pred_ids, perplexities)

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

@@ -1,905 +0,0 @@
-"""Instruct/Chat with LoRA models."""
-
-from collections.abc import Callable, Mapping, Sequence
-import datetime
-import json
-import os
-from typing import Any
-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 evaluate
-from peft import get_peft_model
-from peft import LoraConfig
-import torch
-import transformers
-import trl
-import wandb
-
-from util import dataset_validation_util
-from util import dataset_stats
-from util import device_stats
-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_NAME_OR_PATH = flags.DEFINE_string(
-    'pretrained_model_name_or_path',
-    None,
-    'The pretrained model name or path. 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_name_or_path 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.',
-)
-
-_TRAIN_DATASET = flags.DEFINE_string(
-    'train_dataset',
-    None,
-    'The training dataset name in huggingface or path.',
-)
-
-_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_TRAIN_EPOCHS = flags.DEFINE_float(
-    'num_train_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_train_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.',
-)
-
-_TRAIN_COLUMN = flags.DEFINE_string(
-    'train_column',
-    constants.DEFAULT_TRAIN_COLUMN,
-    '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.',
-)
-
-_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
-    'gradient_checkpointing',
-    False,
-    'Whether to enable gradient checkpointing.',
-)
-
-_ENABLE_PEFT = flags.DEFINE_boolean(
-    'enable_peft',
-    True,
-    'Whether to enable peft.',
-)
-_TRAIN_TEMPLATE = flags.DEFINE_string(
-    'train_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 = flags.DEFINE_string(
-    'train_split',
-    'train',
-    'The train split name.',
-)
-
-_PER_DEVICE_EVAL_BATCH_SIZE = flags.DEFINE_integer(
-    'per_device_eval_batch_size',
-    1,
-    'The per device batch size for model evaluation.',
-)
-
-
-_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_list(
-    'eval_metric_name',
-    ['loss'],
-    'A comma-separated list of metric names to aggregate during model'
-    ' evaluation. The supported metrics are: '
-    + ', '.join(constants.SUPPORTED_EVAL_METRICS),
-)
-
-_EVAL_DATASET = flags.DEFINE_string(
-    'eval_dataset',
-    None,
-    'The Hugging Face dataset name or path to use for evaluation.',
-)
-
-# 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.',
-)
-
-_EVAL_TEMPLATE = flags.DEFINE_string(
-    'eval_template',
-    None,
-    'Template for formatting language model evaluation data.'
-    ' 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.',
-)
-
-_METRIC_FOR_BEST_MODEL = flags.DEFINE_string(
-    'metric_for_best_model',
-    None,
-    'If set, the best model is saved at the end of training based on the'
-    ' metric',
-)
-
-_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.',
-)
-
-_EXAMPLE_PACKING = flags.DEFINE_boolean(
-    'example_packing',
-    False,
-    'Enables example packing during training, which uses '
-    '`ConstantLengthDataset` under the hood.',
-)
-
-_INPUT_MASKING = flags.DEFINE_boolean(
-    'input_masking',
-    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.'
-)
-
-_MAX_GPU_MEMORY_FRACTION = flags.DEFINE_float(
-    'max_gpu_memory_fraction',
-    '0.9',
-    'Maximum GPU memory a caching allocator is allowed to use per GPU.',
-)
-
-
-@flags.multi_flags_validator(
-    [
-        _INPUT_MASKING.name,
-        _EXAMPLE_PACKING.name,
-    ],
-    message='`example_packing=True` does not work with `input_masking=True`',
-)
-def check_example_packing(flags_dict: Mapping[str, Any]) -> bool:
-  """Check to make sure example packing is enabled properly.
-
-  Args:
-    flags_dict: Dictionary containing flags to check.
-
-  Returns:
-    If `example_packing` is set properly.
-  """
-  if flags_dict[_INPUT_MASKING.name] and flags_dict[_EXAMPLE_PACKING.name]:
-    return False
-  return True
-
-
-@flags.multi_flags_validator(
-    [
-        _INPUT_MASKING.name,
-        _TRAIN_TEMPLATE.name,
-    ],
-    message='`train_template` should be provided if using `input_masking=True`',
-)
-def check_input_masking(flags_dict: Mapping[str, Any]) -> bool:
-  """Check to make sure input_masking is enabled properly.
-
-  Args:
-    flags_dict: Dictionary containing flags to check.
-
-  Returns:
-    If `input_masking` is set properly
-  """
-  if (
-      flags_dict[_INPUT_MASKING.name]
-      and flags_dict[_TRAIN_TEMPLATE.name] is None
-  ):
-    return False
-  return True
-
-
-@flags.multi_flags_validator(
-    [
-        _EVAL_DATASET.name,
-        _EVAL_METRIC_NAME.name,
-    ],
-    message=(
-        '`eval_metric_name` should be a valid metric name and present when'
-        ' eval_dataset is provided.'
-    ),
-)
-def _validate_eval_metrics(flags_dict: Mapping[str, Any]) -> bool:
-  """Validates the eval metric name.
-
-  Args:
-    flags_dict: Dictionary containing flags to check.
-
-  Returns:
-    If the eval metrics are valid.
-  """
-  if flags_dict[_EVAL_DATASET.name] is None:
-    return True
-  eval_metrics = flags_dict[_EVAL_METRIC_NAME.name]
-  for eval_metric in eval_metrics:
-    if eval_metric not in constants.SUPPORTED_EVAL_METRICS:
-      raise flags.ValidationError(f'Invalid eval metric: {eval_metric}')
-  if 'perplexity' in eval_metrics and 'loss' not in eval_metrics:
-    _EVAL_METRIC_NAME.value.append('loss')
-    logging.warning(
-        'Adding `loss` to eval_metric_name because `perplexity` is present.'
-    )
-  return True
-
-
-@flags.multi_flags_validator(
-    [
-        _METRIC_FOR_BEST_MODEL.name,
-        _EVAL_METRIC_NAME.name,
-    ],
-    message='`metric_for_best_model` should be in `eval_metric_name`.',
-)
-def _validate_metric_for_best_model(flags_dict: Mapping[str, Any]) -> bool:
-  """Validates the metric for best model.
-
-  Args:
-    flags_dict: Dictionary containing flags to check.
-
-  Returns:
-    If the metric for best model is valid.
-  """
-  if flags_dict[_METRIC_FOR_BEST_MODEL.name] is None:
-    return True
-
-  metric_for_best_model = flags_dict[_METRIC_FOR_BEST_MODEL.name]
-  eval_metric_name = flags_dict[_EVAL_METRIC_NAME.name]
-
-  if metric_for_best_model not in eval_metric_name:
-    raise flags.ValidationError(
-        'Invalid metric for picking the best model:'
-        f' {metric_for_best_model}. The metric should be one'
-        f' of the {eval_metric_name}.'
-    )
-  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.
-
-
-def _calculate_hf_eval_metrics(
-    tokenizer: transformers.PreTrainedTokenizerBase,
-    eval_config: eval_lib.EvalConfig | None,
-) -> tuple[
-    Callable[[transformers.EvalPrediction], Mapping[str, float]], torch.Tensor
-]:
-  """Calculates the HF evaluation metrics.
-
-  Args:
-    tokenizer: The tokenizer to use for evaluation.
-    eval_config: The evaluation config to use.
-
-  Returns:
-    The compute metrics and preprocess logits for metrics.
-  """
-  if eval_config is None:
-    return None, None
-  hf_eval_metrics = {}
-  for metric in eval_config.metric_name:
-    if metric in constants.SUPPORTED_HF_EVAL_METRICS:
-      if metric in constants.ROUGE_VARIANTS:
-        hf_eval_metrics[metric] = evaluate.load('rouge')
-      else:
-        hf_eval_metrics[metric] = evaluate.load(metric)
-
-  if not hf_eval_metrics:
-    return None, None
-  return (
-      eval_lib.create_compute_metrics(tokenizer, hf_eval_metrics),
-      eval_lib.preprocess_logits_for_metrics,
-  )
-
-
-# Copied from https://github.com/artidoro/qlora/blob/main/qlora.py.
-def find_all_linear_names(
-    model: transformers.AutoModelForCausalLM, precision_mode: str
-) -> Sequence[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_name_or_path: str,
-    train_dataset: 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_train_epochs: float | None = None,
-    max_steps: int | None = None,
-    warmup_steps: int = 10,
-    max_seq_length: int = 512,
-    learning_rate: float = 2e-4,
-    precision_mode: str = None,
-    train_column: str = constants.DEFAULT_TRAIN_COLUMN,
-    per_device_train_batch_size: int = 4,
-    gradient_accumulation_steps: int = 4,
-    optim: str = 'paged_adamw_32bit',
-    weight_decay: float = 0.001,
-    gradient_checkpointing: bool = False,
-    enable_peft: bool = True,
-    train_template: str = None,
-    lr_scheduler_type: str = 'constant',
-    save_steps: int = 10,
-    logging_steps: int = 10,
-    train_split: str = 'train',
-    eval_config: eval_lib.EvalConfig | None = None,
-    report_to: str = constants.REPORT_TO_NONE,
-    access_token: str | None = None,
-    train_precision: str = constants.PRECISION_MODE_16B,
-    example_packing: bool = False,
-    attn_implementation: str | None = None,
-    max_grad_norm: float = 0.3,
-    input_masking: bool = False,
-    logger_level: str = 'passive',
-    benchmark_out_file: str | None = None,
-    tuning_data_stats_file: str | None = None,
-    target_modules: str | None = None,
-) -> None:
-  """Finetunes instruct."""
-  logging.info(
-      'on entering instruct_lora, %s,\n%s',
-      device_stats.gpu_stats_str(),
-      device_stats.cpu_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 = dataset_validation_util.load_tokenizer(
-      pretrained_model_name_or_path,
-      'right',
-      access_token=access_token,
-  )
-
-  train_dataset, train_dataset_with_template = (
-      dataset_validation_util.load_dataset_with_template(
-          train_dataset,
-          split=train_split,
-          input_column=train_column,
-          template=train_template,
-          tokenizer=tokenizer,
-      )
-  )
-  train_dataset_with_template = dataset_validation_util.get_filtered_dataset(
-      dataset=train_dataset_with_template,
-      input_column=train_column,
-      max_seq_length=max_seq_length,
-      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 = dataset_stats.get_dataset_stats(
-          raw=train_dataset,
-          templated=train_dataset_with_template,
-          template=train_template,
-          tokenizer=tokenizer,
-          column=train_column,
-          effective_batch_size=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(train_dataset_stats, out_f)
-
-  model = utils.load_model(
-      pretrained_model_name_or_path=pretrained_model_name_or_path,
-      tokenizer=tokenizer,
-      precision_mode=precision_mode,
-      gradient_checkpointing=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)
-    adapter_for_eval_dir = os.path.join(output_dir, 'adapter_for_eval')
-    logging.info('saving adapter for evaluation to %s...', adapter_for_eval_dir)
-    peft_config.save_pretrained(adapter_for_eval_dir)
-    # 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 = transformers.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_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_train_epochs if num_train_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=gradient_checkpointing,
-      gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
-      weight_decay=weight_decay,
-      log_level=logger_level,
-      accelerator_config=accelerator_config,
-      include_num_input_tokens_seen=True,
-  )
-  trainer_kwargs = {}
-  if input_masking and train_template:
-    template_json = dataset_validation_util.get_template(
-        template_path=train_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'] = trl.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
-  )
-  compute_metrics, preprocess_logits = _calculate_hf_eval_metrics(
-      tokenizer, eval_config
-  )
-
-  trainer = eval_lib.create_trainer(
-      cls=trl.SFTTrainer,
-      eval_config=eval_config,
-      model=model,
-      train_dataset=train_dataset_with_template,
-      dataset_text_field=train_column,
-      max_seq_length=max_seq_length,
-      tokenizer=tokenizer,
-      args=training_arguments,
-      packing=example_packing,
-      callbacks=[trainer_stats_callback],
-      compute_metrics=compute_metrics,
-      preprocess_logits_for_metrics=preprocess_logits,
-      **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
-    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
-    )
-  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 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)
-  )
-
-  torch.cuda.set_per_process_memory_fraction(
-      _MAX_GPU_MEMORY_FRACTION.value, device=PartialState().local_process_index
-  )
-
-  utils.print_library_versions()
-  warnings.simplefilter(_WARNINGS_FILTER.value)
-
-  pretrained_model_name_or_path = fileutils.force_gcs_path(
-      _PRETRAINED_MODEL_NAME_OR_PATH.value
-  )
-  if dataset_validation_util.is_gcs_path(pretrained_model_name_or_path):
-    pretrained_model_name_or_path = (
-        dataset_validation_util.download_gcs_uri_to_local(
-            pretrained_model_name_or_path
-        )
-    )
-
-  # 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_DATASET.value:
-    eval_config = eval_lib.EvalConfig(
-        per_device_batch_size=_PER_DEVICE_EVAL_BATCH_SIZE.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.value
-        ),
-        split=_EVAL_SPLIT.value,
-        template=_EVAL_TEMPLATE.value,
-        column=_EVAL_COLUMN.value,
-        tokenize_dataset=False,
-        metric_for_best_model=_METRIC_FOR_BEST_MODEL.value,
-    )
-  else:
-    eval_config = None
-
-  if _REPORT_TO.value == constants.REPORT_TO_WANDB:
-    wandb.login()
-
-  finetune_instruct(
-      pretrained_model_name_or_path=pretrained_model_name_or_path,
-      train_dataset=_TRAIN_DATASET.value,
-      output_dir=_OUTPUT_DIR.value,
-      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_train_epochs=_NUM_TRAIN_EPOCHS.value,
-      warmup_steps=_WARMUP_STEPS.value,
-      max_steps=_MAX_STEPS.value,
-      max_seq_length=_MAX_SEQ_LENGTH.value,
-      learning_rate=_LEARNING_RATE.value,
-      train_column=_TRAIN_COLUMN.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,
-      gradient_checkpointing=_GRADIENT_CHECKPOINTING.value,
-      enable_peft=_ENABLE_PEFT.value,
-      train_template=_TRAIN_TEMPLATE.value,
-      lr_scheduler_type=_LR_SCHEDULER_TYPE.value,
-      save_steps=_SAVE_STEPS.value,
-      logging_steps=_LOGGING_STEPS.value,
-      train_split=_TRAIN_SPLIT.value,
-      eval_config=eval_config,
-      report_to=_REPORT_TO.value,
-      access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
-      train_precision=_TRAIN_PRECISION.value,
-      example_packing=_EXAMPLE_PACKING.value,
-      attn_implementation=_ATTN_IMPLEMENTATION.value,
-      max_grad_norm=_MAX_GRAD_NORM.value,
-      input_masking=_INPUT_MASKING.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()
-
-
-if __name__ == '__main__':
-  app.run(main)

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

@@ -1,108 +0,0 @@
-"""Script to merge PEFT adapter with base model."""
-
-from collections.abc import Mapping, Sequence
-from typing import Any
-
-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_NAME_OR_PATH = flags.DEFINE_string(
-    'pretrained_model_name_or_path',
-    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_name_or_path 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_16B,
-    [
-        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.',
-)
-
-_HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
-    'huggingface_access_token',
-    None,
-    'The access token for loading huggingface gated models.',
-)
-
-
-@flags.multi_flags_validator(
-    [
-        _PRETRAINED_MODEL_NAME_OR_PATH.name,
-        _FINETUNED_LORA_MODEL_DIR.name,
-        _MERGE_BASE_AND_LORA_OUTPUT_DIR.name,
-    ],
-)
-def check_merge_lora_model_flags(flags_dict: Mapping[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_name_or_path = fileutils.force_gcs_path(
-      _PRETRAINED_MODEL_NAME_OR_PATH.value
-  )
-  if dataset_validation_util.is_gcs_path(pretrained_model_name_or_path):
-    pretrained_model_name_or_path = (
-        dataset_validation_util.download_gcs_uri_to_local(
-            pretrained_model_name_or_path
-        )
-    )
-
-  finetuned_lora_model_dir = fileutils.force_gcs_path(
-      _FINETUNED_LORA_MODEL_DIR.value
-  )
-  if dataset_validation_util.is_gcs_path(finetuned_lora_model_dir):
-    finetuned_lora_model_dir = (
-        dataset_validation_util.download_gcs_uri_to_local(
-            finetuned_lora_model_dir
-        )
-    )
-  utils.merge_causal_language_model_with_lora(
-      pretrained_model_name_or_path=pretrained_model_name_or_path,
-      precision_mode=_MERGE_MODEL_PRECISION_MODE.value,
-      finetuned_lora_model_dir=finetuned_lora_model_dir,
-      merged_model_output_dir=_MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
-      access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
-  )
-
-
-if __name__ == '__main__':
-  app.run(main)

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

@@ -1,21 +0,0 @@
-#!/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/templates/README.md

@@ -1,98 +0,0 @@
-# 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

@@ -1,7 +0,0 @@
-{
-  "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

@@ -1,7 +0,0 @@
-{
-  "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

@@ -1,7 +0,0 @@
-{
-  "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

@@ -1,7 +0,0 @@
-{
-  "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

@@ -1,7 +0,0 @@
-{
-  "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 = '\n\n<|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 %}{{ '\n\n<|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

@@ -1,7 +0,0 @@
-{
-  "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-chat.json

@@ -1,7 +0,0 @@
-{
-  "description": "Template used by openai chat.",
-  "source": "https://platform.openai.com/docs/api-reference/fine-tuning/chat-input",
-  "chat_template": "{% set loop_messages = messages %}{% set content = '' %}{% for message in loop_messages %}{% set content = content ~ '\n\n<|start_header_id|>' ~ message.role ~ '<|end_header_id|>\n\n' %}{% if message.content is string %}{% set content = content ~ message.content|trim ~ '<|eot_id|>' %}{% else %}{% set content = content ~ message.content|join(' ', attribute='text')|trim ~ '<|eot_id|>' %}{% endif %}{% 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/openai-completion.json

@@ -1,7 +0,0 @@
-{
-  "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

@@ -1,7 +0,0 @@
-{
-  "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/qwen2_5.json

@@ -1,7 +0,0 @@
-{
-  "description": "Chat template used by Qwen 2.5.",
-  "source": "https://huggingface.co/Qwen/Qwen2.5-72B-Instruct/blob/main/tokenizer_config.json#L198",
-  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
-  "instruction_separator": "<|im_start|>user\n",
-  "response_separator": "<|im_start|>assistant\n"
-}

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

@@ -1,7 +0,0 @@
-{
-  "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 = '\n\n<|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 %}{{ '\n\n<|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/tests/instruct_lora_command_builder.py

@@ -1,453 +0,0 @@
-# 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._gcs_rsync_interval_secs = None
-    self._pretrained_model_name_or_path = None
-    self._train_dataset = None
-    self._train_split = None
-    self._train_template = None
-    self._train_column = 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_train_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._gradient_checkpointing = None
-    self._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 = None
-    self._eval_column = None
-    self._eval_template = None
-    self._eval_split = None
-    self._eval_steps = None
-    self._eval_metric_name = None
-    self._metric_for_best_model = None
-    self._input_masking = 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
-    self._unnamed_args = 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 gcs_rsync_interval_secs(self):
-    return self._gcs_rsync_interval_secs
-
-  @gcs_rsync_interval_secs.setter
-  def gcs_rsync_interval_secs(self, val: str):
-    self._gcs_rsync_interval_secs = val
-
-  @property
-  def pretrained_model_name_or_path(self):
-    return self._pretrained_model_name_or_path
-
-  @pretrained_model_name_or_path.setter
-  def pretrained_model_name_or_path(self, val: str):
-    self._pretrained_model_name_or_path = val
-
-  @property
-  def train_dataset(self):
-    return self._train_dataset
-
-  @train_dataset.setter
-  def train_dataset(self, val: str):
-    self._train_dataset = val
-
-  @property
-  def train_split(self):
-    return self._train_split
-
-  @train_split.setter
-  def train_split(self, val: str):
-    self._train_split = val
-
-  @property
-  def train_template(self):
-    return self._train_template
-
-  @train_template.setter
-  def train_template(self, val: str):
-    self._train_template = val
-
-  @property
-  def train_column(self):
-    return self._train_column
-
-  @train_column.setter
-  def train_column(self, val: str):
-    self._train_column = 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_train_epochs(self):
-    return self._num_train_epochs
-
-  @num_train_epochs.setter
-  def num_train_epochs(self, val: float):
-    self._num_train_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._gradient_checkpointing
-
-  @gradient_checkpointing.setter
-  def gradient_checkpointing(self, val: bool):
-    self._gradient_checkpointing = val
-
-  @property
-  def example_packing(self):
-    return self._example_packing
-
-  @example_packing.setter
-  def example_packing(self, val: bool):
-    self._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
-
-  @eval_dataset.setter
-  def eval_dataset(self, val: str):
-    self._eval_dataset = val
-
-  @property
-  def eval_column(self):
-    return self._eval_column
-
-  @eval_column.setter
-  def eval_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(self):
-    return self._eval_split
-
-  @eval_split.setter
-  def eval_split(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_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 metric_for_best_model(self):
-    return self._metric_for_best_model
-
-  @metric_for_best_model.setter
-  def metric_for_best_model(self, val: str):
-    self._metric_for_best_model = val
-
-  @property
-  def input_masking(self):
-    return self._input_masking
-
-  @input_masking.setter
-  def input_masking(self, val: bool):
-    self._input_masking = 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
-
-  @property
-  def unnamed_args(self):
-    return self._unnamed_args
-
-  @unnamed_args.setter
-  def unnamed_args(self, val: list):
-    self._unnamed_args = val
-
-  def build_cmd(self) -> list[str]:
-    cmd = []
-    args = ''
-    for k, v in self.__dict__.items():
-      if k == '_unnamed_args' and v is not None:
-        args += ' '.join(v)
-        continue
-      if v is not None:
-        cmd.append(f'--{k[1:]}={v}')
-    cmd.append(f'{args}')
-    return cmd

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

@@ -1,142 +0,0 @@
-# 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_name_or_path = 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_name_or_path(self):
-    return self._pretrained_model_name_or_path
-
-  @pretrained_model_name_or_path.setter
-  def pretrained_model_name_or_path(self, val: str):
-    self._pretrained_model_name_or_path = 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/vmg/tests/run_tests.py

@@ -1,61 +0,0 @@
-"""Run the tests from docker command line."""
-
-import subprocess
-import sys
-from typing import Sequence
-
-from absl import app
-from absl import flags
-
-
-_ALLOWED_TEST_FILE_PATHS = (
-    "test_instruct_lora_adapters",
-    "test_instruct_lora_features",
-    "test_instruct_lora_throughput",
-    "test_instruct_lora_trained_model_quality",
-    "test_validate_dataset_with_template",
-)
-
-_TEST_FILE_PATH = flags.DEFINE_multi_enum(
-    "test_file_path",
-    None,
-    _ALLOWED_TEST_FILE_PATHS + ("all",),
-    "The test file path.",
-    required=True,
-)
-
-_IS_AUTOMATED_TEST = flags.DEFINE_bool(
-    "is_automated_test",
-    True,
-    "Whether the test is an automated test.",
-)
-
-
-def main(argv: Sequence[str]) -> None:
-  if len(argv) > 1:
-    raise app.UsageError("Too many command-line arguments.")
-  test_file_path = _TEST_FILE_PATH.value
-  if "all" in test_file_path:
-    test_file_path = _ALLOWED_TEST_FILE_PATHS
-
-  for test_file in test_file_path:
-    cmd = [
-        "python3",
-        f"vertex_vision_model_garden_peft/tests/{test_file}.py",
-    ]
-    if (
-        test_file == "test_instruct_lora_throughput"
-        and _IS_AUTOMATED_TEST.value
-    ):
-      subprocess.run(
-          cmd + ["--", "-k", "peft_train_image_automated_test"],
-          stdout=sys.stdout,
-          stderr=sys.stdout,
-          check=True,
-      )
-    else:
-      subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
-
-
-if __name__ == "__main__":
-  app.run(main)

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

@@ -1,106 +0,0 @@
-# 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.command_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 = 'train'
-    self.task_cmd_builder.train_column = 'input_text'
-    self.task_cmd_builder.train_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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/vmg/tests/test_instruct_lora_features.py

@@ -1,367 +0,0 @@
-# pylint: disable=missing-function-docstring
-# pylint: disable=missing-class-docstring
-"""Tests various features of PEFT train docker."""
-
-import os
-import time
-import unittest
-from absl.testing import absltest
-from absl.testing import parameterized
-import instruct_lora_command_builder as task_cmd_builder
-import test_util
-
-
-class EvalConfigTest(test_util.TestBase):
-
-  def setUp(self):
-    super().setUp()
-
-    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
-    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'
-    self.task_cmd_builder.train_split = 'train'
-    self.task_cmd_builder.train_column = 'text'
-    self.task_cmd_builder.train_template = 'openassistant-guanaco'
-    self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
-    self.task_cmd_builder.logging_dir = '/tmp/logs'
-    self.task_cmd_builder.eval_steps = 10
-    self.task_cmd_builder.eval_dataset = 'mlabonne/guanaco-llama2'
-    self.task_cmd_builder.eval_split = 'test'
-    self.task_cmd_builder.eval_column = 'text'
-    self.task_cmd_builder.eval_template = 'openassistant-guanaco'
-
-  @parameterized.named_parameters(
-      ('all_eval_metric', 'loss,perplexity,bleu,google_bleu,rouge1', 0),
-      ('invalid_metric', 'invalid_metric', 1),
-      ('only_loss', 'loss', 0),
-      ('perplexity_without_loss', 'perplexity,bleu', 0),
-      ('unsupported_eval_metric', 'f1', 1),
-  )
-  def test_hf_eval_metrics(self, eval_metric_name, expected_return_code):
-    self.task_cmd_builder.eval_metric_name = eval_metric_name
-    self.assertEqual(self.run_cmd(), expected_return_code)
-
-  @parameterized.named_parameters(
-      ('valid_best_model_metric', 'loss,perplexity', 'perplexity', 0),
-      ('only_loss', None, 'loss', 0),
-      ('invalid_best_model_metric', 'loss', 'invalid_metric', 1),
-  )
-  def test_metric_for_best_model(
-      self, eval_metric_name, metric_for_best_model, expected_return_code
-  ):
-    self.task_cmd_builder.eval_metric_name = eval_metric_name
-    self.task_cmd_builder.metric_for_best_model = metric_for_best_model
-    self.assertEqual(self.run_cmd(), expected_return_code)
-
-
-class GcsUploadDownloadTest(test_util.TestBase):
-
-  def setUp(self):
-    super().setUp()
-
-    self.command_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 = 'train'
-    self.task_cmd_builder.train_column = 'input_text'
-    self.task_cmd_builder.train_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_name_or_path):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(
-            pretrained_model_name_or_path
-        )
-    )
-    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_name_or_path):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(
-            pretrained_model_name_or_path
-        )
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
-
-    self.command_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_8b_model_download(self):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        'gs://vertex-model-garden-public-us/llama3/llama3-8b-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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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)
-
-  @unittest.skipIf(
-      not test_util.is_gpu_h100(),
-      'Skipping because this test is only for H100',
-  )
-  def test_model_fp8_conversion(self):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
-    )
-    self.task_cmd_builder.merged_model_dir = os.path.join(
-        merged_model_dir, f'merged-{test_util.get_timestamp()}'
-    )
-
-    self.command_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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.command_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.command_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'
-
-  @parameterized.named_parameters(
-      ('multi-chat-string-content', 'openai-multi-chat-example-data.jsonl'),
-      (
-          'multi-chat-array-content',
-          'openai-multi-chat-example-data-array-content.jsonl',
-      ),
-  )
-  def test_openai_chat_template(self, example_dataset):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        example_dataset
-    )
-    self.task_cmd_builder.train_split = 'train'
-    self.task_cmd_builder.train_column = 'messages'
-    self.task_cmd_builder.train_template = 'openai-chat'
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  def test_openai_completion_template(self):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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 = 'train'
-    self.task_cmd_builder.train_column = 'prompt'
-    self.task_cmd_builder.train_template = 'openai-completion'
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  def test_data_stats_chat_template(self):
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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 = 'train'
-    self.task_cmd_builder.train_column = 'messages'
-    self.task_cmd_builder.train_template = 'llama3'
-    self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
-    self.command_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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 = 'train'
-    self.task_cmd_builder.train_column = 'prompt'
-    self.task_cmd_builder.train_template = 'openai-completion'
-    self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
-
-    self.command_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.command_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 = 'train'
-    self.task_cmd_builder.train_column = 'input_text'
-    self.task_cmd_builder.train_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path('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/vmg/tests/test_instruct_lora_throughput.py

@@ -1,276 +0,0 @@
-# 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 = 'train'
-    self.task_cmd_builder.train_column = 'text'
-    self.task_cmd_builder.train_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} | failed | n/a\n')
-    return ret
-
-  @parameterized.product(
-      model_name=[
-          'llama3.1-8b-hf',
-          'llama3.1-70b-hf',
-          'Mistral-7B-v0.1',
-          'Mixtral-8x7B-v0.1',
-          'gemma-2-9b-it',
-          'Qwen2.5-32B-Instruct',
-      ],
-      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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
-
-    self.assertEqual(self.run_cmd_and_handle_failure(), 0)
-
-  @parameterized.product(
-      model_name=[
-          'llama3.1-8b-hf',
-          'llama3.1-70b-hf',
-          'Mistral-7B-v0.1',
-          'Mixtral-8x7B-v0.1',
-          'gemma-2-9b-it',
-          'Qwen2.5-32B-Instruct',
-      ],
-      precision=['4bit', '8bit', 'bfloat16'],
-      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
-      num_gpus=[8],
-      config=['deepspeed_zero2'],
-  )
-  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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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.command_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-8b-hf',
-          'llama3.1-70b-hf',
-          'Qwen2.5-32B-Instruct',
-      ],
-      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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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',
-    )
-    if 'llama' in model_name.lower():
-      self.task_cmd_builder.config_file = (
-          'vertex_vision_model_garden_peft/llama2_fsdp_8gpu.yaml'
-      )
-    elif 'qwen' in model_name.lower():
-      self.task_cmd_builder.config_file = (
-          'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
-      )
-    else:
-      self.fail(f'Unsupported model: {model_name}')
-
-    self.command_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-8b-hf', 'llama3.1-70b-hf'],
-      precision=['4bit', 'bfloat16'],
-      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
-      config=['deepspeed_zero2', 'fsdp'],
-  )
-  def test_peft_train_image_automated_test_llama(
-      self, model_name, precision, max_seq_length, config
-  ):
-    num_gpus = 8
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(model_name)
-    )
-    self.task_cmd_builder.max_seq_length = max_seq_length
-    self.task_cmd_builder.load_precision = precision
-    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.benchmark_out_file = benchmark_out_file
-    if config == 'fsdp':
-      self.task_cmd_builder.config_file = (
-          f'vertex_vision_model_garden_peft/llama_{config}_{num_gpus}gpu.yaml'
-      )
-    else:
-      self.task_cmd_builder.config_file = (
-          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
-      )
-
-    self.command_builder.add_env_var(
-        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
-    )
-    self.run_cmd_and_handle_failure()
-    if test_util.is_gpu_h100():
-      self.assertEqual(
-          test_util.check_benchmark_results(
-              benchmark_out_file, 'llama', 10.0, max_seq_length
-          ),
-          True,
-      )
-
-  @parameterized.product(
-      model_name=['gemma-2-2b-it', 'gemma-2-9b-it', 'gemma-2-27b-it'],
-      precision=['4bit', 'bfloat16'],
-      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
-      config=['deepspeed_zero2', 'deepspeed_zero3', 'fsdp'],
-  )
-  def test_peft_train_image_automated_test_gemma(
-      self, model_name, precision, max_seq_length, config
-  ):
-    num_gpus = 8
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(model_name)
-    )
-    self.task_cmd_builder.attn_implementation = 'eager'
-    self.task_cmd_builder.max_seq_length = max_seq_length
-    self.task_cmd_builder.load_precision = precision
-    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.benchmark_out_file = benchmark_out_file
-    if config == 'fsdp':
-      self.task_cmd_builder.config_file = (
-          f'vertex_vision_model_garden_peft/gemma2_{config}_{num_gpus}gpu.yaml'
-      )
-    else:
-      self.task_cmd_builder.config_file = (
-          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
-      )
-
-    self.command_builder.add_env_var(
-        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
-    )
-    self.run_cmd_and_handle_failure()
-    if test_util.is_gpu_h100():
-      self.assertEqual(
-          test_util.check_benchmark_results(
-              benchmark_out_file, 'gemma', 10.0, max_seq_length
-          ),
-          True,
-      )
-
-  @parameterized.product(
-      model_name=['llama3.1-8b-hf', '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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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.command_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/vmg/tests/test_instruct_lora_trained_model_quality.py

@@ -1,209 +0,0 @@
-# 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_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_train_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.input_masking = 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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(model_name)
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'peft_train_sample.jsonl'
-    )
-    self.task_cmd_builder.train_split = 'train'
-    self.task_cmd_builder.train_column = 'input_text'
-    self.task_cmd_builder.train_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 = 'train'
-    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
-
-    self.command_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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 = 'train'
-    self.task_cmd_builder.train_column = 'text'
-    self.task_cmd_builder.train_template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split = 'test'
-    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
-
-    self.command_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_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(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 = 'train'
-    self.task_cmd_builder.train_column = 'text'
-    self.task_cmd_builder.train_template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split = 'test'
-    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
-
-    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
-  )
-  def test_qwen_model_deepspeed(self, model_name):
-    self.setup_output_dir(f'test_deepspeed_{model_name}')
-    self.task_cmd_builder.pretrained_model_name_or_path = model_name
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'llama-tuning-test/opposite-examples-train.jsonl'
-    )
-    self.task_cmd_builder.train_split = 'train'
-    self.task_cmd_builder.train_column = 'messages'
-    self.task_cmd_builder.train_template = 'qwen2_5'
-    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
-        'llama-tuning-test/opposite-examples-eval.jsonl'
-    )
-    self.task_cmd_builder.eval_split = 'train'
-    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
-
-    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-    # Note(lavrai): The following parameters are needed for the opposite-word
-    # dataset to converge properly.
-    self.task_cmd_builder.gradient_accumulation_steps = 1
-    self.task_cmd_builder.num_train_epochs = 10.0
-    self.task_cmd_builder.logging_steps = 1
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
-  )
-  def test_qwen_model_fsdp(self, model_name):
-    self.setup_output_dir(f'test_fsdp_{model_name}')
-    self.task_cmd_builder.pretrained_model_name_or_path = (
-        test_util.get_pretrained_model_name_or_path(model_name)
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'llama-tuning-test/opposite-examples-train.jsonl'
-    )
-    self.task_cmd_builder.train_split = 'train'
-    self.task_cmd_builder.train_column = 'messages'
-    self.task_cmd_builder.train_template = 'qwen2_5'
-    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
-        'llama-tuning-test/opposite-examples-eval.jsonl'
-    )
-    self.task_cmd_builder.eval_split = 'train'
-    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
-
-    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-    # Note(lavrai): The following parameters are needed for the opposite-word
-    # dataset to converge properly.
-    self.task_cmd_builder.gradient_accumulation_steps = 1
-    self.task_cmd_builder.num_train_epochs = 10.0
-    self.task_cmd_builder.logging_steps = 1
-
-    self.assertEqual(self.run_cmd(), 0)
-
-
-if __name__ == '__main__':
-  absltest.main()

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

@@ -1,353 +0,0 @@
-"""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=0,
-      ),
-  )
-  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 = "train"
-    self.task_cmd_builder.train_column = "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 = "train"
-    self.task_cmd_builder.train_column = "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="small_max_seq_length",
-          max_seq_length=10,
-      ),
-      dict(
-          testcase_name="large_max_seq_length",
-          max_seq_length=1024,
-      ),
-  )
-  def test_validate_dataset_with_template_max_seq_length(
-      self,
-      max_seq_length,
-  ):
-    self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split = "train"
-    self.task_cmd_builder.train_column = "text"
-    self.task_cmd_builder.template = (
-        "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
-    )
-    self.task_cmd_builder.max_seq_length = max_seq_length
-    self.task_cmd_builder.validate_k_rows_of_dataset = None
-    self.task_cmd_builder.use_multiprocessing = True
-    result = self.run_cmd()
-    self.assertEqual(result, 0)
-
-  @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 = split
-    self.task_cmd_builder.train_column = 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 = "train"
-    self.task_cmd_builder.train_column = "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/vmg/tests/validate_dataset_with_template_command_builder.py

@@ -1,97 +0,0 @@
-# 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 = None
-    self._train_column = None
-    self._max_seq_length = 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(self):
-    return self._train_split
-
-  @train_split.setter
-  def train_split(self, val: str):
-    self._train_split = val
-
-  @property
-  def train_column(self):
-    return self._train_column
-
-  @train_column.setter
-  def train_column(self, val: str):
-    self._train_column = 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 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/train_entrypoint.py

@@ -1,292 +0,0 @@
-"""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
-from collections.abc import MutableSequence, Sequence
-import multiprocessing
-import subprocess
-import sys
-from absl import app
-from absl import flags
-from absl import logging
-from util import dataset_validation_util
-from util import cluster_spec
-from vertex_vision_model_garden_peft.train.vmg import utils
-from util import constants
-from util import gcs_syncer
-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.VALIDATE_DATASET_WITH_TEMPLATE: (
-        'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py'
-    ),
-    constants.RUN_TESTS: 'vertex_vision_model_garden_peft/tests/run_tests.py',
-}
-
-
-def launch_script_cmd(
-    script: str,
-    config_file: str | None,
-    accelerate_args: argparse.Namespace = argparse.Namespace(),
-) -> MutableSequence[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."""
-  primary_node_addr, primary_node_port, node_rank, num_nodes = (
-      cluster_spec.get_cluster_spec()
-  )
-  accelerate_args = argparse.Namespace()
-  if num_nodes > 1:
-    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: MutableSequence[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_and_maybe_merge_cmd_and_dirs_to_sync(
-    task_type: str, config_file: str, unknown: Sequence[str]
-) -> Sequence[Sequence[str]]:
-  """Returns the training and merge command(if applicable) and dirs to sync.
-
-  Args:
-    task_type: The task type.
-    config_file: The accelerate config file path.
-    unknown: The unknown args which are not recognised by the parser.
-
-  Returns:
-    The bash commands to execute and the directories to sync.
-  """
-  dirs_to_sync = []
-  # Only populated when multi-node is used.
-  accelerate_args = _get_accelerate_args()
-  node_rank = getattr(accelerate_args, 'machine_rank', 0)
-  training_cmd = launch_script_cmd(
-      _TASK_TO_SCRIPT[task_type],
-      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)
-  )
-
-  local_output_dir, gcs_output_dir = gcs_syncer.manage_sync_path(
-      training_args.output_dir, node_rank
-  )
-  training_args.output_dir = local_output_dir
-  if gcs_syncer.is_gcs_or_gcsfuse_path(gcs_output_dir):
-    dirs_to_sync.append((local_output_dir, gcs_output_dir))
-
-  # Merge only flags.
-  merge_parser = argparse.ArgumentParser()
-  merge_parser.add_argument('--merge_model_precision_mode')
-  merge_parser.add_argument('--merge_base_and_lora_output_dir')
-  merge_args, unknown = merge_parser.parse_known_args(unknown)
-
-  if merge_args.merge_base_and_lora_output_dir:
-    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
-        merge_args.merge_base_and_lora_output_dir, None
-    )
-    merge_args.merge_base_and_lora_output_dir = merge_local_dir
-    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
-      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
-
-  # Common flags shared by merging and training.
-  common_parser = argparse.ArgumentParser()
-  common_parser.add_argument('--pretrained_model_name_or_path', 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 node_rank == 0:
-    lora_dir = utils.get_final_checkpoint_path(training_args.output_dir)
-    lora_local_dir, lora_gcs_dir = gcs_syncer.manage_sync_path(
-        lora_dir, node_rank
-    )
-    if gcs_syncer.is_gcs_or_gcsfuse_path(lora_gcs_dir):
-      dirs_to_sync.append((lora_local_dir, lora_gcs_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_local_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, dirs_to_sync
-
-
-def _get_merge_cmd_and_dirs_to_sync(
-    task_type: str, config_file: str, unknown: Sequence[str]
-) -> Sequence[Sequence[str]]:
-  """Returns the merge command and dirs to sync.
-
-  Args:
-    task_type: The task type.
-    config_file: The accelerate config file path.
-    unknown: The unknown args which are not recognised by the parser.
-
-  Returns:
-    The bash commands to execute and the directories to sync.
-  """
-  # Merge only flags.
-  merge_parser = argparse.ArgumentParser()
-  merge_parser.add_argument('--merge_base_and_lora_output_dir')
-  merge_args, unknown = merge_parser.parse_known_args(unknown)
-
-  dirs_to_sync = []
-  if merge_args.merge_base_and_lora_output_dir:
-    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
-        merge_args.merge_base_and_lora_output_dir, None
-    )
-    merge_args.merge_base_and_lora_output_dir = merge_local_dir
-    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
-      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
-
-  cmd = launch_script_cmd(_TASK_TO_SCRIPT[task_type], config_file)
-  _append_args_to_command_in_place(merge_args, cmd)
-  cmd.extend(unknown)
-  return [cmd], dirs_to_sync
-
-
-def main(unused_argv: Sequence[str]) -> None:
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--config_file')
-  parser.add_argument('--task')
-  parser.add_argument('--gcs_rsync_interval_secs', type=int, default=60)
-  args, unknown = parser.parse_known_args()
-
-  task = args.task
-  dirs_to_sync = None
-
-  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, dirs_to_sync = _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
-        task_type=task, config_file=args.config_file, unknown=unknown
-    )
-  elif task in [constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA]:
-    commands, dirs_to_sync = _get_merge_cmd_and_dirs_to_sync(
-        task_type=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]
-
-  rsync_process = None
-  mp_queue = multiprocessing.Queue(maxsize=1)
-  if dirs_to_sync:
-    rsync_process = gcs_syncer.setup_gcs_rsync(
-        dirs_to_sync, mp_queue, args.gcs_rsync_interval_secs
-    )
-
-  for cmd in commands:
-    logging.info('launching task=%s with cmd: \n%s', task, ' \\\n'.join(cmd))
-    # Both absl logging and python's logging module writes to stderr by default.
-    # Redirect output to stdout on purpose, such that log entries do not get
-    # marked as `Error` in Cloud's Log Explorer.
-    try:
-      subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
-    except subprocess.CalledProcessError as e:
-      if rsync_process is not None and rsync_process.is_alive():
-        logging.info('Terminating GCS rsync process.')
-        rsync_process.terminate()
-      raise e
-  if rsync_process is not None:
-    gcs_syncer.cleanup_gcs_rsync(rsync_process, mp_queue)
-
-
-if __name__ == '__main__':
-  logging.get_absl_handler().python_handler.stream = sys.stdout
-  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

@@ -1,360 +0,0 @@
-"""Common libraries for PEFT."""
-
-from collections.abc import Mapping, Sequence
-import datetime
-import gc
-import os
-from typing import Any
-
-from absl import logging
-import accelerate
-from accelerate import DistributedType
-from accelerate import PartialState
-import peft
-from peft import PeftModel
-from peft import prepare_model_for_kbit_training
-import torch
-import transformers
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
-from transformers import BitsAndBytesConfig
-from transformers import FbgemmFp8Config
-import trl
-
-from util import dataset_validation_util
-from util import constants
-
-_LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
-_LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
-_GEMMA2_MODEL = "gemma-2"
-
-
-def load_model(
-    pretrained_model_name_or_path: str,
-    tokenizer: AutoTokenizer,
-    precision_mode: str = None,
-    gradient_checkpointing: bool = False,
-    gradient_checkpointing_kwargs: Mapping[str, Any] | None = None,
-    access_token: str | None = None,
-    attn_implementation: str | None = None,
-    train_precision: str | None = None,
-    device_map: str | None = 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_kwargs = {
-      "use_cache": not gradient_checkpointing,
-      "device_map": device_map,
-      "torch_dtype": torch_dtype,
-      "quantization_config": quantization_config,
-      "trust_remote_code": False,
-      "token": access_token,
-      "attn_implementation": attn_implementation,
-  }
-  if _GEMMA2_MODEL in pretrained_model_name_or_path:
-    # The cache_implementation for Gemma 2 is set to hybrid by default. This
-    # param is only supported by Gemma 2. The default 'hybrid' value causes an
-    # issue when use_cache is set to False. So we have to use 'None' in such
-    # cases.
-    # https://github.com/huggingface/transformers/commit/238b13478df209ab534f2195a397dc64a3930883
-    model_kwargs["cache_implementation"] = (
-        None if gradient_checkpointing else "hybrid"
-    )
-
-  model = AutoModelForCausalLM.from_pretrained(
-      pretrained_model_name_or_path, **model_kwargs
-  )
-
-  if precision_mode in (constants.PRECISION_MODE_4, constants.PRECISION_MODE_8):
-    model = prepare_model_for_kbit_training(
-        model,
-        use_gradient_checkpointing=gradient_checkpointing,
-        gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
-    )
-
-  if 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 dataset_validation_util.should_add_pad_token(
-      pretrained_model_name_or_path
-  ):
-    model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
-    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_name_or_path: str,
-    merge_precision_mode: str,
-    finetuned_lora_model_dir: str,
-    merged_model_output_dir: str,
-    access_token: str | None = None,
-) -> None:
-  """Internal function to merges the base model with the lora adapter."""
-  logging.info("loading tokenizer...")
-  tokenizer = dataset_validation_util.load_tokenizer(
-      pretrained_model_name_or_path
-  )
-
-  # 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_name_or_path)
-  device_map = "cpu"
-  model = load_model(
-      pretrained_model_name_or_path=pretrained_model_name_or_path,
-      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(
-    pretrained_model_name_or_path: str,
-    precision_mode: str,
-    finetuned_lora_model_dir: str,
-    merged_model_output_dir: str,
-    access_token: str | None = 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...")
-    _merge_causal_language_model_with_lora_internal(
-        pretrained_model_name_or_path=pretrained_model_name_or_path,
-        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_name_or_path,
-        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: str | None = 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()
-
-
-def force_gc():
-  """Collects garbage immediately to release unused CPU/GPU resources."""
-  gc.collect()
-  torch.cuda.empty_cache()
-
-
-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,
-) -> Sequence[str] | None:
-  """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

@@ -1,84 +0,0 @@
-"""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 = flags.DEFINE_string(
-    'train_split',
-    'train',
-    'The train split name.',
-)
-
-_TRAIN_COLUMN = flags.DEFINE_string(
-    'train_column',
-    constants.DEFAULT_TRAIN_COLUMN,
-    '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,
-)
-
-_MAX_SEQ_LENGTH = flags.DEFINE_integer(
-    'max_seq_length',
-    None,
-    'The maximum sequence length.',
-)
-
-_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.value,
-      input_column=_TRAIN_COLUMN.value,
-      template=_TEMPLATE.value,
-      max_seq_length=_MAX_SEQ_LENGTH.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, Tuple
+from typing import Any, Dict, List
 
 import numpy as np
 from PIL import Image
@@ -10,28 +10,6 @@ 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.
 
@@ -62,14 +40,14 @@ def get_label_map(label_map_yaml_filepath: str) -> Dict[str, Any]:
 
 
 def get_object_detection_endpoint_predictions(
-    detector_endpoint: ...,
+    detection_endpoint: ...,
     input_image: np.ndarray,
     detection_thresh: float = 0.2,
 ) -> np.ndarray:
   """Gets endpoint predictions.
 
   Args:
-      detector_endpoint: image object detection endpoint.
+      detection_endpoint: image object detection endpoint.
       input_image: Input image.
       detection_thresh: Detection threshold.
 
@@ -77,10 +55,9 @@ def get_object_detection_endpoint_predictions(
       Object detection predictions from endpoints.
   """
   height, width, _ = input_image.shape
-  predictions = detector_endpoint.predict(
+  predictions = detection_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(
@@ -89,29 +66,6 @@ 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
   ]
@@ -122,5 +76,4 @@ def merge_boxes_and_classes(
   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,12 +36,6 @@ 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_'
@@ -51,7 +45,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'
@@ -66,91 +60,26 @@ 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'
-MERGE_CAUSAL_LANGUAGE_MODEL_LORA = 'merge-causal-language-model-lora'
+CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora'
 INSTRUCT_LORA = 'instruct-lora'
-VALIDATE_DATASET_WITH_TEMPLATE = 'validate-dataset-with-template'
-RUN_TESTS = 'test'
-DEFAULT_TEXT_COLUMN_IN_DATASET = 'quote'
-DEFAULT_TRAIN_COLUMN = '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'
-
-ROUGE_VARIANTS = ('rouge1', 'rouge2', 'rougeL', 'rougeLsum')
-
-# Supported HF evaluation metrics.
-SUPPORTED_HF_EVAL_METRICS = (
-    'perplexity',
-    'bleu',
-    'google_bleu',
-) + ROUGE_VARIANTS
-
-# Supported evaluation metrics.
-SUPPORTED_EVAL_METRICS = ('loss',) + SUPPORTED_HF_EVAL_METRICS
-
-# 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,13 +1,10 @@
 """Fileutil lib to copy files between gcs and local."""
 
-import filecmp
-import fnmatch
+import glob
 import os
 import pathlib
 import shutil
-import subprocess
-import time
-from typing import List, Optional, Tuple
+from typing import Tuple
 import uuid
 
 from absl import logging
@@ -16,17 +13,6 @@ 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.
 
@@ -50,123 +36,6 @@ 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 is_file_available(
-    file_path: str, retry_interval_secs: int = 60, timeout_secs: int = 3600
-) -> bool:
-  """Checks and waits for a file to be available in GCS.
-
-  Args:
-    file_path: The file path to check.
-    retry_interval_secs: The interval in seconds to check the file.
-    timeout_secs: The timeout in seconds to wait for the file.
-
-  Returns:
-    True if the file is available, False otherwise.
-  """
-  start_time = time.time()
-  while True:
-    try:
-      file_check_cmd = ['gcloud', 'storage', 'ls', file_path]
-      result = subprocess.run(
-          file_check_cmd, capture_output=True, text=True, check=True
-      )
-      if file_path in result.stdout:
-        logging.info('File %s exists.', file_path)
-        return True
-    except subprocess.CalledProcessError as e:
-      elapsed_time = time.time() - start_time
-      if elapsed_time > timeout_secs:
-        logging.info(
-            "Timeout: File '%s' not found after %d seconds. Error: %s",
-            file_path,
-            elapsed_time,
-            e,
-        )
-        return False
-
-      logging.info(
-          "File '%s' not found yet. Checking again in %d seconds. Error: %s",
-          file_path,
-          retry_interval_secs,
-          e,
-      )
-      time.sleep(retry_interval_secs)
-
-
-def compare_dirs(
-    local_dir: str,
-    gcsfuse_dir: str,
-    retry_interval_secs: int = 30,
-    timeout_secs: int = 3600,
-) -> bool:
-  """Compares two directories and returns True if they are the same.
-
-  Args:
-    local_dir: The local directory.
-    gcsfuse_dir: The gcsfuse directory.
-    retry_interval_secs: The interval in seconds to check the directories.
-    timeout_secs: The timeout in seconds to wait for the directories.
-
-  Returns:
-    True if the directories are the same, False otherwise.
-  """
-  start_time = time.time()
-  while True:
-    if os.path.exists(local_dir) and os.path.exists(gcsfuse_dir):
-      comparison = filecmp.dircmp(local_dir, gcsfuse_dir)
-      if (
-          not comparison.left_only
-          and not comparison.right_only
-          and not comparison.diff_files
-      ):
-        return True
-    elapsed_time = time.time() - start_time
-    if elapsed_time > timeout_secs:
-      logging.info(
-          "Timeout: Directories '%s' and '%s' do not match after %d seconds.",
-          local_dir,
-          gcsfuse_dir,
-          elapsed_time,
-      )
-      return False
-
-    logging.info(
-        "Directories '%s' and '%s' do not match yet. Checking again in %d"
-        ' seconds.',
-        local_dir,
-        gcsfuse_dir,
-        retry_interval_secs,
-    )
-    time.sleep(retry_interval_secs)
-
-
-def download_gcs_file_to_memory(gcs_uri: str) -> bytes:
-  """Downloads a gcs file to in memory.
-
-  Args:
-    gcs_uri: A string of GCS uri.
-
-  Returns:
-    The content of the gcs file in byte format.
-  """
-  bucket = gcs_uri.split('/')[2]
-  file_path = gcs_uri[len(constants.GCS_URI_PREFIX + bucket + '/') :]
-  client = _get_gcs_client()
-  bucket = client.bucket(bucket)
-  blob = bucket.blob(file_path)
-  return blob.download_as_bytes()
-
-
 def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
   """Download a gcs file to a local dir.
 
@@ -193,47 +62,15 @@ 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 = _get_gcs_client()
+  client = storage.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,
-    allow_patterns: Optional[List[str]] = None,
-    log: bool = True,
-) -> None:
+    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
+):
   """Downloads files in a GCS directory to a local directory.
 
   For example:
@@ -241,21 +78,16 @@ def download_gcs_dir_to_local(
     gs://bucket/foo/a -> /tmp/bar/a
     gs://bucket/foo/b/c -> /tmp/bar/b/c
 
-  Args:
+  Arguments:
     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 = _get_gcs_client()
+  prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
+  client = storage.Client()
   blobs = client.list_blobs(bucket_name, prefix=prefix)
   for blob in blobs:
     if blob.name[-1] == '/':
@@ -263,63 +95,43 @@ 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
     ):
-      if log:
-        logging.info('Skip downloading model bin %s', file_path)
+      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:
-      if log:
-        logging.info('Downloading %s to %s', file_path, local_file_path)
+      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)
-    /tmp/bar/a -> gs://bucket/foo/a
-    /tmp/bar/b/c -> gs://bucket/foo/b/c
+    gs://bucket/foo/a -> /tmp/bar/a
+    gs://bucket/foo/b/c -> /tmp/bar/b/c
 
   Arguments:
     local_dir: A string of local directory path.
     gcs_dir: A string of directory path on GCS.
   """
-  # 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)
+  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)
 
 
 def upload_file_to_gcs_path(
@@ -343,7 +155,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 = _get_gcs_client()
+  storage_client = storage.Client()
   source_file_path = source_path
   destination_file_uri = destination_uri
   logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
@@ -362,9 +174,7 @@ def is_gcs_path(input_path: str) -> bool:
   Returns:
       True if the input path is a GCS path, False otherwise.
   """
-  return input_path is not None and input_path.startswith(
-      constants.GCS_URI_PREFIX
-  )
+  return input_path.startswith(constants.GCS_URI_PREFIX)
 
 
 def release_text_assets(
@@ -422,10 +232,13 @@ 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]
-  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])
+  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)
   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
@@ -441,19 +254,10 @@ 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]
-  out_local_video_file_name = video_output_file_path.replace(
-      file_extension, '_overlay' + file_extension
-  )
+  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'
   return out_local_video_file_name
-
-
-def delete_local_file(local_file_path: str) -> None:
-  """Deletes a local file."""
-  if os.path.exists(local_file_path):
-    os.remove(local_file_path)
-
-
-def delete_local_dir(local_dir: str) -> None:
-  """Deletes a local directory recursively."""
-  if os.path.exists(local_dir):
-    shutil.rmtree(local_dir)

community-content/vertex_model_garden/model_oss/util/gcs_download_launcher.sh

@@ -1,119 +0,0 @@
-#!/bin/bash
-#
-# This launcher downloads model files from GCS to local model directory before
-# launching the actual command.
-#
-# If GCS URI is passed as an environment variable, set GCS_URI_ENV_KEY to the
-# environment variable name.
-# If GCS URI is passed as an argument, set GCS_URI_ARG_KEY to the argument name.
-# The argument must be in the format of '--$GCS_URI_ARG_KEY=gs://*'. Do not
-# separate argument name and value with spaces.
-# This script will also try reading from AIP_STORAGE_URI or AIP_STORAGE_DIR.
-# Note that AIP_STORAGE_DIR is expected to be a local path, so it bypasses the
-# download process.
-#
-# Input priority: AIP_STORAGE_DIR > AIP_STORAGE_URI > GCS_URI_ENV_KEY > GCS_URI_ARG_KEY.
-# Will output the local model directory to GCS_URI_ENV_KEY and GCS_URI_ARG_KEY
-# if they are set. Both will be updated if both set.
-#
-# Requires google-cloud-sdk as a dependency (for gcloud storage CLI).
-
-set -e
-
-readonly LOCAL_MODEL_DIR=${LOCAL_MODEL_DIR:-"/tmp/model_dir"}
-readonly LOCAL_ARGS_FILE=${LOCAL_ARGS_FILE:-"/tmp/args.txt"}
-
-update_model_id() {
-  if [[ ! -z "$GCS_URI_ENV_KEY" ]]; then
-    echo "Updating env var $GCS_URI_ENV_KEY to $AIP_STORAGE_DIR."
-    export "$GCS_URI_ENV_KEY"="$AIP_STORAGE_DIR"
-  fi
-
-  if [[ ! -z "$GCS_URI_ARG_KEY" ]]; then
-    echo "Updating args $GCS_URI_ARG_KEY to $AIP_STORAGE_DIR."
-    updated=0
-    for (( i=1; i <= $#; i++)); do
-      arg="${!i}"
-      if [[ "$arg" == "--$GCS_URI_ARG_KEY="* ]]; then
-        echo "Found $arg, updating to $AIP_STORAGE_DIR."
-        set -- "${@:1:(($i-1))}" "--$GCS_URI_ARG_KEY=$AIP_STORAGE_DIR" "${@:$(($i+1))}";
-        updated=1
-        break
-      fi
-    done
-    if [[ $updated -eq 0 ]]; then
-      echo "Appending args $GCS_URI_ARG_KEY to $AIP_STORAGE_DIR."
-      set -- "$@" "--$GCS_URI_ARG_KEY=$AIP_STORAGE_DIR";
-    fi
-  fi
-  echo "$*" > "$LOCAL_ARGS_FILE"
-}
-
-maybe_download_model() {
-  if [[ -z "$GCS_URI_ENV_KEY" ]] && [[ -z "$GCS_URI_ARG_KEY" ]]; then
-    echo "Internal error: Required GCS_URI_ENV_KEY or GCS_URI_ARG_KEY."
-    exit 1
-  fi
-
-  echo "$*" > "$LOCAL_ARGS_FILE"
-  gcs_uri=""
-  if [[ ! -z "$AIP_STORAGE_DIR" ]]; then
-    # AIP_STORAGE_DIR is expected to be a local path.
-    echo "AIP_STORAGE_DIR set, proceeding to run the launcher."
-    update_model_id "$@"
-    return
-  elif [[ $AIP_STORAGE_URI == gs://* ]]; then
-    # Check AIP_STORAGE_URI environment variable.
-    echo "AIP_STORAGE_URI set and starts with 'gs://', proceeding to download from GCS."
-    gcs_uri="$AIP_STORAGE_URI"
-  elif [[ ! -z "$GCS_URI_ENV_KEY" ]] && [[ ${!GCS_URI_ENV_KEY} == gs://* ]]; then
-    # Check custom environment variable.
-    echo "Custom environment variable ${GCS_URI_ENV_KEY} set and starts with 'gs://', proceeding to download from GCS."
-    gcs_uri="${!GCS_URI_ENV_KEY}"
-  elif [[ ! -z "$GCS_URI_ARG_KEY" ]]; then
-    # Check custom args.
-    for arg in "$@"; do
-      if [[ "$arg" == "--$GCS_URI_ARG_KEY=gs://"* ]]; then
-        gcs_uri="${arg#*=}"
-        echo "Custom args ${GCS_URI_ARG_KEY} set and starts with 'gs://', proceeding to download from GCS."
-        break
-      elif [[ "$arg" == "--$GCS_URI_ARG_KEY" ]]; then
-        echo "Found $GCS_URI_ARG_KEY, but it's not in the format of '--$GCS_URI_ARG_KEY=gs://*'."
-        echo "Ensure the value of $GCS_URI_ARG_KEY is within the same arg, separated by '='."
-        exit 1
-      fi
-    done
-  fi
-
-  if [[ -z "$gcs_uri" ]]; then
-    echo "No GCS URI found, proceeding to run the launcher."
-    return
-  fi
-
-  # Remove trailing '/' if any.
-  gcs_uri="${gcs_uri%%/}"
-  export AIP_STORAGE_DIR="$LOCAL_MODEL_DIR/${gcs_uri##gs://}"
-
-  # Create the target directory.
-  mkdir -p "$AIP_STORAGE_DIR"
-  echo "Downloading model from ${gcs_uri} to ${AIP_STORAGE_DIR}."
-
-  # Use gcloud storage CLI to copy the content from GCS to the target directory.
-  if gcloud storage cp -r "$gcs_uri/*" "$AIP_STORAGE_DIR"; then
-    echo "Model downloaded successfully to ${AIP_STORAGE_DIR}."
-    update_model_id "$@"
-  else
-    echo "Failed to download model from GCS."
-    exit 1
-  fi
-}
-
-run_local_command() {
-  command=$(cat "$LOCAL_ARGS_FILE")
-  rm -f "$LOCAL_ARGS_FILE"
-  echo "Launch command: $command"
-  eval "$command"
-}
-
-maybe_download_model "$@"
-run_local_command

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

@@ -1,176 +0,0 @@
-"""Sync local directory to GCS directory using rsync."""
-
-import multiprocessing
-import os
-import subprocess
-import time
-from typing import Optional, Sequence, Tuple
-
-from absl import logging
-
-from util import constants
-from util import fileutils
-
-_GCS_COMMAND_RETRIES = 3
-_RSYNC_RETRY_INTERVAL_SECS = 30
-
-
-def is_gcs_or_gcsfuse_path(path: str) -> bool:
-  """Returns if the path is a GCS or gcsfuse path.
-
-  Args:
-    path: The path to check.
-
-  Returns:
-    True if the path is a GCS or gcsfuse path.
-  """
-  return path.startswith(
-      (constants.GCS_URI_PREFIX, constants.GCSFUSE_URI_PREFIX)
-  )
-
-
-def manage_sync_path(
-    path: str, node_rank: Optional[int] = None
-) -> Tuple[str, str]:
-  """Returns local dir and GCS location for the given path if the given path is a GCS or gcsfuse path.
-
-  It will also create a local directory if it does not exist. Otherwise, it
-  returns the same path.
-
-  Args:
-    path: The local or GCS path to manage.
-    node_rank: The node rank to be appended to the GCS path.
-
-  Returns:
-    The local and GCS paths.
-  """
-  local_dir = path
-  gcs_dir = path
-  if is_gcs_or_gcsfuse_path(path):
-    local_dir = os.path.join(
-        constants.LOCAL_OUTPUT_DIR,
-        fileutils.force_gcs_fuse_path(path)[1:],
-    )
-    gcs_dir = fileutils.force_gcs_path(path)
-  if not os.path.exists(local_dir):
-    os.makedirs(local_dir, exist_ok=True)
-
-  if node_rank is None:
-    return local_dir, gcs_dir
-  return local_dir, os.path.join(gcs_dir, f"node-{node_rank}")
-
-
-def setup_gcs_rsync(
-    dirs_to_sync: Sequence[Tuple[str, str]],
-    mp_queue: multiprocessing.Queue,
-    gcs_rsync_interval_secs: int,
-) -> multiprocessing.Process:
-  """Sets up the GCS rsync process.
-
-  Args:
-    dirs_to_sync: The absolute directory paths which will be synced to GCS.
-    mp_queue: The multiprocessing queue to check if the training is finished.
-    gcs_rsync_interval_secs: Integer, interval in seconds to run gcs rsync.
-
-  Returns:
-    The GCS rsync process.
-  """
-  rsync_process = multiprocessing.Process(
-      target=start_gcs_rsync,
-      args=(dirs_to_sync, mp_queue, gcs_rsync_interval_secs),
-  )
-  rsync_process.start()
-  return rsync_process
-
-
-def cleanup_gcs_rsync(
-    rsync_process: multiprocessing.Process, mp_queue: multiprocessing.Queue
-) -> None:
-  """Cleans up the GCS rsync process.
-
-  Args:
-    rsync_process: The GCS rsync process.
-    mp_queue: The multiprocessing queue.
-  """
-  mp_queue.put("finish rsync process")
-  rsync_process.join()
-  if rsync_process.exitcode == 0:
-    logging.info("Artifacts have been uploaded to GCS.")
-  else:
-    logging.error(
-        "GCS rsync process failed with exit code %d.", rsync_process.exitcode
-    )
-
-
-def _rsync_local_to_gcs(local_dir: str, gcs_dir: str) -> None:
-  """Syncs the local directory to GCS.
-
-  Args:
-    local_dir: The local directory to sync.
-    gcs_dir: The GCS directory to sync to.
-  """
-  if not os.listdir(local_dir):
-    logging.info("Not rsyncing to GCS since %s is empty.", local_dir)
-    return
-
-  logging.info("Rsyncing %s <--> %s...", local_dir, gcs_dir)
-  cmd = [
-      "gcloud",
-      "storage",
-      "rsync",
-      "-r",
-      "--delete-unmatched-destination-objects",
-  ]
-  cmd.extend([local_dir, gcs_dir])
-
-  attempt = 0
-  while attempt < _GCS_COMMAND_RETRIES:
-    try:
-      subprocess.check_output(cmd)
-      break
-    except subprocess.CalledProcessError as e:
-      attempt += 1
-      if attempt < _GCS_COMMAND_RETRIES:
-        logging.exception(
-            "Attempt %d: Command failed: %s. Retrying in %d seconds...",
-            attempt,
-            e,
-            _RSYNC_RETRY_INTERVAL_SECS,
-        )
-        time.sleep(_RSYNC_RETRY_INTERVAL_SECS)
-      else:
-        logging.exception(
-            "Command failed after %d attempts: %s.", e, _GCS_COMMAND_RETRIES
-        )
-
-  logging.info("%s rsynced to %s.", local_dir, gcs_dir)
-
-
-def start_gcs_rsync(
-    dirs_to_sync: Sequence[Tuple[str, str]],
-    mp_queue: multiprocessing.Queue,
-    gcs_rsync_interval_secs: int,
-) -> None:
-  """Starts a rsync process to sync local directories to GCS directories.
-
-  Args:
-    dirs_to_sync: A list of tuples, where each tuple contains local directory
-      which will be synced to GCS. For example: [('/tmp/local_dir_1',
-      'gs://bucket/gcs_dir_1'), ('/tmp/local_dir_2', 'gs://bucket/gcs_dir_2')]
-    mp_queue: The multiprocessing queue to check if the training is finished.
-    gcs_rsync_interval_secs: Integer, interval in seconds to run gcs rsync.
-  """
-  while True:
-    for local_dir, gcs_dir in dirs_to_sync:
-      _rsync_local_to_gcs(local_dir, gcs_dir)
-    if not mp_queue.empty():
-      break
-    time.sleep(gcs_rsync_interval_secs)
-
-  # Sync up the directory one more time to avoid a race condition.
-  # There can be a case when we are doing an rsync and receive a signal that
-  # the training has been done. The final checkpoint will be skipped in such
-  # case. So we do a final sync to make sure that the all directories
-  # are synced.
-  for local_dir, gcs_dir in dirs_to_sync:
-    _rsync_local_to_gcs(local_dir, gcs_dir)

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

@@ -20,11 +20,3 @@ 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_model_garden/model_oss/util/pytorch_startup_prober.sh

@@ -1,35 +0,0 @@
-#!/bin/bash
-
-# !/bin/bash
-# The Startup prober built to check whether models listed in local disk are
-# loaded in memory and are ready to serve traffic. The script returns 0 if
-# succeed. Any other returned value are consider as an error. More detail could be
-# found from [shell script Exit codes](http://shellscript.sh/exitcodes.html).
-#
-# TorchServe: The Management API listens on port 8081 and is only accessible
-# from localhost by default.
-
-if [[ -z "${MNG_PORT}" ]]; then
-  MNG_PORT=7081  # We default the management_port to 7081.
-else
-  MNG_PORT="${MNG_PORT}"
-fi
-
-check_model_availability(){
-  local MODEL_NAME=$1
-  # Returns whether "READY" is found in the model status.
-  # Reference: https://pytorch.org/serve/management_api.html#describe-model.
-  curl -s "http://localhost:${MNG_PORT}/models/${MODEL_NAME}" | grep "READY" -q
-}
-
-main(){
-  check_model_availability "$MODEL"  # Assume Dockerfile sets MODEL environment parameter.
-  local available=$?
-  if [[ $available -gt 0 ]]
-  then
-    echo "Warning: Model(${MODEL}) is not yet available."
-    return 1
-  fi
-  return 0
-}
-main