Update run_linter.sh

Fix spelling error `notebooked` → `notebooks`

.cloud-build/execute_changed_notebooks_helper.py

@@ -238,7 +238,7 @@ def _get_notebook_python_version(notebook_path: str) -> str:
 
             # Look for the python version specification pattern
             re_match = re.search(
-                r"python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
+                "python version = (\d+\.\d+)", markdown, flags=re.IGNORECASE
             )
             if re_match:
                 # get the version number
@@ -365,7 +365,7 @@ def process_and_execute_notebook(
             # Use gcloud to get tail
             try:
                 result.error_message = subprocess.check_output(
-                    ["gcloud", "storage", "cat", "--range", "-1000", log_file_uri], encoding="UTF-8"
+                    ["gsutil", "cat", "-r", "-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 gcloud storage. Run the following command in your terminal.")
-    print(f'\tgcloud storage cp "{output_file_or_uri}" .')
+    print("Option 1. Using gsutil. Run the following command in your terminal.")
+    print(f'\tgsutil 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,3 +1,5 @@
+notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
+notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
 notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
 notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
 .cloud-build/tests/python_version_test.ipynb

.cloud-build/utils/NotebookProcessors.py

@@ -108,7 +108,7 @@ class VertexAIInstallProprocessor(Preprocessor):
         if "google-cloud-aiplatform" not in content:
             return content
         return (
-            f"gcloud storage cp {self.vertex_ai_wheel} google-cloud-aiplatform.whl\n" + 
+            f"gsutil 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(
-        ["gcloud", "storage", "cp", remote_file_path, destination_file], encoding="UTF-8"
+        ["gsutil", "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(
-        ["gcloud", "storage", "cp", local_file_path, remote_file_path], encoding="UTF-8"
+        ["gsutil", "cp", local_file_path, remote_file_path], encoding="UTF-8"
     )
 
     return remote_file_path

.github/dependabot.yml

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

.github/workflows/ci.yaml

@@ -7,11 +7,11 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - name: Set up Python
-      uses: actions/setup-python@v6
+      uses: actions/setup-python@v5
       with:
-        python-version: '3.12'      
+        python-version: '3.x'      
     - name: Fetch pull request branch
-      uses: actions/checkout@v6
+      uses: actions/checkout@v4
       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.14
+FROM python:3.12
 
 WORKDIR setup
 

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==26.1.0
-pyupgrade==3.21.2
-isort==7.0.0
-flake8==7.3.0
-nbqa==1.9.1
+black==24.4.2
+pyupgrade==3.16.0
+isort==5.13.2
+flake8==7.1.0
+nbqa==1.8.5
 

community-content/CODEOWNERS

@@ -20,7 +20,6 @@
 /vertex_model_garden/model_oss/movinet @KCFindstr
 /vertex_model_garden/model_oss/data_converter @KCFindstr
 /vertex_model_garden/model_oss/peft @weigary
-/vertex_model_garden/model_oss/peft/templates @rayandasoriya
 /vertex_model_garden/model_oss/lm-evaluation-harness @kathyyu-google
 /vertex_model_garden/model_oss/tfvision @dstnluong-google
 /vertex_model_garden/model_oss/fvlm @minwoo33park
@@ -29,5 +28,4 @@
 /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.
-      # gcloud storage cp "$model_archive_uri" "$model_archive_local_path"
+      # gsutil 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" != */ ]] && (gcloud storage ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
+            if [[ "$uri" != */ ]] && (gsutil ls "$uri" | grep --fixed-strings --line-regexp "$uri"); then
                 mkdir -p "$(dirname "$output_path")"
-                gcloud storage cp --recursive "$uri" "$output_path"
+                gsutil -m cp -r "$uri" "$output_path"
             else
                 mkdir -p "$output_path" # When source path is a directory, gsutil requires the destination to also be a directory
-                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.
+                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.
             fi
         - inputValue: GCS path
         - outputPath: Data

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile

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

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt

@@ -1,3 +1,3 @@
-torch==2.8.0
+torch==1.13.1
 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.7.0
+torch==2.2.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": [
-        "! gcloud storage ls $gcs_output_uri_prefix"
+        "! gsutil ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -192,7 +192,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage cp --recursive $gcs_output_uri_prefix/model ./model_server/"
+        "! gsutil cp -r $gcs_output_uri_prefix/model ./model_server/"
       ]
     },
     {
@@ -556,7 +556,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
+        "! gsutil rm -rf $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": [
-        "! gcloud storage ls $gcs_output_uri_prefix"
+        "! gsutil 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:"
-# gcloud storage ls ${JOB_DIR}/
+# gsutil 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}", ";",     "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"]
+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"]

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:"
-# gcloud storage ls ${JOB_DIR}/
+# gsutil 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,7 +478,8 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage buckets create --location $REGION $BUCKET_NAME"      ]
+        "! gsutil mb -l $REGION $BUCKET_NAME"
+      ]
     },
     {
       "cell_type": "markdown",
@@ -497,7 +498,8 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage ls --all-versions --long $BUCKET_NAME"      ]
+        "! gsutil ls -al $BUCKET_NAME"
+      ]
     },
     {
       "cell_type": "markdown",
@@ -580,7 +582,8 @@
       "outputs": [],
       "source": [
         "# Download the sample data into your RAW_DATA_PATH\n",
-        "! 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"      ]
+        "! 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"
+      ]
     },
     {
       "cell_type": "code",
@@ -1618,7 +1621,9 @@
         "! gcloud scheduler jobs delete $SIMULATOR_SCHEDULER_JOB --quiet\n",
         "\n",
         "# Delete Cloud Storage objects that were created.\n",
-        "! gcloud storage rm --recursive $PIPELINE_ROOT\n",        "! gcloud storage rm --recursive $TRAINING_ARTIFACTS_DIR"      ]
+        "! gsutil -m rm -r $PIPELINE_ROOT\n",
+        "! gsutil -m rm -r $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
+tensorflow==2.7.2
 pillow==10.3.0
 tf-agents==0.8.0

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

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

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

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

