Delete deprecated LoRA fine-tuning notebooks and related tutorials.

PiperOrigin-RevId: 976392412

.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(
-                "python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
+                r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
             )
             if re_match:
                 # get the version number
@@ -365,7 +365,7 @@ def process_and_execute_notebook(
             # Use gcloud to get tail
             try:
                 result.error_message = subprocess.check_output(
-                    ["gsutil", "cat", "-r", "-1000", log_file_uri], encoding="UTF-8"
+                    ["gcloud", "storage", "cat", "--range", "-1000", log_file_uri], encoding="UTF-8"
                 )
             except Exception as error:
                 result.error_message = str(error)

.cloud-build/execute_notebook_helper.py

@@ -56,8 +56,8 @@ def execute_notebook(
     print("\n=== DOWNLOAD EXECUTED NOTEBOOK ===\n")
     print(f"Please debug the executed notebook by downloading the executed notebook:")
 
-    print("Option 1. Using gsutil. Run the following command in your terminal.")
-    print(f'\tgsutil cp "{output_file_or_uri}" .')
+    print("Option 1. Using gcloud storage. Run the following command in your terminal.")
+    print(f'\tgcloud storage cp "{output_file_or_uri}" .')
 
     print("Option 2. Using this link.")
     print(f"\thttps://storage.googleapis.com/{output_file_or_uri[5:]}")

.cloud-build/test_notebook_vm.txt

@@ -1,5 +1,3 @@
-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

.cloud-build/utils/NotebookProcessors.py

@@ -108,7 +108,7 @@ class VertexAIInstallProprocessor(Preprocessor):
         if "google-cloud-aiplatform" not in content:
             return content
         return (
-            f"gsutil cp {self.vertex_ai_wheel} google-cloud-aiplatform.whl\n" + 
+            f"gcloud storage cp {self.vertex_ai_wheel} google-cloud-aiplatform.whl\n" + 
             content.replace("google-cloud-aiplatform\n", "google-cloud-aiplatform.whl\n")
             .replace("google-cloud-aiplatform ", "google-cloud-aiplatform.whl ")
         )

.cloud-build/utils/util.py

@@ -15,7 +15,7 @@ def download_file(bucket_name: str, blob_name: str, destination_file: str) -> st
     remote_file_path = "".join(["gs://", "/".join([bucket_name, blob_name])])
 
     subprocess.check_output(
-        ["gsutil", "cp", remote_file_path, destination_file], encoding="UTF-8"
+        ["gcloud", "storage", "cp", remote_file_path, destination_file], encoding="UTF-8"
     )
 
     return destination_file
@@ -27,7 +27,7 @@ def upload_file(
 ) -> str:
     """Copies a local file to a GCS path"""
     subprocess.check_output(
-        ["gsutil", "cp", local_file_path, remote_file_path], encoding="UTF-8"
+        ["gcloud", "storage", "cp", local_file_path, remote_file_path], encoding="UTF-8"
     )
 
     return remote_file_path

.github/workflows/ci.yaml

@@ -7,11 +7,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Set up Python
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
-        python-version: '3.x'      
+        python-version: '3.12'      
     - name: Fetch pull request branch
-      uses: actions/checkout@v4
+      uses: actions/checkout@v6
       with:
         fetch-depth: 0      
     - name: Fetch base main branch

.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.14
 
 WORKDIR setup
 

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==24.10.0
-pyupgrade==3.19.0
-isort==5.13.2
-flake8==7.1.1
+black==26.5.1
+pyupgrade==3.21.2
+isort==8.0.1
+flake8==7.3.0
 nbqa==1.9.1
 

README.md

@@ -1,12 +1,12 @@
 # ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples
 
-This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
+This repository contains notebooks, code samples, sample apps, skills, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
 
 ## Overview
 
 [Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
 
-For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
+⚠️ For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
 
 ## Explore, learn and contribute
 
@@ -16,11 +16,11 @@ You can explore, learn, and contribute to this repository to unleash the full po
 
 Explore this repository, follow the links in the header section of each of the notebooks to -
 
-![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png)  Open and run the notebook in [Colab](https://colab.google/)\
-![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png)  Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
-![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32)  Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
-![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png)  View the notebook on Github
-
+- Open and run the notebook in [Colab](https://colab.google/)
+- Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)
+- Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)
+- View the notebook on Github
+  
 ### Contribute
 
 See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
@@ -35,7 +35,7 @@ To get started using Vertex AI, you must have a Google Cloud project.
 
 ## Repository structure
 
-```bash
+```text
 ├── notebooks
 │   ├── official - Notebooks demonstrating use of each Vertex AI service
 │   │   ├── automl
@@ -45,7 +45,23 @@ To get started using Vertex AI, you must have a Google Cloud project.
 │   │   ├── model_garden
 │   │   ├── ...
 ├── community-content - Sample code and tutorials contributed by the community
-
+├── docs - Deep-dive documentation and advanced setup guides
+└── skills - Suite of AI Agent "Skills" for Vertex AI
+    ├── README.md               # Developer guide for Vertex AI skills
+    ├── vertex-ai/              # Primary router for Vertex AI tasks
+    │   └── SKILL.md            # Entry point that routes across capabilities
+    ├── genai-sdk/              # Gemini API usage with Gen AI SDK
+    │   └── SKILL.md            # Guides for Python, JS/TS, Go, Java, C#
+    ├── vertex-deploy/          # Deploying models to Endpoints
+    │   └── SKILL.md            # Commands for open models & custom weights
+    ├── vertex-inference/       # Inferencing with GenAI models
+    │   └── SKILL.md            # Code samples for Gemini and OpenMaaS
+    └── vertex-tuning/          # Secondary router for model fine-tuning
+        ├── SKILL.md            # Router for tuning tasks
+        ├── gemini/             # Fine-tuning first-party Gemini models
+        │   └── SKILL.md
+        └── open-model/         # Fine-tuning third-party open models
+            └── SKILL.md
 ```
 ## Examples
 

community-content/CODEOWNERS

@@ -29,4 +29,5 @@
 /vertex_model_garden/model_oss/vllm @kathyyu-google
 /vertex_model_garden/benchmarking_reports @lavraicse
 /vertex_model_garden/model_oss/autogluon @lavraicse
+/vertex_distributed_training/a3mega/llama-3-8b-nemo-pretraining @mstyer-google @erwinh85 @mchrestkha
 

community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_PyTorch_model_archive/component.yaml

@@ -148,7 +148,7 @@ implementation:
 
       # Downloading the model archive from GCS
       # TODO: Fix gsutil bugs (requires project ID, has auth issues) and use gsutil instead.
-      # gsutil cp "$model_archive_uri" "$model_archive_local_path"
+      # gcloud storage cp "$model_archive_uri" "$model_archive_local_path"
       pip install google-cloud-storage
       python -c '
       import sys

community-content/pipeline_components/google-cloud/storage/download/component.yaml

@@ -24,12 +24,12 @@ implementation:
 
             # Checking whether the URI points to a single blob, a directory or a URI pattern
             # URI points to a blob when that URI does not end with slash and listing that URI only yields the same URI
-            if [[ "$uri" != */ ]] && (gsutil ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
+            if [[ "$uri" != */ ]] && (gcloud storage ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
                 mkdir -p "$(dirname "$output_path")"
-                gsutil -m cp -r "$uri" "$output_path"
+                gcloud storage cp --recursive "$uri" "$output_path"
             else
                 mkdir -p "$output_path" # When source path is a directory, gsutil requires the destination to also be a directory
-                gsutil -m rsync -r "$uri" "$output_path" # gsutil cp has different path handling than Linux cp. It always puts the source directory (name) inside the destination directory. gsutil rsync does not have that problem.
+                gcloud storage rsync --recursive "$uri" "$output_path" # gsutil cp has different path handling than Linux cp. It always puts the source directory (name) inside the destination directory. gsutil rsync does not have that problem.
             fi
         - inputValue: GCS path
         - outputPath: Data

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt

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

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/trainer/requirements.txt

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

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/vertex_prediction_with_custom_torchserve_container.ipynb

@@ -110,7 +110,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gcloud storage ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -192,7 +192,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil cp -r $gcs_output_uri_prefix/model ./model_server/"
+        "! gcloud storage cp --recursive $gcs_output_uri_prefix/model ./model_server/"
       ]
     },
     {
@@ -556,7 +556,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
+        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
       ]
     },
     {

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/vertex_training_with_custom_container.ipynb

@@ -412,7 +412,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gcloud storage ls $gcs_output_uri_prefix"
       ]
     }
   ],

community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/custom_container/scripts/train-cloud.sh

@@ -77,4 +77,4 @@ echo "After the job is completed successfully, model files will be saved at $JOB
 
 # # Verify the model was exported
 # echo "Verify the model was exported:"
-# gsutil ls ${JOB_DIR}/
+# gcloud storage ls ${JOB_DIR}/

community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/predictor/Dockerfile.serve

@@ -34,4 +34,4 @@ RUN echo "service_envelope=json\n"     "inference_address=http://0.0.0.0:${AIP_H
 USER model-server
 
 # run Torchserve HTTP serve to respond to prediction requests
-CMD ["echo", "AIP_STORAGE_URI=${AIP_STORAGE_URI}", ";",     "gsutil", "cp", "-r", "${AIP_STORAGE_URI}/${MODEL_NAME}.mar", "/home/model-server/model-store/", ";",     "ls", "-ltr", "/home/model-server/model-store/", ";",     "torchserve", "--start", "--ts-config=/home/model-server/config.properties",     "--models", "${MODEL_NAME}=${MODEL_NAME}.mar",     "--model-store", "/home/model-server/model-store"]
+CMD ["echo", "AIP_STORAGE_URI=${AIP_STORAGE_URI}", ";",     "gcloud", "storage", "cp", "--recursive", "${AIP_STORAGE_URI}/${MODEL_NAME}.mar", "/home/model-server/model-store/", ";",     "ls", "-ltr", "/home/model-server/model-store/", ";",     "torchserve", "--start", "--ts-config=/home/model-server/config.properties",     "--models", "${MODEL_NAME}=${MODEL_NAME}.mar",     "--model-store", "/home/model-server/model-store"]

community-content/pytorch_text_classification_using_vertex_sdk_and_gcloud/python_package/scripts/train-cloud.sh

@@ -67,4 +67,4 @@ echo "After the job is completed successfully, model files will be saved at $JOB
 
 # # Verify the model was exported
 # echo "Verify the model was exported:"
-# gsutil ls ${JOB_DIR}/
+# gcloud storage ls ${JOB_DIR}/

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/mlops_pipeline_tf_agents_bandits_movie_recommendation/mlops_pipeline_tf_agents_bandits_movie_recommendation.ipynb

@@ -478,8 +478,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
-      ]
+        "! gcloud storage buckets create --location $REGION $BUCKET_NAME"      ]
     },
     {
       "cell_type": "markdown",
@@ -498,8 +497,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
-      ]
+        "! gcloud storage ls --all-versions --long $BUCKET_NAME"      ]
     },
     {
       "cell_type": "markdown",
@@ -582,8 +580,7 @@
       "outputs": [],
       "source": [
         "# Download the sample data into your RAW_DATA_PATH\n",
-        "! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $RAW_DATA_PATH"
-      ]
+        "! gcloud storage cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\" $RAW_DATA_PATH"      ]
     },
     {
       "cell_type": "code",
@@ -1621,9 +1618,7 @@
         "! gcloud scheduler jobs delete $SIMULATOR_SCHEDULER_JOB --quiet\n",
         "\n",
         "# Delete Cloud Storage objects that were created.\n",
-        "! gsutil -m rm -r $PIPELINE_ROOT\n",
-        "! gsutil -m rm -r $TRAINING_ARTIFACTS_DIR"
-      ]
+        "! gcloud storage rm --recursive $PIPELINE_ROOT\n",        "! gcloud storage rm --recursive $TRAINING_ARTIFACTS_DIR"      ]
     }
   ],
   "metadata": {

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
-pillow==10.3.0
+pillow==12.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
+pillow==12.3.0
 tf-agents==0.8.0
 tensorflow==2.12.1

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
-pillow==10.3.0
+pillow==12.3.0
 tf-agents==0.8.0

community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/step_by_step_sdk_tf_agents_bandits_movie_recommendation/step_by_step_sdk_tf_agents_bandits_movie_recommendation.ipynb

@@ -398,6 +398,7 @@
         "if not IS_GOOGLE_CLOUD_NOTEBOOK:\n",
         "    if \"google.colab\" in sys.modules:\n",
         "        from google.colab import auth as google_auth\n",
+        "\n",
         "        google_auth.authenticate_user()\n",
         "\n",
         "    # If you are running this notebook locally, replace the string below with the\n",
@@ -472,7 +473,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil mb -l $REGION $BUCKET_NAME"
+        "! gcloud storage buckets create --location $REGION $BUCKET_NAME"
       ]
     },
     {
@@ -492,7 +493,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls -al $BUCKET_NAME"
+        "! gcloud storage ls --all-versions --long $BUCKET_NAME"
       ]
     },
     {
@@ -565,7 +566,7 @@
       "outputs": [],
       "source": [
         "# Copy the sample data into your DATA_PATH\n",
-        "! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\"  $DATA_PATH"
+        "! gcloud storage cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\"  $DATA_PATH"
       ]
     },
     {
@@ -579,11 +580,15 @@
         "# Set hyperparameters.\n",
         "BATCH_SIZE = 8  # @param {type:\"integer\"} Training and prediction batch size.\n",
         "TRAINING_LOOPS = 5  # @param {type:\"integer\"} Number of training iterations.\n",
-        "STEPS_PER_LOOP = 2  # @param {type:\"integer\"} Number of driver steps per training iteration.\n",
+        "STEPS_PER_LOOP = (\n",
+        "    2  # @param {type:\"integer\"} Number of driver steps per training iteration.\n",
+        ")\n",
         "\n",
         "# Set MovieLens simulation environment parameters.\n",
         "RANK_K = 20  # @param {type:\"integer\"} Rank for matrix factorization in the MovieLens environment; also the observation dimension.\n",
-        "NUM_ACTIONS = 20  # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
+        "NUM_ACTIONS = (\n",
+        "    20  # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
+        ")\n",
         "PER_ARM = False  # Use the non-per-arm version of the MovieLens environment.\n",
         "\n",
         "# Set agent parameters.\n",
@@ -621,7 +626,8 @@
       "source": [
         "# Define RL environment.\n",
         "env = movielens_py_environment.MovieLensPyEnvironment(\n",
-        "    DATA_PATH, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS, csv_delimiter=\"\\t\")\n",
+        "    DATA_PATH, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS, csv_delimiter=\"\\t\"\n",
+        ")\n",
         "environment = tf_py_environment.TFPyEnvironment(env)\n",
         "\n",
         "# Define RL agent/algorithm.\n",
@@ -631,7 +637,8 @@
         "    tikhonov_weight=TIKHONOV_WEIGHT,\n",
         "    alpha=AGENT_ALPHA,\n",
         "    dtype=tf.float32,\n",
-        "    accepts_per_arm_features=PER_ARM)\n",
+        "    accepts_per_arm_features=PER_ARM,\n",
+        ")\n",
         "print(\"TimeStep Spec (for each batch):\\n\", agent.time_step_spec, \"\\n\")\n",
         "print(\"Action Spec (for each batch):\\n\", agent.action_spec, \"\\n\")\n",
         "print(\"Reward Spec (for each batch):\\n\", environment.reward_spec(), \"\\n\")\n",
@@ -639,7 +646,8 @@
         "# Define RL metric.\n",
         "optimal_reward_fn = functools.partial(\n",
         "    environment_utilities.compute_optimal_reward_with_movielens_environment,\n",
-        "    environment=environment)\n",
+        "    environment=environment,\n",
+        ")\n",
         "regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)\n",
         "metrics = [regret_metric]"
       ]
@@ -704,35 +712,38 @@
         "    if training_data_spec_transformation_fn is None:\n",
         "        data_spec = agent.policy.trajectory_spec\n",
         "    else:\n",
-        "        data_spec = training_data_spec_transformation_fn(\n",
-        "            agent.policy.trajectory_spec)\n",
-        "    replay_buffer = trainer.get_replay_buffer(data_spec, environment.batch_size,\n",
-        "                                              steps_per_loop)\n",
+        "        data_spec = training_data_spec_transformation_fn(agent.policy.trajectory_spec)\n",
+        "    replay_buffer = trainer.get_replay_buffer(\n",
+        "        data_spec, environment.batch_size, steps_per_loop\n",
+        "    )\n",
         "\n",
         "    # `step_metric` records the number of individual rounds of bandit interaction;\n",
         "    # that is, (number of trajectories) * batch_size.\n",
         "    step_metric = tf_metrics.EnvironmentSteps()\n",
         "    metrics = [\n",
         "        tf_metrics.NumberOfEpisodes(),\n",
-        "        tf_metrics.AverageEpisodeLengthMetric(batch_size=environment.batch_size)\n",
+        "        tf_metrics.AverageEpisodeLengthMetric(batch_size=environment.batch_size),\n",
         "    ]\n",
         "    if additional_metrics:\n",
         "        metrics += additional_metrics\n",
         "\n",
         "    if isinstance(environment.reward_spec(), dict):\n",
-        "        metrics += [tf_metrics.AverageReturnMultiMetric(\n",
-        "            reward_spec=environment.reward_spec(),\n",
-        "            batch_size=environment.batch_size)]\n",
-        "    else:\n",
         "        metrics += [\n",
-        "            tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]\n",
+        "            tf_metrics.AverageReturnMultiMetric(\n",
+        "                reward_spec=environment.reward_spec(), batch_size=environment.batch_size\n",
+        "            )\n",
+        "        ]\n",
+        "    else:\n",
+        "        metrics += [tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]\n",
         "\n",
         "    # Store intermediate metric results, indexed by metric names.\n",
         "    metric_results = defaultdict(list)\n",
         "\n",
         "    if training_data_spec_transformation_fn is not None:\n",
-        "        def add_batch_fn(data): return replay_buffer.add_batch(training_data_spec_transformation_fn(data)) \n",
-        "        \n",
+        "\n",
+        "        def add_batch_fn(data):\n",
+        "            return replay_buffer.add_batch(training_data_spec_transformation_fn(data))\n",
+        "\n",
         "    else:\n",
         "        add_batch_fn = replay_buffer.add_batch\n",
         "\n",
@@ -742,10 +753,12 @@
         "        env=environment,\n",
         "        policy=agent.collect_policy,\n",
         "        num_steps=steps_per_loop * environment.batch_size,\n",
-        "        observers=observers)\n",
+        "        observers=observers,\n",
+        "    )\n",
         "\n",
         "    training_loop = trainer.get_training_loop_fn(\n",
-        "        driver, replay_buffer, agent, steps_per_loop)\n",
+        "        driver, replay_buffer, agent, steps_per_loop\n",
+        "    )\n",
         "    saver = policy_saver.PolicySaver(agent.policy)\n",
         "\n",
         "    for _ in range(training_loops):\n",
@@ -783,7 +796,8 @@
         "    environment=environment,\n",
         "    training_loops=TRAINING_LOOPS,\n",
         "    steps_per_loop=STEPS_PER_LOOP,\n",
-        "    additional_metrics=metrics)\n",
+        "    additional_metrics=metrics,\n",
+        ")\n",
         "\n",
         "tf.profiler.experimental.stop()"
       ]
@@ -1092,11 +1106,15 @@
       },
       "outputs": [],
       "source": [
-        "RUN_HYPERPARAMETER_TUNING = True  # Execute hyperparameter tuning instead of regular training.\n",
+        "RUN_HYPERPARAMETER_TUNING = (\n",
+        "    True  # Execute hyperparameter tuning instead of regular training.\n",
+        ")\n",
         "TRAIN_WITH_BEST_HYPERPARAMETERS = False  # Do not train.\n",
         "\n",
         "HPTUNING_RESULT_DIR = \"hptuning/\"  # @param {type: \"string\"} Directory to store the best hyperparameter(s) in `BUCKET_NAME` and locally (temporarily).\n",
-        "HPTUNING_RESULT_PATH = os.path.join(HPTUNING_RESULT_DIR, \"result.json\")  # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
+        "HPTUNING_RESULT_PATH = os.path.join(\n",
+        "    HPTUNING_RESULT_DIR, \"result.json\"\n",
+        ")  # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
       ]
     },
     {
@@ -1124,7 +1142,7 @@
         "    image_uri: str,\n",
         "    args: List[str],\n",
         "    location: str = \"us-central1\",\n",
-        "    api_endpoint: str = \"us-central1-aiplatform.googleapis.com\"\n",
+        "    api_endpoint: str = \"us-central1-aiplatform.googleapis.com\",\n",
         ") -> None:\n",
         "    \"\"\"Creates a hyperparameter tuning job using a custom container.\n",
         "\n",
@@ -1197,8 +1215,8 @@
         "\n",
         "    # Create job\n",
         "    response = client.create_hyperparameter_tuning_job(\n",
-        "        parent=parent,\n",
-        "        hyperparameter_tuning_job=hyperparameter_tuning_job)\n",
+        "        parent=parent, hyperparameter_tuning_job=hyperparameter_tuning_job\n",
+        "    )\n",
         "    job_id = response.name.split(\"/\")[-1]\n",
         "    print(\"Job ID:\", job_id)\n",
         "    print(\"Job config:\", response)\n",
@@ -1242,7 +1260,8 @@
         "    image_uri=f\"gcr.io/{PROJECT_ID}/{HPTUNING_TRAINING_CONTAINER}:latest\",\n",
         "    args=args,\n",
         "    location=REGION,\n",
-        "    api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")"
+        "    api_endpoint=f\"{REGION}-aiplatform.googleapis.com\",\n",
+        ")"
       ]
     },
     {
@@ -1292,7 +1311,8 @@
         "    name = client.hyperparameter_tuning_job_path(\n",
         "        project=project,\n",
         "        location=location,\n",
-        "        hyperparameter_tuning_job=hyperparameter_tuning_job_id)\n",
+        "        hyperparameter_tuning_job=hyperparameter_tuning_job_id,\n",
+        "    )\n",
         "    response = client.get_hyperparameter_tuning_job(name=name)\n",
         "    return response"
       ]
@@ -1313,7 +1333,8 @@
         "        location=REGION,\n",
         "        api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")\n",
         "    if response.state.name == 'JOB_STATE_SUCCEEDED':\n",
-        "        print(\"Job succeeded.\\nJob Time:\", response.update_time - response.create_time)\n",
+        "        print(\"Job succeeded.\n",
+        "Job Time:\", response.update_time - response.create_time)\n",
         "        trials = response.trials\n",
         "        print(\"Trials:\", trials)\n",
         "        break\n",
@@ -1348,8 +1369,8 @@
         "if trials:\n",
         "    # Dict mapping from metric names to the best metric values seen so far\n",
         "    best_objective_values = dict.fromkeys(\n",
-        "        [metric.metric_id for metric in trials[0].final_measurement.metrics],\n",
-        "        -np.inf)\n",
+        "        [metric.metric_id for metric in trials[0].final_measurement.metrics], -np.inf\n",
+        "    )\n",
         "    # Dict mapping from metric names to a list of the best combination(s) of\n",
         "    # hyperparameter(s). Each combination is a dict mapping from hyperparameter\n",
         "    # names to their values.\n",
@@ -1358,12 +1379,13 @@
         "        # `final_measurement` and `parameters` are `RepeatedComposite` objects.\n",
         "        # Reference the structure above to extract the value of your interest.\n",
         "        for metric in trial.final_measurement.metrics:\n",
-        "            params = {\n",
-        "                param.parameter_id: param.value for param in trial.parameters}\n",
+        "            params = {param.parameter_id: param.value for param in trial.parameters}\n",
         "            if metric.value > best_objective_values[metric.metric_id]:\n",
         "                best_params[metric.metric_id] = [params]\n",
         "            elif metric.value == best_objective_values[metric.metric_id]:\n",
-        "                best_params[param.parameter_id].append(params)  # Handle cases where multiple hyperparameter values lead to the same performance.\n",
+        "                best_params[param.parameter_id].append(\n",
+        "                    params\n",
+        "                )  # Handle cases where multiple hyperparameter values lead to the same performance.\n",
         "    print(\"Best hyperparameter value(s):\")\n",
         "    for metric, params in best_params.items():\n",
         "        print(f\"Metric={metric}: {sorted(params)}\")\n",
@@ -1443,7 +1465,9 @@
       },
       "outputs": [],
       "source": [
-        "PREDICTION_CONTAINER = \"prediction-custom-container\"  # @param {type:\"string\"} Name of the container image."
+        "PREDICTION_CONTAINER = (\n",
+        "    \"prediction-custom-container\"  # @param {type:\"string\"} Name of the container image.\n",
+        ")"
       ]
     },
     {
@@ -1475,7 +1499,7 @@
         "  machineType: 'E2_HIGHCPU_8'\"\"\".format(\n",
         "    PROJECT_ID=PROJECT_ID,\n",
         "    PREDICTION_CONTAINER=PREDICTION_CONTAINER,\n",
-        "    ARTIFACTS_DIR=ARTIFACTS_DIR\n",
+        "    ARTIFACTS_DIR=ARTIFACTS_DIR,\n",
         ")\n",
         "\n",
         "with open(\"cloudbuild.yaml\", \"w\") as fp:\n",
@@ -1592,8 +1616,12 @@
       },
       "outputs": [],
       "source": [
-        "RUN_HYPERPARAMETER_TUNING = False  # Execute regular training instead of hyperparameter tuning.\n",
-        "TRAIN_WITH_BEST_HYPERPARAMETERS = True  # @param {type:\"bool\"} Whether to use learned hyperparameters in training."
+        "RUN_HYPERPARAMETER_TUNING = (\n",
+        "    False  # Execute regular training instead of hyperparameter tuning.\n",
+        ")\n",
+        "TRAIN_WITH_BEST_HYPERPARAMETERS = (\n",
+        "    True  # @param {type:\"bool\"} Whether to use learned hyperparameters in training.\n",
+        ")"
       ]
     },
     {
@@ -1633,10 +1661,12 @@
         "job = aiplatform.CustomContainerTrainingJob(\n",
         "    display_name=\"train-movielens\",\n",
         "    container_uri=f\"gcr.io/{PROJECT_ID}/{HPTUNING_TRAINING_CONTAINER}:latest\",\n",
-        "    command=[\"python3\", \"-m\", \"src.training.task\"] + args,  # Pass in training arguments, including hyperparameters.\n",
+        "    command=[\"python3\", \"-m\", \"src.training.task\"]\n",
+        "    + args,  # Pass in training arguments, including hyperparameters.\n",
         "    model_serving_container_image_uri=f\"gcr.io/{PROJECT_ID}/{PREDICTION_CONTAINER}:latest\",\n",
         "    model_serving_container_predict_route=\"/predict\",\n",
-        "    model_serving_container_health_route=\"/health\")\n",
+        "    model_serving_container_health_route=\"/health\",\n",
+        ")\n",
         "\n",
         "print(\"Training Spec:\", job._managed_model)\n",
         "\n",
@@ -1645,7 +1675,8 @@
         "    replica_count=1,\n",
         "    machine_type=\"n1-standard-4\",\n",
         "    accelerator_type=\"ACCELERATOR_TYPE_UNSPECIFIED\",\n",
-        "    accelerator_count=0)"
+        "    accelerator_count=0,\n",
+        ")"
       ]
     },
     {
@@ -1784,7 +1815,7 @@
         "! gcloud ai models delete $model.name --quiet\n",
         "\n",
         "# Delete Cloud Storage objects that were created\n",
-        "! gsutil -m rm -r $ARTIFACTS_DIR"
+        "! gcloud storage rm --recursive $ARTIFACTS_DIR"
       ]
     }
   ],

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_cpu_with_custom_container.ipynb

@@ -324,7 +324,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gcloud storage ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -344,7 +344,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
+        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
       ]
     }
   ],

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_gpu_with_custom_container.ipynb

@@ -328,7 +328,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gcloud storage ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -348,7 +348,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
+        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
       ]
     }
   ],

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/vertex_training_with_local_mode_container.ipynb

@@ -341,7 +341,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil ls $gcs_output_uri_prefix"
+        "! gcloud storage ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -361,7 +361,7 @@
       },
       "outputs": [],
       "source": [
-        "! gsutil rm -rf $gcs_output_uri_prefix"
+        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
       ]
     }
   ],

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/README.md

@@ -0,0 +1,126 @@
+# Vertex AI Training: Llama 3.1 8B pre-training using Nvidia A3 Mega VMs (H100)
+This document provides a step-by-step guide for pre-training a Llama 3.1 8B model on the `en-wiki` dataset using multiple [Vertex AI Custom Training](https://cloud.google.com/vertex-ai/docs/training/overview) `a3-megagpu-8g` nodes.
+
+We will use a custom container based on NVIDIA's [NeMo Framework](https://docs.nvidia.com/nemo-framework/user-guide/24.07/overview.html) to demonstrate a scalable, multi-node training workflow. All required artifacts and commands are included.
+
+## 1. Prerequisites
+
+### 1.1. Google Cloud Project setup
+- **Enable APIs:** Ensure the Vertex AI API is [enabled for your project](http://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).
+- **H100 Mega Quota:** A3 Mega VMs are powered by H100 GPUs. Request quota for `custom_model_training_nvidia_h100_mega_gpus` in one of the [supported regions](https://cloud.google.com/vertex-ai/docs/general/locations#accelerator_support). If using Spot VMs, request `custom_model_training_preemptible_nvidia_h100_mega_gpus` quota instead.
+- **Reservations (Optional but recommended):** For guaranteed capacity, [create a reservation](https://cloud.google.com/compute/docs/instances/reservations-shared) and ensure the reservation is shared with the Vertex AI service account. This guide requires a minimum of **16 H100 GPUs** (2 full A3 Mega nodes).
+
+### 1.2. GCS bucket
+Create a [Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) in the same region where you have quota. If you're using Hierarchical Namespace for your bucket, you may need to update permissions of the Vertex AI Custom Code Service Agent .
+
+This bucket is used for:
+- Staging the training application.
+- Storing model checkpoints and logs.
+- Storing data if you use your own data.
+
+
+## 2. Setup & configuration
+
+### 2.1. Clone the repo
+First clone the repo into your development environment.
+
+```bash
+git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
+```
+
+Navigate to the root folder for this sample.
+
+### 2.2. Environment Setup
+First, configure your local environment. These variables are used in subsequent commands.
+
+```bash
+# Required: Update with your values
+export PROJECT_ID="<your-project-id>"
+export REPOSITORY="<your-artifact-registry-repo-name>" # e.g., "my-containers"
+export BUCKET="<your-gcs-bucket-name>"
+
+# Optional: Change if needed
+export REGION="us-central1"
+
+# --- Do not change the lines below ---
+export ARTIFACT_REGISTRY="${REGION}-docker.pkg.dev/${PROJECT_ID}/${REPOSITORY}"
+export REPO_ROOT=$(git rev-parse --show-toplevel)
+```
+
+## 3. Build and push a docker container image to Artifact Registry
+Normally, you can use any custom training container on Vertex AI Training. In this example you build a NeMo Docker image that is based on the [Nvidia’s NeMo 24.09](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags) image. Use Cloud Build to build and push the container image.
+
+This document picked NeMo as the demonstrating container since it’s a widely adopted GPU LLM training framework providing high performance and versatile training functionalities.
+
+In addition to the base image, some customizations are included to form the final prebuilt image:
+- Some dependencies are installed to integrate with Vertex AI Training.
+- An entrypoint script that sets up required environments and calls the training job.
+- Some patches are applied to the NeMo code to let it  load the dataset from a GCS bucket.
+
+Run this command to build the container and push the container into the Google Artifact Registry.
+
+```bash
+cd "${REPO_ROOT}/community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining"
+export IMAGE_NAME="vertex-nemo-llama"
+gcloud builds submit . \
+    --project="${PROJECT_ID}" \
+    --region="${REGION}" \
+    --config=docker/cloudbuild.yml \
+    --substitutions="_ARTIFACT_REGISTRY=${ARTIFACT_REGISTRY},_IMAGE_NAME=${IMAGE_NAME}" \
+    --timeout="2h" \
+    --machine-type="e2-highcpu-32"
+```
+
+## 4. Launch the Training Job
+
+
+###  4.1. Job Configuration File
+Once the container is built, update the job_config.json to set up the training job.
+File: job_config.json
+```json
+{
+  "project_id": "<project-id>",
+  "region": "<region>",
+  "zone": "<zone if using reservation>",
+  "bucket": "<bucket>",
+  "dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
+  "image_uri": "<docker image uri from artifact registry>",
+  "strategy": "spot",
+  "nodes": "2",
+  "machine_type": "a3-megagpu-8g",
+  "gpu_type": "NVIDIA_H100_MEGA_80GB",
+  "gpus_per_node": "8",
+  "recipe_name": "llama3_1_8b_pretrain_a3mega",
+  "job_prefix": "vertex-spot-",
+  "reservation_name": ""
+}
+```
+
+### 4.2 Launch the Training Job
+
+First, create a Python virtual environment using your tool of choice, then install
+the requirements specified in `requirements.txt`. Using `pip`, the command would be:
+```bash
+pip install -r requirements.txt
+```
+
+Now launch the Vertex AI training job using the provided Python script.
+
+```bash
+python3 scripts/launch.py --config_file=job_config.json
+```
+
+This script reads job_config.json, defines the cluster specification (2 nodes, 8 GPUs each), and submits the custom training job to Vertex AI.
+
+## 5. Monitor and Clean Up
+
+### 5.1. Monitoring
+Vertex AI Console: Track the job's status in the Google Cloud Console under Vertex AI > Training > Custom Jobs.
+Logs: View detailed logs in Cloud Logging by filtering for your job name.
+Checkpoints: Model checkpoints are saved to your GCS bucket at the path specified in your training script's configuration.
+
+### 5.2. Cleaning Up
+To avoid ongoing charges, delete the resources you created:
+- The Artifact Registry image.
+- The contents of the GCS bucket (checkpoints, logs).
+- The Vertex AI Custom Job will eventually complete or fail, incurring no further cost.

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/configs/llama3_1_8b_pretrain_a3mega.yaml

@@ -0,0 +1,265 @@
+# Reference:
+# https://github.com/NVIDIA/NeMo-Framework-Launcher/blob/24.07/launcher_scripts/conf/training/llama/llama3_1_8b.yaml
+name: llama3_1_8b_pretrain_a3mega
+restore_from_path: null  # used when starting from a .nemo file
+
+trainer:
+  devices: 8
+  num_nodes: 1
+  accelerator: gpu
+  precision: bf16
+  logger: false  # logger provided by exp_manager
+  enable_checkpointing: false
+  use_distributed_sampler: false
+  max_epochs: -1  # PTL default. In practice, max_steps will be reached first.
+  max_steps: 30  # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches
+  log_every_n_steps: 1
+  val_check_interval: null
+  limit_val_batches: 1
+  limit_test_batches: 1
+  accumulate_grad_batches: 1  # do not modify, grad acc is automatic for training megatron models
+  gradient_clip_val: 1.0
+  benchmark: false
+  enable_model_summary: false  # default PTL callback for this does not support model parallelism, instead we log manually
+
+exp_manager:
+  explicit_log_dir: null
+  exp_dir: /data
+  name: ${name}
+  create_dllogger_logger: true
+  dllogger_logger_kwargs:
+    verbose: true
+    stdout: true
+    json_file: "/data/dllogger.json"
+  create_wandb_logger: false
+  wandb_logger_kwargs:
+    project: null
+    name: null
+  resume_if_exists: true
+  resume_ignore_no_checkpoint: true
+  create_checkpoint_callback: false
+  checkpoint_callback_params:
+    monitor: val_loss
+    save_top_k: 3
+    mode: min
+    always_save_nemo: false  # saves nemo file during validation, not implemented for model parallel
+    save_nemo_on_train_end: false  # not recommended when training large models on clusters with short time limits
+    filename: 'megatron_gpt--{val_loss:.2f}-{step}-{consumed_samples}'
+    model_parallel_size: ${multiply:${model.tensor_model_parallel_size}, ${model.pipeline_model_parallel_size}}
+  seconds_to_sleep: 5  # Allows node_rank!=0 to sleep and let node0 to init, like preparing data
+
+model:
+  mcore_gpt: true
+  # specify micro_batch_size, global_batch_size, and model parallelism
+  # gradient accumulation will be done automatically based on data_parallel_size
+  micro_batch_size: 1  # limited by GPU memory
+  global_batch_size: 1024  # will use more micro batches to reach global batch size
+  tensor_model_parallel_size: 1  # intra-layer model parallelism
+  pipeline_model_parallel_size: 2  # inter-layer model parallelism
+  context_parallel_size: 1
+  virtual_pipeline_model_parallel_size: null  # interleaved pipeline
+  ## Sequence Parallelism
+  # Makes tensor parallelism more memory efficient for LLMs (20B+) by parallelizing layer norms and dropout sequentially
+  # See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
+  sequence_parallel: false
+
+  fsdp: false
+  fsdp_cpu_offload: true
+  fsdp_sharding_strategy: "full"  # Method to shard model states. Available options are 'full', 'hybrid', and 'grad'.
+  fsdp_grad_reduce_dtype: "16"  # Gradient reduction data type.
+  fsdp_sharded_checkpoint: false  # Store and load FSDP shared checkpoint.
+  fsdp_use_orig_params: false  # Set to True to use FSDP for specific peft scheme.
+
+  # Distributed checkpoint setup
+  dist_ckpt_format: "torch_dist"  # Set to 'torch_dist' to use PyTorch distributed checkpoint format.
+  dist_ckpt_load_on_device: true  # whether to load checkpoint weights directly on GPU or to CPU
+  dist_ckpt_parallel_save: true  # if true, each worker will write its own part of the dist checkpoint
+  dist_ckpt_parallel_save_within_dp: false  # if true, save will be parallelized only within a DP group (whole world otherwise), which might slightly reduce the save overhead
+  dist_ckpt_parallel_load: false  # if true, each worker will load part of the dist checkpoint and exchange with NCCL. Might use some extra GPU memory
+  dist_ckpt_torch_dist_multiproc: 2  # number of extra processes per rank used during ckpt save with PyTorch distributed format
+  dist_ckpt_assume_constant_structure: false  # set to True only if the state dict structure doesn't change within a single job. Allows caching some computation across checkpoint saves.
+  dist_ckpt_parallel_dist_opt: true  # parallel save/load of a DistributedOptimizer. 'True' allows performant save and reshardable checkpoints. Set to 'False' only in order to minimize the number of checkpoint files.
+  dist_ckpt_load_strictness: null  # defines checkpoint keys mismatch behavior (only during dist-ckpt load). Choices: assume_ok_unexpected (default - try loading without any check), log_all (log mismatches), raise_all (raise mismatches)
+
+  # model architecture
+  encoder_seq_length: 8192
+  max_position_embeddings: ${.encoder_seq_length}
+  num_layers: 32  # 8b: 32 | 70b: 80 | 405b: 126
+  hidden_size: 4096  # 8b: 4096 | 70b: 8192 | 405b: 16384
+  ffn_hidden_size: 14336  # 8b: 14336 | 70b: 28672 | 405b: 53248
+  num_attention_heads: 32  # 8b: 32 | 70b: 64 | 405b: 128
+  num_query_groups: 8  # Number of query groups for group query attention. If None, normal attention is used. 8b: 8 | 70b: 8 | 405b: 16
+  init_method_std: 0.01  # Standard deviation of the zero mean normal distribution used for weight initialization.  8b: 0.01 | 70b: 0.008944 | 405b: 0.02
+  use_scaled_init_method: true  # use scaled residuals initialization
+  hidden_dropout: 0.0  # Dropout probability for hidden state transformer.
+  attention_dropout: 0.0  # Dropout probability for attention
+  ffn_dropout: 0.0  # Dropout probability in the feed-forward layer.
+  kv_channels: null  # Projection weights dimension in multi-head attention. Set to hidden_size // num_attention_heads if null
+  apply_query_key_layer_scaling: true  # scale Q * K^T by 1 / layer-number.
+  normalization: 'rmsnorm'  # Normalization layer to use. Options are 'layernorm', 'rmsnorm'
+  layernorm_epsilon: 1e-5
+  do_layer_norm_weight_decay: false  # True means weight decay on all params
+  make_vocab_size_divisible_by: 128  # Pad the vocab size to be divisible by this value for computation efficiency.
+  pre_process: true  # add embedding
+  post_process: true  # add pooler
+  persist_layer_norm: true  # Use of persistent fused layer norm kernel.
+  bias: false  # Whether to use bias terms in all weight matrices.
+  activation: 'fast-swiglu'  # Options ['gelu', 'geglu', 'swiglu', 'reglu', 'squared-relu', 'fast-geglu', 'fast-swiglu', 'fast-reglu']
+  headscale: false  # Whether to learn extra parameters that scale the output of the each self-attention head.
+  transformer_block_type: 'pre_ln'  # Options ['pre_ln', 'post_ln', 'normformer']
+  openai_gelu: false  # Use OpenAI's GELU instead of the default GeLU
+  normalize_attention_scores: true  # Whether to scale the output Q * K^T by 1 / sqrt(hidden_size_per_head). This arg is provided as a configuration option mostly for compatibility with models that have been weight-converted from HF. You almost always want to se this to True.
+  position_embedding_type: 'rope'  # Position embedding type. Options ['learned_absolute', 'rope']
+  rotary_percentage: 1.0  # If using position_embedding_type=rope, then the per head dim is multiplied by this.
+  attention_type: 'multihead'  # Attention type. Options ['multihead']
+  share_embeddings_and_output_weights: false  # Share embedding and output layer weights.
+  scale_positional_embedding: true  # This is false for llama3 models. Only used for >= llama3.1.
+
+  # Use GPT2BPETokenizer for test, because the testing dataset is tokenized by this tokenizer.
+  # https://docs.nvidia.com/nemo-framework/user-guide/24.07/playbooks/singlenodepretrain.html#data-download-and-pre-processing
+  tokenizer:
+    library: megatron
+    type: GPT2BPETokenizer
+    model: null  # /path/to/tokenizer.model
+    vocab_file: null
+    merge_file: null
+    delimiter: null  # only used for tabular tokenizer
+    sentencepiece_legacy: false  # Legacy=True allows you to add special tokens to sentencepiece tokenizers.
+
+  # Mixed precision
+  native_amp_init_scale: 4294967296  # 2 ** 32
+  native_amp_growth_interval: 1000
+  hysteresis: 2  # Gradient scale hysteresis
+  fp32_residual_connection: false  # Move residual connections to fp32
+  fp16_lm_cross_entropy: false  # Move the cross entropy unreduced loss calculation for lm head to fp16
+
+  # Megatron O2-style half-precision
+  megatron_amp_O2: true  # Enable O2-level automatic mixed precision using main parameters
+  grad_allreduce_chunk_size_mb: 125
+
+  # Fusion
+  grad_div_ar_fusion: true  # Fuse grad division into torch.distributed.all_reduce. Only used with O2 and no pipeline parallelism..
+  gradient_accumulation_fusion: true  # Fuse weight gradient accumulation to GEMMs. Only used with pipeline parallelism and O2.
+  bias_activation_fusion: true  # Use a kernel that fuses the bias addition from weight matrices with the subsequent activation function.
+  bias_dropout_add_fusion: true  # Use a kernel that fuses the bias addition, dropout and residual connection addition.
+  masked_softmax_fusion: true  # Use a kernel that fuses the attention softmax with it's mask.
+  apply_rope_fusion: true  # Use a kernel to add rotary positional embeddings. Only used if position_embedding_type=rope
+  cross_entropy_loss_fusion: true
+
+  # Miscellaneous
+  seed: 1234
+  resume_from_checkpoint: null  # manually set the checkpoint file to load from
+  use_cpu_initialization: false  # Init weights on the CPU (slow for large models)
+  onnx_safe: false  # Use work-arounds for known problems with Torch ONNX exporter.
+  apex_transformer_log_level: 30  # Python logging level displays logs with severity greater than or equal to this
+  gradient_as_bucket_view: true  # PyTorch DDP argument. Allocate gradients in a contiguous bucket to save memory (less fragmentation and buffer memory)
+  sync_batch_comm: false  # Enable stream synchronization after each p2p communication between pipeline stages
+
+  ## Activation Checkpointing
+  # NeMo Megatron supports 'selective' activation checkpointing where only the memory intensive part of attention is checkpointed.
+  # These memory intensive activations are also less compute intensive which makes activation checkpointing more efficient for LLMs (20B+).
+  # See Reducing Activation Recomputation in Large Transformer Models: https://arxiv.org/abs/2205.05198 for more details.
+  # 'full' will checkpoint the entire transformer layer.
+  activations_checkpoint_granularity: null  # 'selective' or 'full'
+  activations_checkpoint_method: null  # 'uniform', 'block'
+  # 'uniform' divides the total number of transformer layers and checkpoints the input activation
+  # of each chunk at the specified granularity. When used with 'selective', 'uniform' checkpoints all attention blocks in the model.
+  # 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity
+  activations_checkpoint_num_layers: null
+  # when using 'uniform' this creates groups of transformer layers to checkpoint. Usually set to 1. Increase to save more memory.
+  # when using 'block' this this will checkpoint the first activations_checkpoint_num_layers per pipeline stage.
+  num_micro_batches_with_partial_activation_checkpoints: null
+  # This feature is valid only when used with pipeline-model-parallelism.
+  # When an integer value is provided, it sets the number of micro-batches where only a partial number of Transformer layers get checkpointed
+  # and recomputed within a window of micro-batches. The rest of micro-batches in the window checkpoint all Transformer layers. The size of window is
+  # set by the maximum outstanding micro-batch backpropagations, which varies at different pipeline stages. The number of partial layers to checkpoint
+  # per micro-batch is set by 'activations_checkpoint_num_layers' with 'activations_checkpoint_method' of 'block'.
+  # This feature enables using activation checkpoint at a fraction of micro-batches up to the point of full GPU memory usage.
+  activations_checkpoint_layers_per_pipeline: null
+  # This feature is valid only when used with pipeline-model-parallelism.
+  # When an integer value (rounded down when float is given) is provided, it sets the number of Transformer layers to skip checkpointing at later
+  # pipeline stages. For example, 'activations_checkpoint_layers_per_pipeline' of 3 makes pipeline stage 1 to checkpoint 3 layers less than
+  # stage 0 and stage 2 to checkpoint 6 layers less stage 0, and so on. This is possible because later pipeline stage
+  # uses less GPU memory with fewer outstanding micro-batch backpropagations. Used with 'num_micro_batches_with_partial_activation_checkpoints',
+  # this feature removes most of activation checkpoints at the last pipeline stage, which is the critical execution path.
+
+  ## Transformer Engine
+  transformer_engine: true
+  fp8: false  # enables fp8 in TransformerLayer forward
+  fp8_e4m3: false  # sets fp8_format = recipe.Format.E4M3
+  fp8_hybrid: false  # sets fp8_format = recipe.Format.HYBRID
+  fp8_margin: 0  # scaling margin
+  fp8_interval: 1  # scaling update interval
+  fp8_amax_history_len: 1024  # Number of steps for which amax history is recorded per tensor
+  fp8_amax_compute_algo: 'max'  # 'most_recent' or 'max'. Algorithm for computing amax from history
+  ub_tp_comm_overlap: false  # do not turn on because of b/397797926
+  use_flash_attention: true
+  gc_interval: 100
+
+  ## Offloading Activations/Weights to CPU
+  cpu_offloading: false
+  cpu_offloading_num_layers: ${sum:${.num_layers},-1}  # This value should be between [1,num_layers-1] as we don't want to offload the final layer's activations and expose any offloading duration for the final layer
+  cpu_offloading_activations: true
+  cpu_offloading_weights: true
+
+  data:
+    # Path to data must be specified by the user.
+    # Supports List, String and Dictionary
+    # List : can override from the CLI: "model.data.data_prefix=[.5,/raid/data/pile/my-gpt3_00_text_document,.5,/raid/data/pile/my-gpt3_01_text_document]",
+    # Or see example below:
+    # data_prefix:
+    # - .5
+    # - /raid/data/pile/my-gpt3_00_text_document
+    # - .5
+    # - /raid/data/pile/my-gpt3_01_text_document
+    # Dictionary: can override from CLI "model.data.data_prefix"={"train":[1.0, /path/to/data], "validation":/path/to/data, "test":/path/to/test}
+    # Or see example below:
+    # "model.data.data_prefix: {train:[1.0,/path/to/data], validation:[/path/to/data], test:[/path/to/test]}"
+    data_prefix: [1.0, /data/hfbpe_gpt_training_data_text_document]
+    index_mapping_dir: null  # path to save index mapping .npy files, by default will save in the same location as data_prefix
+    data_impl: mmap
+    splits_string: 900,50,50
+    seq_length: ${model.encoder_seq_length}
+    skip_warmup: true
+    num_workers: 2
+    dataloader_type: single  # cyclic
+    reset_position_ids: false  # Reset position ids after end-of-document token
+    reset_attention_mask: false  # Reset attention mask after end-of-document token
+    eod_mask_loss: false  # Mask loss for the end of document tokens
+    validation_drop_last: true  # Set to false if the last partial validation samples is to be consumed
+    no_seqlen_plus_one_input_tokens: false  # Set to True to disable fetching (sequence length + 1) input tokens, instead get (sequence length) input tokens and mask the last token
+    pad_samples_to_global_batch_size: false  # Set to True if you want to pad the last partial batch with -1's to equal global batch size
+    shuffle_documents: true  # Set to False to disable documents shuffling. Sample index will still be shuffled
+
+  # Nsys profiling options
+  nsys_profile:
+    enabled: false
+    start_step: 0  # Global batch to start profiling
+    end_step: 1  # Global batch to end profiling
+    ranks: [0]  # Global rank IDs to profile
+    gen_shape: false  # Generate model and kernel details including input shapes
+
+  memory_profile:
+    enabled: false
+    start_step: 0
+    end_step: 1
+    ranks: [0]
+    output_path: /data  # Must be a dir
+
+  optim:
+    name: distributed_fused_adam  # E.g., fused_adam or set _target_: torch.optim.AdamW field
+    lr: 2e-5
+    weight_decay: 0.01
+    betas:
+    - 0.9
+    - 0.98
+    bucket_cap_mb: 125
+    overlap_grad_sync: true
+    overlap_param_sync: true
+    contiguous_grad_buffer: true
+    contiguous_param_buffer: true
+    sched:
+      name: CosineAnnealing
+      warmup_steps: 400
+      constant_steps: 0
+      min_lr: 2e-6

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/cloudbuild.yml

@@ -0,0 +1,26 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+steps:
+- name: 'gcr.io/cloud-builders/docker'
+  args:
+  - 'build'
+  - '--tag=${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'
+  - '--file=docker/vertex-dist-recipes.Dockerfile'
+  - '.'
+  automapSubstitutions: true
+  env:
+  - 'DOCKER_BUILDKIT=1'
+images:
+- '${_ARTIFACT_REGISTRY}/${_IMAGE_NAME}'

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/patches/24.09/gpu_stats.patch

@@ -0,0 +1,41 @@
+diff --git a/nemo/collections/nlp/parts/megatron_trainer_builder.py b/nemo/collections/nlp/parts/megatron_trainer_builder.py
+index b2c85cde4..a3a9670c3 100644
+--- a/nemo/collections/nlp/parts/megatron_trainer_builder.py
++++ b/nemo/collections/nlp/parts/megatron_trainer_builder.py
+@@ -19,6 +19,7 @@ from lightning_fabric.utilities.exceptions import MisconfigurationException
+ from omegaconf import DictConfig
+ from pytorch_lightning import Trainer
+ from pytorch_lightning.callbacks import ModelSummary
++from pytorch_lightning.callbacks import Callback
+ from pytorch_lightning.plugins.environments import TorchElasticEnvironment
+
+ from nemo.collections.common.metrics.perf_metrics import FLOPsMeasurementCallback
+@@ -38,6 +39,23 @@ from nemo.utils.callbacks.dist_ckpt_io import (
+     AsyncFinalizerCallback,
+     DistributedCheckpointIO,
+ )
++from vmg.util.device_stats import gpu_stats_str
++
++class GpuStatsMon(Callback):
++    def on_train_start(self, trainer, pl_module) -> None:
++        rank=pl_module.global_rank
++        print(f'train_start: {rank=} {gpu_stats_str()}', flush=True)
++
++    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx) -> None:
++        rank=pl_module.global_rank
++        print(f'batch_start: {rank=} {gpu_stats_str()}', flush=True)
++
++    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx) -> None:
++        rank=pl_module.global_rank
++        print(f'batch_end: {rank=} {gpu_stats_str()}', flush=True)
+
+
+ class MegatronTrainerBuilder:
+@@ -178,6 +196,7 @@ class MegatronTrainerBuilder:
+         if self.cfg.get('exp_manager', {}).get('log_tflops_per_sec_per_gpu', True):
+             callbacks.append(FLOPsMeasurementCallback(self.cfg))
+
++        callbacks.append(GpuStatsMon())
+         return callbacks
+
+     def create_trainer(self, callbacks=None) -> Trainer:

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/patches/24.09/local_rank.patch

@@ -0,0 +1,41 @@
+diff -ruN old-datasets/blended_megatron_dataset_builder.py datasets/blended_megatron_dataset_builder.py
+--- old-datasets/blended_megatron_dataset_builder.py	2025-05-02 04:08:45.369199665 +0000
++++ datasets/blended_megatron_dataset_builder.py	2025-05-02 04:10:47.369119891 +0000
+@@ -2,6 +2,7 @@
+
+ import logging
+ import math
++import os
+ from concurrent.futures import ThreadPoolExecutor
+ from typing import Any, Callable, Iterable, List, Optional, Type, Union
+
+@@ -353,7 +354,7 @@
+         num_dataset_builder_threads = self.config.num_dataset_builder_threads
+
+         if torch.distributed.is_initialized():
+-            rank = torch.distributed.get_rank()
++            rank = int(os.getenv("LOCAL_RANK", "0"))
+             # First, build on rank 0
+             if rank == 0:
+                 num_workers = num_dataset_builder_threads
+@@ -475,7 +476,7 @@
+             Optional[Union[DistributedDataset, Iterable]]: The DistributedDataset instantion, the Iterable instantiation, or None
+         """
+         if torch.distributed.is_initialized():
+-            rank = torch.distributed.get_rank()
++            rank = int(os.getenv("LOCAL_RANK", "0"))
+
+             dataset = None
+
+diff -ruN old-datasets/gpt_dataset.py datasets/gpt_dataset.py
+--- old-datasets/gpt_dataset.py	2025-05-02 04:08:45.369199665 +0000
++++ datasets/gpt_dataset.py	2025-05-02 04:09:30.309170278 +0000
+@@ -351,7 +351,7 @@
+
+         if not path_to_cache or (
+             not cache_hit
+-            and (not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0)
++            and (not torch.distributed.is_initialized() or int(os.getenv("LOCAL_RANK", "0")) == 0)
+         ):
+
+             log_single_rank(

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/patches/24.09/nemo2hf.patch

@@ -0,0 +1,13 @@
+diff --git a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
+index 8da15148d..005cae6c9 100644
+--- a/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
++++ b/scripts/checkpoint_converters/convert_llama_nemo_to_hf.py
+@@ -104,6 +104,8 @@ def convert(input_nemo_file, output_hf_file, precision=None, cpu_only=False) ->
+     dummy_trainer = Trainer(devices=1, accelerator='cpu', strategy=NLPDDPStrategy())
+     model_config = MegatronGPTModel.restore_from(input_nemo_file, trainer=dummy_trainer, return_config=True)
+     model_config.tensor_model_parallel_size = 1
++    model_config.virtual_pipeline_model_parallel_size = None
++    model_config.sequence_parallel = False
+     model_config.pipeline_model_parallel_size = 1
+     if cpu_only:
+         map_location = torch.device('cpu')

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/patches/24.09/sigabort.patch

@@ -0,0 +1,24 @@
+diff --git a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
+index bfe8ea359..dfeaf93b5 100644
+--- a/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
++++ b/examples/nlp/language_modeling/tuning/megatron_gpt_finetuning.py
+@@ -13,6 +13,8 @@
+ # limitations under the License.
+
+ import torch.multiprocessing as mp
++import torch.distributed as dist
++
+ from omegaconf.omegaconf import OmegaConf
+
+ from nemo.collections.nlp.models.language_modeling.megatron_gpt_sft_model import MegatronGPTSFTModel
+@@ -76,6 +78,10 @@ def main(cfg) -> None:
+
+     trainer.fit(model)
+
++    if dist.is_available() and dist.is_initialized():
++        dist.barrier()
++        dist.destroy_process_group()
++
+
+ if __name__ == '__main__':
+     main()

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/patches/24.09/throughput_calc.patch

@@ -0,0 +1,13 @@
+diff --git a/src/utils/training_metrics/process_training_results.py b/src/utils/training_metrics/process_training_results.py
+index 3e82a66..e61e1d8 100644
+--- a/src/utils/training_metrics/process_training_results.py
++++ b/src/utils/training_metrics/process_training_results.py
+@@ -134,7 +134,7 @@ def get_average_step_time(file: str, start_step: int, end_step: int) -> float:
+     for line in datajson:
+         if line.get("step") != "PARAMETER":
+             step = line.get("step")
+-            if step >= start_step and step <= end_step:
++            if step >= start_step and step <= end_step and "train_step_timing in s" in line["data"]:
+                 time_step_accumulator += line["data"].get("train_step_timing in s")
+                 num_steps += 1
+     if num_steps == 0:

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/requirements.txt

@@ -0,0 +1,10 @@
+dllogger@git+https://github.com/NVIDIA/dllogger@v1.0.0
+
+# Fixing these libraries versions to avoid conflicting or broken packages.
+immutabledict==4.2.1
+protobuf==5.29.6
+opencv-python-headless==4.11.0.86
+docutils==0.16
+urllib3==2.7.0
+google-cloud-storage==3.0.0
+retrying

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/uninstall.txt

@@ -0,0 +1,18 @@
+# cuml-cu12==24.8.0 was installed in nemo:24.09
+# Removing cuml=24.4.0 to avoid conflicting packages.
+cudf==24.4.0
+cugraph==24.4.0
+cugraph-service-server==24.4.0
+cuml==24.4.0
+dask-cudf==24.4.0
+raft-dask==24.4.0
+cugraph-dgl==24.4.0
+cugraph-pyg==24.4.0
+# The following packages are removed temporarily to avoid conflicting packages
+# and can be brought back if needed.
+tensorrt-llm==0.12.0
+img2dataset==1.45.0
+Sphinx==8.1.3
+sphinxcontrib-bibtex==2.6.3
+torchx==0.7.0
+nemo-run

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/docker/vertex-dist-recipes.Dockerfile

@@ -0,0 +1,66 @@
+# Dockerfile wrapping NeMo.
+#
+# To workaround base nemo docker image using too many layers, we use Multi-stage
+# build to first collect the additional files we'll need.
+FROM alpine:latest AS prep_files
+WORKDIR /workspace
+RUN mkdir -p configs vdt vdt/util
+COPY scripts/*.py vdt/
+COPY scripts/util/*.py vdt/util/
+COPY configs/* configs/
+COPY docker/patches/24.09/* vdt/patches/
+RUN chmod a+rwX -R vdt
+# Copy license.
+RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
+
+# Available tags
+# https://catalog.ngc.nvidia.com/orgs/nvidia/containers/nemo/tags
+# It installs NeMo source code in /opt/NeMo folder, with tag=r2.0.0
+FROM nvcr.io/nvidia/nemo:24.09
+
+RUN apt-get update && apt-get install -y sudo zsh tmux && \
+    rm -rf /var/lib/apt/lists*
+
+RUN 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 google-cloud-sdk -y && \
+    rm -rf /var/lib/apt/lists*
+
+# Install libraries with pip
+ENV PIP_ROOT_USER_ACTION=ignore
+
+# We expect this will be run in the root directory of the vertex-dist-recipes repo
+ARG HOST_SRC_DIR="."
+
+# The pre-installed NeMo introduces a lot of deps conflicts.
+# We uninstall the confilicting libs and reinstall some of them as needed.
+COPY ${HOST_SRC_DIR}/docker/uninstall.txt /tmp/uninstall.txt
+RUN cat /tmp/uninstall.txt | grep -v '#' | xargs pip uninstall -y
+COPY ${HOST_SRC_DIR}/docker/requirements.txt /tmp/requirements.txt
+RUN pip install -r /tmp/requirements.txt
+
+# Make sure there's no inconsistent pip libraries.
+RUN pip check
+
+WORKDIR /workspace
+
+# Copy configs
+COPY ${HOST_SRC_DIR}/configs/* /opt/NeMo/examples/nlp/language_modeling/conf/
+
+# Copy all additional files we need from `prep_files` image.
+COPY --from=prep_files /workspace/ .
+
+# Install for `src/utils/training_metrics/process_training_results.py` to report
+# throughput and MFU numbers.
+RUN git clone https://github.com/AI-Hypercomputer/gpu-recipes.git
+
+# This hack is needed for multi-node training while not using a sharing file system.
+RUN patch --verbose -l -d /opt/megatron-lm/megatron/core/datasets -p1 -i /workspace/vdt/patches/local_rank.patch; \
+    git -C /workspace/gpu-recipes apply /workspace/vdt/patches/throughput_calc.patch; \
+    git -C /opt/NeMo apply /workspace/vdt/patches/nemo2hf.patch; \
+    git -C /opt/NeMo apply /workspace/vdt/patches/sigabort.patch;
+# git -C /opt/NeMo apply /workspace/vdt/patches/gpu_stats.patch;
+
+# Do not put an entrypoint here. Specify the entrypoint in the docker run script.

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/job_config.json

@@ -0,0 +1,16 @@
+{
+  "project_id": "<your_project_id>",
+  "region": "us-central1",
+  "zone": "us-central1-c",
+  "bucket": "<your_bucket",
+  "dataset_bucket": "github-repo/data/third-party/enwiki-latest-pages-articles",
+  "image_uri": "<your_image_uri>",
+  "strategy": "spot",
+  "nodes": "2",
+  "machine_type": "a3-megagpu-8g",
+  "gpu_type": "NVIDIA_H100_MEGA_80GB",
+  "gpus_per_node": "8",
+  "recipe_name": "llama3_1_8b_pretrain_a3mega",
+  "job_prefix": "vertex-ai",
+  "reservation_name": ""
+}

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/requirements.txt

@@ -0,0 +1,49 @@
+absl-py==2.2.2
+annotated-types==0.7.0
+anyio==4.9.0
+black==26.3.1
+cachetools==5.5.2
+certifi==2025.4.26
+charset-normalizer==3.4.2
+click==8.1.8
+docstring_parser==0.16
+google-api-core==2.24.2
+google-auth==2.40.1
+google-cloud-aiplatform==1.133.0
+google-cloud-bigquery==3.31.0
+google-cloud-core==2.4.3
+google-cloud-resource-manager==1.14.2
+google-cloud-storage==2.19.0
+google-crc32c==1.7.1
+google-genai==1.14.0
+google-resumable-media==2.7.2
+googleapis-common-protos==1.70.0
+grpc-google-iam-v1==0.14.2
+grpcio==1.71.0
+grpcio-status==1.71.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+idna==3.15
+mypy_extensions==1.1.0
+numpy==2.2.5
+packaging==25.0
+pathspec==0.12.1
+platformdirs==4.3.8
+proto-plus==1.26.1
+protobuf==5.29.6
+pyasn1==0.6.4
+pyasn1_modules==0.4.2
+pydantic==2.11.4
+pydantic_core==2.33.2
+python-dateutil==2.9.0.post0
+pytz==2025.2
+requests==2.33.0
+rsa==4.9.1
+shapely==2.1.0
+six==1.17.0
+sniffio==1.3.1
+typing-inspection==0.4.0
+typing_extensions==4.13.2
+urllib3==2.7.0
+websockets==15.0.1

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/scripts/launch.py

@@ -0,0 +1,173 @@
+"""Launch script for Vertex distributed training"""
+
+# Copy the sample_job_config.json file to job_config.json
+# to define the job parameters.
+#
+# Run like this:
+#
+#   python3 vertex_dist_train/launch.py --config_file=job_config.json
+#
+
+import datetime
+import json
+import os
+import pprint
+from collections.abc import Sequence
+from typing import Any, List
+
+from absl import app, flags
+from google.cloud import aiplatform
+from google.cloud.aiplatform_v1.types.custom_job import Scheduling
+from pytz import timezone
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("config_file", None, "Path to JSON config file")
+flags.DEFINE_boolean(
+    "debug", False, "Debug mode: just print the command, don't run it."
+)
+
+
+def launch_job(
+    job_name: str,
+    project: str,
+    region: str,
+    gcs_bucket: str,
+    image_uri: str,
+    entrypoint_cmd: List[str],
+    trainer_args: List[Any],
+    num_nodes: int,
+    machine_type: str,
+    num_gpus_per_node: int,
+    gpu_type: str,
+    strategy: str,
+    reservation_name: str = "",
+):
+    assert strategy in ("dws", "spot", "reservation")
+    aiplatform.init(
+        project=project, location=region, staging_bucket=gcs_bucket
+    )
+
+    train_job = aiplatform.CustomContainerTrainingJob(
+        display_name=job_name,
+        container_uri=image_uri,
+        command=entrypoint_cmd,
+    )
+
+    job_args = dict(
+        args=trainer_args,
+        enable_web_access=True,
+        replica_count=num_nodes,
+        machine_type=machine_type,
+        accelerator_type=gpu_type,
+        accelerator_count=num_gpus_per_node,
+        boot_disk_size_gb=1000,
+        restart_job_on_worker_restart=True,
+        #restart_job_on_worker_restart=False,
+    )
+
+    if strategy == "spot":
+        job_args.update({"scheduling_strategy": Scheduling.Strategy.SPOT.name})
+    elif strategy == "dws":
+        job_args.update(
+            {"scheduling_strategy": Scheduling.Strategy.FLEX_START.name}
+        )
+    elif strategy == "reservation":
+        assert reservation_name != "", (
+            "If using a reservation, provide the reservation_name in the "
+            "format `projects/{project_id_or_number}/zones/{zone}/"
+            "reservations/{reservation_name}`"
+        )
+        job_args.update(
+            {
+                "reservation_affinity_type": "SPECIFIC_RESERVATION",
+                "reservation_affinity_key": "compute.googleapis.com/reservation-name",
+                "reservation_affinity_values": [reservation_name],
+            }
+        )
+
+    pprint.pprint(job_args)
+    if not FLAGS.debug:
+        train_job.submit(**job_args)
+
+
+def main(argv: Sequence[str]) -> None:
+    config_file_path = FLAGS.config_file
+    print(f"Reading job config from {config_file_path}")
+    with open(config_file_path, encoding="utf-8") as config_file:
+        config = json.load(config_file)
+
+    project_id = config["project_id"]
+    region = config["region"]
+    zone = config["zone"]
+    bucket = config["bucket"]
+    dataset_bucket = config["dataset_bucket"]
+    n_nodes = int(config["nodes"])
+    machine_type = config["machine_type"]
+    num_gpus_per_node = int(config["gpus_per_node"])
+    gpu_type = config["gpu_type"]
+    reservation_name = config.get("reservation_name")
+    reservation_full_name = (
+        f"projects/{project_id}/zones/{zone}/reservations/{reservation_name}"
+        if "reservation_name" in config
+        else ""
+    )
+
+    strategy = config["strategy"]
+    recipe_name = config["recipe_name"]
+    job_prefix = config["job_prefix"]
+    image_uri = config["image_uri"]
+
+    # Job name
+    timestamp = (
+        datetime.datetime.now()
+        .astimezone(timezone("US/Pacific"))
+        .strftime("%Y%m%d_%H%M%S")
+    )
+    job_name = f"{recipe_name}-{timestamp}"
+    if job_prefix:
+        job_name = f"{job_prefix}-{job_name}"
+
+    base_output_dir = os.path.join("/gcs", bucket, job_name)
+
+    # Training command and args
+    entrypoint_cmd = ["python3", "vdt/run.py"]
+
+    dataset_bucket = f"gs://{config['dataset_bucket']}"
+
+    trainer_args = [
+        f"--train_data_gcs={dataset_bucket}",
+        "/opt/NeMo/examples/nlp/language_modeling/megatron_gpt_pretraining.py",
+        "--config-path=conf/",
+        f"--config-name={recipe_name}.yaml",
+        f"exp_manager.explicit_log_dir={base_output_dir}",
+        f"exp_manager.dllogger_logger_kwargs.json_file={base_output_dir}/dllogger.json",
+        "+exp_manager.create_tensorboard_logger=true",
+        "exp_manager.create_checkpoint_callback=false",
+        f"trainer.num_nodes={n_nodes}",
+        f"trainer.devices={num_gpus_per_node}",
+        "trainer.max_steps=10",
+        "trainer.log_every_n_steps=1",
+        "model.tokenizer.vocab_file=/data/gpt2-vocab.json",
+        "model.tokenizer.merge_file=/data/gpt2-merges.txt",
+        "model.data.data_prefix=[1.0,/data/hfbpe_gpt_training_data_text_document]",
+    ]
+
+    launch_job(
+        job_name=job_name,
+        project=project_id,
+        region=region,
+        gcs_bucket=bucket,
+        image_uri=image_uri,
+        entrypoint_cmd=entrypoint_cmd,
+        trainer_args=trainer_args,
+        num_nodes=n_nodes,
+        machine_type=machine_type,
+        num_gpus_per_node=num_gpus_per_node,
+        gpu_type=gpu_type,
+        strategy=strategy,
+        reservation_name=reservation_full_name,
+    )
+
+
+if __name__ == "__main__":
+    app.run(main)

community-content/vertex-distributed-training/a3mega/llama-3-8b-nemo-pretraining/scripts/run.py

@@ -0,0 +1,85 @@
+"""Entrypoint for Vertex Distributed Training container."""
+
+import argparse
+import os
+import sys
+from collections.abc import Sequence
+from subprocess import STDOUT, check_output, run
+
+from absl import app, flags, logging
+from util import cluster_spec
+
+from retrying import retry
+
+# PyTorch barrier call which synchronizes all of the nodes before launching the training process.
+# This makes sure that processes will block until all processes are ready.
+# Improves the reliability of spot VM usage for multi-node training jobs
+
+@retry(stop_max_attempt_number=100, wait_exponential_multiplier=1000)
+def barrier_with_retry() -> None:
+    import torch
+    logging.info("Starting barrier on RANK {}".format(os.environ["RANK"]))
+    torch.distributed.init_process_group()
+    torch.distributed.barrier()
+    torch.distributed.destroy_process_group()
+    logging.info("Finished barrier on RANK {}".format(os.environ["RANK"]))
+
+def main(unused_argv: Sequence[str]) -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--train_data_gcs",
+        type=str,
+        help="Download training data from gcs path",
+    )
+    args, unknown = parser.parse_known_args()
+
+    for key, val in os.environ.items():
+        logging.info("ENV %s=%s", key, val)
+
+    if args.train_data_gcs:
+        local_dir = "/data"
+        if not os.path.exists(local_dir):
+            os.mkdir(local_dir)
+        logging.info("downloading %s to %s...", args.train_data_gcs, local_dir)
+        check_output(
+            [
+                "gcloud",
+                "storage",
+                "cp",
+                "-r",
+                f"{args.train_data_gcs}/*",
+                local_dir,
+            ],
+            stderr=STDOUT,
+        )
+        logging.info("%s downloaded.", args.train_data_gcs)
+
+    primary_node_addr, primary_node_port, node_rank, num_nodes = (
+        cluster_spec.get_cluster_spec()
+    )
+
+    cmd = [
+        "torchrun",
+        "--nproc-per-node=8",
+        f"--nnodes={num_nodes}",
+        f"--node_rank={node_rank}",
+    ]
+    if num_nodes > 1:
+        cmd += [
+            "--max-restarts=3",
+            "--rdzv-backend=static",
+            f'--rdzv_id={os.getenv("CLOUD_ML_JOB_ID", primary_node_port)}',
+            f"--rdzv-endpoint={primary_node_addr}:{primary_node_port}",
+        ]
+    cmd += unknown
+
+    logging.info("launching with cmd: \n%s", " \\\n".join(cmd))
+    barrier_with_retry()
+    run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
+
+
+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-distributed-training/a3mega/llama-3-8b-nemo-pretraining/scripts/util/cluster_spec.py

@@ -0,0 +1,81 @@
+"""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-distributed-training/a3mega/llama-3-8b-nemo-pretraining/scripts/util/cluster_spec_test.py

@@ -0,0 +1,59 @@
+"""Add tests for cluster_spec.py."""
+
+import os
+
+from . import cluster_spec
+
+
+# TODO(styer): Use pytest instead
+class ClusterSpecTest(googletest.TestCase):
+
+    def setUp(self):
+        super().setUp()
+        self.curr_env_var = os.environ.copy()
+
+    def tearDown(self):
+        super().tearDown()
+        os.environ = self.curr_env_var
+
+    def test_get_cluster_spec_from_env_vars(self):
+        os.environ["CLUSTER_SPEC"] = ""
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "8080"
+        os.environ["RANK"] = "0"
+        os.environ["NNODES"] = "2"
+        cluster_info = cluster_spec.get_cluster_spec()
+        self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
+        self.assertEqual(cluster_info.primary_node_port, "8080")
+        self.assertEqual(cluster_info.node_rank, 0)
+        self.assertEqual(cluster_info.num_nodes, 2)
+
+    def test_get_cluster_spec_from_cluster_spec(self):
+        os.environ[
+            "CLUSTER_SPEC"
+        ] = """
+    {
+      "cluster": {
+        "workerpool0": [
+          "127.0.0.1:8080"
+        ],
+        "workerpool1": [
+          "127.0.0.2:8080",
+          "127.0.0.3:8080"
+        ]
+      },
+      "task": {
+        "type": "workerpool1",
+        "index": 0
+      }
+    }
+    """
+        cluster_info = cluster_spec.get_cluster_spec()
+        self.assertEqual(cluster_info.primary_node_addr, "127.0.0.1")
+        self.assertEqual(cluster_info.primary_node_port, "8080")
+        self.assertEqual(cluster_info.node_rank, 1)
+        self.assertEqual(cluster_info.num_nodes, 3)
+
+
+if __name__ == "__main__":
+    googletest.main()

community-content/vertex_cpr_samples/sklearn/predictor_SGDClassifier.py

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

community-content/vertex_cpr_samples/torch/predictor_resnet.py

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

community-content/vertex_cpr_samples/xgboost/predictor_xgbranker.py

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

community-content/vertex_model_garden/model_oss/cloudnerf/local_colmap_and_resize.sh

@@ -66,7 +66,7 @@ mkdir -p "$local_folder"
 mkdir -p "$output_folder"
 
 # Download the content from the GCS URI
-gsutil -m cp -r "$gcs_dataset_path"/* "$local_folder/"
+gcloud storage cp --recursive "$gcs_dataset_path"/* "$local_folder/"
 
 # Process files in the local folder
 for file in "$local_folder"/*; do
@@ -122,23 +122,23 @@ cp -r "$output_folder" "$images_folder"/images_2
 pushd "$images_folder"/images_2
 ls | xargs -P 8 -I {} mogrify -resize 50% {}
 popd
-gsutil -m cp -r "$images_folder"/images_2/*  "$gcs_experiment_path"/data/images_2
+gcloud storage cp --recursive "$images_folder"/images_2/*  "$gcs_experiment_path"/data/images_2
 
 cp -r "$output_folder" "$images_folder"/images_4
 pushd "$images_folder"/images_4
 ls | xargs -P 8 -I {} mogrify -resize 25% {}
 popd
-gsutil -m cp -r "$images_folder"/images_4/*  "$gcs_experiment_path"/data/images_4
+gcloud storage cp --recursive "$images_folder"/images_4/*  "$gcs_experiment_path"/data/images_4
 
 cp -r "$output_folder" "$images_folder"/images_8
 pushd "$images_folder"/images_8
 ls | xargs -P 8 -I {} mogrify -resize 12.5% {}
 popd
-gsutil -m cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
+gcloud storage cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
 
 # Copy images and sparse reconstruction files to gcs experiment folder.
-gsutil -m cp "$images_folder"/images/*  "$gcs_experiment_path"/data/images
-gsutil -m cp -r "$local_folder"/sparse  "$gcs_experiment_path"/data
-gsutil -m cp "$local_folder"/database.db  "$gcs_experiment_path"/data
+gcloud storage cp "$images_folder"/images/*  "$gcs_experiment_path"/data/images
+gcloud storage cp --recursive "$local_folder"/sparse  "$gcs_experiment_path"/data
+gcloud storage cp "$local_folder"/database.db  "$gcs_experiment_path"/data
 
 echo "Processing complete."

community-content/vertex_model_garden/model_oss/cloudnerf/render.sh

@@ -99,14 +99,14 @@ create_dir_if_not_exists "$CHECKPOINTS_PATH"
 touch "$local_experiment_path/$exp_folder_name/log_render.txt"
 
 # Copy experiment from GCS bucket to local
-gsutil -m cp -r "${args[-gcs_experiment_path]}/data" "$local_experiment_path/$exp_folder_name" || exit 1
-gsutil -m cp -r "${args[-gcs_experiment_path]}/checkpoints/${training_job_name}/*" "$CHECKPOINTS_PATH" || exit 1
+gcloud storage cp --recursive "${args[-gcs_experiment_path]}/data" "$local_experiment_path/$exp_folder_name" || exit 1
+gcloud storage cp --recursive "${args[-gcs_experiment_path]}/checkpoints/${training_job_name}/*" "$CHECKPOINTS_PATH" || exit 1
 
 # Check and copy keyframes file.
 if [[ -n ${args[-gcs_keyframes_file]} ]]; then
   keyframes_file_basename=$(basename "${args[-gcs_keyframes_file]}")
   local_keyframes_file="$local_dataset_path/$keyframes_file_basename"
-  gsutil cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
+  gcloud storage cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
   echo "Local keyframe file: $local_keyframes_file"
   launch_rendering "$local_keyframes_file"
 else
@@ -114,4 +114,4 @@ else
 fi
 
 # Copy rendered data back to GCS.
-gsutil -m cp -r "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"
+gcloud storage cp --recursive "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"

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

@@ -74,7 +74,7 @@ create_dir_if_not_exists "$local_experiment_path"
 create_dir_if_not_exists "$local_experiment_path/$scene_folder_name"
 
 # Copy experiment from GCS bucket to local.
-gsutil -m cp -r "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
+gcloud storage cp --recursive "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
 
 echo "GCS Experiment: $gcs_experiment_path"
 echo "Gin Config File: $gin_config_file"
@@ -89,6 +89,6 @@ accelerate launch train.py --gin_configs="$gin_config_file" \
   --gin_bindings="Config.factor = ${factor}" \
   --gin_bindings="Config.max_steps = ${max_training_steps}"
 
-gsutil -m rm -r "${gcs_experiment_path}/checkpoints/${training_job_name}"
-gsutil -m cp -r "$local_experiment_path/$scene_folder_name/config.gin" "${gcs_experiment_path}/${training_job_name}_config.gin"
-gsutil -m cp -r "$local_experiment_path/$scene_folder_name/checkpoints/*/*" "${gcs_experiment_path}/checkpoints/${training_job_name}"
+gcloud storage rm --recursive "${gcs_experiment_path}/checkpoints/${training_job_name}"
+gcloud storage cp --recursive "$local_experiment_path/$scene_folder_name/config.gin" "${gcs_experiment_path}/${training_job_name}_config.gin"
+gcloud storage cp --recursive "$local_experiment_path/$scene_folder_name/checkpoints/*/*" "${gcs_experiment_path}/checkpoints/${training_job_name}"

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

@@ -1,13 +1,16 @@
 """Common util functions for notebook."""
 
 import base64
+from collections.abc import Sequence
 import datetime
 import io
 import json
 import os
 import subprocess
-from typing import Any, Dict, Sequence
+import time
+from typing import Any
 
+from google import auth
 from google.cloud import storage
 import matplotlib.pyplot as plt
 import numpy as np
@@ -281,7 +284,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:
@@ -331,6 +334,7 @@ def vqa_predict(
     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.
@@ -354,7 +358,9 @@ def vqa_predict(
         "image": resized_image_base64,
     })
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return [pred.get("response") for pred in response.predictions]
 
 
@@ -364,6 +370,7 @@ def caption_predict(
     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.
@@ -378,7 +385,9 @@ def caption_predict(
     instance["prompt"] = caption_prompt_format.format(language_code)
 
   instances = [instance]
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -387,6 +396,7 @@ def ocr_predict(
     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.
@@ -398,7 +408,9 @@ def ocr_predict(
     instance["prompt"] = ocr_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -407,6 +419,7 @@ def detect_predict(
     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.
@@ -418,7 +431,9 @@ def detect_predict(
     instance["prompt"] = detect_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -495,6 +510,17 @@ 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")
@@ -504,12 +530,13 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
         return int(region_data["effectiveLimit"])
       else:
         return 0
-  return -1
+  return global_quota
 
 
 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:
@@ -519,6 +546,7 @@ def get_resource_id(
     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.
@@ -528,11 +556,15 @@ def get_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",
   }
@@ -547,6 +579,10 @@ def get_resource_id(
   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
+  }
   serving_accelerator_map = {
       key: f"custom_model_serving_{accelerator_suffix_map[key]}"
       for key in accelerator_suffix_map
@@ -575,8 +611,11 @@ def get_resource_id(
   else:
     if is_dynamic_workload_scheduler:
       raise ValueError("Dynamic Workload Scheduler does not work for serving.")
-    if accelerator_type in serving_accelerator_map:
-      return serving_accelerator_map[accelerator_type]
+    accelerator_map = (
+        spot_serving_accelerator_map if is_spot else serving_accelerator_map
+    )
+    if accelerator_type in accelerator_map:
+      return accelerator_map[accelerator_type]
     else:
       raise ValueError(
           f"Could not find accelerator type: {accelerator_type} for serving."
@@ -589,13 +628,28 @@ 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,
-):
-  """Checks if the project and the region has the required quota."""
+) -> 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,
   )
@@ -618,3 +672,76 @@ 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

@@ -7,7 +7,7 @@ import json
 import multiprocessing
 import os
 import subprocess
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict, Tuple, Union
 from absl import logging
 import accelerate
 import datasets
@@ -20,8 +20,10 @@ 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/templates"
+    "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"
@@ -68,7 +70,9 @@ def force_gcs_fuse_path(gcs_uri: str) -> str:
 
 
 def download_gcs_uri_to_local(
-    gcs_uri: str, destination_dir: str = LOCAL_BASE_MODEL_DIR
+    gcs_uri: str,
+    destination_dir: str = LOCAL_BASE_MODEL_DIR,
+    check_path_exists: bool = True,
 ) -> str:
   """Downloads GCS URI to local.
 
@@ -79,6 +83,7 @@ def download_gcs_uri_to_local(
   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.
@@ -87,7 +92,7 @@ def download_gcs_uri_to_local(
       destination_dir,
       os.path.basename(os.path.normpath(gcs_uri)),
   )
-  if os.path.exists(target):
+  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:
@@ -97,10 +102,10 @@ def download_gcs_uri_to_local(
     if not os.path.exists(destination_dir):
       os.mkdir(destination_dir)
     subprocess.check_output([
-        "gsutil",
-        "-m",
+        "gcloud",
+        "storage",
         "cp",
-        "-r",
+        "--recursive",
         gcs_uri,
         destination_dir,
     ])
@@ -171,57 +176,6 @@ def _format_template_fn(
   """
   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()}
-      format_str = (
-          template_json[_PROMPT_INPUT_KEY]
-          if format_dict.get(input_column)
-          else template_json[_PROMPT_NO_INPUT_KEY]
-      )
-      return {input_column: format_str.format(**format_dict)}
-
-    return format_fn
-  elif (
-      _PROMPT_INPUT_KEY in template_json
-      or _PROMPT_NO_INPUT_KEY in template_json
-  ):
-    raise ValueError(
-        "chat_template templates do not support input/no_input 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]
-    return lambda example: {
-        input_column: tokenizer.apply_chat_template(
-            example[input_column], tokenize=False, add_generation_prompt=False
-        )
-    }
-
-
-def _format_template_fn_notebook(
-    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]:
@@ -366,13 +320,140 @@ def _get_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,
-) -> Any:
+) -> Tuple[Any, Any]:
   """Loads dataset with templates.
 
   Args:
@@ -386,19 +467,15 @@ def load_dataset_with_template(
     tokenizer: The tokenizer to use for chat_template templates.
 
   Returns:
-    A dataset compatible with the template.
+    The raw dataset and the dataset compatible with the template.
   """
-  dataset = _get_dataset(dataset_name, split=split)
+  raw = _get_dataset(dataset_name, split=split)
   if template:
-    dataset = dataset.map(
-        _format_template_fn(
-            template,
-            input_column=input_column,
-            tokenizer=tokenizer,
-        )
-    )
+    templated = format_dataset(raw, input_column, template, tokenizer)
+  else:
+    templated = None
 
-  return dataset
+  return raw, templated
 
 
 def validate_dataset_with_template(
@@ -407,6 +484,7 @@ def validate_dataset_with_template(
     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,
@@ -428,6 +506,7 @@ def validate_dataset_with_template(
     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.
@@ -470,17 +549,22 @@ def validate_dataset_with_template(
         f" https://github.com/GoogleCloudPlatform/{_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME}/tree/main/{_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR}."
     )
 
-  _get_dataset(dataset_name, split, num_proc).map(
-      _format_template_fn_notebook(
-          template_path,
-          input_column=input_column,
-          tokenizer=tokenizer,
-      )
+  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/train/test_utils/command_builder.py

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

@@ -0,0 +1,471 @@
+"""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/tests/docker_command_builder.py

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

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

@@ -1,152 +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_tasks = 'builtin_eval'
-    self.task_cmd_builder.eval_metric_name = 'loss'
-    self.task_cmd_builder.per_device_batch_size = 1
-    self.task_cmd_builder.gradient_accumulation_steps = 8
-    self.task_cmd_builder.lora_rank = 16
-    self.task_cmd_builder.lora_alpha = 32
-    self.task_cmd_builder.lora_dropout = 0.05
-    self.task_cmd_builder.learning_rate = 5e-5
-    self.task_cmd_builder.num_epochs = 2.0
-    self.task_cmd_builder.warmup_ratio = 0.01
-    self.task_cmd_builder.max_steps = -1
-    self.task_cmd_builder.save_steps = 10
-    self.task_cmd_builder.eval_steps = 10
-    self.task_cmd_builder.max_seq_length = 4096
-    self.task_cmd_builder.load_precision = '4bit'
-    self.task_cmd_builder.gradient_checkpointing = True
-    self.task_cmd_builder.completion_only = True
-    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
-    self.task_cmd_builder.report_to = 'tensorboard'
-
-  def setup_output_dir(self, testcase_name: str):
-    testcase_output_dir = os.path.join(
-        self.test_suite_output_dir, testcase_name
-    )
-    self.task_cmd_builder.ckpt_dir = os.path.join(
-        testcase_output_dir, 'adapter'
-    )
-    self.task_cmd_builder.logging_dir = os.path.join(
-        testcase_output_dir, 'logs'
-    )
-    self.task_cmd_builder.merged_model_dir = os.path.join(
-        testcase_output_dir, 'merged'
-    )
-
-  @parameterized.named_parameters(
-      ('llama3-8b', 'llama3-8b-hf'),
-      ('llama3.1-8b', 'llama3.1-8b-hf'),
-  )
-  def test_8b_model_deepspeed(self, model_name):
-    self.setup_output_dir(f'test_deepspeed_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'peft_train_sample.jsonl'
-    )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
-    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
-        'peft_eval_sample.jsonl'
-    )
-    self.task_cmd_builder.eval_split_name = 'train'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('llama3-70b', 'llama3-70b-hf'),
-      ('llama3.1-70b', 'llama3.1-70b-hf'),
-  )
-  def test_70b_model_deepspeed(self, model_name):
-    self.setup_output_dir(f'test_deepspeed_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split_name = 'test'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-  @parameterized.named_parameters(
-      ('llama3-70b', 'llama3-70b-hf'),
-      ('llama3.1-70b', 'llama3.1-70b-hf'),
-  )
-  def test_70b_model_fsdp(self, model_name):
-    self.setup_output_dir(f'test_fsdp_{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
-    )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
-    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
-    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
-    self.task_cmd_builder.eval_split_name = 'test'
-    self.task_cmd_builder.eval_instruct_column = (
-        self.task_cmd_builder.instruct_column
-    )
-    self.task_cmd_builder.eval_template = self.task_cmd_builder.template
-
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
-
-    self.assertEqual(self.run_cmd(), 0)
-
-
-if __name__ == '__main__':
-  absltest.main()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,140 @@
+"""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,5 +1,7 @@
 """Different trainer callbacks for PEFT Trainer."""
 
+from collections.abc import MutableMapping
+import math
 import time
 
 from absl import logging
@@ -9,7 +11,7 @@ from transformers.trainer_callback import TrainerCallback
 from transformers.trainer_callback import TrainerControl
 from transformers.trainer_callback import TrainerState
 
-from vertex_vision_model_garden_peft.train.vmg import utils
+from util import device_stats
 
 
 class TrainerStatsCallback(TrainerCallback):
@@ -25,6 +27,30 @@ class TrainerStatsCallback(TrainerCallback):
     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,
@@ -35,21 +61,29 @@ class TrainerStatsCallback(TrainerCallback):
     if self._partial_state.is_main_process:
       if state.global_step == 1:
         self._prev_time = time.time()
-        delta_t = float('nan')
+        self._prev_num_token = state.num_input_tokens_seen
+        throughput = 0.0
       else:
         cur_time = time.time()
-        delta_t = cur_time - self._prev_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._avg_throughput += (delta_t - self._avg_throughput) / (
+        self._prev_num_token = cur_num_token
+        self._avg_throughput += (throughput - self._avg_throughput) / (
             state.global_step - 1
         )
 
-      gpu_stats = utils.gpu_stats()
-      self._peak_mem = max(gpu_stats.total_mem, self._peak_mem)
+      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: %s, throughput: %.2f s/it',
-          utils.gpu_stats_str(gpu_stats),
-          delta_t,
+          '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(
@@ -61,7 +95,11 @@ class TrainerStatsCallback(TrainerCallback):
   ):
     if self._partial_state.is_main_process:
       self._start_time = time.time()
-      logging.info('on_train_begin: %s', utils.gpu_stats_str())
+      logging.info(
+          'on_train_begin: %s, %s',
+          device_stats.gpu_stats_str(),
+          device_stats.cpu_stats_str(),
+      )
 
   def on_train_end(
       self,
@@ -72,15 +110,17 @@ class TrainerStatsCallback(TrainerCallback):
   ):
     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: %.2f s/it, peak_mem: %.2f GB',
+          'training time %.2f s, throughput (including overhead, e.g., ckpt'
+          ' saving): %.2f token/s, peak_mem: %.2f GB',
           train_time,
-          self._avg_throughput,
+          throughput,
           self._peak_mem,
       )
       if self._filename:
         with open(self._filename, 'a') as out_f:
           out_f.write(
-              f'{self._max_seq_length/1024.0:.1f}k | {self._peak_mem:.2f} |'
+              f'{self._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/gemma2_fsdp_8gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: 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/qwen2_fsdp_8gpu.yaml

@@ -0,0 +1,28 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: 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/dockerfile/merge_env.yaml

@@ -31,7 +31,7 @@ dependencies:
 - pip:
   - --extra-index-url https://download.pytorch.org/whl/cu121
   - absl-py==2.1.0
-  - accelerate==0.33.0  # Needed for fp8
+  - accelerate==0.34.2  # Needed for fp8
   - datasets==2.19.2
   - fbgemm-gpu==0.8.0+cu121  # Needed for fp8
   - kfp==2.5.0
@@ -39,5 +39,5 @@ dependencies:
   - protobuf==3.20.3
   - pynvml==11.5.3
   - torch==2.4.0+cu121  # Needed for fp8
-  - transformers==4.43.1
-  - trl==0.9.6
+  - transformers==4.47.1
+  - trl==0.11.2

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

@@ -5,23 +5,28 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 
 # keep sorted
-accelerate==0.31.0
+accelerate==0.34.2
 auto_gptq==0.7.1+cu118
-autoawq==0.2.5
+autoawq==0.2.8
 bitsandbytes==0.43.2
 cloudml-hypertune==0.1.0.dev6
-datasets==2.19.2
-deepspeed==0.14.4
-diffusers==0.25.1
+datasets==2.20.0
+deepspeed==0.15.2
+diffusers==0.38.0
+evaluate==0.4.3
 fsspec==2024.3.1
 gcsfs==2024.3.1
-lm_eval==0.4.3
+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.43.1
-trl==0.9.6
+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

@@ -23,6 +23,14 @@ RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples
 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
@@ -58,15 +66,14 @@ WORKDIR /diffusers/examples
 
 RUN mkdir -p ./vertex_vision_model_garden_peft/
 COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
-# custom `lm_eval` task.
-ARG LM_EVAL_DIR=$(python -c 'import site; print(site.getsitepackages()[0])')/lm_eval
-RUN mkdir -p $LM_EVAL_DIR/tasks/vertex && \
-  mv ./vertex_vision_model_garden_peft/custom_loglikelihood.yaml $LM_EVAL_DIR/tasks/vertex/
 COPY model_oss/peft/train/vmg/*.py ./vertex_vision_model_garden_peft/train/vmg/
 COPY model_oss/peft/train/vmg/templates /diffusers/examples/util/templates
-COPY model_oss/util /diffusers/examples/util
+COPY model_oss/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/tests/*.py ./vertex_vision_model_garden_peft/tests/
+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/

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

@@ -1,183 +1,284 @@
 """Library for running evaluations during training."""
 
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 import dataclasses
-from typing import Any, Optional, Type
+import string
+from typing import Type
 
 from absl import logging
-import datasets
-from lm_eval import evaluator
-from lm_eval import tasks
-from lm_eval import utils
-from lm_eval.api import model as lm_model
-from lm_eval.api import registry
-from lm_eval.models import huggingface
-from peft import peft_model
+import evaluate
+import numpy as np
+import torch
 import transformers
-from transformers import trainer
 
 from util import dataset_validation_util
 from util import constants
 
 
-_DESCRIPTION_EVALUATION = "evaluation"
-_BUILTIN_EVAL_TASK = "builtin_eval"
+_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
-  tasks: list[str]
   per_device_batch_size: int
-  num_fewshot: Optional[int]
-  limit: Optional[float]
-  metric_name: str
+  limit: float | None
+  metric_name: Sequence[str]
   tokenize_dataset: bool
   dataset_path: str = ""
   split: str = "test"
   template: str = ""
-  column: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET
-
-
-class PeftCausalLMModel(huggingface.HFLM):
-  """PeftCausalLMModel that supports loading an in-memory model."""
-
-  AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
-
-  def __init__(
-      self,
-      model: peft_model.PeftModelForCausalLM,
-      tokenizer: transformers.PreTrainedTokenizerBase,
-      batch_size_per_gpu: int,
-  ):
-    lm_model.LM.__init__(self)
-    self._model = model
-    self.tokenizer = tokenizer
-    self.vocab_size = tokenizer.vocab_size
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-    self._config = model.config
-    self.batch_size_per_gpu = batch_size_per_gpu
-    self._device = model.device
-    self._max_length = None  # Will be automatically determined from config.
-    self._add_special_tokens = (
-        None  # Will be automatically determined from AUTO_MODEL_CLASS.
-    )
+  column: str = constants.DEFAULT_TRAIN_COLUMN
+  metric_for_best_model: str | None = None
 
 
 def create_trainer(
     cls: Type[transformers.Trainer],
-    eval_config: Optional[EvalConfig],
-    tokenizer: Optional[transformers.PreTrainedTokenizerBase],
-    args: trainer.TrainingArguments,
+    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."""
+  """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
 
-  if eval_config.tasks == [_BUILTIN_EVAL_TASK]:
-    try:
-      eval_dataset = dataset_validation_util.load_dataset_with_template(
-          dataset_name=eval_config.dataset_path,
-          split=eval_config.split,
+  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,
-          template=eval_config.template,
+          max_seq_length=kwargs["max_seq_length"],
           tokenizer=tokenizer,
       )
-      if eval_config.limit is not None:
-        if eval_config.limit >= 1:
-          limit = int(eval_config.limit)
-        else:
-          limit = int(eval_config.limit * len(eval_dataset))
-        eval_dataset = eval_dataset.select(range(limit))
-      if eval_config.tokenize_dataset:
-        eval_dataset = eval_dataset.map(
-            lambda samples: tokenizer(samples[eval_config.column])
-        )
-      kwargs["eval_dataset"] = eval_dataset
-    except (OSError, ValueError, IndexError) as e:
-      logging.warning(
-          "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
-          eval_config.dataset_path,
-          e,
+    if eval_config.tokenize_dataset:
+      eval_dataset = eval_dataset.map(
+          lambda samples: tokenizer(samples[eval_config.column])
       )
-      del args.evaluation_strategy
-      del args.eval_steps
-      del args.per_device_eval_batch_size
-    return cls(args=args, **kwargs)
+    kwargs["eval_dataset"] = eval_dataset
+  except (OSError, ValueError, IndexError) as e:
+    logging.warning(
+        "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
+        eval_config.dataset_path,
+        e,
+    )
+    del args.evaluation_strategy
+    del args.eval_steps
+    del args.per_device_eval_batch_size
+  return cls(args=args, **kwargs)
 
-  class LMEvalTrainer(cls):
-    """Trainer with lm_eval injected as the eval library."""
 
-    def __init__(self, **kwargs):
-      super().__init__(**kwargs)
-      task_names = utils.pattern_match(eval_config.tasks, registry.ALL_TASKS)
-      logging.info("Selected Eval Tasks: %s", task_names)
-      task_args = {}
-      if eval_config.num_fewshot is not None:
-        task_args["num_fewshot"] = eval_config.num_fewshot
-      if eval_config.dataset_path:
-        task_args["dataset_path"] = "json"
-        task_args["dataset_kwargs"] = {
-            "data_files": {"test": eval_config.dataset_path},
-        }
-      self._eval_task_dict = tasks.get_task_dict(task_names, **task_args)
+def _cleanup_text(text: str) -> str:
+  """Cleans up the prediction and references text.
 
-    def evaluation_loop(
-        self,
-        dataloader: trainer.DataLoader,
-        description: str,
-        prediction_loss_only: Optional[bool] = None,
-        ignore_keys: Optional[list[str]] = None,
-        metric_key_prefix: str = "eval",
-    ) -> trainer.EvalLoopOutput:
-      """Custom evaluation loop that invokes lm_eval."""
-      if description.lower() != _DESCRIPTION_EVALUATION:
-        return super().evaluation_loop(
-            dataloader,
-            description,
-            prediction_loss_only,
-            ignore_keys,
-            metric_key_prefix,
-        )
+  Args:
+    text: The text to clean up.
 
-      model = self._wrap_model(self.model, training=False)
-      lm = PeftCausalLMModel(
-          model,
-          self.tokenizer or self.data_collator.tokenizer,
-          eval_config.per_device_batch_size,
-      )
-      results: dict[str, Any] = evaluator.evaluate(
-          lm=lm,
-          task_dict=self._eval_task_dict,
-          limit=eval_config.limit,
-      )["results"]
-      metric_name = eval_config.metric_name
-      # Compute average value if there are multiple tasks.
-      metric_values: list[float] = []
-      for result in results.values():
-        for key, value in result.items():
-          if key.split(",")[0] == metric_name:
-            metric_values.append(value)
-      if not metric_values:
-        raise ValueError(
-            f"Metric {metric_name} not found in eval response: {results}"
-        )
-      metric_average = sum(metric_values) / len(metric_values)
-      logging.info("%s value: %f\n%s", metric_name, metric_average, results)
-      return trainer.EvalLoopOutput(
-          # Only metrics field is set. Other fields are dummy values.
-          predictions=None,
-          label_ids=None,
-          metrics={f"{metric_key_prefix}_{metric_name}": metric_average},
-          num_samples=0,
-      )
+  Returns:
+    Cleaned up text.
+  """
+  text = text.translate(_STRING_TRANSLATOR)
+  text = text.strip()
+  text = " ".join(text.split())
+  return text.lower()
 
-  # Use empty eval dataset as a placeholder.
-  return LMEvalTrainer(
-      args=args, eval_dataset=datasets.Dataset.from_dict({"test": []}), **kwargs
+
+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,10 +1,10 @@
 """Instruct/Chat with LoRA models."""
 
-import dataclasses
+from collections.abc import Callable, Mapping, Sequence
 import datetime
 import json
 import os
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 import warnings
 
 from absl import app
@@ -13,17 +13,17 @@ from absl import logging
 from accelerate import DistributedType
 from accelerate import PartialState
 import bitsandbytes as bnb
-import hypertune
+import evaluate
 from peft import get_peft_model
 from peft import LoraConfig
 import torch
-from transformers import AutoModelForCausalLM
-from transformers import TrainingArguments
-from trl import DataCollatorForCompletionOnlyLM
-from trl import SFTTrainer
+import 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
@@ -31,14 +31,15 @@ from util import constants
 from util import fileutils
 
 
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
+_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_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    '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,
 )
 
@@ -48,10 +49,10 @@ _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
     'The access token for loading huggingface gated models.',
 )
 
-_DATASET_NAME = flags.DEFINE_string(
-    'dataset_name',
+_TRAIN_DATASET = flags.DEFINE_string(
+    'train_dataset',
     None,
-    'The dataset name in huggingface.',
+    'The training dataset name in huggingface or path.',
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
@@ -119,8 +120,8 @@ _WEIGHT_DECAY = flags.DEFINE_float(
     'The weight decay in the learning rate scheduler.',
 )
 
-_NUM_EPOCHS = flags.DEFINE_float(
-    'num_epochs',
+_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'
@@ -130,8 +131,8 @@ _NUM_EPOCHS = flags.DEFINE_float(
 _MAX_STEPS = flags.DEFINE_integer(
     'max_steps',
     None,
-    'Total number of training steps. Overrides num_epochs if set. Only used for'
-    ' "instruct-lora."',
+    'Total number of training steps. Overrides num_train_epochs if set. Only'
+    ' used for "instruct-lora."',
 )
 
 _MAX_SEQ_LENGTH = flags.DEFINE_integer(
@@ -146,9 +147,9 @@ _LEARNING_RATE = flags.DEFINE_float(
     'The learning rate after the potential warmup period.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -170,8 +171,8 @@ _GRADIENT_ACCUMULATION_STEPS = flags.DEFINE_integer(
     'The gradient accumulation steps.',
 )
 
-_ENABLE_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
-    'enable_gradient_checkpointing',
+_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
+    'gradient_checkpointing',
     False,
     'Whether to enable gradient checkpointing.',
 )
@@ -181,8 +182,8 @@ _ENABLE_PEFT = flags.DEFINE_boolean(
     True,
     'Whether to enable peft.',
 )
-_TEMPLATE = flags.DEFINE_string(
-    'template',
+_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'
@@ -219,32 +220,18 @@ _EVAL_STEPS = flags.DEFINE_integer(
     'The number of training steps between evaluations.',
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_EVAL_TASKS = flags.DEFINE_list(
-    'eval_tasks',
-    None,
-    'List of eval task names (can have wildcards) as in'
-    ' https://github.com/EleutherAI/lm-evaluation-harness. Will not run'
-    ' evaluation if not set. Runs the built-in trainer evaluation loop if set'
-    ' to `builtin_eval`.',
-)
-
-_EVAL_PER_DEVICE_BATCH_SIZE = flags.DEFINE_integer(
-    'eval_per_device_batch_size',
+_PER_DEVICE_EVAL_BATCH_SIZE = flags.DEFINE_integer(
+    'per_device_eval_batch_size',
     1,
     'The per device batch size for model evaluation.',
 )
 
-_EVAL_NUM_FEWSHOT = flags.DEFINE_integer(
-    'eval_num_fewshot',
-    None,
-    'Run N-shot language model evaluation. Not implemented in `builtin_eval`.',
-)
 
 _EVAL_LIMIT = flags.DEFINE_float(
     'eval_limit',
@@ -253,17 +240,18 @@ _EVAL_LIMIT = flags.DEFINE_float(
     ' total number of examples.',
 )
 
-_EVAL_METRIC_NAME = flags.DEFINE_string(
+_EVAL_METRIC_NAME = flags.DEFINE_list(
     'eval_metric_name',
-    'acc',
-    'The metric name to aggregate during model evaluation.',
+    ['loss'],
+    'A comma-separated list of metric names to aggregate during model'
+    ' evaluation. The supported metrics are: '
+    + ', '.join(constants.SUPPORTED_EVAL_METRICS),
 )
 
-_EVAL_DATASET_PATH = flags.DEFINE_string(
-    'eval_dataset_path',
+_EVAL_DATASET = flags.DEFINE_string(
+    'eval_dataset',
     None,
-    'Overrides the default evaluation dataset path. In `builtin_eval` mode,'
-    ' this can be any Hugging Face dataset name.',
+    'The Hugging Face dataset name or path to use for evaluation.',
 )
 
 # We set the default eval split as `test`, based on observation from
@@ -271,13 +259,13 @@ _EVAL_DATASET_PATH = flags.DEFINE_string(
 _EVAL_SPLIT = flags.DEFINE_string(
     'eval_split',
     'test',
-    'Eval split name in the eval dataset for `builtin_eval`.',
+    'Eval split name in the eval dataset.',
 )
 
 _EVAL_TEMPLATE = flags.DEFINE_string(
     'eval_template',
     None,
-    'Template for formatting language model evaluation data for `builtin_eval`.'
+    '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.',
 )
@@ -285,7 +273,14 @@ _EVAL_TEMPLATE = flags.DEFINE_string(
 _EVAL_COLUMN = flags.DEFINE_string(
     'eval_column',
     None,
-    'Eval column name in the eval dataset for `builtin_eval`.',
+    '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(
@@ -299,15 +294,15 @@ _TRAIN_PRECISION = flags.DEFINE_enum(
     'Precision to train the model.',
 )
 
-_USE_EXAMPLE_PACKING = flags.DEFINE_boolean(
-    'use_example_packing',
+_EXAMPLE_PACKING = flags.DEFINE_boolean(
+    'example_packing',
     False,
     'Enables example packing during training, which uses '
     '`ConstantLengthDataset` under the hood.',
 )
 
-_COMPLETION_ONLY = flags.DEFINE_boolean(
-    'completion_only',
+_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',
@@ -356,28 +351,53 @@ _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(
     [
-        _COMPLETION_ONLY.name,
-        _USE_EXAMPLE_PACKING.name,
+        _INPUT_MASKING.name,
+        _EXAMPLE_PACKING.name,
     ],
-    message=(
-        '`use_example_packing=True` does not work with `completion_only=True`'
-    ),
+    message='`example_packing=True` does not work with `input_masking=True`',
 )
-def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
+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 `use_example_packing` is set properly.
+    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[_COMPLETION_ONLY.name]
-      and flags_dict[_USE_EXAMPLE_PACKING.name]
+      flags_dict[_INPUT_MASKING.name]
+      and flags_dict[_TRAIN_TEMPLATE.name] is None
   ):
     return False
   return True
@@ -385,22 +405,65 @@ def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
 
 @flags.multi_flags_validator(
     [
-        _COMPLETION_ONLY.name,
-        _TEMPLATE.name,
+        _EVAL_DATASET.name,
+        _EVAL_METRIC_NAME.name,
     ],
-    message='`template` should be provided if using `completion_only=True`',
+    message=(
+        '`eval_metric_name` should be a valid metric name and present when'
+        ' eval_dataset is provided.'
+    ),
 )
-def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
-  """Check to make sure completion_only is enabled properly.
+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 `completion_only` is set properly
+    If the eval metrics are valid.
   """
-  if flags_dict[_COMPLETION_ONLY.name] and flags_dict[_TEMPLATE.name] is None:
-    return False
+  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
 
 
@@ -411,10 +474,43 @@ def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
 # 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: AutoModelForCausalLM, precision_mode: str
-) -> list[str]:
+    model: transformers.AutoModelForCausalLM, precision_mode: str
+) -> Sequence[str]:
   """Finds all linear module names."""
   if precision_mode == constants.PRECISION_MODE_4:
     cls = bnb.nn.Linear4bit
@@ -433,47 +529,51 @@ def find_all_linear_names(
 
 
 def finetune_instruct(
-    pretrained_model_id: str,
-    dataset_name: str,
+    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_epochs: Optional[float] = None,
-    max_steps: Optional[int] = None,
+    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,
-    instruct_column_in_dataset: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    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,
-    enable_gradient_checkpointing: bool = False,
+    gradient_checkpointing: bool = False,
     enable_peft: bool = True,
-    template: str = None,
+    train_template: str = None,
     lr_scheduler_type: str = 'constant',
     save_steps: int = 10,
     logging_steps: int = 10,
-    train_split_name: str = 'train',
-    eval_config: Optional[eval_lib.EvalConfig] = None,
+    train_split: str = 'train',
+    eval_config: eval_lib.EvalConfig | None = None,
     report_to: str = constants.REPORT_TO_NONE,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
     train_precision: str = constants.PRECISION_MODE_16B,
-    use_example_packing: bool = False,
-    attn_implementation: Optional[str] = None,
+    example_packing: bool = False,
+    attn_implementation: str | None = None,
     max_grad_norm: float = 0.3,
-    completion_only: bool = False,
+    input_masking: bool = False,
     logger_level: str = 'passive',
-    benchmark_out_file: Optional[str] = None,
-    tuning_data_stats_file: Optional[str] = None,
-    target_modules: Optional[str] = None,
+    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', utils.gpu_stats_str())
+  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
@@ -483,17 +583,25 @@ def finetune_instruct(
   if PartialState().distributed_type == DistributedType.MULTI_GPU:
     gradient_checkpointing_kwargs['use_reentrant'] = False
 
-  tokenizer = utils.load_tokenizer(
-      pretrained_model_id,
+  tokenizer = dataset_validation_util.load_tokenizer(
+      pretrained_model_name_or_path,
       'right',
       access_token=access_token,
   )
 
-  train_dataset = dataset_validation_util.load_dataset_with_template(
-      dataset_name,
-      split=train_split_name,
-      input_column=instruct_column_in_dataset,
-      template=template,
+  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,
   )
 
@@ -508,24 +616,26 @@ def finetune_instruct(
           'getting tuning data stats with effective batch size %s',
           effective_batch_size,
       )
-      train_dataset_stats = utils.get_dataset_stats(
-          train_dataset,
-          tokenizer,
-          instruct_column_in_dataset,
-          effective_batch_size,
+      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(dataclasses.asdict(train_dataset_stats), out_f)
+        json.dump(train_dataset_stats, out_f)
 
   model = utils.load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=precision_mode,
-      enable_gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing=gradient_checkpointing,
       access_token=access_token,
       attn_implementation=attn_implementation,
       train_precision=train_precision,
@@ -550,7 +660,9 @@ def finetune_instruct(
     # `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.
@@ -568,14 +680,13 @@ def finetune_instruct(
   # b/357970482#comment3
   accelerator_config = {'use_configured_state': True}
 
-  training_arguments = TrainingArguments(
+  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_strategy='steps',
       save_total_limit=3,
       logging_dir=os.path.join(logging_output_dir, 'logs'),
       logging_steps=logging_steps,
@@ -583,21 +694,24 @@ def finetune_instruct(
       fp16=(train_precision == constants.PRECISION_MODE_16),
       bf16=(train_precision == constants.PRECISION_MODE_16B),
       max_grad_norm=max_grad_norm,
-      num_train_epochs=num_epochs if num_epochs else -1,
+      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=enable_gradient_checkpointing,
+      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 completion_only and template:
-    template_json = dataset_validation_util.get_template(template_path=template)
+  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
     )
@@ -608,7 +722,7 @@ def finetune_instruct(
           ' `DataCollatorForCompletionOnlyLM`'
       )
 
-    trainer_kwargs['data_collator'] = DataCollatorForCompletionOnlyLM(
+    trainer_kwargs['data_collator'] = trl.DataCollatorForCompletionOnlyLM(
         instruction_template=instruction_sep,
         response_template=response_sep,
         tokenizer=tokenizer,
@@ -618,17 +732,23 @@ def finetune_instruct(
   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=SFTTrainer,
+      cls=trl.SFTTrainer,
       eval_config=eval_config,
       model=model,
-      train_dataset=train_dataset,
-      dataset_text_field=instruct_column_in_dataset,
+      train_dataset=train_dataset_with_template,
+      dataset_text_field=train_column,
       max_seq_length=max_seq_length,
       tokenizer=tokenizer,
       args=training_arguments,
-      packing=use_example_packing,
+      packing=example_packing,
       callbacks=[trainer_stats_callback],
+      compute_metrics=compute_metrics,
+      preprocess_logits_for_metrics=preprocess_logits,
       **trainer_kwargs,
   )
 
@@ -658,7 +778,7 @@ def finetune_instruct(
     # This method saves the sharded weights like `accelerator.save_state`, see
     # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
     trainer.save_model(output_dir)
-    model = trainer.model.cpu()  # Avoids GPU OOM
+    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`.
@@ -668,7 +788,6 @@ def finetune_instruct(
         is_main_process=PartialState().is_main_process,
         save_embedding_layers=False,  # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
     )
-    model.cuda()  # Move back to GPU to do eval.
   else:
     trainer.model.save_pretrained(
         final_checkpoint,
@@ -683,15 +802,6 @@ def finetune_instruct(
     # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L1001 #pylint: disable=line-too-long
     trainer.log_metrics('eval', metrics)
     trainer.save_metrics('eval', metrics)
-    if PartialState().is_main_process:
-      hp_metric = metrics[f'eval_{eval_config.metric_name}']
-      hpt = hypertune.HyperTune()
-      hpt.report_hyperparameter_tuning_metric(
-          hyperparameter_metric_tag=constants.HP_METRIC_TAG,
-          metric_value=hp_metric,
-      )
-      logging.info('Send HP metric: %f to hyperparameter tuning.', hp_metric)
-    PartialState().wait_for_everyone()
 
   if not enable_peft:
     tokenizer.save_pretrained(
@@ -705,38 +815,41 @@ def main(unused_argv: Sequence[str]) -> None:
       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_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
-    )
-
-  output_dir = utils.GcsOrLocalDirectory(
-      _OUTPUT_DIR.value, check_empty=True, upload_from_all_nodes=True
+  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_TASKS.value:
+  if _EVAL_DATASET.value:
     eval_config = eval_lib.EvalConfig(
-        tasks=_EVAL_TASKS.value,
-        per_device_batch_size=_EVAL_PER_DEVICE_BATCH_SIZE.value,
-        num_fewshot=_EVAL_NUM_FEWSHOT.value,
+        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_PATH.value
+            _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
@@ -745,40 +858,40 @@ def main(unused_argv: Sequence[str]) -> None:
     wandb.login()
 
   finetune_instruct(
-      pretrained_model_id=pretrained_model_id,
-      dataset_name=_DATASET_NAME.value,
-      output_dir=output_dir.local_dir,
+      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_epochs=_NUM_EPOCHS.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,
-      instruct_column_in_dataset=_INSTRUCT_COLUMN_IN_DATASET.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,
-      enable_gradient_checkpointing=_ENABLE_GRADIENT_CHECKPOINTING.value,
+      gradient_checkpointing=_GRADIENT_CHECKPOINTING.value,
       enable_peft=_ENABLE_PEFT.value,
-      template=_TEMPLATE.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_name=_TRAIN_SPLIT_NAME.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,
-      use_example_packing=_USE_EXAMPLE_PACKING.value,
+      example_packing=_EXAMPLE_PACKING.value,
       attn_implementation=_ATTN_IMPLEMENTATION.value,
       max_grad_norm=_MAX_GRAD_NORM.value,
-      completion_only=_COMPLETION_ONLY.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,
@@ -787,8 +900,6 @@ def main(unused_argv: Sequence[str]) -> None:
   # Frees the model from GPU.
   utils.force_gc()
 
-  output_dir.upload_to_gcs(skip_if_exists=True)
-
 
 if __name__ == '__main__':
   app.run(main)

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

@@ -1,6 +1,7 @@
 """Script to merge PEFT adapter with base model."""
 
-from typing import Any, Dict, Sequence
+from collections.abc import Mapping, Sequence
+from typing import Any
 
 from absl import app
 from absl import flags
@@ -11,14 +12,14 @@ from util import constants
 from util import fileutils
 
 
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
+_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_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    ' 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,
 )
 
@@ -30,7 +31,7 @@ _MERGE_BASE_AND_LORA_OUTPUT_DIR = flags.DEFINE_string(
 
 _MERGE_MODEL_PRECISION_MODE = flags.DEFINE_enum(
     'merge_model_precision_mode',
-    constants.PRECISION_MODE_16,
+    constants.PRECISION_MODE_16B,
     [
         constants.PRECISION_MODE_4,
         constants.PRECISION_MODE_8,
@@ -48,20 +49,6 @@ _FINETUNED_LORA_MODEL_DIR = flags.DEFINE_string(
     'The directory storing finetuned LoRA model weights.',
 )
 
-_RESTRICT_MODEL_UPLOAD_DOCKER_URI = flags.DEFINE_string(
-    'restrict_model_upload_docker_uri',
-    '',
-    'If set, mark output model as only uploadable to Model Registry with the'
-    ' specified Docker URI.',
-)
-
-_EXECUTOR_INPUT = flags.DEFINE_string(
-    'executor_input',
-    '',
-    'For internal use. Kubeflow pipeline context when running trainer as part'
-    ' of an internal pipeline.',
-)
-
 _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
     'huggingface_access_token',
     None,
@@ -71,12 +58,12 @@ _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
 
 @flags.multi_flags_validator(
     [
-        _PRETRAINED_MODEL_ID.name,
+        _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: Dict[str, Any]) -> bool:
+def check_merge_lora_model_flags(flags_dict: Mapping[str, Any]) -> bool:
   """Check if required flags are set on merge model LoRA task.
 
   Args:
@@ -89,44 +76,33 @@ def check_merge_lora_model_flags(flags_dict: Dict[str, Any]) -> bool:
 
 
 def main(unused_argv: Sequence[str]) -> None:
-  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
+  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 = utils.GcsOrLocalDirectory(
+  finetuned_lora_model_dir = fileutils.force_gcs_path(
       _FINETUNED_LORA_MODEL_DIR.value
   )
-
-  merge_base_and_lora_output_dir = utils.GcsOrLocalDirectory(
-      _MERGE_BASE_AND_LORA_OUTPUT_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_id=pretrained_model_id,
+      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.local_dir,
-      merged_model_output_dir=merge_base_and_lora_output_dir.local_dir,
+      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 _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value:
-    utils.write_first_party_model_metadata(
-        merge_base_and_lora_output_dir.local_dir,
-        _RESTRICT_MODEL_UPLOAD_DOCKER_URI.value,
-    )
-
-  if _EXECUTOR_INPUT.value:
-    utils.write_kfp_outputs(
-        _EXECUTOR_INPUT.value,
-        {
-            'saved_model': _MERGE_BASE_AND_LORA_OUTPUT_DIR.value,
-        },
-    )
-
-  merge_base_and_lora_output_dir.upload_to_gcs(skip_if_exists=True)
-
 
 if __name__ == '__main__':
   app.run(main)

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

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

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

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

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

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

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

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

@@ -3,16 +3,18 @@
 # DO NOT MODIFY: this file is auto-generated
 # See go/vmg-oss-peft-tests#command-builder-genpy
 
+
 class InstructLoraCommandBuilder:
 
   def __init__(self):
     self._config_file = None
     self._task = None
-    self._pretrained_model_id = None
-    self._dataset_name = None
-    self._train_split_name = None
-    self._template = None
-    self._instruct_column_in_dataset = None
+    self._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
@@ -22,14 +24,14 @@ class InstructLoraCommandBuilder:
     self._lora_alpha = None
     self._lora_dropout = None
     self._max_steps = None
-    self._num_epochs = 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._enable_gradient_checkpointing = None
-    self._use_example_packing = None
+    self._gradient_checkpointing = None
+    self._example_packing = None
     self._attn_implementation = None
     self._optimizer = None
     self._warmup_ratio = None
@@ -37,14 +39,14 @@ class InstructLoraCommandBuilder:
     self._save_steps = None
     self._logging_steps = None
     self._huggingface_access_token = None
-    self._eval_dataset_path = None
+    self._eval_dataset = None
     self._eval_column = None
     self._eval_template = None
     self._eval_split = None
     self._eval_steps = None
-    self._eval_tasks = None
     self._eval_metric_name = None
-    self._completion_only = None
+    self._metric_for_best_model = None
+    self._input_masking = None
     self._max_grad_norm = None
     self._logger_level = None
     self._benchmark_out_file = None
@@ -52,6 +54,7 @@ class InstructLoraCommandBuilder:
     self._enable_peft = None
     self._merge_model_precision_mode = None
     self._target_modules = None
+    self._unnamed_args = None
 
   @property
   def config_file(self):
@@ -70,44 +73,52 @@ class InstructLoraCommandBuilder:
     self._task = val
 
   @property
-  def pretrained_model_id(self):
-    return self._pretrained_model_id
+  def gcs_rsync_interval_secs(self):
+    return self._gcs_rsync_interval_secs
 
-  @pretrained_model_id.setter
-  def pretrained_model_id(self, val: str):
-    self._pretrained_model_id = val
+  @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._dataset_name
+    return self._train_dataset
 
   @train_dataset.setter
   def train_dataset(self, val: str):
-    self._dataset_name = val
+    self._train_dataset = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def template(self):
-    return self._template
+  def train_template(self):
+    return self._train_template
 
-  @template.setter
-  def template(self, val: str):
-    self._template = val
+  @train_template.setter
+  def train_template(self, val: str):
+    self._train_template = val
 
   @property
-  def instruct_column(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column.setter
-  def instruct_column(self, val: str):
-    self._instruct_column_in_dataset = val
+  @train_column.setter
+  def train_column(self, val: str):
+    self._train_column = val
 
   @property
   def ckpt_dir(self):
@@ -182,12 +193,12 @@ class InstructLoraCommandBuilder:
     self._max_steps = val
 
   @property
-  def num_epochs(self):
-    return self._num_epochs
+  def num_train_epochs(self):
+    return self._num_train_epochs
 
-  @num_epochs.setter
-  def num_epochs(self, val: float):
-    self._num_epochs = val
+  @num_train_epochs.setter
+  def num_train_epochs(self, val: float):
+    self._num_train_epochs = val
 
   @property
   def max_seq_length(self):
@@ -231,19 +242,19 @@ class InstructLoraCommandBuilder:
 
   @property
   def gradient_checkpointing(self):
-    return self._enable_gradient_checkpointing
+    return self._gradient_checkpointing
 
   @gradient_checkpointing.setter
   def gradient_checkpointing(self, val: bool):
-    self._enable_gradient_checkpointing = val
+    self._gradient_checkpointing = val
 
   @property
   def example_packing(self):
-    return self._use_example_packing
+    return self._example_packing
 
   @example_packing.setter
   def example_packing(self, val: bool):
-    self._use_example_packing = val
+    self._example_packing = val
 
   @property
   def attn_implementation(self):
@@ -303,18 +314,18 @@ class InstructLoraCommandBuilder:
 
   @property
   def eval_dataset(self):
-    return self._eval_dataset_path
+    return self._eval_dataset
 
   @eval_dataset.setter
   def eval_dataset(self, val: str):
-    self._eval_dataset_path = val
+    self._eval_dataset = val
 
   @property
-  def eval_instruct_column(self):
+  def eval_column(self):
     return self._eval_column
 
-  @eval_instruct_column.setter
-  def eval_instruct_column(self, val: str):
+  @eval_column.setter
+  def eval_column(self, val: str):
     self._eval_column = val
 
   @property
@@ -326,11 +337,11 @@ class InstructLoraCommandBuilder:
     self._eval_template = val
 
   @property
-  def eval_split_name(self):
+  def eval_split(self):
     return self._eval_split
 
-  @eval_split_name.setter
-  def eval_split_name(self, val: str):
+  @eval_split.setter
+  def eval_split(self, val: str):
     self._eval_split = val
 
   @property
@@ -341,14 +352,6 @@ class InstructLoraCommandBuilder:
   def eval_steps(self, val: int):
     self._eval_steps = val
 
-  @property
-  def eval_tasks(self):
-    return self._eval_tasks
-
-  @eval_tasks.setter
-  def eval_tasks(self, val: str):
-    self._eval_tasks = val
-
   @property
   def eval_metric_name(self):
     return self._eval_metric_name
@@ -358,12 +361,20 @@ class InstructLoraCommandBuilder:
     self._eval_metric_name = val
 
   @property
-  def completion_only(self):
-    return self._completion_only
+  def metric_for_best_model(self):
+    return self._metric_for_best_model
 
-  @completion_only.setter
-  def completion_only(self, val: bool):
-    self._completion_only = val
+  @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):
@@ -421,9 +432,22 @@ class InstructLoraCommandBuilder:
   def target_modules(self, val: str):
     self._target_modules = val
 
-  def build_cmd(self) -> str:
+  @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

@@ -8,7 +8,7 @@ class QuantizeModelCommandBuilder:
 
   def __init__(self):
     self._task = None
-    self._pretrained_model_id = None
+    self._pretrained_model_name_or_path = None
     self._quantization_method = None
     self._quantization_precision_mode = None
     self._quantization_dataset_name = None
@@ -31,12 +31,12 @@ class QuantizeModelCommandBuilder:
     self._task = val
 
   @property
-  def pretrained_model_id(self):
-    return self._pretrained_model_id
+  def pretrained_model_name_or_path(self):
+    return self._pretrained_model_name_or_path
 
-  @pretrained_model_id.setter
-  def pretrained_model_id(self, val: str):
-    self._pretrained_model_id = val
+  @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):

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

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

@@ -31,7 +31,7 @@ class AdapterTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    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'
@@ -39,9 +39,9 @@ class AdapterTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.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
@@ -87,8 +87,8 @@ class AdapterTest(test_util.TestBase):
   def test_llama_adapters(self, model_name):
     test_function_name = inspect.stack()[0][3]
     self.setup_output_dir(f'{test_function_name}-{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    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)

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

@@ -4,19 +4,81 @@
 
 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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    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'
@@ -24,9 +86,9 @@ class GcsUploadDownloadTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.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
@@ -45,9 +107,11 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_single_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  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)
@@ -61,24 +125,26 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_multi_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    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_70b_model_download(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        'gs://vertex-model-garden-public-us/llama3/llama3-70b-hf'
+  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)
@@ -90,8 +156,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge(self, output_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    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(
@@ -107,9 +173,13 @@ class GcsUploadDownloadTest(test_util.TestBase):
     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_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    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()}'
@@ -128,17 +198,17 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge_and_upload_deepspeed(self, merged_model_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.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_zero3_8gpu.yaml'
+        '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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    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()
@@ -149,8 +219,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('save-multiple-times', 1),
   )
   def test_llama3_8b_save_and_merge_8_gpus_fsdp(self, save_steps):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.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
@@ -159,7 +229,7 @@ class GcsUploadDownloadTest(test_util.TestBase):
     )
     self.task_cmd_builder.merged_model_dir = '/tmp/merged'
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    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()
@@ -171,7 +241,7 @@ class TemplateAndDataStatsTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    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'
@@ -187,35 +257,42 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
-  def test_openai_chat_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+  @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(
-        'openai-multi-chat-example-data.jsonl'
+        example_dataset
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.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_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    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_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    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_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    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'
@@ -223,17 +300,17 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-multi-chat-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    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_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    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'
@@ -241,12 +318,12 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-completion-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     self.assertEqual(self.run_cmd(), 0)
 
 
@@ -255,7 +332,7 @@ class TargetModulesTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    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'
@@ -263,9 +340,9 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.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
@@ -278,8 +355,8 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
   def test_target_modules(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.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'
 

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

@@ -43,9 +43,9 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.example_packing = True
     self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2-1k'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.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'
 
@@ -54,23 +54,24 @@ class TrainerThroughputTest(test_util.TestBase):
     if ret != 0:
       with open(self.task_cmd_builder.benchmark_out_file, 'a') as f:
         max_seq_length = self.task_cmd_builder.max_seq_length
-        f.write(f'{max_seq_length/1024.0:.1f}k | failed | n/a\n')
+        f.write(f'{max_seq_length/1024.0:.1f} | failed | n/a\n')
     return ret
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          '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_id = (
-        test_util.get_pretrained_model_id(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.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -78,29 +79,30 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir, f'bm_{model_name}_{precision}.txt'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          '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', 'deepspeed_zero3'],
+      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_id = (
-        test_util.get_pretrained_model_id(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.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -112,14 +114,18 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.config_file = (
         f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
     )
-    self.docker_builder.add_env_var(
+    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-70b-hf'],
+      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],
@@ -127,8 +133,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_lora(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.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
@@ -136,18 +142,108 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir,
         f'bm_fsdp_{num_gpus}gpu_{model_name}_{precision}.txt',
     )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
+    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.docker_builder.add_env_var(
+    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-70b-hf'],
+      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],
@@ -155,8 +251,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_full_finetuning(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.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
@@ -169,7 +265,7 @@ class TrainerThroughputTest(test_util.TestBase):
     )
     self.task_cmd_builder.enable_peft = False
 
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 

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

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

@@ -31,7 +31,7 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       dict(
           testcase_name="out_of_range_rows",
           validate_top_k_rows=100000,
-          expected_result=1,
+          expected_result=0,
       ),
   )
   def test_validate_dataset_with_template_top_k_rows(
@@ -40,8 +40,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -79,8 +79,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -92,6 +92,32 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
     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",
@@ -204,8 +230,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = split
-    self.task_cmd_builder.instruct_column_in_dataset = input_column
+    self.task_cmd_builder.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
@@ -311,8 +337,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.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

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

@@ -10,8 +10,9 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._task = None
     self._template = None
     self._dataset_name = None
-    self._train_split_name = None
-    self._instruct_column_in_dataset = 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
@@ -41,20 +42,28 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._dataset_name = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def instruct_column_in_dataset(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column_in_dataset.setter
-  def instruct_column_in_dataset(self, val: str):
-    self._instruct_column_in_dataset = val
+  @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):

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

@@ -8,16 +8,18 @@ environment. Otherwise, `python3` is used.
 """
 
 import argparse
-import json
-import os
+from collections.abc import MutableSequence, Sequence
+import multiprocessing
 import subprocess
-from typing import List, Optional, Sequence
+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
 
 
@@ -37,20 +39,21 @@ _TASK_TO_SCRIPT = {
     constants.INSTRUCT_LORA: (
         'vertex_vision_model_garden_peft/train/vmg/instruct_lora.py'
     ),
-    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: 'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py',
-    constants.QUANTIZE_MODEL: (
-        'vertex_vision_model_garden_peft/train/vmg/quantize_model.py'
+    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: (
+        'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py'
     ),
-    constants.SEQUENCE_CLASSIFICATION_LORA: 'vertex_vision_model_garden_peft/train/vmg/sequence_classification_lora.py',
-    constants.VALIDATE_DATASET_WITH_TEMPLATE: 'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py',
+    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: Optional[str],
+    config_file: str | None,
     accelerate_args: argparse.Namespace = argparse.Namespace(),
-) -> List[str]:
+) -> MutableSequence[str]:
   """Returns the command to launch the script."""
   if config_file:
     cmd = [
@@ -69,59 +72,23 @@ def launch_script_cmd(
 
 def _get_accelerate_args() -> argparse.Namespace:
   """Returns the accelerate args."""
-  # For the format of the cluster spec, see
-  # https://cloud.google.com/vertex-ai/docs/training/distributed-training#cluster-spec-format # pylint: disable=line-too-long
-  cluster_spec = os.getenv('CLUSTER_SPEC', default=None)
-  if not cluster_spec:
-    return argparse.Namespace()
-  logging.info('CLUSTER_SPEC: %s', cluster_spec)
-
-  cluster_data = json.loads(cluster_spec)
-  if (
-      'workerpool1' not in cluster_data['cluster']
-      or not cluster_data['cluster']['workerpool1']
-  ):
-    return argparse.Namespace()
-
-  # Get primary node info
-  primary_node = cluster_data['cluster']['workerpool0'][0]
-  logging.info('primary node: %s', primary_node)
-  primary_node_addr, primary_node_port = primary_node.split(':')
-  logging.info('primary node address: %s', primary_node_addr)
-  logging.info('primary node port: %s', primary_node_port)
-
-  # Determine node rank of this machine
-  workerpool = cluster_data['task']['type']
-  if workerpool == 'workerpool0':
-    node_rank = 0
-  elif workerpool == 'workerpool1':
-    # Add 1 for the primary node, since `index` is the index of workerpool1.
-    node_rank = cluster_data['task']['index'] + 1
-  else:
-    raise ValueError(
-        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
-        f' {workerpool}'
-    )
-  logging.info('node rank: %s', node_rank)
-
-  # Calculate total nodes
-  num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
-  num_nodes = num_worker_nodes + 1  # Add 1 for the primary node
-  logging.info('num nodes: %s', num_nodes)
-
+  primary_node_addr, primary_node_port, node_rank, num_nodes = (
+      cluster_spec.get_cluster_spec()
+  )
   accelerate_args = argparse.Namespace()
-  accelerate_args.machine_rank = node_rank
-  accelerate_args.num_machines = num_nodes
-  accelerate_args.main_process_ip = primary_node_addr
-  accelerate_args.main_process_port = primary_node_port
-  accelerate_args.max_restarts = 0
-  accelerate_args.monitor_interval = 120
+  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: List[str]
+    args: argparse.Namespace, command: MutableSequence[str]
 ):
   for key, value in vars(args).items():
     # If not specified, skip.
@@ -129,15 +96,25 @@ def _append_args_to_command_in_place(
       command.append(f'--{key}={value}')
 
 
-def _get_train_cmd_and_maybe_merge_cmd(
-    task: str, config_file: str, unknown: Sequence[str]
+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 command and maybe the merge command if applicable."""
+  """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],
+      _TASK_TO_SCRIPT[task_type],
       config_file,
       accelerate_args=accelerate_args,
   )
@@ -152,17 +129,30 @@ def _get_train_cmd_and_maybe_merge_cmd(
       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('--executor_input')
-  merge_parser.add_argument('--restrict_model_upload_docker_uri')
   merge_parser.add_argument('--merge_base_and_lora_output_dir')
   merge_args, unknown = merge_parser.parse_known_args(unknown)
 
+  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_id', required=True)
+  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)
 
@@ -173,17 +163,19 @@ def _get_train_cmd_and_maybe_merge_cmd(
   commands = [training_cmd]
 
   # Only the main node runs merging.
-  if (
-      merge_args.merge_base_and_lora_output_dir
-      and getattr(accelerate_args, 'machine_rank', 0) == 0
-  ):
+  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_dir}',
+        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)
@@ -196,16 +188,51 @@ def _get_train_cmd_and_maybe_merge_cmd(
     ]
     commands.append(conda_run_cmd)
 
-  return commands
+  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.
@@ -224,8 +251,12 @@ def main(unused_argv: Sequence[str]) -> None:
     ] + list(map(dataset_validation_util.force_gcs_fuse_path, unknown))
     commands = [config_gen_cmd, task_cmd]
   elif task in [constants.INSTRUCT_LORA]:
-    commands = _get_train_cmd_and_maybe_merge_cmd(
-        task=task, config_file=args.config_file, unknown=unknown
+    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
@@ -233,10 +264,29 @@ def main(unused_argv: Sequence[str]) -> None:
     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))
-    subprocess.run(cmd, check=True)
+    # 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,232 +1,44 @@
 """Common libraries for PEFT."""
 
-import dataclasses
+from collections.abc import Mapping, Sequence
 import datetime
 import gc
-import multiprocessing as mp
 import os
-import subprocess
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 
 from absl import logging
 import accelerate
 from accelerate import DistributedType
 from accelerate import PartialState
-from google.protobuf import json_format
-from kfp.pipeline_spec import pipeline_spec_pb2
-import numpy as np
 import peft
 from peft import PeftModel
 from peft import prepare_model_for_kbit_training
-import pynvml
 import torch
 import transformers
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 from transformers import BitsAndBytesConfig
 from transformers import FbgemmFp8Config
-from transformers.integrations import is_deepspeed_zero3_enabled
 import trl
 
 from util import dataset_validation_util
 from util import constants
-from util import fileutils
 
-
-_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
-_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
 _LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
 _LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
-
-
-
-class GcsOrLocalDirectory(os.PathLike):
-  """A class to represent a directory with upload support if GCS path is given.
-
-  This class is used to represent a directory. It can be used for a temporary
-  local directory and for uploading files to the GCS directory later if the
-  given path is a GCS directory. If the given path is a local directory, a call
-  to gcs_dir attribute will raise an error. This class has multi-node and
-  multi-process support with accelerate.
-
-  Attributes:
-    local_dir: The local directory to store the files.
-    gcs_dir: The path to the GCS directory.
-  """
-
-  def __init__(
-      self,
-      path: str,
-      check_empty: bool = False,
-      upload_from_all_nodes: bool = False,
-  ):
-    """Initializes the GcsOrLocalDirectory.
-
-    Args:
-      path: The path to the directory.
-      check_empty: If True, check if the GCS directory is empty. No-op for local
-        directory.
-      upload_from_all_nodes: If True, upload the local directory to GCS from all
-        nodes.
-    """
-    if len(path) > 1:
-      path = path.rstrip("/")
-
-    self._upload_from_all_nodes = upload_from_all_nodes
-
-    if path.startswith(constants.GCS_URI_PREFIX) or path.startswith(
-        constants.GCSFUSE_URI_PREFIX
-    ):
-      self._is_gcs_path = True
-      self._local_dir = _get_local_dir_from_gcs_dir(path)
-      self._gcs_dir = fileutils.force_gcs_path(path)
-      os.makedirs(self.local_dir, exist_ok=True)
-
-      with PartialState().main_process_first():
-        if (
-            check_empty
-            and PartialState().is_main_process
-            and not _is_gcs_dir_empty(self._gcs_dir)
-        ):
-          raise ValueError(f"{self._gcs_dir} needs to be empty.")
-    else:
-      self._is_gcs_path = False
-      self._local_dir = path
-      self._gcs_dir = path
-
-  def __fspath__(self) -> str:
-    return self.local_dir
-
-  @property
-  def local_dir(self) -> str:
-    return self._local_dir
-
-  @property
-  def gcs_dir(self) -> str:
-    """Returns the GCS directory path.
-
-    Returns:
-      The GCS directory path.
-
-    Raises:
-      ValueError: If the path is not a GCS path.
-    """
-    if not self._is_gcs_path:
-      raise ValueError(f"{self._gcs_dir} is not a GCS path.")
-    return self._gcs_dir
-
-  def upload_to_gcs(
-      self,
-      skip_if_exists: bool = True,
-      force_upload: bool = False,
-  ):
-    """Uploads the local directory to GCS."""
-    if not self._is_gcs_path:
-      logging.info(
-          "Not uploading to GCS since %s is not a GCS path.", self.local_dir
-      )
-      return
-
-    if not os.listdir(self.local_dir):
-      logging.info("Not uploading to GCS since %s is empty.", self.local_dir)
-      return
-
-    target = os.path.dirname(self.gcs_dir) + "/"
-    # Avoid race condition uploading the same file from multiple processes.
-    with PartialState().main_process_first():
-      if not PartialState().is_local_main_process:
-        # Non local main processes don't upload.
-        pass
-      elif self._upload_from_all_nodes or PartialState().is_main_process:
-        logging.info("Uploading %s to %s...", self.local_dir, target)
-        cmd = [
-            "gsutil",
-            "-m",
-            "cp",
-            "-r",
-        ]
-        if skip_if_exists:
-          cmd.append("-n")
-        if force_upload:
-          cmd.append("-f")
-        cmd.extend([self.local_dir, target])
-        subprocess.check_output(cmd)
-        logging.info("%s uploaded.", self.local_dir)
-
-
-def _get_local_dir_from_gcs_dir(path: str) -> str:
-  return os.path.join(
-      constants.LOCAL_OUTPUT_DIR,
-      dataset_validation_util.force_gcs_fuse_path(path)[1:],
-  )
-
-
-def _is_gcs_dir_empty(path: str) -> bool:
-  """Checks if a GCS directory is empty.
-
-  Args:
-    path: The GCS directory path.
-
-  Returns:
-    True if the directory is empty.
-
-  Raises:
-    subprocess.CalledProcessError: If the gsutil command failure reason is not
-      because the dir is empty.
-  """
-  path = path.rstrip("/") + "/"
-  try:
-    subprocess.check_output(["gsutil", "ls", path], stderr=subprocess.STDOUT)
-  except subprocess.CalledProcessError as e:
-    if (
-        str(e.output, encoding="utf-8")
-        == "CommandException: One or more URLs matched no objects.\n"
-    ):
-      return True
-    else:
-      logging.info(str(e.output, encoding="utf-8"))
-      raise
-  else:
-    return False
-
-
-def load_tokenizer(
-    pretrained_model_id: str,
-    padding_side: Optional[str] = None,
-    access_token: Optional[str] = None,
-) -> AutoTokenizer:
-  """Loads tokenizer based on `pretrained_model_id`."""
-  tokenizer_kwargs = {}
-  if should_add_eos_token(pretrained_model_id):
-    tokenizer_kwargs["add_eos_token"] = True
-  if padding_side:
-    tokenizer_kwargs["padding_side"] = padding_side
-
-  with PartialState().local_main_process_first():
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_id,
-        trust_remote_code=False,
-        use_fast=True,
-        token=access_token,
-        **tokenizer_kwargs,
-    )
-
-  if should_add_pad_token(pretrained_model_id):
-    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-
-  return tokenizer
+_GEMMA2_MODEL = "gemma-2"
 
 
 def load_model(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     tokenizer: AutoTokenizer,
     precision_mode: str = None,
-    enable_gradient_checkpointing: bool = False,
-    gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
-    access_token: Optional[str] = None,
-    attn_implementation: Optional[str] = None,
-    train_precision: Optional[str] = None,
-    device_map: Optional[str] = None,
+    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."""
@@ -304,25 +116,37 @@ def load_model(
     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_id,
-      use_cache=not enable_gradient_checkpointing,
-      device_map=device_map,
-      torch_dtype=torch_dtype,
-      quantization_config=quantization_config,
-      trust_remote_code=True,
-      token=access_token,
-      attn_implementation=attn_implementation,
+      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=enable_gradient_checkpointing,
+        use_gradient_checkpointing=gradient_checkpointing,
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
     )
 
-  if enable_gradient_checkpointing:
+  if gradient_checkpointing:
     model.gradient_checkpointing_enable(
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
     )
@@ -345,8 +169,10 @@ def load_model(
     # https://stackoverflow.com/a/77408076
     model.config.use_cache = False
 
-  if should_add_pad_token(pretrained_model_id):
-    model.resize_token_embeddings(len(tokenizer))
+  if 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.
@@ -357,22 +183,24 @@ def load_model(
 
 
 def _merge_causal_language_model_with_lora_internal(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     merge_precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Internal function to merges the base model with the lora adapter."""
   logging.info("loading tokenizer...")
-  tokenizer = load_tokenizer(pretrained_model_id)
+  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_id)
+  logging.info("loading model %s...", pretrained_model_name_or_path)
   device_map = "cpu"
   model = load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=merge_precision_mode,
       access_token=access_token,
@@ -402,42 +230,12 @@ def _merge_causal_language_model_with_lora_internal(
   )
 
 
-def merge_causal_language_model_with_lora_fsdp(
-    pretrained_model_id: str,
-    merge_precision_mode: str,
-    finetuned_lora_model_dir: str,
-    merged_model_output_dir: str,
-    access_token: Optional[str] = None,
-) -> None:
-  """Merges the base model with the lora adapter for FSDP.
-
-  Only the main process should call this function.
-
-  Args:
-    pretrained_model_id: Predefined base model name or path to directory
-      containing model checkpoints.
-    merge_precision_mode: Precision mode for saving model weights.
-    finetuned_lora_model_dir: Path to directory containing PEFT-finetuned model
-      weights.
-    merged_model_output_dir: Path to directory to save the merged model.
-    access_token: Access token for accessing the model.
-  """
-  assert PartialState().is_main_process
-  _merge_causal_language_model_with_lora_internal(
-      pretrained_model_id=pretrained_model_id,
-      merge_precision_mode=merge_precision_mode,
-      finetuned_lora_model_dir=finetuned_lora_model_dir,
-      merged_model_output_dir=merged_model_output_dir,
-      access_token=access_token,
-  )
-
-
 def merge_causal_language_model_with_lora(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Merges the base model with the lora adapter."""
 
@@ -452,39 +250,13 @@ def merge_causal_language_model_with_lora(
 
   if PartialState().is_main_process:
     logging.info("Starting merging job...")
-    # When deepspeed Zero3 is enabled, users are not allowed to specify
-    # `device_map` when loading the model (even on CPU).
-    #
-    # To work-around this, we kick off another process (from the
-    # is_main_process) and set up the environment to avoid using Deepspeed when
-    # doing the merging.
-    if is_deepspeed_zero3_enabled():
-      ctx = mp.get_context("spawn")
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "false"
-      merge_job = ctx.Process(
-          target=_merge_causal_language_model_with_lora_internal,
-          args=(
-              pretrained_model_id,
-              merge_precision_mode,
-              finetuned_lora_model_dir,
-              local_merged_model_dir,
-          ),
-          kwargs={
-              "access_token": access_token,
-          },
-      )
-      merge_job.start()
-      merge_job.join()
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
-    else:
-      _merge_causal_language_model_with_lora_internal(
-          pretrained_model_id=pretrained_model_id,
-          merge_precision_mode=merge_precision_mode,
-          finetuned_lora_model_dir=finetuned_lora_model_dir,
-          merged_model_output_dir=local_merged_model_dir,
-          access_token=access_token,
-      )
-
+    _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.
@@ -492,7 +264,7 @@ def merge_causal_language_model_with_lora(
 
   if precision_mode == constants.PRECISION_MODE_FP8:
     convert_model_to_fp8(
-        pretrained_model_name_or_path=pretrained_model_id,
+        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,
@@ -503,7 +275,7 @@ def convert_model_to_fp8(
     pretrained_model_name_or_path: str,
     merged_model_output_dir: str,
     quantized_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Converts the model to fp8.
 
@@ -533,166 +305,12 @@ def convert_model_to_fp8(
   PartialState().wait_for_everyone()
 
 
-@dataclasses.dataclass
-class TuningDataStats:
-  tuning_dataset_example_count: int
-  total_billable_token_count: int
-  tuning_step_count: int
-
-
-def get_dataset_stats(
-    dataset: Any,
-    tokenizer: transformers.PreTrainedTokenizer,
-    column: str,
-    effective_batch_size: int,
-) -> TuningDataStats:
-  """Calculates dataset statistics, e.g., total number of tokens."""
-  tokenized_dataset = dataset.map(lambda x: tokenizer(x[column]))
-  inputs = tokenized_dataset["input_ids"]
-  tuning_dataset_example_count = int(len(inputs))
-  total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
-  tuning_step_count = (
-      tuning_dataset_example_count + effective_batch_size - 1
-  ) // effective_batch_size
-  return TuningDataStats(
-      tuning_dataset_example_count,
-      total_billable_token_count,
-      tuning_step_count,
-  )
-
-
 def force_gc():
   """Collects garbage immediately to release unused CPU/GPU resources."""
   gc.collect()
   torch.cuda.empty_cache()
 
 
-def should_add_pad_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special pad token."""
-  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
-
-
-def should_add_eos_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special eos token."""
-  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
-
-
-def write_kfp_outputs(
-    executor_input: str, output_artifacts: Dict[str, str]
-) -> None:
-  """Writes KFP outputs given a dict of output artifact names and URIs."""
-  # Only the main process writes to avoid race condition.
-  if PartialState().is_main_process:
-    executor_input = json_format.Parse(
-        executor_input, pipeline_spec_pb2.ExecutorInput()
-    )
-    outputs = executor_input.outputs
-    # set all artifacts
-    for name, uri in output_artifacts.items():
-      artifact_list = outputs.artifacts.get(name)
-      if not artifact_list or not artifact_list.artifacts:
-        raise ValueError(f"Artifact name={name} does not exist.")
-      artifact_list.artifacts[0].uri = uri
-
-    # write output file
-    executor_output = pipeline_spec_pb2.ExecutorOutput(
-        artifacts=outputs.artifacts
-    )
-    os.makedirs(os.path.dirname(outputs.output_file), exist_ok=True)
-    with open(outputs.output_file, "w") as f:
-      f.write(json_format.MessageToJson(executor_output, indent=None))
-
-  # Wait for the main process to finish before moving on to the next task.
-  PartialState().wait_for_everyone()
-
-
-def upload_local_dir_to_gcs(local_dir: str, gcs_path: str):
-  """Uploads local dir to GCS."""
-
-  if PartialState().is_main_process:
-    logging.info("uploading %s to %s...", local_dir, gcs_path)
-    subprocess.check_output([
-        "gsutil",
-        "-m",
-        "cp",
-        "-r",
-        local_dir,
-        gcs_path,
-    ])
-    logging.info("%s uploaded.", local_dir)
-
-  PartialState().wait_for_everyone()
-
-
-def write_first_party_model_metadata(output_dir: str, docker_uri: str) -> None:
-  """Multi-process friendly version of fileutils.write_first_party_model_metadata."""
-  if PartialState().is_main_process:
-    fileutils.write_first_party_model_metadata(output_dir, docker_uri)
-  PartialState().wait_for_everyone()
-
-
-@dataclasses.dataclass
-class GpuStats:
-  """Holds information about GPU usage stats.
-
-  For memory related, see
-  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
-  """
-
-  # total memory
-  total_mem: float
-  # memory occupied.
-  occupied: float
-  # memory reserved, but not used.
-  unused: float
-  # nvidia-smi usually reports more memory usages than pytorch (for driver,
-  # kernel and etc). `smi_diff` tracks this difference.
-  smi_diff: float
-  # Gpu utilization.
-  util: float
-
-  # Allows unpacking operation like
-  # total_mem, occupied, unused, smi_diff, util = GpuStats(...)
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def gpu_stats() -> GpuStats:
-  """Reports GPU memory usage and utilization."""
-  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
-  bytes_per_gb = 1024.0**3
-  device = torch.cuda.current_device()
-  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
-  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
-  unused = reserved - occupied
-
-  def smi_mem(device):
-    try:
-      pynvml.nvmlInit()
-      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
-      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
-      return info.used / bytes_per_gb
-    except pynvml.NVMLError:
-      return 0.0
-
-  mem_used_smi = smi_mem(device)
-  smi_diff = mem_used_smi - reserved
-
-  util = torch.cuda.utilization(device)
-  return GpuStats(mem_used_smi, occupied, unused, smi_diff, util)
-
-
-def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
-  if stats is None:
-    stats = gpu_stats()
-  total, occupied, unused, smi_diff, util = stats
-  return (
-      f"GPU memory: {total:.2f}({occupied=:.2f}, {unused=:.2f},"
-      f" {smi_diff=:.2f}) GB. Utilization: {util:.2f}%"
-  )
-
-
 def init_partial_state(
     timeout: datetime.timedelta = datetime.timedelta(seconds=600),
 ) -> None:
@@ -722,7 +340,7 @@ def get_final_checkpoint_path(output_dir: str) -> str:
 
 def _maybe_get_modules_to_not_convert_by_model_id(
     pretrained_model_name_or_path: str,
-) -> Optional[Sequence[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()

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

@@ -17,15 +17,15 @@ _DATASET_NAME = flags.DEFINE_string(
     required=True,
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -38,6 +38,12 @@ _TEMPLATE = flags.DEFINE_string(
     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,
@@ -64,9 +70,10 @@ def main(unused_argv: Sequence[str]) -> None:
 
   dataset_validation_util.validate_dataset_with_template(
       dataset_name=_DATASET_NAME.value,
-      split=_TRAIN_SPLIT_NAME.value,
-      input_column=_INSTRUCT_COLUMN_IN_DATASET.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,

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

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

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

@@ -36,6 +36,12 @@ BEST_CKPT_METRIC_COMP = 'higher'
 
 # Reported hyperparameter tuning metric tag.
 HP_METRIC_TAG = 'model_performance'
+HP_LOSS_TAG = 'model_loss'
+
+# Reported places.
+REPORT_TO_NONE = 'none'
+REPORT_TO_WANDB = 'wandb'
+REPORT_TO_TENSORBOARD = 'tensorboard'
 
 # HPT trial prefix.
 TRIAL_PREFIX = 'trial_'
@@ -45,7 +51,7 @@ ML_USE_TRAINING = 'training'
 ML_USE_VALIDATION = 'validation'
 ML_USE_TEST = 'test'
 
-# COCO json keys
+# COCO json keys.
 COCO_JSON_ANNOTATIONS = 'annotations'
 COCO_JSON_ANNOTATION_IMAGE_ID = 'image_id'
 COCO_JSON_ANNOTATION_CATEGORY_ID = 'category_id'
@@ -60,36 +66,91 @@ COCO_JSON_IMAGE_HEIGHT = 'height'
 COCO_JSON_IMAGE_COCO_URL = 'coco_url'
 COCO_ANNOTATION_BBOX = 'bbox'
 
-# GCS prefixes
+# GCS prefixes.
 GCS_URI_PREFIX = 'gs://'
 GCSFUSE_URI_PREFIX = '/gcs/'
 
 LOCAL_EVALUATION_RESULT_DIR = '/tmp/evaluation_result_dir'
 LOCAL_MODEL_DIR = '/tmp/model_dir'
+LOCAL_LORA_DIR = '/tmp/lora_dir'
 LOCAL_BASE_MODEL_DIR = '/tmp/base_model_dir'
 LOCAL_DATA_DIR = '/tmp/data'
+LOCAL_OUTPUT_DIR = '/tmp/output_dir'
+LOCAL_PREDICTION_RESULT_DIR = '/tmp/prediction_result_dir'
+SHARED_MEM_DIR = '/dev/shm'
 
 # Huggingface files.
 HF_MODEL_WEIGHTS_SUFFIX = '.bin'
 
 # PEFT finetuning constants.
+TEXT_TO_IMAGE = 'text-to-image'
 TEXT_TO_IMAGE_LORA = 'text-to-image-lora'
+TEXT_TO_IMAGE_DREAMBOOTH = 'text-to-image-dreambooth'
+TEXT_TO_IMAGE_DREAMBOOTH_LORA = 'text-to-image-dreambooth-lora'
+TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL = 'text-to-image-dreambooth-lora-sdxl'
 SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
-CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora'
+MERGE_CAUSAL_LANGUAGE_MODEL_LORA = 'merge-causal-language-model-lora'
 INSTRUCT_LORA = 'instruct-lora'
-CAUSAL_LANGUAGE_MODELING_LORA_TARGET_MODULES = [
-  "q_proj",
-  "v_proj",
-]
-INSTRUCT_LORA_TARGET_MODULES = [
-    "query_key_value",
-    "dense",
-    "dense_h_to_4h",
-    "dense_4h_to_h",
-]
+VALIDATE_DATASET_WITH_TEMPLATE = 'validate-dataset-with-template'
+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,10 +1,13 @@
 """Fileutil lib to copy files between gcs and local."""
 
-import glob
+import filecmp
+import fnmatch
 import os
 import pathlib
 import shutil
-from typing import Tuple
+import subprocess
+import time
+from typing import List, Optional, Tuple
 import uuid
 
 from absl import logging
@@ -13,6 +16,17 @@ from google.cloud import storage
 from util import constants
 
 
+_GCS_CLIENT = None
+
+
+def _get_gcs_client() -> storage.Client:
+  """Gets the default GCS client."""
+  global _GCS_CLIENT
+  if _GCS_CLIENT is None:
+    _GCS_CLIENT = storage.Client()
+  return _GCS_CLIENT
+
+
 def generate_tmp_path(extension: str = '') -> str:
   """Generates a temporary file path with UUID.
 
@@ -36,6 +50,123 @@ 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.
 
@@ -62,15 +193,47 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str):
     raise ValueError(
         f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
     )
-  client = storage.Client()
+  client = _get_gcs_client()
   os.makedirs(os.path.dirname(local_path), exist_ok=True)
   with open(local_path, 'wb') as f:
     client.download_blob_to_file(gcs_uri, f)
 
 
+def download_gcs_file_list_to_local(
+    gcs_uri_list: List[str], local_dir: str
+) -> List[str]:
+  """Downloads a list of GCS files to a local directory.
+
+  Args:
+    gcs_uri_list: A list of GCS file paths.
+    local_dir: Local directory in which the GCS files are saved.
+
+  Returns:
+    The local file paths corresponding to the input GCS file paths.
+
+  Raises:
+    ValueError: An input file path is not a GCS path.
+  """
+  local_paths = []
+  for gcs_uri in gcs_uri_list:
+    if not is_gcs_path(gcs_uri):
+      raise ValueError(
+          f'{gcs_uri} is not a GCS path starting with'
+          f' {constants.GCS_URI_PREFIX}.'
+      )
+    local_path = os.path.join(local_dir, gcs_uri.replace('gs://', ''))
+    download_gcs_file_to_local(gcs_uri, local_path)
+    local_paths.append(local_path)
+  return local_paths
+
+
 def download_gcs_dir_to_local(
-    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
-):
+    gcs_dir: str,
+    local_dir: str,
+    skip_hf_model_bin: bool = False,
+    allow_patterns: Optional[List[str]] = None,
+    log: bool = True,
+) -> None:
   """Downloads files in a GCS directory to a local directory.
 
   For example:
@@ -78,16 +241,21 @@ def download_gcs_dir_to_local(
     gs://bucket/foo/a -> /tmp/bar/a
     gs://bucket/foo/b/c -> /tmp/bar/b/c
 
-  Arguments:
+  Args:
     gcs_dir: A string of directory path on GCS.
     local_dir: A string of local directory path.
     skip_hf_model_bin: True to skip downloading HF model bin files.
+    allow_patterns: A list of allowed patterns. If provided, only files matching
+      one or more patterns are downloaded.
+    log: True to log each downloaded file.
   """
   if not is_gcs_path(gcs_dir):
     raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.')
   bucket_name = gcs_dir.split('/')[2]
-  prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
-  client = storage.Client()
+  prefix = (
+      gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/') + '/'
+  )
+  client = _get_gcs_client()
   blobs = client.list_blobs(bucket_name, prefix=prefix)
   for blob in blobs:
     if blob.name[-1] == '/':
@@ -95,43 +263,63 @@ def download_gcs_dir_to_local(
     file_path = blob.name[len(prefix) :].strip('/')
     local_file_path = os.path.join(local_dir, file_path)
     os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+    if allow_patterns and all(
+        [not fnmatch.fnmatch(file_path, p) for p in allow_patterns]
+    ):
+      continue
     if (
         file_path.endswith(constants.HF_MODEL_WEIGHTS_SUFFIX)
         and skip_hf_model_bin
     ):
-      logging.info('Skip downloading model bin %s', file_path)
+      if log:
+        logging.info('Skip downloading model bin %s', file_path)
       with open(local_file_path, 'w') as f:
-        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}/{file_path}')
+        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}{file_path}')
     else:
-      logging.info('Downloading %s to %s', file_path, local_file_path)
+      if log:
+        logging.info('Downloading %s to %s', file_path, local_file_path)
       blob.download_to_filename(local_file_path)
 
 
+def _get_relative_paths(base_dir: str) -> List[str]:
+  """Gets relative paths of all files in a local base directory."""
+  path = pathlib.Path(base_dir)
+  relative_paths = []
+  for local_file in path.rglob('*'):
+    if os.path.isfile(local_file):
+      relative_path = os.path.relpath(local_file, base_dir)
+      relative_paths.append(relative_path)
+  return relative_paths
+
+
+def _upload_local_files_to_gcs(
+    relative_paths: List[str], local_dir: str, gcs_dir: str
+):
+  """Uploads local files to gcs."""
+  bucket_name = gcs_dir.split('/')[2]
+  blob_dir = '/'.join(gcs_dir.split('/')[3:])
+  client = _get_gcs_client()
+  bucket = client.bucket(bucket_name)
+  for relative_path in relative_paths:
+    blob = bucket.blob(os.path.join(blob_dir, relative_path))
+    blob.upload_from_filename(os.path.join(local_dir, relative_path))
+
+
 def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
   """Uploads local dir to gcs.
 
   For example:
     upload_local_dir_to_gcs(/tmp/bar, gs://bucket/foo)
-    gs://bucket/foo/a -> /tmp/bar/a
-    gs://bucket/foo/b/c -> /tmp/bar/b/c
+    /tmp/bar/a -> gs://bucket/foo/a
+    /tmp/bar/b/c -> gs://bucket/foo/b/c
 
   Arguments:
     local_dir: A string of local directory path.
     gcs_dir: A string of directory path on GCS.
   """
-  bucket_name = gcs_dir.split('/')[2]
-  blob_dir = '/'.join(gcs_dir.split('/')[3:])
-  client = storage.Client()
-  bucket = client.bucket(bucket_name)
-  for local_file in glob.glob(local_dir + '/**'):
-    if os.path.isfile(local_file):
-      logging.info(
-          'Uploading %s to %s',
-          local_file,
-          os.path.join(constants.GCS_URI_PREFIX, bucket_name, blob_dir),
-      )
-      blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file)))
-      blob.upload_from_filename(local_file)
+  # Relative paths of all files in local_dir.
+  relative_paths = _get_relative_paths(local_dir)
+  _upload_local_files_to_gcs(relative_paths, local_dir, gcs_dir)
 
 
 def upload_file_to_gcs_path(
@@ -155,7 +343,7 @@ def upload_file_to_gcs_path(
   if not source_path_obj.exists():
     raise RuntimeError(f'Source path does not exist: {source_path}')
 
-  storage_client = storage.Client()
+  storage_client = _get_gcs_client()
   source_file_path = source_path
   destination_file_uri = destination_uri
   logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
@@ -174,7 +362,9 @@ def is_gcs_path(input_path: str) -> bool:
   Returns:
       True if the input path is a GCS path, False otherwise.
   """
-  return input_path.startswith(constants.GCS_URI_PREFIX)
+  return input_path is not None and input_path.startswith(
+      constants.GCS_URI_PREFIX
+  )
 
 
 def release_text_assets(
@@ -232,13 +422,10 @@ def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
   """
   _, local_video_file_name = os.path.split(video_file_path)
   file_extension = os.path.splitext(video_file_path)[1]
-  if file_extension:
-    remote_video_file_name = local_video_file_name.replace(
-        file_extension, '_overlay.mp4'
-    )
-  else:
-    remote_video_file_name = local_video_file_name + '_overlay.mp4'
-  local_file_path = generate_tmp_path(file_extension)
+  remote_video_file_name = local_video_file_name.replace(
+      file_extension, '_overlay.mp4'
+  )
+  local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1])
   logging.info('Downloading %s to %s...', video_file_path, local_file_path)
   download_gcs_file_to_local(video_file_path, local_file_path)
   return local_file_path, remote_video_file_name
@@ -254,10 +441,19 @@ def get_output_video_file(video_output_file_path: str) -> str:
     str: Local video output file path.
   """
   file_extension = os.path.splitext(video_output_file_path)[1]
-  if file_extension:
-    out_local_video_file_name = video_output_file_path.replace(
-        file_extension, '_overlay' + file_extension
-    )
-  else:
-    out_local_video_file_name = video_output_file_path + '_overlay'
+  out_local_video_file_name = video_output_file_path.replace(
+      file_extension, '_overlay' + file_extension
+  )
   return out_local_video_file_name
+
+
+def 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

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

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