fix: official standard

.cloud-build/cleanup/cleanup.py

@@ -16,7 +16,6 @@ from resource_cleanup_manager import (
     ResourceCleanupManager,
     MatchingEngineIndexEndpointResourceCleanupManager,
     MatchingEngineIndexResourceCleanupManager,
-    FeatureStoreLegacyCleanupManager,
     FeatureStoreCleanupManager,
     PipelineJobCleanupManager,
     TrainingJobCleanupManager,
@@ -36,10 +35,7 @@ def run_cleanup_managers(managers: List[ResourceCleanupManager], is_dry_run: boo
 
         print(f"Fetching {type_name}'s...")
         resources = manager.list()
-        try:
-            print(f"Found {len(resources)} {type_name}'s")
-        except Exception as e:
-            print(f"{type_name} {e}")
+        print(f"Found {len(resources)} {type_name}'s")
         for resource in resources:
             try:
                 if not manager.is_deletable(resource):
@@ -66,7 +62,6 @@ managers: List[ResourceCleanupManager] = [
     ModelResourceCleanupManager(),  # ModelResourceCleanupManager must follow EndpointResourceCleanupManager due to deployed models blocking model deletion.
     MatchingEngineIndexEndpointResourceCleanupManager(),
     MatchingEngineIndexResourceCleanupManager(),
-    FeatureStoreLegacyCleanupManager(),
     FeatureStoreCleanupManager(),
     PipelineJobCleanupManager(),
     TrainingJobCleanupManager(),

.cloud-build/cleanup/resource_cleanup_manager.py

@@ -12,15 +12,9 @@ from typing import Any, Type
 
 from google.cloud import aiplatform
 from google.cloud.aiplatform import base
-from google.cloud.aiplatform_v1beta1 import (FeatureOnlineStoreAdminServiceClient,
-                                             FeatureOnlineStore)
 from google.cloud import storage
 from proto.datetime_helpers import DatetimeWithNanoseconds
 
-PROJECT_ID = "python-docs-samples-tests"
-REGION = "us-central1"
-API_ENDPOINT = f"{REGION}-aiplatform.googleapis.com"
-
 # If a resource was updated within this number of seconds, do not delete.
 RESOURCE_UPDATE_BUFFER_IN_SECONDS = 60 * 60 * 8
 
@@ -135,51 +129,12 @@ class MatchingEngineIndexEndpointResourceCleanupManager(VertexAIResourceCleanupM
         resource.undeploy_all()
         resource.delete(force=True)
 
-class FeatureStoreLegacyCleanupManager(VertexAIResourceCleanupManager):
-    # TODO: only deleting legacy
-    #    not deleting ingestions jobs
-    #       ingest_from_xxx methods do not return a job ID, there is no list command, aka no python way to delete
-    #    not deleting batch serving jobs
-    #       batch_serve_to_xxx methods do not return a job ID, there is no list command, aka no python way to delete
+class FeatureStoreCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.Featurestore
 
     def resource_name(self, resource: Any) -> str:
         return resource.name
 
-    def delete(self, resource):
-        resource.delete(force=True)
-
-
-class FeatureStoreCleanupManager(VertexAIResourceCleanupManager):
-    # for FS 2.0
-    # TODO: use _v1beta1, and gapic clients
-    #    delete features, feature groups, feature views, feature online stores
-    vertex_ai_resource = FeatureOnlineStore
-
-    admin_client = FeatureOnlineStoreAdminServiceClient(
-        client_options={"api_endpoint": API_ENDPOINT}
-    )
-
-    def resource_name(self, resource: Any) -> str:
-        return resource.name
-
-    def type_name(self) -> str:
-        return "FeatureOnlineStore"
-
-    def list(self) -> Any:
-        try:
-            return self.admin_client.list_feature_online_stores(parent=f"projects/{PROJECT_ID}/locations/{REGION}")
-        except Exception as e:
-            print(e)
-            return []
-
-    def delete(self, resource):
-        try:
-            self.admin_client.delete_feature_online_store(name=resource.name, force=True)
-        except Exception as e:
-            print(e)
-
-
 class PipelineJobCleanupManager(VertexAIResourceCleanupManager):
     vertex_ai_resource = aiplatform.PipelineJob
 

.cloud-build/execute_changed_notebooks_cli.py

@@ -18,7 +18,6 @@
 import argparse
 import pathlib
 import os
-import csv
 
 import execute_changed_notebooks_helper
 
@@ -38,7 +37,7 @@ parser = argparse.ArgumentParser(description="Run changed notebooks.")
 parser.add_argument(
     "--test_paths_file",
     type=pathlib.Path,
-    help="The path to the file that has newline-delimited folders of notebooks that should be tested.",
+    help="The path to the file that has newline-limited folders of notebooks that should be tested.",
     required=True,
 )
 parser.add_argument(
@@ -129,20 +128,6 @@ parser.add_argument(
     default=10,
     required=False,
 )
-parser.add_argument(
-    "--run_first_file",
-    type=pathlib.Path,
-    help="The path to the file that has newline-delimited of notebooks to run in the first batch",
-    default=None,
-    required=False,
-)
-parser.add_argument(
-    "--aiplatform_whl",
-    type=str,
-    help="The GCS path to a whl version google-cloud-aiplatform",
-    default=None,
-    required=False,
-)
 parser.add_argument(
     "--dry_run",
     type=str2bool,
@@ -172,29 +157,6 @@ else:
     accumulative_results = execute_changed_notebooks_helper.load_results(results_bucket, results_file)
 
     notebooks = [changed_notebook for changed_notebook in changed_notebooks if execute_changed_notebooks_helper.select_notebook(changed_notebook, accumulative_results, args.test_percent)]
-    # cap the number of notebooks to the specified percentage
-    max_notebooks = int((len(changed_notebooks) * (args.test_percent/100)))
-    if (len(notebooks) > max_notebooks):
-        notebooks = notebooks[:max_notebooks]
-
-run_first = []
-if args.run_first_file:
-    if not os.path.isfile(args.run_first_file):
-        print("Error: file does not exist", args.run_first_file)
-    else:
-        with open(args.run_first_file, 'r') as csvfile:
-            reader = csv.reader(csvfile)
-            for row in reader:
-                notebook = row[0]
-                run_first.append(notebook)
-
-    for notebook in run_first:
-        if notebook in notebooks:
-            # remove from existing list
-            notebooks.remove(notebook)
-            # add back to the front of the list
-            notebooks.insert(0, notebook)
-            print(f"Run first: {notebook}")
 
 if args.dry_run:
     print("Dry run ...\n")
@@ -215,5 +177,4 @@ else:
         variable_vpc_network=args.variable_vpc_network,
         private_pool_id=args.private_pool_id,
         concurrent_notebooks=args.concurrent_notebooks,
-        aiplatform_whl=args.aiplatform_whl
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -41,13 +41,10 @@ from utils import NotebookProcessors, util
 
 # A buffer so that workers finish before the orchestrating job
 WORKER_TIMEOUT_BUFFER_IN_SECONDS: int = 60 * 60
-
 PYTHON_VERSION = "3.9"  # Set default python version
 
 # rolling time window for accumulating build results for selecting notebooks
 MAX_RESULTS_AGE_SECONDS: int = (60 * 60) * 24 * 60  # 60 days
-# maximum time since last run to force a run on the current build
-MAX_AGE_BEFORE_FORCE_RUN: int = (60 * 60) * 24 * 30
 
 
 def format_timedelta(delta: datetime.timedelta) -> str:
@@ -121,12 +118,8 @@ def load_results(results_bucket: str,
                 if notebook in accumulative_results:
                     accumulative_results[notebook]['passed'] += build_results[notebook]['passed']
                     accumulative_results[notebook]['failed'] += build_results[notebook]['failed']
-                    if accumulative_results[notebook]['last_time_ran'] < time_created:
-                        accumulative_results[notebook]['last_time_ran'] = time_created
                 else:
                     accumulative_results[notebook] = build_results[notebook]
-                    accumulative_results[notebook]['failed_on_latest_run'] = build_results[notebook]['failed']
-                    accumulative_results[notebook]['last_time_ran'] = time_created
 
         print(accumulative_results)
     except Exception as e:
@@ -145,37 +138,19 @@ def select_notebook(changed_notebook: str,
     if changed_notebook in accumulative_results:
         pass_count = accumulative_results[changed_notebook]['passed']
         fail_count = accumulative_results[changed_notebook]['failed']
-        failed_on_latest_run = accumulative_results[changed_notebook]['failed_on_latest_run']
-        last_time_ran = accumulative_results[changed_notebook]['last_time_ran']
     else:
         pass_count = 1
         fail_count = 0
-        failed_on_latest_run = 0
-        last_time_ran = datetime.datetime.now().replace(tzinfo=None)
 
-    # If notebook has not been ran in a long time, force running it
-    if (datetime.datetime.now().replace(tzinfo=None) - last_time_ran).total_seconds() > MAX_AGE_BEFORE_FORCE_RUN:
-        should_test_do_to_age = True
-    else:
-        should_test_do_to_age  = False
-
-
-    # if failed on the last time it was ran, select the notebook
-    if failed_on_latest_run:
-        inferred_failure_rate = 1
-    # otherwise, calculate the frequency of failure
-    else:
-        inferred_failure_rate = fail_count / (pass_count + fail_count)
+    inferred_failure_rate = fail_count / (pass_count + fail_count)
 
     # If failure rate is high, the chance of testing should be higher
     should_test_due_to_failure = random.uniform(0, 1) <= inferred_failure_rate
 
-    #if accumulative_resultsi[changed_notebook]['latest_date_ran']
-
     # Additionally, only test a percentage of these
     should_test_due_to_random_subset = random.uniform(0, 1) <= (test_percent / 100)
 
-    if should_test_due_to_failure or should_test_due_to_random_subset or should_test_do_to_age:
+    if should_test_due_to_failure or should_test_due_to_random_subset:
         print(f"Selected: {changed_notebook}, {should_test_due_to_failure}, {should_test_due_to_random_subset}")
         return True
     else:
@@ -238,7 +213,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
@@ -271,7 +246,7 @@ def process_and_execute_notebook(
     private_pool_id: Optional[str],
     deadline: datetime.datetime,
     notebook: str,
-    should_get_tail_logs: bool = True,
+    should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
 
     print(f"Running notebook: {notebook}")
@@ -453,39 +428,15 @@ def _save_results(results: List[NotebookExecutionResult],
         else:
             pass_count = 0
             fail_count = 1
-        if result.error_message is None:
-            error_type = ''
-        elif '500 Internal' in result.error_message or 'INTERNAL' in result.error_message or 'internal error' in result.error_message:
-            error_type = 'INTERNAL'
-        elif 'context deadline exceeded' in result.error_message or 'TIMEOUT' in result.error_message:
-            error_type = 'TIMEOUT'
-        elif 'Quota' in result.error_message or 'quotas are exceeded' in result.error_message:
-            error_type = 'QUOTA'
-        elif 'ServiceUnavailable' in result.error_message:
-            error_type = 'SERVICEUNAVAILABLE'
-        elif 'ModuleNotFoundError' in result.error_message:
-            error_type = 'IMPORT'
-        elif result.is_pass:
-            error_type = ''
-        else:
-            error_type = 'undetermined'
-
-        if error_type != '':
-            log_url = result.log_url
-        else:
-            log_url = ''
-
         build_results[result.path] = {
                 'duration': result.duration.total_seconds(),
                 'start_time': str(result.start_time),
                 'passed': pass_count,
-                'failed': fail_count,
-                'error_type': error_type,
-                'log_url': log_url
+                'failed': fail_count
         }
         print(f"adding {result.path}")
 
-    print(f"Saving accumulative results to {results_file}, nentries {len(build_results)}")
+    print("Saving accumulative results ...")
     content = json.dumps(build_results)
 
     client = storage.Client()
@@ -508,7 +459,6 @@ def process_and_execute_notebooks(
     variable_vpc_network: Optional[str] = None,
     private_pool_id: Optional[str] = None,
     concurrent_notebooks: Optional[int] = 10,
-    aiplatform_whl: Optional[str] = None,
 ):
     """
     Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -540,7 +490,6 @@ def process_and_execute_notebooks(
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
         concurrent_notebooks (int): Max number of notebooks per minute to run in parallel.
-        aiplatform_whl: alternate whl version of Vertex AI SDK to install
     """
 
     # Calculate deadline

.cloud-build/notebook-execution-test-cloudbuild.yaml

@@ -36,7 +36,7 @@ steps:
     - -c
     - |
       . workspace/env/bin/activate &&
-      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GCP_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT} --concurrent_notebooks=${_CONCURRENT_NOTEBOOKS} --run_first_file=${_RUN_FIRST_FILE}
+      python3 .cloud-build/execute_changed_notebooks_cli.py --test_paths_file "${_TEST_PATHS_FILE}" --base_branch "${_FORCED_BASE_BRANCH}" --container_uri ${_PYTHON_IMAGE} --staging_bucket ${_GCS_STAGING_BUCKET}  --artifacts_bucket ${_GCS_STAGING_BUCKET}/executed_notebooks/PR_${_PR_NUMBER}/BUILD_${BUILD_ID} --variable_project_id ${PROJECT_ID} --variable_region ${_GCP_REGION} --variable_service_account ${_GCP_SERVICE_ACCOUNT} --variable_vpc_network "${_GPC_VPC_NETWORK_NAME}" `if [ ! -z "${_PRIVATE_POOL_NAME}" ]; then echo "--private_pool_id ${_PRIVATE_POOL_NAME}"; fi` --build_id ${BUILD_ID} --test_percent=${_TEST_PERCENT} --concurrent_notebooks=${_CONCURRENT_NOTEBOOKS}
     env:
     - 'IS_TESTING=1'
 timeout: 86400s

.cloud-build/run_first.csv

@@ -1,8 +0,0 @@
-notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb
-notebooks/official/generative_ai/rlhf_tune_llm.ipynb
-notebooks/official/generative_ai/tune_peft.ipynb
-notebooks/official/prediction/llm_streaming_prediction.ipynb
-notebooks/official/migration/sdk-automl-text-classification-batch-prediction.ipynb
-notebooks/official/vizier/get_started_vertex_vizier.ipynb
-notebooks/official/workbench/sentiment_analysis/Sentiment_Analysis.ipynb
-notebooks/official/model_monitoring/get_started_with_model_monitoring_automl.ipynb

.cloud-build/test_failing.txt

@@ -0,0 +1,40 @@
+notebooks/official/training/pytorch_gcs_data_training.ipynb
+notebooks/official/custom/custom_training_tensorboard_profiler.ipynb
+notebooks/official/custom/custom-tabular-bq-managed-dataset.ipynb
+notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-locally-trained-models.ipynb
+notebooks/official/ml_metadata/sdk-metric-parameter-tracking-for-custom-jobs.ipynb
+notebooks/official/tabnet/tabnet_vertex_tutorial.ipynb
+notebooks/official/tabnet/get_started_with_tabnet.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb
+notebooks/official/pipelines/multicontender_vs_champion_deployment_method.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_images.ipynb
+notebooks/official/pipelines/rapid_prototyping_bqml_automl.ipynb
+notebooks/official/pipelines/challenger_vs_blessed_deployment_method.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_create_stack_overflow_embeddings.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_for_indexing.ipynb
+notebooks/official/matching_engine/sdk_matching_engine_create_text_to_image_embeddings.ipynb
+notebooks/official/explainable_ai/sdk_custom_image_classification_online_explain.ipynb
+notebooks/official/explainable_ai/xai_image_classification_feature_attributions.ipynb
+notebooks/official/explainable_ai/sdk_custom_image_classification_batch_explain.ipynb
+notebooks/official/tabular_workflows/tabnet_on_vertex_pipelines.ipynb
+notebooks/official/model_registry/get_started_with_model_registry.ipynb
+notebooks/official/model_registry/bqml_vertexai_model_registry.ipynb
+notebooks/official/sdk/SDK_Custom_Training_Python_Package_Managed_Text_Dataset_Tensorflow_Serving_Container.ipynb
+notebooks/official/model_monitoring/batch_prediction_model_monitoring.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_setup.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_custom.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_custom_tf_serving.ipynb
+notebooks/official/model_monitoring/model_monitoring.ipynb
+notebooks/official/tensorboard/tensorboard_profiler_custom_training_with_prebuilt_container.ipynb
+notebooks/official/tensorboard/tensorboard_hyperparameter_tuning_with_hparams.ipynb
+notebooks/official/tensorboard/tensorboard_profiler_custom_training.ipynb
+notebooks/official/model_evaluation/custom_tabular_regression_model_evaluation.ipynb
+notebooks/official/model_evaluation/custom_tabular_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/automl_video_classification_model_evaluation.ipynb
+notebooks/official/experiments/comparing_local_trained_models.ipynb
+notebooks/official/automl/automl_image_classification_online_online_prediction.ipynb
+notebooks/official/automl/automl-text-classification.ipynb
+notebooks/official/automl/sdk_automl_video_object_tracking_batch.ipynb
+notebooks/official/feature_store/sdk-feature-store-pandas.ipynb
+notebooks/official/prediction/custom_batch_prediction_feature_filter.ipynb
+notebooks/official/prediction/pytorch_image_classification_with_prebuilt_serving_containers.ipynb

.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/test_passing.txt

@@ -0,0 +1,80 @@
+notebooks/official/training/hyperparameter_tuning_tensorflow.ipynb
+notebooks/official/training/get_started_with_vertex_distributed_training.ipynb
+notebooks/official/training/hyperparameter_tuning_xgboost.ipynb
+notebooks/official/training/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb
+notebooks/official/training/distributed_hyperparameter_tuning.ipynb
+notebooks/official/training/pytorch-text-sentiment-classification-custom-train-deploy.ipynb
+notebooks/official/training/xgboost_data_parallel_training_on_cpu_using_dask.ipynb
+notebooks/official/training/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb
+notebooks/official/bigquery_ml/get_started_with_bqml_training.ipynb
+notebooks/official/bigquery_ml/bqml-online-prediction.ipynb
+notebooks/official/custom/custom_training_container_and_model_registry.ipynb
+notebooks/official/custom/sdk-custom-image-classification-online.ipynb
+notebooks/official/custom/sdk-custom-image-classification-batch.ipynb
+notebooks/official/custom/SDK_FBProphet_Forecasting_Online.ipynb
+notebooks/official/custom/get_started_vertex_training_xgboost.ipynb
+notebooks/official/custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb
+notebooks/official/custom/SDK_Custom_Container_Prediction.ipynb
+notebooks/official/reduction_server/pytorch_distributed_training_reduction_server.ipynb
+notebooks/official/tabnet/ai-explanations-tabnet-algorithm.ipynb
+notebooks/official/vizier/get_started_vertex_vizier.ipynb
+notebooks/official/vizier/gapic-vizier-multi-objective-optimization.ipynb
+notebooks/official/pipelines/get_started_with_hpt_pipeline_components.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_automl_tabular.ipynb
+notebooks/official/pipelines/custom_tabular_train_batch_pred_bq_pipeline.ipynb
+notebooks/official/pipelines/metrics_viz_run_compare_kfp.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_model_upload_predict_evaluate.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_model_train_upload_deploy.ipynb
+notebooks/official/pipelines/get_started_with_machine_management.ipynb
+notebooks/official/pipelines/custom_model_training_and_batch_prediction.ipynb
+notebooks/official/pipelines/control_flow_kfp.ipynb
+notebooks/official/pipelines/lightweight_functions_component_io_kfp.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_bqml_text.ipynb
+notebooks/official/pipelines/pipelines_intro_kfp.ipynb
+notebooks/official/pipelines/automl_tabular_classification_beans.ipynb
+notebooks/official/pipelines/google_cloud_pipeline_components_dataproc_tabular.ipynb
+notebooks/official/explainable_ai/sdk_automl_tabular_classification_online_explain.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_online_explain.ipynb
+notebooks/official/explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_online_explain_get_metadata.ipynb
+notebooks/official/explainable_ai/sdk_custom_tabular_regression_batch_explain.ipynb
+notebooks/official/tabular_workflows/prophet_on_vertex_pipelines.ipynb
+notebooks/official/tabular_workflows/wide_and_deep_on_vertex_pipelines.ipynb
+notebooks/official/sdk/SDK_AutoML_Video_Classification.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl_image_batch.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_automl_image_online.ipynb
+notebooks/official/model_monitoring/get_started_with_model_monitoring_xgboost.ipynb
+notebooks/official/tensorboard/tensorboard_custom_training_with_custom_container.ipynb
+notebooks/official/tensorboard/tensorboard_custom_training_with_prebuilt_container.ipynb
+notebooks/official/tensorboard/tensorboard_vertex_ai_pipelines_integration.ipynb
+notebooks/official/model_evaluation/automl_text_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/get_started_with_custom_model_evaluation_import.ipynb
+notebooks/official/model_evaluation/automl_tabular_classification_model_evaluation.ipynb
+notebooks/official/model_evaluation/automl_tabular_regression_model_evaluation.ipynb
+notebooks/official/experiments/get_started_with_vertex_experiments.ipynb
+notebooks/official/experiments/comparing_pipeline_runs.ipynb
+notebooks/official/experiments/get_started_with_vertex_experiments_autologging.ipynb
+notebooks/official/experiments/build_model_experimentation_lineage_with_prebuild_code.ipynb
+notebooks/official/experiments/delete_outdated_tensorboard_experiments.ipynb
+notebooks/official/automl/sdk_automl_tabular_regression_batch_bq.ipynb
+notebooks/official/automl/sdk_automl_text_sentiment_analysis_online.ipynb
+notebooks/official/automl/sdk_automl_text_entity_extraction_online.ipynb
+notebooks/official/automl/sdk_automl_forecasting_hierarchical_batch.ipynb
+notebooks/official/automl/automl_text_entity_extraction_batch_prediction.ipynb
+notebooks/official/automl/automl_image_classification_batch_prediction.ipynb
+notebooks/official/automl/automl_text_sentiment_analysis_batch_prediction.ipynb
+notebooks/official/automl/sdk_automl_tabular_regression_online_bq.ipynb
+notebooks/official/automl/get_started_automl_training.ipynb
+notebooks/official/automl/automl-tabular-classification.ipynb
+notebooks/official/automl/automl_image_object_detection_export_edge.ipynb
+notebooks/official/automl/sdk_automl_image_object_detection_batch.ipynb
+notebooks/official/automl/automl_tabular_on_vertex_pipelines.ipynb
+notebooks/official/automl/sdk_automl_video_classification_batch.ipynb
+notebooks/official/automl/sdk_automl_video_action_recognition_batch.ipynb
+notebooks/official/automl/sdk_automl_tabular_forecasting_batch.ipynb
+notebooks/official/automl/automl_image_object_detection_online_prediction.ipynb
+notebooks/official/automl/automl_forecasting_bqml_arima_plus_comparison.ipynb
+notebooks/official/datasets/get_started_bq_datasets.ipynb
+notebooks/official/datasets/get_started_with_data_labeling.ipynb
+notebooks/official/feature_store/feature_store_streaming_ingestion_sdk.ipynb

.cloud-build/utils/NotebookProcessors.py

@@ -35,7 +35,7 @@ class RemoveNoExecuteCells(Preprocessor):
 
 
 class UpdateVariablesPreprocessor(Preprocessor):
-    def __init__(self, replacement_map: Dict[str, str]):
+    def __init__(self, replacement_map: Dict):
         self._replacement_map = replacement_map
 
     @staticmethod
@@ -98,28 +98,3 @@ class UniqueStringsPreprocessor(Preprocessor):
             executable_cells.append(cell)
         notebook.cells = executable_cells
         return notebook, resources
-
-class VertexAIInstallProprocessor(Preprocessor):
-    def __init__(self, vertex_ai_wheel):
-        self.vertex_ai_wheel = vertex_ai_wheel
-
-    @staticmethod
-    def update_vertex_ai_install(content: str):
-        if "google-cloud-aiplatform" not in content:
-            return content
-        return (
-            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 ")
-        )
-
-    def preprocess(self, notebook, resources=None):
-        executable_cells = []
-        for cell in notebook.cells:
-            if cell.cell_type == "code":
-                cell.source = self.update_vertex_ai_install(
-                    content=cell.source,
-                )
-
-            executable_cells.append(cell)
-        notebook.cells = executable_cells

.cloud-build/utils/build_results_viewer.py

@@ -5,69 +5,18 @@ Cloud Storage location: gs://cloud-build-notebooks-presubmit/build_results/
 '''
 import argparse
 import json
-from  util import download_file
-import csv
-import datetime
-from google.cloud import storage
-
-BUILD_BUCKET = "cloud-build-notebooks-presubmit"
-BUILD_FOLDER = "build_results"
 
 
 parser = argparse.ArgumentParser()
 parser.add_argument('--file', dest='file',
-                    default=None, type=str, help='build results filei (local or GCS)')
-args = parser.parse_args()
+                    default='build.json', type=str, help='build results file')
+import json
 
-investigate = {}
-with open('investigate.csv', 'r') as csvfile:
-    reader = csv.reader(csvfile)
-    for row in reader:
-        investigate[row[0][:-6]] = row[1]
-
-if not args.file:
-    client = storage.Client()
-    blobs = client.list_blobs(BUILD_BUCKET, prefix=BUILD_FOLDER)
-    newest_time = datetime.datetime(2000, 1, 1)
-    for blob in blobs:
-        # individual PR
-        if blob.size < 2000:
-            continue
-        time_created = blob.time_created.replace(tzinfo=None)
-        if time_created > newest_time:
-            newest_time = time_created
-            args.file = f"gs://{BUILD_BUCKET}/{blob.name}"
-
-if args.file.startswith("gs://"):
-    path = args.file[5:]
-    bucket = path.split('/')[0]
-    file   = path[len(bucket)+1:]
-    download_file(bucket, file, "build.json")
-    args.file = "build.json"
-
-with open(args.file, 'r') as f:
+with open('build.json', 'r') as f:
     results = json.load(f)
 
 for item in results.items():
-    notebook = item[0][len("/notebooks/official/")-1:-6]
     if item[1]['passed']:
-        passed = "PASS"
+        print(f"{item[0]},PASSED")
     else:
-        if notebook in investigate:
-            passed = "INVG"
-        else:
-            passed = "FAIL"
-      
-    error = item[1]['error_type']
-  
-    if passed == "FAIL":
-        if error == '':
-            error = "undetermined"
-        if 'log_url' in item[1]:
-            log_url = item[1]['log_url']
-        else:
-            log_url = ''
-    else:
-        log_url = ''
-
-    print(f"{notebook:75} {passed} {error:10} {log_url}")
+        print(f"{item[0]},FAILED")

.cloud-build/utils/investigate.csv

@@ -1,19 +0,0 @@
-notebook,status
-prediction/llm_streaming_prediction.ipynb,wait_for_fix
-custom/get_started_with_vertex_endpoint_and_shared_vm.ipynb,issue 2527
-feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store.ipynb,wait_for_reaper
-feature_store/online_feature_serving_and_vector_retrieval_bigquery_data_with_feature_store.ipynb,wait_for_reaper
-pipelines/google_cloud_pipeline_components_TPU_model_train_upload_deploy.ipynb,wait_for_fix
-explainable_ai/sdk_custom_image_classification_batch_explain.ipynb,issue 2528
-explainable_ai/sdk_custom_image_classification_online_explain.ipynb,issue 2528
-explainable_ai/sdk_custom_tabular_regression_online_explain.ipynb,issue 2528
-explainable_ai/sdk_automl_tabular_binary_classification_batch_explain.ipynb,issue 2528
-explainable_ai/xai_image_classification_feature_attributions.ipynb,issue 2528
-matching_engine,sdk_matching_engine_create_stack_overflow_embeddings.ipynb,issue 2530
-automl/automl_forecasting_bqml_arima_plus_comparison.ipynb,flaky
-model_evaluation/custom_tabular_regression_model_evaluation.ipynb,regr
-experiments/get_started_with_vertex_experiments.ipynb,regr
-experiments/comparing_local_trained_models.ipynb,regr
-generative_ai/tune_peft.ipynb,internal
-pipelines/custom_model_training_and_batch_prediction.ipynb,regr
-feature_store/online_feature_serving_and_fetching_bigquery_data_with_feature_store_optimized.ipynb,wait_for_reaper

.cloud-build/web-index/do_not_index.txt

@@ -1,11 +0,0 @@
-sdk2_remote_tabnet_training.ipynb
-remote_hyperparameter_tuning.ipynb
-remote_prediction.ipynb
-remote_training_bigframes_pytorch.ipynb
-remote_training_bigframes_sklearn.ipynb
-remote_training_bigframes_tensorflow.ipynb
-remote_training_lightning.ipynb
-remote_training_pytorch.ipynb
-remote_training_sklearn.ipynb
-remote_training_tensorflow_with_autologging.ipynb
-

.cloud-build/web-index/notebook-index-cloudbuild.yaml

@@ -15,7 +15,7 @@ steps:
     - -c
     - |
       . workspace/env/bin/activate &&
-      python3 notebooks/notebook_template_review.py --web --title --steps --desc --linkback --notebook-dir=notebooks/official --skip-file=${_DO_NOT_INDEX_FILE} >web.html
+      python3 notebooks/notebook_template_review.py --web --title --steps --desc --linkback --notebook-dir=notebooks/official >web.html
 artifacts:
   objects:
           location: gs://${_GCS_ARTIFACTS_BUCKET}/webdoc

.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@v5
+      uses: actions/setup-python@v4
       with:
         python-version: '3.x'      
     - name: Fetch pull request branch
-      uses: actions/checkout@v4
+      uses: actions/checkout@v3
       with:
         fetch-depth: 0      
     - name: Fetch base main branch

.github/workflows/linter/Dockerfile

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

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==25.1.0
-pyupgrade==3.20.0
-isort==6.0.1
-flake8==7.3.0
-nbqa==1.9.1
+black==23.3.0
+pyupgrade==3.7.0
+isort==5.12.0
+flake8==6.0.0
+nbqa==1.7.0
 

.github/workflows/linter/run_linter.sh

@@ -58,7 +58,7 @@ done
 # Only check notebooks in test folders modified in this pull request.
 # Note: Use process substitution to persist the data in the array
 if [ ${#notebooks[@]} -eq 0 ]; then
-    echo "Checking for changed notebooks using git"
+    echo "Checking for changed notebooked using git"
     while read -r file || [ -n "$line" ]; do
         notebooks+=("$file")
     done < <(git diff --name-only main... | grep '\.ipynb$')

README.md

@@ -1,176 +1,37 @@
-# ![Google Cloud](https://avatars.githubusercontent.com/u/2810941?s=60&v=4) Google Cloud Vertex AI Samples
+# Google Cloud Vertex AI Samples
 
-This repository contains notebooks, code samples, sample apps, and other resources that demonstrate how to use, develop and manage machine learning and generative AI workflows using Google Cloud Vertex AI.
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](LICENSE)
+
+Welcome to the Google Cloud [Vertex AI](https://cloud.google.com/vertex-ai/docs/) sample repository.
 
 ## Overview
 
-[Vertex AI](https://cloud.google.com/vertex-ai) is a fully-managed, unified AI development platform for building and using generative AI. This repository is designed to help you get started with Vertex AI. Whether you're new to Vertex AI or an experienced ML practitioner, you'll find valuable resources here.
-
-For more Vertex AI Generative AI notebook samples, please visit the Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository.
-
-## Explore, learn and contribute
-
-You can explore, learn, and contribute to this repository to unleash the full potential of machine learning on Vertex AI! 
-
-### Explore and learn
-
-Explore this repository, follow the links in the header section of each of the notebooks to -
-
-![Colab](https://cloud.google.com/ml-engine/images/colab-logo-32px.png)  Open and run the notebook in [Colab](https://colab.google/)\
-![Colab Enterprise](https://cloud.google.com/ml-engine/images/colab-enterprise-logo-32px.png)  Open and run the notebook in [Colab Enterprise](https://cloud.google.com/colab/docs/introduction)\
-![Workbench](https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32)  Open and run the notebook in [Vertex AI Workbench](https://cloud.google.com/vertex-ai/docs/workbench/introduction)\
-![Github](https://cloud.google.com/ml-engine/images/github-logo-32px.png)  View the notebook on Github
-
-### Contribute
-
-See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
-
-## Get started
-
-To get started using Vertex AI, you must have a Google Cloud project. 
-
-- If you don't have a Google Cloud project, you can learn and build on GCP for free using [Free Trail](https://cloud.google.com/free).
-- Once you have a Google Cloud project, you can learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).
-
+The repository contains [notebooks](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/notebooks) and [community content](https://github.com/GoogleCloudPlatform/vertex-ai-samples/tree/master/community-content) that demonstrate how to develop and manage ML workflows using Google Cloud Vertex AI.
 
 ## Repository structure
 
 ```bash
+├── community-content - Sample code and tutorials contributed by the community
 ├── notebooks
+│   ├── community - Notebooks contributed by the community
 │   ├── official - Notebooks demonstrating use of each Vertex AI service
 │   │   ├── automl
 │   │   ├── custom
 │   │   ├── ...
-│   ├── community - Notebooks contributed by the community
-│   │   ├── model_garden
-│   │   ├── ...
-├── community-content - Sample code and tutorials contributed by the community
-
 ```
-## Examples
 
-<!-- markdownlint-disable MD033 -->
-<table>
+## Contributing
 
-  <tr>
-    <th style="text-align: center;">Category</th>
-    <th style="text-align: center;">Product</th>
-    <th style="text-align: center;">Description</th>
-  </tr>
-  <tr>
-    <td>Model</td>
-    <td>
-      <a href="notebooks/community/model_garden"><code>Model Garden/</code></a>
-    </td>
-    <td>
-      Curated collection of first-party, open-source, and third-party models available on Vertex AI including Gemini, Gemma, Llama 3, Claude 3 and many more.
-    </td>
-  </tr>
-  <tr>
-    <td>Data</td>
-    <td>
-      <a href="notebooks/official/feature_store"><code>Feature Store/</code></a>
-    </td>
-    <td>
-      Set up and manage online serving using Vertex AI Feature Store.
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/datasets"><code>datasets/</code></a>
-    </td>
-    <td>
-      Use BigQuery and Data Labeling service with Vertex AI.
-    </td>
-  </tr>
-  <tr>
-    <td>Model development</td>
-    <td>
-      <a href="notebooks/official/automl"><code>automl/</code></a>
-    </td>
-    <td>
-      Train and make predictions on AutoML models
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/custom"><code>custom/</code></a>
-    </td>
-    <td>
-      Create, deploy and serve custom  models on Vertex AI
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/ray_on_vertex_ai"><code>ray_on_vertex_ai/</code></a>
-    </td>
-    <td>
-      Use Colab Enterprise and Vertex AI SDK for Python to connect to the Ray Cluster.
-    </td>
-  </tr>
-  <tr>
-    <td>Deploy and use</td>
-    <td>
-      <a href="notebooks/official/prediction"><code>prediction/</code></a>
-    </td>
-    <td>
-      Build, train and deploy models using prebuilt containers for custom training and prediction.
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/model_registry"><code>model_registry/</code></a>
-    </td>
-    <td>
-      Use Model Registry to create and register a model.
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/explainable_ai"><code>Explainable AI/</code></a>
-    </td>
-    <td>
-      Use Vertex Explainable AI's feature-based and example-based explanations to explain how or why a model produced a specific prediction.
-    </td>
-  </tr>
-  <tr>
-    <td></td>
-    <td>
-      <a href="notebooks/official/ml_metadata"><code>ml_metadata/</code></a>
-    </td>
-    <td>
-      Record the metadata and artifacts and query that metadata to help analyze, debug, and audit the performance of your ML system.
-    </td>
-  </tr>
-  <tr>
-    <td>Tools</td>
-    <td>
-      <a href="notebooks/official/pipelines"><code>Pipelines/</code></a>
-    </td>
-    <td>
-      Use `Vertex AI Pipelines` and `Google Cloud Pipeline Components` to build, tune, or deploy a custom model.
-    </td>
-  </tr>
-</table>
-<!-- markdownlint-enable MD033 -->
+Contributions welcome! See the [Contributing Guide](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/CONTRIBUTING.md).
 
+## Getting help
 
-## Get help
-
-Please use the [Issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
+Please use the [issues page](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues) to provide feedback or submit a bug report.
 
 ## Disclaimer
 
 This is not an officially supported Google product. The code in this repository is for demonstrative purposes only.
 
+## Feedback
 
-## References
-- [Vertex AI Jupyter Notebook tutorials](https://cloud.google.com/vertex-ai/docs/tutorials/jupyter-notebooks)
-- Vertex AI [Generative AI](https://github.com/GoogleCloudPlatform/generative-ai) GitHub repository
-- [Vertex AI documentaton](https://cloud.google.com/vertex-ai/docs)
-  
+Please feel free to fill out our [survey](https://bit.ly/vertex-ai-samples-survey) to give us feedback on the repo and its content.

community-content/CODEOWNERS

@@ -10,24 +10,7 @@
 /pipeline_components @Ark-kun
 /pipeline_components/image_ml_model_training @lakeyk
 /prediction_featurestore_integration @googleapis/vertex-prediction-team
-/vertex_model_garden/model_oss/notebook_util @minwoo33park
-/vertex_model_garden/model_oss/util @weigary
-/vertex_model_garden/model_oss/diffusers @weigary
-/vertex_model_garden/model_oss/keras @dstnluong-google
-/vertex_model_garden/model_oss/transformers @dstnluong-google
-/vertex_model_garden/model_oss/pic2word @jismailyan-google
-/vertex_model_garden/model_oss/open_clip @lydhr
-/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
-/vertex_model_garden/model_oss/imagebind @kathyyu-google
-/vertex_model_garden/model_oss/llava @py4
-/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
-
+/vertex_vision_model_garden/model_oss/util @weigary
+/vertex_vision_model_garden/model_oss/diffusers @weigary
+/vertex_vision_model_garden/model_oss/keras @dstnluong-google
+/vertex_vision_model_garden/model_oss/transformers @dstnluong-google

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -9,7 +9,7 @@ binarize_column_using_Pandas_on_CSV_data_op = components.load_component_from_url
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_classification_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -23,7 +23,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 #train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_XGBoost_and_import_to_Vertex_AI/pipeline.py

@@ -8,7 +8,7 @@ fill_all_missing_values_using_Pandas_on_CSV_data_op = components.load_component_
 split_rows_into_subsets_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/dataset_manipulation/Split_rows_into_subsets/in_CSV/component.yaml")
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 deploy_model_to_endpoint_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Deploy_to_endpoint/component.yaml")
 
 # %% Pipeline definition

community-content/Train_tabular_models_with_many_frameworks_and_import_to_Vertex_AI_using_Pipelines/Train_tabular_regression_model_using_all_frameworks_and_import_to_Vertex_AI/pipeline.py

@@ -22,7 +22,7 @@ upload_PyTorch_model_archive_to_Google_Cloud_Vertex_AI_op = components.load_comp
 # XGBoost
 train_XGBoost_model_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Train/component.yaml")
 xgboost_predict_on_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/XGBoost/Predict/component.yaml")
-upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/d5c9918850a6cc70004c4269dae066cfe2e664eb/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
+upload_XGBoost_model_to_Google_Cloud_Vertex_AI_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/399405402d95f4a011e2d2e967c96f8508ba5688/community-content/pipeline_components/google-cloud/Vertex_AI/Models/Upload_XGBoost_model/component.yaml")
 
 # Scikit-learn
 train_linear_regression_model_using_scikit_learn_from_CSV_op = components.load_component_from_url("https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/1f5cf6e06409b704064b2086c0a705e4e6b4fcde/community-content/pipeline_components/ML_frameworks/Scikit_learn/Train_linear_regression_model/from_CSV/component.yaml")

community-content/cpr-examples/timm_serving/requirements.txt

@@ -1,5 +1,5 @@
 absl-py==1.1.0
-fastapi==0.109.1
+fastapi==0.75.2
 uvicorn==0.18.2
 timm==0.5.4
 smart_open==6.0.0

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

@@ -64,8 +64,8 @@ implementation:
           labels["component-source"] = "github-com-ark-kun-pipeline-components"
 
           # The serving container decides the model type based on the model file extension.
-          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.bst
-          _, renamed_model_path = tempfile.mkstemp(suffix=".bst")
+          # So we need to rename the mode file (e.g. /tmp/inputs/model/data) to *.pkl
+          _, renamed_model_path = tempfile.mkstemp(suffix=".pkl")
           shutil.copyfile(src=model_path, dst=renamed_model_path)
 
           model = aiplatform.Model.upload_xgboost_model_file(

community-content/pipeline_components/image_ml_model_training/load_image_classification_model/component.yaml

@@ -87,7 +87,7 @@ outputs:
 - {name: image_size_path, type: HeightWidth}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.
@@ -109,4 +109,4 @@ implementation:
       {inputValue: l2_regularization_penalty},
       --image-size-path,
       {outputPath: image_size_path},
-    ]
+    ]

community-content/pipeline_components/image_ml_model_training/preprocess_image_data/component.yaml

@@ -34,7 +34,7 @@ outputs:
     path for the validation data,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/train_image_classification_model/component.yaml

@@ -55,7 +55,7 @@ outputs:
     for the saved model,'}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_csv/component.yaml

@@ -20,7 +20,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pipeline_components/image_ml_model_training/transcode_tfrecord_image_dataset_from_jsonl/component.yaml

@@ -22,7 +22,7 @@ outputs:
     path for the TFRecord image data}
 implementation:
   container:
-    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.2
+    image: us-docker.pkg.dev/vertex-ai/ready-to-go-image-classification/image-components:v0.1
     # command is a list of strings (command-line arguments).
     # The YAML language has two syntaxes for lists and you can use either of them.
     # Here we use the "flow syntax" - comma-separated strings inside square brackets.

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile

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

community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt

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

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/trainer/requirements.txt

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

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

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

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

@@ -1,4 +1,4 @@
 google-cloud-pubsub==2.5.0
-pillow==10.3.0
+pillow==9.0.1
 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
-pillow==10.3.0
+tensorflow==2.7.2
+pillow==9.0.1
 tf-agents==0.8.0

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/requirements.txt

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

community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/trainer/requirements.txt

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

community-content/vertex-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==4.25.8
-opencv-python-headless==4.11.0.86
-docutils==0.16
-urllib3==2.5.0
-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.4
-pyasn1==0.6.1
-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.4.0
-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/README.md

@@ -1,15 +0,0 @@
-# Vertex AI custom prediction routines samples
-
-## Overview
-Vertex Custom Prediction Routines(CPR) simplify the process of building custom containers 
-and make local model testing easy. Here are the sameple codes for different libraries.
-
-
-### Objectives
-The objective is to provide various samples for Vertex Custom Prediction Routine(CPR).
-
-
-### Supporting libraries
-* torch
-* sklearn
-* xgboost

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/torch/predictor_transformers.py

@@ -1,73 +0,0 @@
-import ast
-import json
-import os
-import pickle
-import torch
-
-from google.cloud.aiplatform.utils import prediction_utils
-from google.cloud.aiplatform.prediction.predictor import Predictor
-from transformers import AutoModelForQuestionAnswering
-from typing import Dict, List
-
-class TorchTransformersPredictor(Predictor):
-
-    def __init__(self):
-        return
-
-    def load(self, artifacts_uri: str) -> None:
-        prediction_utils.download_model_artifacts(artifacts_uri)
-        
-        if os.path.isfile("setup_config.json"):
-            with open("setup_config.json") as setup_config_file:
-                self.setup_config = json.load(setup_config_file)
-
-        if os.path.exists("model.pt"):
-            self.model = AutoModelForQuestionAnswering.from_pretrained("model.pt")
-            self.model.eval()
-        else:
-            raise ValueError("One of the following model files must be provided: model.pt.")
-
-    def preprocess(self, prediction_input: dict) -> torch.Tensor:
-        max_length = self.setup_config["max_length"]
-        instances = prediction_input["instances"]
-        question_context = ast.literal_eval(instances)
-        question = question_context["question"]
-        context = question_context["context"]
-        inputs = self.tokenizer.encode_plus(
-            question,
-            context,
-            max_length=int(max_length),
-            pad_to_max_length=True,
-            add_special_tokens=True,
-            return_tensors="pt",
-        )
-        input_ids = inputs["input_ids"]
-        attention_mask = inputs["attention_mask"]
-        return torch.Tensor(input_ids, attention_mask)
-    
-    @torch.inference_mode()
-    def predict(self, instances: torch.Tensor) -> List[str]:
-        input_ids, attention_mask = instances
-        outputs = self._model(input_ids, attention_mask)
-        answer_start_scores = outputs.start_logits
-        answer_end_scores = outputs.end_logits
-        
-        num_rows, num_cols = answer_start_scores.shape
-        inferences = []
-        for i in range(num_rows):
-            answer_start_scores_one_seq = answer_start_scores[i].unsqueeze(0)
-            answer_start = torch.argmax(answer_start_scores_one_seq)
-            answer_end_scores_one_seq = answer_end_scores[i].unsqueeze(0)
-            answer_end = torch.argmax(answer_end_scores_one_seq) + 1
-            prediction = self.tokenizer.convert_tokens_to_string(
-                self.tokenizer.convert_ids_to_tokens(
-                    input_ids[i].tolist()[answer_start:answer_end]
-                )
-            )
-            inferences.append(prediction)
-        return inferences
-
-    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/benchmarking_reports/jax_vit_benchmarking_report.md

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

community-content/vertex_model_garden/benchmarking_reports/pytorch_openllama_hyperparameter_tuning_benchmark_report.md

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

community-content/vertex_model_garden/benchmarking_reports/stable_diffusion_v1-5_benchmarking_report.md

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

community-content/vertex_model_garden/model_oss/autogluon/dockerfile/autogluon_serve.Dockerfile

@@ -1,50 +0,0 @@
-# Dockerfile for serving dockers with AutoGluon.
-#
-# To build:
-# docker build -f model_oss/autogluon/dockerfile/serve.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 pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
-
-USER root
-
-# AutoGluon might require libgomp for some dependencies.
-RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
-        curl \
-        wget \
-        vim \
-        libgomp1
-
-# Install AutoGluon and other dependencies.
-RUN pip install --upgrade pip
-RUN pip install autogluon==1.0.0
-RUN pip install flask==3.0.0
-
-# Dependencies needed to work with GCS.
-RUN pip install absl-py==2.0.0
-RUN pip install google-cloud-storage==2.7.0
-
-# Copy scripts into the container.
-COPY model_oss/autogluon /autogluon
-COPY model_oss/util /autogluon/util
-WORKDIR /autogluon
-
-# Copy license.
-RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
-RUN wget https://github.com/pallets/flask/blob/main/LICENSE.rst
-
-# Expose the port the app runs on.
-EXPOSE 8501
-
-# Set the working directory to a specific path for consistency.
-WORKDIR /autogluon
-
-# Change to a non-root user for security purposes.
-RUN useradd -m autogluonuser
-USER autogluonuser
-
-# Run Flask application.
-CMD ["python", "serve.py"]

community-content/vertex_model_garden/model_oss/autogluon/dockerfile/autogluon_train.Dockerfile

@@ -1,36 +0,0 @@
-# Dockerfile for training dockers with Autogluon.
-#
-# To build:
-# docker build -f model_oss/autogluon/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 pytorch/pytorch:2.1.2-cuda11.8-cudnn8-runtime
-
-# Install tools.
-ENV DEBIAN_FRONTEND=noninteractive
-ENV PIP_ROOT_USER_ACTION=ignore
-
-RUN apt-get update && apt-get -y upgrade && apt-get install -y --no-install-recommends \
-apt-utils \
-curl \
-wget \
-git \
-jq \
-gnupg \
-build-essential \
-tesseract-ocr \
-vim
-
-# Install libraries.
-RUN pip install autogluon==1.0.0
-
-COPY model_oss/autogluon /autogluon
-WORKDIR /autogluon
-
-# Copy license.
-RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE
-
-ENTRYPOINT ["python", "train.py"]