community-content/tf_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,7 +398,6 @@
         "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",
@@ -473,7 +472,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage buckets create --location $REGION $BUCKET_NAME"
+        "! gsutil mb -l $REGION $BUCKET_NAME"
       ]
     },
     {
@@ -493,7 +492,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage ls --all-versions --long $BUCKET_NAME"
+        "! gsutil ls -al $BUCKET_NAME"
       ]
     },
     {
@@ -566,7 +565,7 @@
       "outputs": [],
       "source": [
         "# Copy the sample data into your DATA_PATH\n",
-        "! 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"
+        "! gsutil cp \"gs://cloud-samples-data/vertex-ai/community-content/tf_agents_bandits_movie_recommendation_with_kfp_and_vertex_sdk/u.data\"  $DATA_PATH"
       ]
     },
     {
@@ -580,15 +579,11 @@
         "# 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 = (\n",
-        "    2  # @param {type:\"integer\"} Number of driver steps per training iteration.\n",
-        ")\n",
+        "STEPS_PER_LOOP = 2  # @param {type:\"integer\"} Number of driver steps per training iteration.\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 = (\n",
-        "    20  # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
-        ")\n",
+        "NUM_ACTIONS = 20  # @param {type:\"integer\"} Number of actions (movie items) to choose from.\n",
         "PER_ARM = False  # Use the non-per-arm version of the MovieLens environment.\n",
         "\n",
         "# Set agent parameters.\n",
@@ -626,8 +621,7 @@
       "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",
-        ")\n",
+        "    DATA_PATH, RANK_K, BATCH_SIZE, num_movies=NUM_ACTIONS, csv_delimiter=\"\\t\")\n",
         "environment = tf_py_environment.TFPyEnvironment(env)\n",
         "\n",
         "# Define RL agent/algorithm.\n",
@@ -637,8 +631,7 @@
         "    tikhonov_weight=TIKHONOV_WEIGHT,\n",
         "    alpha=AGENT_ALPHA,\n",
         "    dtype=tf.float32,\n",
-        "    accepts_per_arm_features=PER_ARM,\n",
-        ")\n",
+        "    accepts_per_arm_features=PER_ARM)\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",
@@ -646,8 +639,7 @@
         "# Define RL metric.\n",
         "optimal_reward_fn = functools.partial(\n",
         "    environment_utilities.compute_optimal_reward_with_movielens_environment,\n",
-        "    environment=environment,\n",
-        ")\n",
+        "    environment=environment)\n",
         "regret_metric = tf_bandit_metrics.RegretMetric(optimal_reward_fn)\n",
         "metrics = [regret_metric]"
       ]
@@ -712,38 +704,35 @@
         "    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(agent.policy.trajectory_spec)\n",
-        "    replay_buffer = trainer.get_replay_buffer(\n",
-        "        data_spec, environment.batch_size, steps_per_loop\n",
-        "    )\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",
         "\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 += [\n",
-        "            tf_metrics.AverageReturnMultiMetric(\n",
-        "                reward_spec=environment.reward_spec(), batch_size=environment.batch_size\n",
-        "            )\n",
-        "        ]\n",
+        "        metrics += [tf_metrics.AverageReturnMultiMetric(\n",
+        "            reward_spec=environment.reward_spec(),\n",
+        "            batch_size=environment.batch_size)]\n",
         "    else:\n",
-        "        metrics += [tf_metrics.AverageReturnMetric(batch_size=environment.batch_size)]\n",
+        "        metrics += [\n",
+        "            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",
-        "\n",
-        "        def add_batch_fn(data):\n",
-        "            return replay_buffer.add_batch(training_data_spec_transformation_fn(data))\n",
-        "\n",
+        "        def add_batch_fn(data): return replay_buffer.add_batch(training_data_spec_transformation_fn(data)) \n",
+        "        \n",
         "    else:\n",
         "        add_batch_fn = replay_buffer.add_batch\n",
         "\n",
@@ -753,12 +742,10 @@
         "        env=environment,\n",
         "        policy=agent.collect_policy,\n",
         "        num_steps=steps_per_loop * environment.batch_size,\n",
-        "        observers=observers,\n",
-        "    )\n",
+        "        observers=observers)\n",
         "\n",
         "    training_loop = trainer.get_training_loop_fn(\n",
-        "        driver, replay_buffer, agent, steps_per_loop\n",
-        "    )\n",
+        "        driver, replay_buffer, agent, steps_per_loop)\n",
         "    saver = policy_saver.PolicySaver(agent.policy)\n",
         "\n",
         "    for _ in range(training_loops):\n",
@@ -796,8 +783,7 @@
         "    environment=environment,\n",
         "    training_loops=TRAINING_LOOPS,\n",
         "    steps_per_loop=STEPS_PER_LOOP,\n",
-        "    additional_metrics=metrics,\n",
-        ")\n",
+        "    additional_metrics=metrics)\n",
         "\n",
         "tf.profiler.experimental.stop()"
       ]
@@ -1106,15 +1092,11 @@
       },
       "outputs": [],
       "source": [
-        "RUN_HYPERPARAMETER_TUNING = (\n",
-        "    True  # Execute hyperparameter tuning instead of regular training.\n",
-        ")\n",
+        "RUN_HYPERPARAMETER_TUNING = True  # Execute hyperparameter tuning instead of regular training.\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(\n",
-        "    HPTUNING_RESULT_DIR, \"result.json\"\n",
-        ")  # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
+        "HPTUNING_RESULT_PATH = os.path.join(HPTUNING_RESULT_DIR, \"result.json\")  # @param {type: \"string\"} Path to the file containing the best hyperparameter(s)."
       ]
     },
     {
@@ -1142,7 +1124,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",
@@ -1215,8 +1197,8 @@
         "\n",
         "    # Create job\n",
         "    response = client.create_hyperparameter_tuning_job(\n",
-        "        parent=parent, hyperparameter_tuning_job=hyperparameter_tuning_job\n",
-        "    )\n",
+        "        parent=parent,\n",
+        "        hyperparameter_tuning_job=hyperparameter_tuning_job)\n",
         "    job_id = response.name.split(\"/\")[-1]\n",
         "    print(\"Job ID:\", job_id)\n",
         "    print(\"Job config:\", response)\n",
@@ -1260,8 +1242,7 @@
         "    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\",\n",
-        ")"
+        "    api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")"
       ]
     },
     {
@@ -1311,8 +1292,7 @@
         "    name = client.hyperparameter_tuning_job_path(\n",
         "        project=project,\n",
         "        location=location,\n",
-        "        hyperparameter_tuning_job=hyperparameter_tuning_job_id,\n",
-        "    )\n",
+        "        hyperparameter_tuning_job=hyperparameter_tuning_job_id)\n",
         "    response = client.get_hyperparameter_tuning_job(name=name)\n",
         "    return response"
       ]
@@ -1333,8 +1313,7 @@
         "        location=REGION,\n",
         "        api_endpoint=f\"{REGION}-aiplatform.googleapis.com\")\n",
         "    if response.state.name == 'JOB_STATE_SUCCEEDED':\n",
-        "        print(\"Job succeeded.\n",
-        "Job Time:\", response.update_time - response.create_time)\n",
+        "        print(\"Job succeeded.\\nJob Time:\", response.update_time - response.create_time)\n",
         "        trials = response.trials\n",
         "        print(\"Trials:\", trials)\n",
         "        break\n",
@@ -1369,8 +1348,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], -np.inf\n",
-        "    )\n",
+        "        [metric.metric_id for metric in trials[0].final_measurement.metrics],\n",
+        "        -np.inf)\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",
@@ -1379,13 +1358,12 @@
         "        # `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 = {param.parameter_id: param.value for param in trial.parameters}\n",
+        "            params = {\n",
+        "                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(\n",
-        "                    params\n",
-        "                )  # Handle cases where multiple hyperparameter values lead to the same performance.\n",
+        "                best_params[param.parameter_id].append(params)  # 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",
@@ -1465,9 +1443,7 @@
       },
       "outputs": [],
       "source": [
-        "PREDICTION_CONTAINER = (\n",
-        "    \"prediction-custom-container\"  # @param {type:\"string\"} Name of the container image.\n",
-        ")"
+        "PREDICTION_CONTAINER = \"prediction-custom-container\"  # @param {type:\"string\"} Name of the container image."
       ]
     },
     {
@@ -1499,7 +1475,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",
@@ -1616,12 +1592,8 @@
       },
       "outputs": [],
       "source": [
-        "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",
-        ")"
+        "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."
       ]
     },
     {
@@ -1661,12 +1633,10 @@
         "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\"]\n",
-        "    + args,  # Pass in training arguments, including hyperparameters.\n",
+        "    command=[\"python3\", \"-m\", \"src.training.task\"] + 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",
-        ")\n",
+        "    model_serving_container_health_route=\"/health\")\n",
         "\n",
         "print(\"Training Spec:\", job._managed_model)\n",
         "\n",
@@ -1675,8 +1645,7 @@
         "    replica_count=1,\n",
         "    machine_type=\"n1-standard-4\",\n",
         "    accelerator_type=\"ACCELERATOR_TYPE_UNSPECIFIED\",\n",
-        "    accelerator_count=0,\n",
-        ")"
+        "    accelerator_count=0)"
       ]
     },
     {
@@ -1815,7 +1784,7 @@
         "! gcloud ai models delete $model.name --quiet\n",
         "\n",
         "# Delete Cloud Storage objects that were created\n",
-        "! gcloud storage rm --recursive $ARTIFACTS_DIR"
+        "! gsutil -m rm -r $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": [
-        "! gcloud storage ls $gcs_output_uri_prefix"
+        "! gsutil ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -344,7 +344,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
+        "! gsutil rm -rf $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": [
-        "! gcloud storage ls $gcs_output_uri_prefix"
+        "! gsutil ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -348,7 +348,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
+        "! gsutil rm -rf $gcs_output_uri_prefix"
       ]
     }
   ],

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/requirements.txt

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

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/trainer/requirements.txt

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

community-content/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": [
-        "! gcloud storage ls $gcs_output_uri_prefix"
+        "! gsutil ls $gcs_output_uri_prefix"
       ]
     },
     {
@@ -361,7 +361,7 @@
       },
       "outputs": [],
       "source": [
-        "! gcloud storage rm --recursive --continue-on-error $gcs_output_uri_prefix"
+        "! gsutil rm -rf $gcs_output_uri_prefix"
       ]
     }
   ],

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

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

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

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

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

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

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

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

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

@@ -1,10 +0,0 @@
-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.6.3
-google-cloud-storage==3.0.0
-retrying

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

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

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

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

@@ -1,49 +0,0 @@
-absl-py==2.2.2
-annotated-types==0.7.0
-anyio==4.9.0
-black==25.1.0
-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.92.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.10
-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.2
-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.32.4
-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.6.3
-websockets==15.0.1

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

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

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

@@ -1,81 +0,0 @@
-"""Get cluster info from environment variables."""
-
-import dataclasses
-import json
-import os
-
-from absl import logging
-
-
-@dataclasses.dataclass
-class ClusterInfo:
-    """Contains information about the cluster.
-
-    Attributes:
-      primary_node_addr: The address of the primary node.
-      primary_node_port: The port of the primary node.
-      node_rank: The rank of the node.
-      num_nodes: The number of nodes in the cluster.
-    """
-
-    primary_node_addr: str | None = None
-    primary_node_port: str | None = None
-    node_rank: int = 0
-    num_nodes: int = 1
-
-    # Allows unpacking operation like
-    # primary_node_addr, primary_node_port, _, _ = ClusterInfo()
-    # See https://stackoverflow.com/a/70753113
-    def __iter__(self):
-        return iter(dataclasses.astuple(self))
-
-
-def get_cluster_spec() -> ClusterInfo:
-    """Parses CLUSTER_SPEC environment variable and returns the cluster info.
-
-    Returns:
-      A ClusterInfo object.
-    """
-    cluster_spec = os.getenv("CLUSTER_SPEC", None)
-
-    # If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
-    if not cluster_spec:
-        cluster_info = ClusterInfo(
-            primary_node_addr=os.getenv("MASTER_ADDR", None),
-            primary_node_port=os.getenv("MASTER_PORT", None),
-            node_rank=int(os.getenv("RANK", "0")),
-            num_nodes=int(os.getenv("NNODES", "1")),
-        )
-        return cluster_info
-
-    cluster_data = json.loads(cluster_spec)
-    # Get primary node info
-    primary_node = cluster_data["cluster"]["workerpool0"][0]
-    logging.info("primary node: %s", primary_node)
-    primary_node_addr, primary_node_port = primary_node.split(":")
-    logging.info("primary node address: %s", primary_node_addr)
-    logging.info("primary node port: %s", primary_node_port)
-
-    # Determine node rank of this machine
-    workerpool = cluster_data["task"]["type"]
-    if workerpool == "workerpool0":
-        node_rank = 0
-    elif workerpool == "workerpool1":
-        # Add 1 for the primary node, since `index` is the index of workerpool1.
-        node_rank = cluster_data["task"]["index"] + 1
-    else:
-        raise ValueError(
-            "Only workerpool0 and workerpool1 are supported. Unknown workerpool:"
-            f" {workerpool}"
-        )
-    logging.info("node rank: %s", node_rank)
-
-    # Calculate total nodes.
-    num_nodes = 1  # For the primary node.
-    if "workerpool1" in cluster_data["cluster"]:
-        num_nodes += len(cluster_data["cluster"]["workerpool1"])
-    logging.info("num nodes: %s", num_nodes)
-
-    return ClusterInfo(
-        primary_node_addr, primary_node_port, node_rank, num_nodes
-    )

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

@@ -1,59 +0,0 @@
-"""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_Classifier.py

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

community-content/vertex_cpr_samples/sklearn/predictor_LinearRegression.py

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

community-content/vertex_cpr_samples/sklearn/predictor_SGDClassifier.py

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

community-content/vertex_cpr_samples/torch/predictor_resnet.py

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

community-content/vertex_cpr_samples/xgboost/predictor_Classifier.py

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

community-content/vertex_cpr_samples/xgboost/predictor_xgbranker.py

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

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

@@ -66,7 +66,7 @@ mkdir -p "$local_folder"
 mkdir -p "$output_folder"
 
 # Download the content from the GCS URI
-gcloud storage cp --recursive "$gcs_dataset_path"/* "$local_folder/"
+gsutil -m cp -r "$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
-gcloud storage cp --recursive "$images_folder"/images_2/*  "$gcs_experiment_path"/data/images_2
+gsutil -m cp -r "$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
-gcloud storage cp --recursive "$images_folder"/images_4/*  "$gcs_experiment_path"/data/images_4
+gsutil -m cp -r "$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
-gcloud storage cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
+gsutil -m cp "$images_folder"/images_8/*  "$gcs_experiment_path"/data/images_8
 
 # Copy images and sparse reconstruction files to gcs experiment folder.
-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
+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
 
 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
-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
+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
 
 # 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"
-  gcloud storage cp "${args[-gcs_keyframes_file]}" "$local_keyframes_file" || exit 1
+  gsutil 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.
-gcloud storage cp --recursive "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"
+gsutil -m cp -r "$OUTPUT_RENDER_PATH" "${args[-gcs_experiment_path]}/render/${rendering_job_name}"

community-content/vertex_model_garden/model_oss/cloudnerf/requirements.txt

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

community-content/vertex_model_garden/model_oss/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.
-gcloud storage cp --recursive "${gcs_experiment_path}/data" "$local_experiment_path/$scene_folder_name" || exit 1
+gsutil -m cp -r "${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}"
 
-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}"
+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}"

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -1,95 +0,0 @@
-"""Class that bundles docker related flags."""
-
-import getpass
-import os
-import pwd
-
-
-class CommandBuilder:
-  """Base class for building commands."""
-
-  def __init__(self):
-    self._defaults = []
-    self._env_vars = {}
-
-  def add_env_var(self, var: str, val: str) -> None:
-    """Add environment variable to the command.
-
-    Args:
-      var: environment variable name.
-      val: environment variable value.
-    """
-    self._env_vars[var] = val
-
-  def add_mount_map(self, host_path, docker_path):
-    pass
-
-
-class DockerCommandBuilder(CommandBuilder):
-  """Bundle docker related flags."""
-
-  def __init__(self, docker_uri: str, shm_size: str = '128gb'):
-    super().__init__()
-    self._docker_uri = [docker_uri]
-    self.privilege_mode = []
-    self.entrypoint = []
-
-    self._defaults = [
-        'docker',
-        'run',
-        '--gpus=all',
-        '--net=host',
-        '--rm',
-        f'--shm-size={shm_size}',
-    ]
-
-    self._mount_maps = []
-    user = getpass.getuser()
-    # username ends with `_google_com` is managed by ldap and does not have a
-    # corresponding entry in /etc/passwd or /etc/group file. We cannot enable
-    # non-root docker user with below method.
-    if not user.endswith('_google_com'):
-      uid = os.getuid()
-      gid = pwd.getpwuid(uid).pw_gid
-      self._defaults += [
-          f'--user={uid}:{gid}',
-          '--volume=/etc/group:/etc/group:ro',
-          '--volume=/etc/passwd:/etc/passwd:ro',
-      ]
-
-  def add_mount_map(self, host_path, docker_path):
-    self._mount_maps.append(f'--volume={host_path}:{docker_path}')
-
-  def add_privilege_mode(self):
-    self.privilege_mode = ['--privileged']
-
-  def add_entrypoint(self, entrypoint: list[str]):
-    self.entrypoint = entrypoint
-
-  def build_cmd(self) -> str:
-    return (
-        self._defaults
-        + [f'--env={var}={val}' for var, val in self._env_vars.items()]
-        + self._mount_maps
-        + self.privilege_mode
-        + self._docker_uri
-        + self.entrypoint
-    )
-
-
-class PythonCommandBuilder(CommandBuilder):
-  """Bundle Python test command related flags."""
-
-  def __init__(self):
-    super().__init__()
-    self._defaults = [
-        'python3',
-        './vertex_vision_model_garden_peft/train/vmg/train_entrypoint.py',
-    ]
-
-  def build_cmd(self) -> str:
-    os.environ.update(self._env_vars)
-    return self._defaults
-
-  def add_entrypoint(self, entrypoint: list[str]):
-    self._defaults = entrypoint

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

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

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

@@ -1,79 +0,0 @@
-"""Get cluster info from environment variables."""
-
-import dataclasses
-import json
-import os
-
-from absl import logging
-
-
-@dataclasses.dataclass
-class ClusterInfo:
-  """Contains information about the cluster.
-
-  Attributes:
-    primary_node_addr: The address of the primary node.
-    primary_node_port: The port of the primary node.
-    node_rank: The rank of the node.
-    num_nodes: The number of nodes in the cluster.
-  """
-
-  primary_node_addr: str | None = None
-  primary_node_port: str | None = None
-  node_rank: int = 0
-  num_nodes: int = 1
-
-  # Allows unpacking operation like
-  # primary_node_addr, primary_node_port, _, _ = ClusterInfo()
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def get_cluster_spec() -> ClusterInfo:
-  """Parses CLUSTER_SPEC environment variable and returns the cluster info.
-
-  Returns:
-    A ClusterInfo object.
-  """
-  cluster_spec = os.getenv('CLUSTER_SPEC', None)
-
-  # If CLUSTER_SPEC is not set, use individual vars to construct cluster info.
-  if not cluster_spec:
-    cluster_info = ClusterInfo(
-        primary_node_addr=os.getenv('MASTER_ADDR', None),
-        primary_node_port=os.getenv('MASTER_PORT', None),
-        node_rank=int(os.getenv('RANK', '0')),
-        num_nodes=int(os.getenv('NNODES', '1')),
-    )
-    return cluster_info
-
-  cluster_data = json.loads(cluster_spec)
-  # Get primary node info
-  primary_node = cluster_data['cluster']['workerpool0'][0]
-  logging.info('primary node: %s', primary_node)
-  primary_node_addr, primary_node_port = primary_node.split(':')
-  logging.info('primary node address: %s', primary_node_addr)
-  logging.info('primary node port: %s', primary_node_port)
-
-  # Determine node rank of this machine
-  workerpool = cluster_data['task']['type']
-  if workerpool == 'workerpool0':
-    node_rank = 0
-  elif workerpool == 'workerpool1':
-    # Add 1 for the primary node, since `index` is the index of workerpool1.
-    node_rank = cluster_data['task']['index'] + 1
-  else:
-    raise ValueError(
-        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
-        f' {workerpool}'
-    )
-  logging.info('node rank: %s', node_rank)
-
-  # Calculate total nodes.
-  num_nodes = 1  # For the primary node.
-  if 'workerpool1' in cluster_data['cluster']:
-    num_nodes += len(cluster_data['cluster']['workerpool1'])
-  logging.info('num nodes: %s', num_nodes)
-
-  return ClusterInfo(primary_node_addr, primary_node_port, node_rank, num_nodes)

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

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

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

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

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

@@ -1,140 +0,0 @@
-"""Util functions for reporting device (GPU, CPU) stats."""
-
-import dataclasses
-
-import psutil
-import pynvml
-import torch
-
-
-@dataclasses.dataclass
-class GpuStats:
-  """Holds information about GPU usage stats.
-
-  For memory related, see
-  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
-  """
-
-  # device id
-  device_id: int
-  # memory reserved.
-  reserved: float
-  # memory occupied.
-  occupied: float
-  # memory reserved, but not used.
-  unused: float
-  # nvidia-smi usually reports more memory usages than pytorch (for driver,
-  # kernel and etc). `smi_diff` tracks this difference.
-  smi_diff: float
-  # Gpu utilization.
-  util: float
-
-  # Allows unpacking operation like
-  # device_id, reserved, occupied, unused, smi_diff, util = GpuStats(...)
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def gpu_stats() -> GpuStats:
-  """Reports GPU memory usage and utilization."""
-  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
-  bytes_per_gb = 1024.0**3
-  device = torch.cuda.current_device()
-  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
-  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
-  unused = reserved - occupied
-
-  def smi_mem(device):
-    try:
-      pynvml.nvmlInit()
-      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
-      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
-      return info.used / bytes_per_gb
-    except pynvml.NVMLError:
-      return 0.0
-
-  mem_used_smi = smi_mem(device)
-  smi_diff = mem_used_smi - reserved
-
-  util = torch.cuda.utilization(device)
-  return GpuStats(device, reserved, occupied, unused, smi_diff, util)
-
-
-def gpu_stats_str(stats: GpuStats | None = None) -> str:
-  if stats is None:
-    stats = gpu_stats()
-  device, reserved, occupied, unused, smi_diff, util = stats
-  return (
-      f"GPU ({device=}) memory: {reserved:.2f}({occupied=:.2f}, {unused=:.2f}),"
-      f" {smi_diff=:.2f} GB. Utilization: {util:.2f}%"
-  )
-
-
-@dataclasses.dataclass
-class CpuStats:
-  """Holds information about CPU usage stats."""
-
-  # Total CPU virtual memory i.e. virtual memory allocated + unallocated.
-  total_virtual_mem: float
-  # CPU virtual memory available for use.
-  unallocated_virtual_mem: float
-  # CPU virtual memory already used.
-  allocated_virtual_mem: float
-  # Total CPU swap memory i.e. swap memory allocated + unallocated.
-  total_swap_mem: float
-  # CPU swap memory available for use.
-  unallocated_swap_mem: float
-  # CPU swap memory already used.
-  allocated_swap_mem: float
-  # CPU utilization percentage.
-  utilization: float
-
-
-def cpu_stats() -> CpuStats:
-  """Reports CPU memory usage and utilization."""
-
-  # https://psutil.readthedocs.io/en/latest/#memory
-  gb = 1024.0**3
-  vmem = psutil.virtual_memory()
-  vmem_total = vmem.total / gb
-  vmem_available = vmem.available / gb
-  vmem_used = vmem_total - vmem_available
-  smem = psutil.swap_memory()
-  swap_total = smem.total / gb
-  swap_free = smem.free / gb
-  swap_used = smem.used / gb
-  # https://psutil.readthedocs.io/en/latest/#psutil.cpu_percent
-  cpu_util = psutil.cpu_percent(interval=1e-6)
-  return CpuStats(
-      total_virtual_mem=vmem_total,
-      unallocated_virtual_mem=vmem_available,
-      allocated_virtual_mem=vmem_used,
-      total_swap_mem=swap_total,
-      unallocated_swap_mem=swap_free,
-      allocated_swap_mem=swap_used,
-      utilization=cpu_util,
-  )
-
-
-def cpu_stats_str(stats: CpuStats | None = None) -> str:
-  """Returns a string representation of the CPU stats."""
-
-  if stats is None:
-    stats = cpu_stats()
-  total, occupied, unused = (
-      stats.total_virtual_mem,
-      stats.allocated_virtual_mem,
-      stats.unallocated_virtual_mem,
-  )
-  virtual_mem = (
-      f"CPU virtual memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
-  )
-  total, occupied, unused = (
-      stats.total_swap_mem,
-      stats.allocated_swap_mem,
-      stats.unallocated_swap_mem,
-  )
-  swap_mem = f"CPU swap memory: {total:.2f}({occupied=:.2f}, {unused=:.2f}) GB"
-  percent = stats.utilization
-  return f"{virtual_mem} {swap_mem} CPU Utilization: {percent:.2f}%"

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

@@ -1,126 +0,0 @@
-"""Different trainer callbacks for PEFT Trainer."""
-
-from collections.abc import MutableMapping
-import math
-import time
-
-from absl import logging
-import accelerate
-from transformers import TrainingArguments
-from transformers.trainer_callback import TrainerCallback
-from transformers.trainer_callback import TrainerControl
-from transformers.trainer_callback import TrainerState
-
-from util import device_stats
-
-
-class TrainerStatsCallback(TrainerCallback):
-  """Trainer callback to report trainer stats."""
-
-  def __init__(self, max_seq_length, filename=None):
-    self._max_seq_length = max_seq_length
-    self._filename = filename
-
-    self._partial_state = accelerate.PartialState()
-    self._start_time = float('nan')
-    self._prev_time = float('nan')
-    self._peak_mem = 0.0
-    self._avg_throughput = 0.0
-
-  def on_log(
-      self,
-      args: TrainingArguments,
-      state: TrainerState,
-      control: TrainerControl,
-      logs: MutableMapping[str, float] | None = None,
-      **kwargs,
-  ) -> None:
-    """Calculates perplexity from train loss.
-
-    Args:
-      args: Arguments passed to the trainer.
-      state: State of the trainer.
-      control: Control of the trainer.
-      logs: A dict of logs from the training loop.
-      **kwargs: Additional keyword arguments, not used in this callback.
-    """
-    del kwargs  # Unused.
-    if self._partial_state.is_main_process:
-      train_loss = logs.get('loss') if logs is not None else None
-      if train_loss is not None:
-        perplexity = round(float(math.exp(train_loss)), 4)
-        logs['perplexity'] = perplexity
-
-  def on_step_end(
-      self,
-      args: TrainingArguments,
-      state: TrainerState,
-      control: TrainerControl,
-      **kwargs,
-  ):
-    if self._partial_state.is_main_process:
-      if state.global_step == 1:
-        self._prev_time = time.time()
-        self._prev_num_token = state.num_input_tokens_seen
-        throughput = 0.0
-      else:
-        cur_time = time.time()
-        cur_num_token = state.num_input_tokens_seen
-        throughput = (cur_num_token - self._prev_num_token) / (
-            cur_time - self._prev_time
-        )
-        self._prev_time = cur_time
-        self._prev_num_token = cur_num_token
-        self._avg_throughput += (throughput - self._avg_throughput) / (
-            state.global_step - 1
-        )
-
-      gpu_stats = device_stats.gpu_stats()
-      self._peak_mem = max(
-          gpu_stats.reserved + gpu_stats.smi_diff, self._peak_mem
-      )
-      logging.info(
-          'on_step_end: Throughput: %.2f token/s. %s, %s',
-          throughput,
-          device_stats.gpu_stats_str(gpu_stats),
-          device_stats.cpu_stats_str(),
-      )
-
-  def on_train_begin(
-      self,
-      args: TrainingArguments,
-      state: TrainerState,
-      control: TrainerControl,
-      **kwargs,
-  ):
-    if self._partial_state.is_main_process:
-      self._start_time = time.time()
-      logging.info(
-          'on_train_begin: %s, %s',
-          device_stats.gpu_stats_str(),
-          device_stats.cpu_stats_str(),
-      )
-
-  def on_train_end(
-      self,
-      args: TrainingArguments,
-      state: TrainerState,
-      control: TrainerControl,
-      **kwargs,
-  ):
-    if self._partial_state.is_main_process:
-      train_time = time.time() - self._start_time
-      throughput = state.num_input_tokens_seen / train_time
-      logging.info(
-          'training time %.2f s, throughput (including overhead, e.g., ckpt'
-          ' saving): %.2f token/s, peak_mem: %.2f GB',
-          train_time,
-          throughput,
-          self._peak_mem,
-      )
-      if self._filename:
-        with open(self._filename, 'a') as out_f:
-          out_f.write(
-              f'{self._max_seq_length/1024.0:.1f} | {self._peak_mem:.2f} |'
-              f' {self._avg_throughput:.2f}\n'
-          )

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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