Delete notebooks/official/pipelines/google_cloud_pipeline_components_automl_text.ipynb

b/442906902

.cloud-build/execute_changed_notebooks_helper.py

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

.cloud-build/test_notebook_vm.txt

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

.github/dependabot.yml

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

.github/workflows/linter/Dockerfile

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

.github/workflows/linter/requirements.txt

@@ -2,9 +2,9 @@ git+https://github.com/tensorflow/docs
 ipython
 jupyter
 nbconvert
-black==24.10.0
-pyupgrade==3.17.0
-isort==5.13.2
-flake8==7.1.1
-nbqa==1.9.0
+black==25.1.0
+pyupgrade==3.20.0
+isort==6.0.1
+flake8==7.3.0
+nbqa==1.9.1
 

community-content/CODEOWNERS

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

community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/trainer/requirements.txt

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,10 @@
+dllogger@git+https://github.com/NVIDIA/dllogger@v1.0.0
+
+# Fixing these libraries versions to avoid conflicting or broken packages.
+immutabledict==4.2.1
+protobuf==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

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

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

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

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

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

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

@@ -0,0 +1,49 @@
+absl-py==2.2.2
+annotated-types==0.7.0
+anyio==4.9.0
+black==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

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

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

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

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

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

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

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

community-content/vertex_cpr_samples/sklearn/predictor_SGDClassifier.py

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

community-content/vertex_cpr_samples/torch/predictor_resnet.py

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

community-content/vertex_cpr_samples/xgboost/predictor_xgbranker.py

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

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

@@ -1,13 +1,16 @@
 """Common util functions for notebook."""
 
 import base64
+from collections.abc import Sequence
 import datetime
 import io
 import json
 import os
 import subprocess
-from typing import Any, Dict, Sequence
+import time
+from typing import Any
 
+from google import auth
 from google.cloud import storage
 import matplotlib.pyplot as plt
 import numpy as np
@@ -230,7 +233,7 @@ def download_image(url: str) -> str:
     base64 encoded image.
   """
   response = requests.get(url)
-  return Image.open(io.BytesIO(response.content))
+  return Image.open(io.BytesIO(response.content))  # pytype: disable=bad-return-type  # pillow-102-upgrade
 
 
 def resize_image(image: Any, new_width: int = 1000) -> Any:
@@ -281,7 +284,7 @@ def decode_image(
   return image
 
 
-def get_label_map(label_map_yaml_filepath: str) -> Dict[int, str]:
+def get_label_map(label_map_yaml_filepath: str) -> dict[int, str]:
   """Returns class id to label mapping given a filepath to the label map.
 
   Args:
@@ -331,6 +334,7 @@ def vqa_predict(
     image: Any,
     language_code: str = "en",
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> Sequence[str]:
   """Predicts the answer to a question about an image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -354,7 +358,9 @@ def vqa_predict(
         "image": resized_image_base64,
     })
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return [pred.get("response") for pred in response.predictions]
 
 
@@ -364,6 +370,7 @@ def caption_predict(
     image: Any,
     caption_prompt: bool = False,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Predicts a caption for a given image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -378,7 +385,9 @@ def caption_predict(
     instance["prompt"] = caption_prompt_format.format(language_code)
 
   instances = [instance]
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -387,6 +396,7 @@ def ocr_predict(
     ocr_prompt: str,
     image: Any,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Extracts text from a given image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -398,7 +408,9 @@ def ocr_predict(
     instance["prompt"] = ocr_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
@@ -407,6 +419,7 @@ def detect_predict(
     detect_prompt: str,
     image: Any,
     new_width: int = 1000,
+    use_dedicated_endpoint: bool = False,
 ) -> str:
   """Predicts the answer to a question about an image using an Endpoint."""
   # Resize and convert image to base64 string.
@@ -418,10 +431,47 @@ def detect_predict(
     instance["prompt"] = detect_prompt
   instances = [instance]
 
-  response = endpoint.predict(instances=instances)
+  response = endpoint.predict(
+      instances=instances, use_dedicated_endpoint=use_dedicated_endpoint
+  )
   return response.predictions[0].get("response")
 
 
+def copy_model_artifacts(
+    model_id: str,
+    model_source: str,
+    model_destination: str,
+) -> None:
+  """Copies model artifacts from model_source to model_destination.
+
+  model_source and model_destination should be GCS path.
+
+  Args:
+    model_id: The model id.
+    model_source: The source of the model artifact.
+    model_destination: The destination of the model artifact.
+  """
+  if not model_source.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{model_source} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  if not model_destination.startswith(GCS_URI_PREFIX):
+    raise ValueError(
+        f"{model_destination} is not a GCS path starting with {GCS_URI_PREFIX}."
+    )
+  model_source = f"{model_source}/{model_id}"
+  model_destination = f"{model_destination}/{model_id}"
+  print("Copying model artifact from ", model_source, " to ", model_destination)
+  subprocess.check_output([
+      "gcloud",
+      "storage",
+      "cp",
+      "-r",
+      model_source,
+      model_destination,
+  ])
+
+
 def get_quota(project_id: str, region: str, resource_id: str) -> int:
   """Returns the quota for a resource in a region.
 
@@ -460,6 +510,17 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
   ):
     return -1
   all_regions_data = quota_data[0]["consumerQuotaLimits"][0]["quotaBuckets"]
+
+  # If the quota data does not have dimensions, it is global quota. However,
+  # global quota may be overridden by regional quota. So we need to check the
+  # global quota first.
+  global_quota = -1
+  if (
+      all_regions_data
+      and "dimensions" not in all_regions_data[0]
+      and "effectiveLimit" in all_regions_data[0]
+  ):
+    global_quota = int(all_regions_data[0]["effectiveLimit"])
   for region_data in all_regions_data:
     if (
         region_data.get("dimensions")
@@ -469,12 +530,13 @@ def get_quota(project_id: str, region: str, resource_id: str) -> int:
         return int(region_data["effectiveLimit"])
       else:
         return 0
-  return -1
+  return global_quota
 
 
 def get_resource_id(
     accelerator_type: str,
     is_for_training: bool,
+    is_spot: bool = False,
     is_restricted_image: bool = False,
     is_dynamic_workload_scheduler: bool = False,
 ) -> str:
@@ -484,6 +546,7 @@ def get_resource_id(
     accelerator_type: The accelerator type.
     is_for_training: Whether the resource is used for training. Set false for
       serving use case.
+    is_spot: Whether the resource is used with Spot.
     is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
     is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
       Workload Scheduler.
@@ -493,11 +556,15 @@ def get_resource_id(
   """
   accelerator_suffix_map = {
       "NVIDIA_TESLA_V100": "nvidia_v100_gpus",
+      "NVIDIA_TESLA_P100": "nvidia_p100_gpus",
       "NVIDIA_L4": "nvidia_l4_gpus",
       "NVIDIA_TESLA_A100": "nvidia_a100_gpus",
       "NVIDIA_A100_80GB": "nvidia_a100_80gb_gpus",
       "NVIDIA_H100_80GB": "nvidia_h100_gpus",
+      "NVIDIA_H100_MEGA_80GB": "nvidia_h100_mega_gpus",
+      "NVIDIA_H200_141GB": "nvidia_h200_gpus",
       "NVIDIA_TESLA_T4": "nvidia_t4_gpus",
+      "TPU_V6e": "tpu_v6e",
       "TPU_V5e": "tpu_v5e",
       "TPU_V3": "tpu_v3",
   }
@@ -512,6 +579,10 @@ def get_resource_id(
   restricted_image_training_accelerator_map = {
       "NVIDIA_A100_80GB": "restricted_image_training_nvidia_a100_80gb_gpus",
   }
+  spot_serving_accelerator_map = {
+      key: f"custom_model_serving_preemptible_{accelerator_suffix_map[key]}"
+      for key in accelerator_suffix_map
+  }
   serving_accelerator_map = {
       key: f"custom_model_serving_{accelerator_suffix_map[key]}"
       for key in accelerator_suffix_map
@@ -540,8 +611,11 @@ def get_resource_id(
   else:
     if is_dynamic_workload_scheduler:
       raise ValueError("Dynamic Workload Scheduler does not work for serving.")
-    if accelerator_type in serving_accelerator_map:
-      return serving_accelerator_map[accelerator_type]
+    accelerator_map = (
+        spot_serving_accelerator_map if is_spot else serving_accelerator_map
+    )
+    if accelerator_type in accelerator_map:
+      return accelerator_map[accelerator_type]
     else:
       raise ValueError(
           f"Could not find accelerator type: {accelerator_type} for serving."
@@ -554,13 +628,28 @@ def check_quota(
     accelerator_type: str,
     accelerator_count: int,
     is_for_training: bool,
+    is_spot: bool = False,
     is_restricted_image: bool = False,
     is_dynamic_workload_scheduler: bool = False,
-):
-  """Checks if the project and the region has the required quota."""
+) -> None:
+  """Checks if the project and the region has the required quota.
+
+  Args:
+    project_id: The project id.
+    region: The region.
+    accelerator_type: The accelerator type.
+    accelerator_count: The number of accelerators to check quota for.
+    is_for_training: Whether the resource is used for training. Set false for
+      serving use case.
+    is_spot: Whether the resource is used with Spot.
+    is_restricted_image: Whether the image is hosted in `vertex-ai-restricted`.
+    is_dynamic_workload_scheduler: Whether the resource is used with Dynamic
+      Workload Scheduler.
+  """
   resource_id = get_resource_id(
       accelerator_type,
       is_for_training=is_for_training,
+      is_spot=is_spot,
       is_restricted_image=is_restricted_image,
       is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,
   )
@@ -583,3 +672,76 @@ def check_quota(
         f"Quota not enough for {resource_id} in {region}: {quota} <"
         f" {accelerator_count}. {quota_request_instruction}"
     )
+
+
+def get_deploy_source() -> str:
+  """Gets deploy_source string based on running environment."""
+  vertex_product = os.environ.get("VERTEX_PRODUCT", "")
+  match vertex_product:
+    case "COLAB_ENTERPRISE":
+      return "notebook_colab_enterprise"
+    case "WORKBENCH_INSTANCE":
+      return "notebook_workbench"
+    case _:
+      # Legacy workbench, legacy colab, or other custom environments.
+      return "notebook_environment_unspecified"
+
+
+def _is_operation_done(op_name: str, region: str) -> bool:
+  """Checks if the operation is done.
+
+  Args:
+    op_name: The name of the operation to poll.
+    region: The region of the operation.
+
+  Returns:
+    True if the operation is done, False otherwise.
+
+  Raises:
+    ValueError: If the operation failed.
+  """
+  creds, _ = auth.default()
+  auth_req = auth.transport.requests.Request()
+  creds.refresh(auth_req)
+  headers = {
+      "Authorization": f"Bearer {creds.token}",
+  }
+  url = f"https://{region}-aiplatform.googleapis.com/ui/{op_name}"
+  response = requests.get(url, headers=headers)
+  operation_data = response.json()
+  if "error" in operation_data:
+    raise ValueError(f"Operation failed: {operation_data['error']}")
+  return operation_data.get("done", False)
+
+
+def poll_and_wait(
+    op_name: str, region: str, total_wait: int, interval: int = 60
+) -> None:
+  """Polls the operation and waits for it to complete.
+
+  Args:
+    op_name: The name of the operation to poll.
+    region: The region of the operation.
+    total_wait: The total wait time in seconds.
+    interval: The interval between each poll in seconds.
+
+  Raises:
+    TimeoutError: If the operation times out.
+  """
+  start_time = time.time()
+  while True:
+    if _is_operation_done(op_name, region):
+      break
+    time_elapsed = time.time() - start_time
+    if time_elapsed > total_wait:
+      raise TimeoutError(
+          f"Operation timed out after {int(time_elapsed)} seconds."
+      )
+    print(
+        "\rStill waiting for operation... Elapsed time in seconds:"
+        f" {int(time_elapsed):<6}",
+        end="",
+        flush=True,
+    )
+    time.sleep(interval)
+

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

@@ -7,7 +7,7 @@ import json
 import multiprocessing
 import os
 import subprocess
-from typing import Any, Callable, Dict, Union
+from typing import Any, Callable, Dict, Tuple, Union
 from absl import logging
 import accelerate
 import datasets
@@ -20,8 +20,10 @@ LOCAL_TEMPLATE_DIR = "/tmp/template_dir"
 _TEMPLATE_DIRNAME = "templates"
 _VERTEX_AI_SAMPLES_GITHUB_REPO_NAME = "vertex-ai-samples"
 _VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR = (
-    "community-content/vertex_model_garden/model_oss/peft/templates"
+    "community-content/vertex_model_garden/model_oss/peft/train/vmg/templates"
 )
+_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
+_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
 _DESCRIPTION_KEY = "description"
 _SOURCE_KEY = "source"
 _PROMPT_INPUT_KEY = "prompt_input"
@@ -68,7 +70,9 @@ def force_gcs_fuse_path(gcs_uri: str) -> str:
 
 
 def download_gcs_uri_to_local(
-    gcs_uri: str, destination_dir: str = LOCAL_BASE_MODEL_DIR
+    gcs_uri: str,
+    destination_dir: str = LOCAL_BASE_MODEL_DIR,
+    check_path_exists: bool = True,
 ) -> str:
   """Downloads GCS URI to local.
 
@@ -79,6 +83,7 @@ def download_gcs_uri_to_local(
   Args:
     gcs_uri: GCS URI to download.
     destination_dir: Local directory directory.
+    check_path_exists: Whether to check if the path exists.
 
   Returns:
     Local path to target folder/file.
@@ -87,7 +92,7 @@ def download_gcs_uri_to_local(
       destination_dir,
       os.path.basename(os.path.normpath(gcs_uri)),
   )
-  if os.path.exists(target):
+  if check_path_exists and os.path.exists(target):
     logging.info("File %s already exists.", target)
     return target
   if accelerate.PartialState().is_local_main_process:
@@ -171,57 +176,6 @@ def _format_template_fn(
   """
   template_json = get_template(template)
 
-  if _CHAT_TEMPLATE_KEY not in template_json:
-
-    def format_fn(example: Dict[str, str]) -> Dict[str, str]:
-      format_dict = {key: value for key, value in example.items()}
-      format_str = (
-          template_json[_PROMPT_INPUT_KEY]
-          if format_dict.get(input_column)
-          else template_json[_PROMPT_NO_INPUT_KEY]
-      )
-      return {input_column: format_str.format(**format_dict)}
-
-    return format_fn
-  elif (
-      _PROMPT_INPUT_KEY in template_json
-      or _PROMPT_NO_INPUT_KEY in template_json
-  ):
-    raise ValueError(
-        "chat_template templates do not support input/no_input templates."
-    )
-  else:
-    if tokenizer is None:
-      raise ValueError("A tokenizer is required for chat_template templates.")
-    # Assign HuggingFace jinja template.
-    tokenizer.chat_template = template_json[_CHAT_TEMPLATE_KEY]
-    return lambda example: {
-        input_column: tokenizer.apply_chat_template(
-            example[input_column], tokenize=False, add_generation_prompt=False
-        )
-    }
-
-
-def _format_template_fn_notebook(
-    template: str,
-    input_column: str,
-    tokenizer: transformers.PreTrainedTokenizer | None = None,
-) -> Callable[[Dict[str, str]], Dict[str, str]]:
-  """Formats a dataset example according to a template.
-
-  Args:
-    template: Name of the JSON template file under `templates/` or GCS path to
-      the template file.
-    input_column: The input column in the dataset to be used or updated by the
-      template. If it does not exist, the template's `prompt_no_input` will be
-      used, and the input_column will be created.
-    tokenizer: The tokenizer to use for chat_template templates.
-
-  Returns:
-    A function that formats data according to the template.
-  """
-  template_json = get_template(template)
-
   if _CHAT_TEMPLATE_KEY not in template_json:
 
     def format_fn(example: Dict[str, str]) -> Dict[str, str]:
@@ -366,13 +320,140 @@ def _get_dataset(
   return datasets.load_dataset(dataset_name, split=split, num_proc=num_proc)
 
 
+def should_add_pad_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special pad token.
+
+  Args:
+    model_id: The name of the model.
+
+  Returns:
+    True if the model requires adding a special pad token, False otherwise.
+  """
+  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
+
+
+def should_add_eos_token(model_id: str) -> bool:
+  """Returns whether the model requires adding a special eos token.
+
+  Args:
+    model_id: The name of the model.
+
+  Returns:
+    True if the model requires adding a special eos token, False otherwise.
+  """
+  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
+
+
+def load_tokenizer(
+    pretrained_model_id: str,
+    padding_side: str | None = None,
+    access_token: str | None = None,
+) -> transformers.AutoTokenizer:
+  """Loads tokenizer based on `pretrained_model_id`.
+
+  Args:
+    pretrained_model_id: The name of the pretrained model.
+    padding_side: The side to pad the input on.
+    access_token: The access token to use for the tokenizer.
+
+  Returns:
+    The tokenizer.
+  """
+  tokenizer_kwargs = {}
+  if should_add_eos_token(pretrained_model_id):
+    tokenizer_kwargs["add_eos_token"] = True
+  if padding_side:
+    tokenizer_kwargs["padding_side"] = padding_side
+
+  with accelerate.PartialState().local_main_process_first():
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        pretrained_model_id,
+        trust_remote_code=False,
+        use_fast=True,
+        token=access_token,
+        **tokenizer_kwargs,
+    )
+
+  if should_add_pad_token(pretrained_model_id):
+    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+  return tokenizer
+
+
+def get_filtered_dataset(
+    dataset: Any,
+    input_column: str,
+    max_seq_length: int,
+    tokenizer: transformers.PreTrainedTokenizer,
+) -> Any:
+  """Returns the dataset by removing examples that are longer than max_seq_length.
+
+  Args:
+    dataset: The dataset to filter.
+    input_column: The input column in the dataset to be used.
+    max_seq_length: The maximum sequence length.
+    tokenizer: The tokenizer.
+  """
+  actual_dataset_length = len(dataset)
+  filtered_dataset = dataset.filter(
+      lambda x: len(tokenizer(x[input_column])["input_ids"]) <= max_seq_length
+  )
+  filtered_dataset_length = len(filtered_dataset)
+  if actual_dataset_length != filtered_dataset_length:
+    examples_removed_percent = (
+        (actual_dataset_length - filtered_dataset_length)
+        * 100
+        / actual_dataset_length
+    )
+    logging.info(
+        "(%.2f%%) of examples token length is <= max-seq-length(%d); (%.2f%%) >"
+        " max-seq-length. Filtering out %d example(s) which are longer than"
+        " max-seq-length.",
+        100 - examples_removed_percent,
+        max_seq_length,
+        examples_removed_percent,
+        actual_dataset_length - filtered_dataset_length,
+    )
+
+  return filtered_dataset
+
+
+def format_dataset(
+    dataset: datasets.Dataset,
+    input_column: str,
+    template: str = None,
+    tokenizer: transformers.PreTrainedTokenizer | None = None,
+) -> datasets.Dataset:
+  """Takes a raw dataset and formats it using a template and tokenizer.
+
+  Args:
+    dataset: The raw (unprocessed) dataset to format.
+    input_column: The input column in the dataset to be used or updaded by the
+      template. If it does not exist, the template's `prompt_no_input` will be
+      used, and the input_column will be created.
+    template: Name of the JSON template file under `templates/` or GCS path to
+      the template file.
+    tokenizer: The tokenizer to use for chat_template templates.
+
+  Returns:
+    A dataset compatible with the template.
+  """
+  return dataset.map(
+      _format_template_fn(
+          template,
+          input_column=input_column,
+          tokenizer=tokenizer,
+      )
+  )
+
+
 def load_dataset_with_template(
     dataset_name: str,
     split: str,
     input_column: str,
     template: str = None,
     tokenizer: transformers.PreTrainedTokenizer | None = None,
-) -> Any:
+) -> Tuple[Any, Any]:
   """Loads dataset with templates.
 
   Args:
@@ -386,19 +467,15 @@ def load_dataset_with_template(
     tokenizer: The tokenizer to use for chat_template templates.
 
   Returns:
-    A dataset compatible with the template.
+    The raw dataset and the dataset compatible with the template.
   """
-  dataset = _get_dataset(dataset_name, split=split)
+  raw = _get_dataset(dataset_name, split=split)
   if template:
-    dataset = dataset.map(
-        _format_template_fn(
-            template,
-            input_column=input_column,
-            tokenizer=tokenizer,
-        )
-    )
+    templated = format_dataset(raw, input_column, template, tokenizer)
+  else:
+    templated = None
 
-  return dataset
+  return raw, templated
 
 
 def validate_dataset_with_template(
@@ -407,6 +484,7 @@ def validate_dataset_with_template(
     input_column: str,
     template: str,
     tokenizer: transformers.PreTrainedTokenizer | None = None,
+    max_seq_length: int | None = None,
     use_multiprocessing: bool = False,
     validate_percentage_of_dataset: int | None = None,
     validate_k_rows_of_dataset: int | None = None,
@@ -428,6 +506,7 @@ def validate_dataset_with_template(
     template: Name of the JSON template file under `templates/` or GCS path to
       the template file.
     tokenizer: The tokenizer to use for chat_template templates.
+    max_seq_length: The maximum sequence length.
     use_multiprocessing: If True, it will use multiprocessing to load the
       dataset.
     validate_percentage_of_dataset: The percentage of the dataset to load.
@@ -470,17 +549,22 @@ def validate_dataset_with_template(
         f" https://github.com/GoogleCloudPlatform/{_VERTEX_AI_SAMPLES_GITHUB_REPO_NAME}/tree/main/{_VERTEX_AI_SAMPLES_GITHUB_TEMPLATE_DIR}."
     )
 
-  _get_dataset(dataset_name, split, num_proc).map(
-      _format_template_fn_notebook(
-          template_path,
-          input_column=input_column,
-          tokenizer=tokenizer,
-      )
+  dataset = format_dataset(
+      _get_dataset(dataset_name, split, num_proc),
+      input_column,
+      template_path,
+      tokenizer,
   )
 
+  if tokenizer is not None:
+    get_filtered_dataset(
+        dataset=dataset,
+        input_column=input_column,
+        max_seq_length=max_seq_length,
+        tokenizer=tokenizer,
+    )
   print(
       "Dataset {} is compatible with the {} template.".format(
           os.path.basename(dataset_name), os.path.basename(template)
       )
   )
-

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -5,23 +5,28 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 
 # keep sorted
-accelerate==0.31.0
+accelerate==0.34.2
 auto_gptq==0.7.1+cu118
-autoawq==0.2.5
+autoawq==0.2.8
 bitsandbytes==0.43.2
 cloudml-hypertune==0.1.0.dev6
-datasets==2.19.2
-deepspeed==0.14.4
+datasets==2.20.0
+deepspeed==0.15.2
 diffusers==0.25.1
+evaluate==0.4.3
 fsspec==2024.3.1
 gcsfs==2024.3.1
-lm_eval==0.4.3
+immutabledict==4.2.1
 ninja==1.11.1 # Needed to avoid `ninja 1.11.1.1 is not supported on this platform` error
+nltk==3.9.1
 optimum==1.17.1
 peft==0.12.0
 pynvml==11.5.3
+rouge_score==0.1.2
 torch==2.2.2+cu118
 torchvision==0.17.2+cu118
-transformers==4.43.1
-trl==0.9.6
+transformers==4.47.1
+trl==0.11.2
 wandb==0.17.1
+ydata-profiling==4.7.0 # Upgrade the version from 4.6.0 to 4.7.0 to fix the old `pydantic` package error.
+psutil==6.0.0

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

@@ -23,6 +23,14 @@ RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples
 ENV PIP_ROOT_USER_ACTION=ignore
 RUN pip install --upgrade pip
 
+# Remove packages that are not needed and are causing conflicts.
+# dataproc_jupyter_plugin was installed as a part of pytorch-cu121.2-2.py310
+# container which we don't need. It depends on ibis-framework and bigframes.
+# The package and its dependencies request lower versions of pyarrow/pydantic
+# than deepspeed/datasets. So, dataproc_jupyter_plugin conflicts with
+# deepspeed/datasets.
+RUN pip uninstall -y dataproc_jupyter_plugin ibis-framework bigframes
+
 # Prefer to install with requirement file as much as possible for reasons
 # described in b/355034754.
 COPY model_oss/peft/train/vmg/dockerfile/requirements.txt /tmp/requirements.txt
@@ -58,15 +66,14 @@ WORKDIR /diffusers/examples
 
 RUN mkdir -p ./vertex_vision_model_garden_peft/
 COPY model_oss/peft/train/vmg/configs/* ./vertex_vision_model_garden_peft/
-# custom `lm_eval` task.
-ARG LM_EVAL_DIR=$(python -c 'import site; print(site.getsitepackages()[0])')/lm_eval
-RUN mkdir -p $LM_EVAL_DIR/tasks/vertex && \
-  mv ./vertex_vision_model_garden_peft/custom_loglikelihood.yaml $LM_EVAL_DIR/tasks/vertex/
 COPY model_oss/peft/train/vmg/*.py ./vertex_vision_model_garden_peft/train/vmg/
 COPY model_oss/peft/train/vmg/templates /diffusers/examples/util/templates
-COPY model_oss/util /diffusers/examples/util
+COPY model_oss/peft/train/util/*.py /diffusers/examples/util/
+COPY model_oss/util/* /diffusers/examples/util/
 COPY model_oss/notebook_util/dataset_validation_util.py /diffusers/examples/util
-COPY model_oss/peft/train/tests/*.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/vmg/tests/*.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/test_utils/test_util.py ./vertex_vision_model_garden_peft/tests/
+COPY model_oss/peft/train/test_utils/command_builder.py ./vertex_vision_model_garden_peft/tests/
 
 RUN chmod a+rwX -R /diffusers/examples/
 ENV PYTHONPATH /diffusers/examples/

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

@@ -1,183 +1,284 @@
 """Library for running evaluations during training."""
 
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 import dataclasses
-from typing import Any, Optional, Type
+import string
+from typing import Type
 
 from absl import logging
-import datasets
-from lm_eval import evaluator
-from lm_eval import tasks
-from lm_eval import utils
-from lm_eval.api import model as lm_model
-from lm_eval.api import registry
-from lm_eval.models import huggingface
-from peft import peft_model
+import evaluate
+import numpy as np
+import torch
 import transformers
-from transformers import trainer
 
 from util import dataset_validation_util
 from util import constants
 
 
-_DESCRIPTION_EVALUATION = "evaluation"
-_BUILTIN_EVAL_TASK = "builtin_eval"
+_STRING_TRANSLATOR = str.maketrans("", "", string.punctuation)
+
+_GREATER_IS_BETTER_MAP = {
+    "loss": False,
+    "perplexity": False,
+    "bleu": True,
+    "google_bleu": True,
+    "rouge1": True,
+    "rouge2": True,
+    "rougeL": True,
+    "rougeLsum": True,
+}
 
 
 @dataclasses.dataclass(frozen=True)
 class EvalConfig:
+  """Configuration for running evaluations during training.
+
+  Attributes:
+    steps: The number of steps to run evaluation.
+    tasks: The list of tasks to run evaluation on.
+    per_device_batch_size: The per device batch size for evaluation.
+    limit: The maximum number of examples to evaluate.
+    metric_name: The name of the metric to compute.
+    tokenize_dataset: Whether to tokenize the dataset.
+    dataset_path: The path to the dataset.
+    split: The split of the dataset to evaluate.
+    template: The template to use for the dataset.
+    column: The column name of the dataset.
+    metric_for_best_model: The metric to use for loading the best model.
+  """
+
   steps: int
-  tasks: list[str]
   per_device_batch_size: int
-  num_fewshot: Optional[int]
-  limit: Optional[float]
-  metric_name: str
+  limit: float | None
+  metric_name: Sequence[str]
   tokenize_dataset: bool
   dataset_path: str = ""
   split: str = "test"
   template: str = ""
-  column: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET
-
-
-class PeftCausalLMModel(huggingface.HFLM):
-  """PeftCausalLMModel that supports loading an in-memory model."""
-
-  AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
-
-  def __init__(
-      self,
-      model: peft_model.PeftModelForCausalLM,
-      tokenizer: transformers.PreTrainedTokenizerBase,
-      batch_size_per_gpu: int,
-  ):
-    lm_model.LM.__init__(self)
-    self._model = model
-    self.tokenizer = tokenizer
-    self.vocab_size = tokenizer.vocab_size
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-    self._config = model.config
-    self.batch_size_per_gpu = batch_size_per_gpu
-    self._device = model.device
-    self._max_length = None  # Will be automatically determined from config.
-    self._add_special_tokens = (
-        None  # Will be automatically determined from AUTO_MODEL_CLASS.
-    )
+  column: str = constants.DEFAULT_TRAIN_COLUMN
+  metric_for_best_model: str | None = None
 
 
 def create_trainer(
     cls: Type[transformers.Trainer],
-    eval_config: Optional[EvalConfig],
-    tokenizer: Optional[transformers.PreTrainedTokenizerBase],
-    args: trainer.TrainingArguments,
+    eval_config: EvalConfig | None,
+    tokenizer: transformers.PreTrainedTokenizerBase | None,
+    args: transformers.TrainingArguments,
     **kwargs,
 ) -> transformers.Trainer:
-  """Creates a trainer. If eval config is provided, injects evaluation loop."""
+  """Creates a trainer. If eval config is provided, injects evaluation loop.
+
+  Args:
+    cls: The trainer class.
+    eval_config: The evaluation config.
+    tokenizer: The tokenizer.
+    args: The training arguments.
+    **kwargs: The keyword arguments.
+
+  Returns:
+    A trainer.
+  """
   if not eval_config:
     return cls(args=args, **kwargs)
 
   args.eval_strategy = "steps"
   args.eval_steps = eval_config.steps
   args.per_device_eval_batch_size = eval_config.per_device_batch_size
+  args.metric_for_best_model = eval_config.metric_for_best_model
+  args.greater_is_better = _GREATER_IS_BETTER_MAP.get(
+      eval_config.metric_for_best_model, None
+  )
+  args.save_strategy = (
+      transformers.trainer_utils.SaveStrategy.STEPS
+      if eval_config.metric_for_best_model is None
+      else transformers.trainer_utils.SaveStrategy.BEST
+  )
+
   kwargs["tokenizer"] = tokenizer
 
-  if eval_config.tasks == [_BUILTIN_EVAL_TASK]:
-    try:
-      eval_dataset = dataset_validation_util.load_dataset_with_template(
-          dataset_name=eval_config.dataset_path,
-          split=eval_config.split,
+  try:
+    _, eval_dataset = dataset_validation_util.load_dataset_with_template(
+        dataset_name=eval_config.dataset_path,
+        split=eval_config.split,
+        input_column=eval_config.column,
+        template=eval_config.template,
+        tokenizer=tokenizer,
+    )
+    if eval_config.limit is not None:
+      if eval_config.limit >= 1:
+        limit = int(eval_config.limit)
+      else:
+        limit = int(eval_config.limit * len(eval_dataset))
+      eval_dataset = eval_dataset.select(range(limit))
+    if tokenizer is not None:
+      eval_dataset = dataset_validation_util.get_filtered_dataset(
+          dataset=eval_dataset,
           input_column=eval_config.column,
-          template=eval_config.template,
+          max_seq_length=kwargs["max_seq_length"],
           tokenizer=tokenizer,
       )
-      if eval_config.limit is not None:
-        if eval_config.limit >= 1:
-          limit = int(eval_config.limit)
-        else:
-          limit = int(eval_config.limit * len(eval_dataset))
-        eval_dataset = eval_dataset.select(range(limit))
-      if eval_config.tokenize_dataset:
-        eval_dataset = eval_dataset.map(
-            lambda samples: tokenizer(samples[eval_config.column])
-        )
-      kwargs["eval_dataset"] = eval_dataset
-    except (OSError, ValueError, IndexError) as e:
-      logging.warning(
-          "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
-          eval_config.dataset_path,
-          e,
+    if eval_config.tokenize_dataset:
+      eval_dataset = eval_dataset.map(
+          lambda samples: tokenizer(samples[eval_config.column])
       )
-      del args.evaluation_strategy
-      del args.eval_steps
-      del args.per_device_eval_batch_size
-    return cls(args=args, **kwargs)
+    kwargs["eval_dataset"] = eval_dataset
+  except (OSError, ValueError, IndexError) as e:
+    logging.warning(
+        "Failed to load eval dataset %s. Evaluation will be skipped.\n%s",
+        eval_config.dataset_path,
+        e,
+    )
+    del args.evaluation_strategy
+    del args.eval_steps
+    del args.per_device_eval_batch_size
+  return cls(args=args, **kwargs)
 
-  class LMEvalTrainer(cls):
-    """Trainer with lm_eval injected as the eval library."""
 
-    def __init__(self, **kwargs):
-      super().__init__(**kwargs)
-      task_names = utils.pattern_match(eval_config.tasks, registry.ALL_TASKS)
-      logging.info("Selected Eval Tasks: %s", task_names)
-      task_args = {}
-      if eval_config.num_fewshot is not None:
-        task_args["num_fewshot"] = eval_config.num_fewshot
-      if eval_config.dataset_path:
-        task_args["dataset_path"] = "json"
-        task_args["dataset_kwargs"] = {
-            "data_files": {"test": eval_config.dataset_path},
-        }
-      self._eval_task_dict = tasks.get_task_dict(task_names, **task_args)
+def _cleanup_text(text: str) -> str:
+  """Cleans up the prediction and references text.
 
-    def evaluation_loop(
-        self,
-        dataloader: trainer.DataLoader,
-        description: str,
-        prediction_loss_only: Optional[bool] = None,
-        ignore_keys: Optional[list[str]] = None,
-        metric_key_prefix: str = "eval",
-    ) -> trainer.EvalLoopOutput:
-      """Custom evaluation loop that invokes lm_eval."""
-      if description.lower() != _DESCRIPTION_EVALUATION:
-        return super().evaluation_loop(
-            dataloader,
-            description,
-            prediction_loss_only,
-            ignore_keys,
-            metric_key_prefix,
-        )
+  Args:
+    text: The text to clean up.
 
-      model = self._wrap_model(self.model, training=False)
-      lm = PeftCausalLMModel(
-          model,
-          self.tokenizer or self.data_collator.tokenizer,
-          eval_config.per_device_batch_size,
-      )
-      results: dict[str, Any] = evaluator.evaluate(
-          lm=lm,
-          task_dict=self._eval_task_dict,
-          limit=eval_config.limit,
-      )["results"]
-      metric_name = eval_config.metric_name
-      # Compute average value if there are multiple tasks.
-      metric_values: list[float] = []
-      for result in results.values():
-        for key, value in result.items():
-          if key.split(",")[0] == metric_name:
-            metric_values.append(value)
-      if not metric_values:
-        raise ValueError(
-            f"Metric {metric_name} not found in eval response: {results}"
-        )
-      metric_average = sum(metric_values) / len(metric_values)
-      logging.info("%s value: %f\n%s", metric_name, metric_average, results)
-      return trainer.EvalLoopOutput(
-          # Only metrics field is set. Other fields are dummy values.
-          predictions=None,
-          label_ids=None,
-          metrics={f"{metric_key_prefix}_{metric_name}": metric_average},
-          num_samples=0,
-      )
+  Returns:
+    Cleaned up text.
+  """
+  text = text.translate(_STRING_TRANSLATOR)
+  text = text.strip()
+  text = " ".join(text.split())
+  return text.lower()
 
-  # Use empty eval dataset as a placeholder.
-  return LMEvalTrainer(
-      args=args, eval_dataset=datasets.Dataset.from_dict({"test": []}), **kwargs
+
+def create_compute_metrics(
+    tokenizer: transformers.PreTrainedTokenizerBase,
+    eval_metrics: Mapping[str, evaluate.EvaluationModule],
+) -> Callable[[transformers.EvalPrediction], MutableMapping[str, float]]:
+  """Creates a compute_metrics function using Hugging Face evaluate library.
+
+  Args:
+    tokenizer: The tokenizer for decoding predictions.
+    eval_metrics: The eval metrics to compute.
+
+  Returns:
+    Function that computes comprehensive metrics.
+  """
+
+  def _preprocess_data(
+      predictions: np.ndarray, labels: np.ndarray
+  ) -> tuple[Sequence[str], Sequence[str]]:
+    """Preprocesses predictions and lavels before evaluation.
+
+    Args:
+      predictions: The predictions to preprocess.
+      labels: The labels to preprocess.
+
+    Returns:
+      A tuple (preprocessed predictions, labels).
+    """
+    # Handle padding and special tokens.
+    predictions = np.where(
+        predictions != -100, predictions, tokenizer.pad_token_id
+    )
+    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+
+    # Decode to text.
+    pred_texts = tokenizer.batch_decode(predictions, skip_special_tokens=True)
+    label_texts = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+    # Clean up text.
+    cleaned_pred_texts = [_cleanup_text(text) for text in pred_texts]
+    cleaned_label_texts = [_cleanup_text(text) for text in label_texts]
+
+    return cleaned_pred_texts, cleaned_label_texts
+
+  def _compute_metrics_with_tokenizer(
+      eval_pred: transformers.EvalPrediction,
+  ) -> MutableMapping[str, float]:
+    """Computes metrics using Hugging Face evaluate library.
+
+    Args:
+      eval_pred: The evaluation prediction.
+
+    Returns:
+      A dictionary of metrics.
+    """
+    predictions, perplexities = eval_pred.predictions
+    labels = eval_pred.label_ids
+
+    pred_texts, label_texts = _preprocess_data(predictions, labels)
+
+    metrics = {}
+
+    for eval_metric, computed_eval_metric in eval_metrics.items():
+      match eval_metric:
+        case "perplexity":
+          # We don't use the perplexity from HF Evaluate since it loads the
+          # model again. This causes an increase in the GPU utilization and
+          # hence an OOM. Due to this, we compute the perplexity ourselves
+          # using the eval_loss over the unmasked tokens in
+          # preprocess_logits_for_metrics fn.
+          metrics[eval_metric] = np.mean(perplexities)
+        case "bleu" | "google_bleu":
+          num_valid_labels = len(list(filter(None, label_texts)))
+          if num_valid_labels:
+            eval_score = computed_eval_metric.compute(
+                predictions=pred_texts,
+                references=[[text] for text in label_texts],
+            )
+            metrics[eval_metric] = eval_score[eval_metric]
+          else:
+            metrics[eval_metric] = 0.0
+        case "rouge1" | "rouge2" | "rougeL" | "rougeLsum":
+          rouge_scores = computed_eval_metric.compute(
+              predictions=pred_texts,
+              references=label_texts,
+              use_stemmer=True,
+          )
+          metrics[eval_metric] = rouge_scores[eval_metric]
+
+    pred_lengths = [len(pred.split()) for pred in pred_texts]
+    label_lengths = [len(label.split()) for label in label_texts]
+
+    metrics["gen_len"] = np.mean(pred_lengths)
+    metrics["ref_len"] = np.mean(label_lengths)
+    metrics["length_ratio"] = np.mean(
+        [len(p) / len(r) if r else 0 for p, r in zip(pred_texts, label_texts)]
+    )
+
+    # Round all metrics to 4 decimal places.
+    return {k: round(float(v), 4) for k, v in metrics.items()}
+
+  return _compute_metrics_with_tokenizer
+
+
+def preprocess_logits_for_metrics(
+    logits: torch.Tensor, labels: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+  """Preprocesses the logits before caching them for eval metric calculation.
+
+  Args:
+    logits: Logits predicted by the model.
+    labels: Ground truth labels.
+
+  Returns:
+    A tuple (pred_ids, perplexities).
+  """
+  # Calculate prediction IDs.
+  pred_ids = logits.argmax(dim=-1)
+
+  # This step shifts the logits and labels to align them correctly, where we are
+  # predicting the next token in a sequence. The last logit doesn't have a
+  # corresponding label, and the first label doesn't have a preceding logit to
+  # predict it. This calculation of perplexity is inspired from
+  # https://github.com/huggingface/evaluate/blob/main/metrics/perplexity/perplexity.py.
+  shift_logits = logits[..., :-1, :].contiguous()
+  shift_labels = labels[..., 1:].contiguous()
+  attn_mask = shift_labels != -100
+  loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
+
+  perplexities = torch.exp(
+      (loss_fct(shift_logits.transpose(1, 2), shift_labels) * attn_mask).sum(1)
+      / attn_mask.sum(1)
   )
+
+  return (pred_ids, perplexities)

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

@@ -1,10 +1,10 @@
 """Instruct/Chat with LoRA models."""
 
-import dataclasses
+from collections.abc import Callable, Mapping, Sequence
 import datetime
 import json
 import os
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 import warnings
 
 from absl import app
@@ -13,17 +13,17 @@ from absl import logging
 from accelerate import DistributedType
 from accelerate import PartialState
 import bitsandbytes as bnb
-import hypertune
+import evaluate
 from peft import get_peft_model
 from peft import LoraConfig
 import torch
-from transformers import AutoModelForCausalLM
-from transformers import TrainingArguments
-from trl import DataCollatorForCompletionOnlyLM
-from trl import SFTTrainer
+import transformers
+import trl
 import wandb
 
 from util import dataset_validation_util
+from util import dataset_stats
+from util import device_stats
 from vertex_vision_model_garden_peft.train.vmg import callbacks
 from vertex_vision_model_garden_peft.train.vmg import eval_lib
 from vertex_vision_model_garden_peft.train.vmg import utils
@@ -31,14 +31,15 @@ from util import constants
 from util import fileutils
 
 
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    'pretrained_model_id',
+_PRETRAINED_MODEL_NAME_OR_PATH = flags.DEFINE_string(
+    'pretrained_model_name_or_path',
     None,
-    'The pretrained model id. Supported models can be causal language modeling'
-    ' models from https://github.com/huggingface/peft/tree/main. Note, there'
-    ' might be different paddings for different models. This tool assumes the'
-    ' pretrained_model_id contains model name, and then choose proper padding'
-    ' methods. e.g. it must contain `llama` for `Llama2 models`.',
+    'The pretrained model name or path. Supported models can be causal language'
+    ' modeling models from https://github.com/huggingface/peft/tree/main. Note,'
+    ' there might be different paddings for different models. This tool assumes'
+    ' the pretrained_model_name_or_path contains model name, and then choose'
+    ' proper padding methods. e.g. it must contain `llama` for `Llama2'
+    ' models`.',
     required=True,
 )
 
@@ -48,10 +49,10 @@ _HUGGINGFACE_ACCESS_TOKEN = flags.DEFINE_string(
     'The access token for loading huggingface gated models.',
 )
 
-_DATASET_NAME = flags.DEFINE_string(
-    'dataset_name',
+_TRAIN_DATASET = flags.DEFINE_string(
+    'train_dataset',
     None,
-    'The dataset name in huggingface.',
+    'The training dataset name in huggingface or path.',
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
@@ -119,8 +120,8 @@ _WEIGHT_DECAY = flags.DEFINE_float(
     'The weight decay in the learning rate scheduler.',
 )
 
-_NUM_EPOCHS = flags.DEFINE_float(
-    'num_epochs',
+_NUM_TRAIN_EPOCHS = flags.DEFINE_float(
+    'num_train_epochs',
     None,
     'The number of training epochs. Only used for'
     ' "sequence-classification-lora" with an integer value and for'
@@ -130,8 +131,8 @@ _NUM_EPOCHS = flags.DEFINE_float(
 _MAX_STEPS = flags.DEFINE_integer(
     'max_steps',
     None,
-    'Total number of training steps. Overrides num_epochs if set. Only used for'
-    ' "instruct-lora."',
+    'Total number of training steps. Overrides num_train_epochs if set. Only'
+    ' used for "instruct-lora."',
 )
 
 _MAX_SEQ_LENGTH = flags.DEFINE_integer(
@@ -146,9 +147,9 @@ _LEARNING_RATE = flags.DEFINE_float(
     'The learning rate after the potential warmup period.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -170,8 +171,8 @@ _GRADIENT_ACCUMULATION_STEPS = flags.DEFINE_integer(
     'The gradient accumulation steps.',
 )
 
-_ENABLE_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
-    'enable_gradient_checkpointing',
+_GRADIENT_CHECKPOINTING = flags.DEFINE_boolean(
+    'gradient_checkpointing',
     False,
     'Whether to enable gradient checkpointing.',
 )
@@ -181,8 +182,8 @@ _ENABLE_PEFT = flags.DEFINE_boolean(
     True,
     'Whether to enable peft.',
 )
-_TEMPLATE = flags.DEFINE_string(
-    'template',
+_TRAIN_TEMPLATE = flags.DEFINE_string(
+    'train_template',
     None,
     'Template for formatting language model training data. Must be a filename'
     ' under `templates` folder, without `.json` extension, e.g. `alpaca`, or a'
@@ -219,32 +220,18 @@ _EVAL_STEPS = flags.DEFINE_integer(
     'The number of training steps between evaluations.',
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_EVAL_TASKS = flags.DEFINE_list(
-    'eval_tasks',
-    None,
-    'List of eval task names (can have wildcards) as in'
-    ' https://github.com/EleutherAI/lm-evaluation-harness. Will not run'
-    ' evaluation if not set. Runs the built-in trainer evaluation loop if set'
-    ' to `builtin_eval`.',
-)
-
-_EVAL_PER_DEVICE_BATCH_SIZE = flags.DEFINE_integer(
-    'eval_per_device_batch_size',
+_PER_DEVICE_EVAL_BATCH_SIZE = flags.DEFINE_integer(
+    'per_device_eval_batch_size',
     1,
     'The per device batch size for model evaluation.',
 )
 
-_EVAL_NUM_FEWSHOT = flags.DEFINE_integer(
-    'eval_num_fewshot',
-    None,
-    'Run N-shot language model evaluation. Not implemented in `builtin_eval`.',
-)
 
 _EVAL_LIMIT = flags.DEFINE_float(
     'eval_limit',
@@ -253,17 +240,18 @@ _EVAL_LIMIT = flags.DEFINE_float(
     ' total number of examples.',
 )
 
-_EVAL_METRIC_NAME = flags.DEFINE_string(
+_EVAL_METRIC_NAME = flags.DEFINE_list(
     'eval_metric_name',
-    'acc',
-    'The metric name to aggregate during model evaluation.',
+    ['loss'],
+    'A comma-separated list of metric names to aggregate during model'
+    ' evaluation. The supported metrics are: '
+    + ', '.join(constants.SUPPORTED_EVAL_METRICS),
 )
 
-_EVAL_DATASET_PATH = flags.DEFINE_string(
-    'eval_dataset_path',
+_EVAL_DATASET = flags.DEFINE_string(
+    'eval_dataset',
     None,
-    'Overrides the default evaluation dataset path. In `builtin_eval` mode,'
-    ' this can be any Hugging Face dataset name.',
+    'The Hugging Face dataset name or path to use for evaluation.',
 )
 
 # We set the default eval split as `test`, based on observation from
@@ -271,13 +259,13 @@ _EVAL_DATASET_PATH = flags.DEFINE_string(
 _EVAL_SPLIT = flags.DEFINE_string(
     'eval_split',
     'test',
-    'Eval split name in the eval dataset for `builtin_eval`.',
+    'Eval split name in the eval dataset.',
 )
 
 _EVAL_TEMPLATE = flags.DEFINE_string(
     'eval_template',
     None,
-    'Template for formatting language model evaluation data for `builtin_eval`.'
+    'Template for formatting language model evaluation data.'
     ' Must be a filename under `templates` folder, without `.json` extension,'
     ' e.g. `alpaca`, or a Cloud Storage URI to a JSON file.',
 )
@@ -285,7 +273,14 @@ _EVAL_TEMPLATE = flags.DEFINE_string(
 _EVAL_COLUMN = flags.DEFINE_string(
     'eval_column',
     None,
-    'Eval column name in the eval dataset for `builtin_eval`.',
+    'Eval column name in the eval dataset.',
+)
+
+_METRIC_FOR_BEST_MODEL = flags.DEFINE_string(
+    'metric_for_best_model',
+    None,
+    'If set, the best model is saved at the end of training based on the'
+    ' metric',
 )
 
 _TRAIN_PRECISION = flags.DEFINE_enum(
@@ -299,15 +294,15 @@ _TRAIN_PRECISION = flags.DEFINE_enum(
     'Precision to train the model.',
 )
 
-_USE_EXAMPLE_PACKING = flags.DEFINE_boolean(
-    'use_example_packing',
+_EXAMPLE_PACKING = flags.DEFINE_boolean(
+    'example_packing',
     False,
     'Enables example packing during training, which uses '
     '`ConstantLengthDataset` under the hood.',
 )
 
-_COMPLETION_ONLY = flags.DEFINE_boolean(
-    'completion_only',
+_INPUT_MASKING = flags.DEFINE_boolean(
+    'input_masking',
     False,
     'If set, it uses DataCollatorForCompletionOnlyLM to train the model on the'
     ' generated prompts only, i.e., masking out the input',
@@ -356,28 +351,53 @@ _TARGET_MODULES = flags.DEFINE_list(
     'target_modules', None, 'The names of the modules to apply LoRA adapter to.'
 )
 
+_MAX_GPU_MEMORY_FRACTION = flags.DEFINE_float(
+    'max_gpu_memory_fraction',
+    '0.9',
+    'Maximum GPU memory a caching allocator is allowed to use per GPU.',
+)
+
 
 @flags.multi_flags_validator(
     [
-        _COMPLETION_ONLY.name,
-        _USE_EXAMPLE_PACKING.name,
+        _INPUT_MASKING.name,
+        _EXAMPLE_PACKING.name,
     ],
-    message=(
-        '`use_example_packing=True` does not work with `completion_only=True`'
-    ),
+    message='`example_packing=True` does not work with `input_masking=True`',
 )
-def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
+def check_example_packing(flags_dict: Mapping[str, Any]) -> bool:
   """Check to make sure example packing is enabled properly.
 
   Args:
     flags_dict: Dictionary containing flags to check.
 
   Returns:
-    If `use_example_packing` is set properly.
+    If `example_packing` is set properly.
+  """
+  if flags_dict[_INPUT_MASKING.name] and flags_dict[_EXAMPLE_PACKING.name]:
+    return False
+  return True
+
+
+@flags.multi_flags_validator(
+    [
+        _INPUT_MASKING.name,
+        _TRAIN_TEMPLATE.name,
+    ],
+    message='`train_template` should be provided if using `input_masking=True`',
+)
+def check_input_masking(flags_dict: Mapping[str, Any]) -> bool:
+  """Check to make sure input_masking is enabled properly.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If `input_masking` is set properly
   """
   if (
-      flags_dict[_COMPLETION_ONLY.name]
-      and flags_dict[_USE_EXAMPLE_PACKING.name]
+      flags_dict[_INPUT_MASKING.name]
+      and flags_dict[_TRAIN_TEMPLATE.name] is None
   ):
     return False
   return True
@@ -385,22 +405,65 @@ def check_example_packing(flags_dict: Dict[str, Any]) -> bool:
 
 @flags.multi_flags_validator(
     [
-        _COMPLETION_ONLY.name,
-        _TEMPLATE.name,
+        _EVAL_DATASET.name,
+        _EVAL_METRIC_NAME.name,
     ],
-    message='`template` should be provided if using `completion_only=True`',
+    message=(
+        '`eval_metric_name` should be a valid metric name and present when'
+        ' eval_dataset is provided.'
+    ),
 )
-def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
-  """Check to make sure completion_only is enabled properly.
+def _validate_eval_metrics(flags_dict: Mapping[str, Any]) -> bool:
+  """Validates the eval metric name.
 
   Args:
     flags_dict: Dictionary containing flags to check.
 
   Returns:
-    If `completion_only` is set properly
+    If the eval metrics are valid.
   """
-  if flags_dict[_COMPLETION_ONLY.name] and flags_dict[_TEMPLATE.name] is None:
-    return False
+  if flags_dict[_EVAL_DATASET.name] is None:
+    return True
+  eval_metrics = flags_dict[_EVAL_METRIC_NAME.name]
+  for eval_metric in eval_metrics:
+    if eval_metric not in constants.SUPPORTED_EVAL_METRICS:
+      raise flags.ValidationError(f'Invalid eval metric: {eval_metric}')
+  if 'perplexity' in eval_metrics and 'loss' not in eval_metrics:
+    _EVAL_METRIC_NAME.value.append('loss')
+    logging.warning(
+        'Adding `loss` to eval_metric_name because `perplexity` is present.'
+    )
+  return True
+
+
+@flags.multi_flags_validator(
+    [
+        _METRIC_FOR_BEST_MODEL.name,
+        _EVAL_METRIC_NAME.name,
+    ],
+    message='`metric_for_best_model` should be in `eval_metric_name`.',
+)
+def _validate_metric_for_best_model(flags_dict: Mapping[str, Any]) -> bool:
+  """Validates the metric for best model.
+
+  Args:
+    flags_dict: Dictionary containing flags to check.
+
+  Returns:
+    If the metric for best model is valid.
+  """
+  if flags_dict[_METRIC_FOR_BEST_MODEL.name] is None:
+    return True
+
+  metric_for_best_model = flags_dict[_METRIC_FOR_BEST_MODEL.name]
+  eval_metric_name = flags_dict[_EVAL_METRIC_NAME.name]
+
+  if metric_for_best_model not in eval_metric_name:
+    raise flags.ValidationError(
+        'Invalid metric for picking the best model:'
+        f' {metric_for_best_model}. The metric should be one'
+        f' of the {eval_metric_name}.'
+    )
   return True
 
 
@@ -411,10 +474,43 @@ def check_completion_only(flags_dict: Dict[str, Any]) -> bool:
 # https://github.com/huggingface/notebooks/blob/main/sagemaker/28_train_llms_with_qlora/scripts/run_clm.py.
 
 
+def _calculate_hf_eval_metrics(
+    tokenizer: transformers.PreTrainedTokenizerBase,
+    eval_config: eval_lib.EvalConfig | None,
+) -> tuple[
+    Callable[[transformers.EvalPrediction], Mapping[str, float]], torch.Tensor
+]:
+  """Calculates the HF evaluation metrics.
+
+  Args:
+    tokenizer: The tokenizer to use for evaluation.
+    eval_config: The evaluation config to use.
+
+  Returns:
+    The compute metrics and preprocess logits for metrics.
+  """
+  if eval_config is None:
+    return None, None
+  hf_eval_metrics = {}
+  for metric in eval_config.metric_name:
+    if metric in constants.SUPPORTED_HF_EVAL_METRICS:
+      if metric in constants.ROUGE_VARIANTS:
+        hf_eval_metrics[metric] = evaluate.load('rouge')
+      else:
+        hf_eval_metrics[metric] = evaluate.load(metric)
+
+  if not hf_eval_metrics:
+    return None, None
+  return (
+      eval_lib.create_compute_metrics(tokenizer, hf_eval_metrics),
+      eval_lib.preprocess_logits_for_metrics,
+  )
+
+
 # Copied from https://github.com/artidoro/qlora/blob/main/qlora.py.
 def find_all_linear_names(
-    model: AutoModelForCausalLM, precision_mode: str
-) -> list[str]:
+    model: transformers.AutoModelForCausalLM, precision_mode: str
+) -> Sequence[str]:
   """Finds all linear module names."""
   if precision_mode == constants.PRECISION_MODE_4:
     cls = bnb.nn.Linear4bit
@@ -433,47 +529,51 @@ def find_all_linear_names(
 
 
 def finetune_instruct(
-    pretrained_model_id: str,
-    dataset_name: str,
+    pretrained_model_name_or_path: str,
+    train_dataset: str,
     output_dir: str,
     logging_output_dir: str,
     lora_rank: int = 64,
     lora_alpha: int = 16,
     lora_dropout: float = 0.1,
     warmup_ratio: int = 0.03,
-    num_epochs: Optional[float] = None,
-    max_steps: Optional[int] = None,
+    num_train_epochs: float | None = None,
+    max_steps: int | None = None,
     warmup_steps: int = 10,
     max_seq_length: int = 512,
     learning_rate: float = 2e-4,
     precision_mode: str = None,
-    instruct_column_in_dataset: str = constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+    train_column: str = constants.DEFAULT_TRAIN_COLUMN,
     per_device_train_batch_size: int = 4,
     gradient_accumulation_steps: int = 4,
     optim: str = 'paged_adamw_32bit',
     weight_decay: float = 0.001,
-    enable_gradient_checkpointing: bool = False,
+    gradient_checkpointing: bool = False,
     enable_peft: bool = True,
-    template: str = None,
+    train_template: str = None,
     lr_scheduler_type: str = 'constant',
     save_steps: int = 10,
     logging_steps: int = 10,
-    train_split_name: str = 'train',
-    eval_config: Optional[eval_lib.EvalConfig] = None,
+    train_split: str = 'train',
+    eval_config: eval_lib.EvalConfig | None = None,
     report_to: str = constants.REPORT_TO_NONE,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
     train_precision: str = constants.PRECISION_MODE_16B,
-    use_example_packing: bool = False,
-    attn_implementation: Optional[str] = None,
+    example_packing: bool = False,
+    attn_implementation: str | None = None,
     max_grad_norm: float = 0.3,
-    completion_only: bool = False,
+    input_masking: bool = False,
     logger_level: str = 'passive',
-    benchmark_out_file: Optional[str] = None,
-    tuning_data_stats_file: Optional[str] = None,
-    target_modules: Optional[str] = None,
+    benchmark_out_file: str | None = None,
+    tuning_data_stats_file: str | None = None,
+    target_modules: str | None = None,
 ) -> None:
   """Finetunes instruct."""
-  logging.info('on entering instruct_lora, %s', utils.gpu_stats_str())
+  logging.info(
+      'on entering instruct_lora, %s,\n%s',
+      device_stats.gpu_stats_str(),
+      device_stats.cpu_stats_str(),
+  )
   gradient_checkpointing_kwargs = {}
   # DDP provides limited support with the reentrant variant of gradient
   # checkpoint [1]. Below is an indirect way of checking whether DDP will be
@@ -483,17 +583,25 @@ def finetune_instruct(
   if PartialState().distributed_type == DistributedType.MULTI_GPU:
     gradient_checkpointing_kwargs['use_reentrant'] = False
 
-  tokenizer = utils.load_tokenizer(
-      pretrained_model_id,
+  tokenizer = dataset_validation_util.load_tokenizer(
+      pretrained_model_name_or_path,
       'right',
       access_token=access_token,
   )
 
-  train_dataset = dataset_validation_util.load_dataset_with_template(
-      dataset_name,
-      split=train_split_name,
-      input_column=instruct_column_in_dataset,
-      template=template,
+  train_dataset, train_dataset_with_template = (
+      dataset_validation_util.load_dataset_with_template(
+          train_dataset,
+          split=train_split,
+          input_column=train_column,
+          template=train_template,
+          tokenizer=tokenizer,
+      )
+  )
+  train_dataset_with_template = dataset_validation_util.get_filtered_dataset(
+      dataset=train_dataset_with_template,
+      input_column=train_column,
+      max_seq_length=max_seq_length,
       tokenizer=tokenizer,
   )
 
@@ -508,24 +616,26 @@ def finetune_instruct(
           'getting tuning data stats with effective batch size %s',
           effective_batch_size,
       )
-      train_dataset_stats = utils.get_dataset_stats(
-          train_dataset,
-          tokenizer,
-          instruct_column_in_dataset,
-          effective_batch_size,
+      train_dataset_stats = dataset_stats.get_dataset_stats(
+          raw=train_dataset,
+          templated=train_dataset_with_template,
+          template=train_template,
+          tokenizer=tokenizer,
+          column=train_column,
+          effective_batch_size=effective_batch_size,
       )
       logging.info('stats: %s', train_dataset_stats)
       tuning_data_stats_file = dataset_validation_util.force_gcs_fuse_path(
           tuning_data_stats_file
       )
       with open(tuning_data_stats_file, 'w') as out_f:
-        json.dump(dataclasses.asdict(train_dataset_stats), out_f)
+        json.dump(train_dataset_stats, out_f)
 
   model = utils.load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=precision_mode,
-      enable_gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing=gradient_checkpointing,
       access_token=access_token,
       attn_implementation=attn_implementation,
       train_precision=train_precision,
@@ -550,7 +660,9 @@ def finetune_instruct(
     # `get_peft_model`, which may revert other changes we did before. That's why
     # we are calling `get_peft_model` explicitly here.
     model = get_peft_model(model, peft_config)
-
+    adapter_for_eval_dir = os.path.join(output_dir, 'adapter_for_eval')
+    logging.info('saving adapter for evaluation to %s...', adapter_for_eval_dir)
+    peft_config.save_pretrained(adapter_for_eval_dir)
     # This is to work-around mix-precision training. This issue is not fixed as
     # of transformers==4.41.2.
     # See b/332760883#comment30 for more details.
@@ -568,14 +680,13 @@ def finetune_instruct(
   # b/357970482#comment3
   accelerator_config = {'use_configured_state': True}
 
-  training_arguments = TrainingArguments(
+  training_arguments = transformers.TrainingArguments(
       report_to=report_to,
       output_dir=output_dir,
       per_device_train_batch_size=per_device_train_batch_size,
       gradient_accumulation_steps=gradient_accumulation_steps,
       optim=optim,
       save_steps=save_steps,
-      save_strategy='steps',
       save_total_limit=3,
       logging_dir=os.path.join(logging_output_dir, 'logs'),
       logging_steps=logging_steps,
@@ -583,21 +694,24 @@ def finetune_instruct(
       fp16=(train_precision == constants.PRECISION_MODE_16),
       bf16=(train_precision == constants.PRECISION_MODE_16B),
       max_grad_norm=max_grad_norm,
-      num_train_epochs=num_epochs if num_epochs else -1,
+      num_train_epochs=num_train_epochs if num_train_epochs else -1,
       max_steps=max_steps if max_steps else -1,
       warmup_ratio=warmup_ratio,
       warmup_steps=warmup_steps,
       group_by_length=False,
       lr_scheduler_type=lr_scheduler_type,
-      gradient_checkpointing=enable_gradient_checkpointing,
+      gradient_checkpointing=gradient_checkpointing,
       gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
       weight_decay=weight_decay,
       log_level=logger_level,
       accelerator_config=accelerator_config,
+      include_num_input_tokens_seen=True,
   )
   trainer_kwargs = {}
-  if completion_only and template:
-    template_json = dataset_validation_util.get_template(template_path=template)
+  if input_masking and train_template:
+    template_json = dataset_validation_util.get_template(
+        template_path=train_template
+    )
     instruction_sep = dataset_validation_util.get_instruction_separator(
         template_json
     )
@@ -608,7 +722,7 @@ def finetune_instruct(
           ' `DataCollatorForCompletionOnlyLM`'
       )
 
-    trainer_kwargs['data_collator'] = DataCollatorForCompletionOnlyLM(
+    trainer_kwargs['data_collator'] = trl.DataCollatorForCompletionOnlyLM(
         instruction_template=instruction_sep,
         response_template=response_sep,
         tokenizer=tokenizer,
@@ -618,17 +732,23 @@ def finetune_instruct(
   trainer_stats_callback = callbacks.TrainerStatsCallback(
       max_seq_length, benchmark_out_file
   )
+  compute_metrics, preprocess_logits = _calculate_hf_eval_metrics(
+      tokenizer, eval_config
+  )
+
   trainer = eval_lib.create_trainer(
-      cls=SFTTrainer,
+      cls=trl.SFTTrainer,
       eval_config=eval_config,
       model=model,
-      train_dataset=train_dataset,
-      dataset_text_field=instruct_column_in_dataset,
+      train_dataset=train_dataset_with_template,
+      dataset_text_field=train_column,
       max_seq_length=max_seq_length,
       tokenizer=tokenizer,
       args=training_arguments,
-      packing=use_example_packing,
+      packing=example_packing,
       callbacks=[trainer_stats_callback],
+      compute_metrics=compute_metrics,
+      preprocess_logits_for_metrics=preprocess_logits,
       **trainer_kwargs,
   )
 
@@ -658,7 +778,7 @@ def finetune_instruct(
     # This method saves the sharded weights like `accelerator.save_state`, see
     # https://huggingface.co/docs/accelerate/en/usage_guides/fsdp#saving-and-loading
     trainer.save_model(output_dir)
-    model = trainer.model.cpu()  # Avoids GPU OOM
+    model = trainer.model
     state_dict = trainer.accelerator.get_state_dict(model)
     # To aggregate the weights from all the devices, we need to use
     # `state_dict=state_dict`.
@@ -668,7 +788,6 @@ def finetune_instruct(
         is_main_process=PartialState().is_main_process,
         save_embedding_layers=False,  # Only pad token is added. See go/lora-adapter-pad-token #pylint: disable=line-too-long
     )
-    model.cuda()  # Move back to GPU to do eval.
   else:
     trainer.model.save_pretrained(
         final_checkpoint,
@@ -683,15 +802,6 @@ def finetune_instruct(
     # https://github.com/huggingface/transformers/blob/v4.38.2/src/transformers/trainer_pt_utils.py#L1001 #pylint: disable=line-too-long
     trainer.log_metrics('eval', metrics)
     trainer.save_metrics('eval', metrics)
-    if PartialState().is_main_process:
-      hp_metric = metrics[f'eval_{eval_config.metric_name}']
-      hpt = hypertune.HyperTune()
-      hpt.report_hyperparameter_tuning_metric(
-          hyperparameter_metric_tag=constants.HP_METRIC_TAG,
-          metric_value=hp_metric,
-      )
-      logging.info('Send HP metric: %f to hyperparameter tuning.', hp_metric)
-    PartialState().wait_for_everyone()
 
   if not enable_peft:
     tokenizer.save_pretrained(
@@ -705,38 +815,41 @@ def main(unused_argv: Sequence[str]) -> None:
       timeout=datetime.timedelta(seconds=_NCCL_TIMEOUT.value)
   )
 
+  torch.cuda.set_per_process_memory_fraction(
+      _MAX_GPU_MEMORY_FRACTION.value, device=PartialState().local_process_index
+  )
+
   utils.print_library_versions()
   warnings.simplefilter(_WARNINGS_FILTER.value)
 
-  pretrained_model_id = fileutils.force_gcs_path(_PRETRAINED_MODEL_ID.value)
-  if dataset_validation_util.is_gcs_path(pretrained_model_id):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        pretrained_model_id
-    )
-
-  output_dir = utils.GcsOrLocalDirectory(
-      _OUTPUT_DIR.value, check_empty=True, upload_from_all_nodes=True
+  pretrained_model_name_or_path = fileutils.force_gcs_path(
+      _PRETRAINED_MODEL_NAME_OR_PATH.value
   )
+  if dataset_validation_util.is_gcs_path(pretrained_model_name_or_path):
+    pretrained_model_name_or_path = (
+        dataset_validation_util.download_gcs_uri_to_local(
+            pretrained_model_name_or_path
+        )
+    )
 
   # GCS Fuse does not sync flushed files if not closed. See b/361771727.
   logging_output_dir = fileutils.force_gcs_path(_LOGGING_OUTPUT_DIR.value)
 
   # Creates evaluation config.
-  if _EVAL_TASKS.value:
+  if _EVAL_DATASET.value:
     eval_config = eval_lib.EvalConfig(
-        tasks=_EVAL_TASKS.value,
-        per_device_batch_size=_EVAL_PER_DEVICE_BATCH_SIZE.value,
-        num_fewshot=_EVAL_NUM_FEWSHOT.value,
+        per_device_batch_size=_PER_DEVICE_EVAL_BATCH_SIZE.value,
         limit=_EVAL_LIMIT.value,
         metric_name=_EVAL_METRIC_NAME.value,
         steps=_EVAL_STEPS.value,
         dataset_path=dataset_validation_util.force_gcs_fuse_path(
-            _EVAL_DATASET_PATH.value
+            _EVAL_DATASET.value
         ),
         split=_EVAL_SPLIT.value,
         template=_EVAL_TEMPLATE.value,
         column=_EVAL_COLUMN.value,
         tokenize_dataset=False,
+        metric_for_best_model=_METRIC_FOR_BEST_MODEL.value,
     )
   else:
     eval_config = None
@@ -745,40 +858,40 @@ def main(unused_argv: Sequence[str]) -> None:
     wandb.login()
 
   finetune_instruct(
-      pretrained_model_id=pretrained_model_id,
-      dataset_name=_DATASET_NAME.value,
-      output_dir=output_dir.local_dir,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
+      train_dataset=_TRAIN_DATASET.value,
+      output_dir=_OUTPUT_DIR.value,
       logging_output_dir=logging_output_dir,
       precision_mode=_PRECISION_MODE.value,
       lora_rank=_LORA_RANK.value,
       lora_alpha=_LORA_ALPHA.value,
       lora_dropout=_LORA_DROPOUT.value,
       warmup_ratio=_WARMUP_RATIO.value,
-      num_epochs=_NUM_EPOCHS.value,
+      num_train_epochs=_NUM_TRAIN_EPOCHS.value,
       warmup_steps=_WARMUP_STEPS.value,
       max_steps=_MAX_STEPS.value,
       max_seq_length=_MAX_SEQ_LENGTH.value,
       learning_rate=_LEARNING_RATE.value,
-      instruct_column_in_dataset=_INSTRUCT_COLUMN_IN_DATASET.value,
+      train_column=_TRAIN_COLUMN.value,
       per_device_train_batch_size=_PER_DEVICE_TRAIN_BATCH_SIZE.value,
       optim=_OPTIMIZER.value,
       weight_decay=_WEIGHT_DECAY.value,
       gradient_accumulation_steps=_GRADIENT_ACCUMULATION_STEPS.value,
-      enable_gradient_checkpointing=_ENABLE_GRADIENT_CHECKPOINTING.value,
+      gradient_checkpointing=_GRADIENT_CHECKPOINTING.value,
       enable_peft=_ENABLE_PEFT.value,
-      template=_TEMPLATE.value,
+      train_template=_TRAIN_TEMPLATE.value,
       lr_scheduler_type=_LR_SCHEDULER_TYPE.value,
       save_steps=_SAVE_STEPS.value,
       logging_steps=_LOGGING_STEPS.value,
-      train_split_name=_TRAIN_SPLIT_NAME.value,
+      train_split=_TRAIN_SPLIT.value,
       eval_config=eval_config,
       report_to=_REPORT_TO.value,
       access_token=_HUGGINGFACE_ACCESS_TOKEN.value,
       train_precision=_TRAIN_PRECISION.value,
-      use_example_packing=_USE_EXAMPLE_PACKING.value,
+      example_packing=_EXAMPLE_PACKING.value,
       attn_implementation=_ATTN_IMPLEMENTATION.value,
       max_grad_norm=_MAX_GRAD_NORM.value,
-      completion_only=_COMPLETION_ONLY.value,
+      input_masking=_INPUT_MASKING.value,
       logger_level=_LOGGER_LEVEL.value,
       benchmark_out_file=_BENCHMARK_OUT_FILE.value,
       tuning_data_stats_file=_TUNING_DATA_STATS_FILE.value,
@@ -787,8 +900,6 @@ def main(unused_argv: Sequence[str]) -> None:
   # Frees the model from GPU.
   utils.force_gc()
 
-  output_dir.upload_to_gcs(skip_if_exists=True)
-
 
 if __name__ == '__main__':
   app.run(main)

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

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

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

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

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

@@ -1,229 +0,0 @@
-"""Sequence classification with LoRA models."""
-
-from typing import Sequence
-
-from absl import app
-from absl import flags
-from datasets import load_dataset
-import evaluate
-from peft import get_peft_model
-from peft import LoraConfig
-import torch
-from torch.optim import AdamW
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-from transformers import AutoModelForSequenceClassification
-from transformers import AutoTokenizer
-from transformers import get_linear_schedule_with_warmup
-
-from util import dataset_validation_util
-
-
-_PRETRAINED_MODEL_ID = flags.DEFINE_string(
-    "pretrained_model_id",
-    None,
-    "The pretrained model id. Supported models can be causal language modeling"
-    " models from https://github.com/huggingface/peft/tree/main. Note, there"
-    " might be different paddings for different models. This tool assumes the"
-    " pretrained_model_id contains model name, and then choose proper padding"
-    " methods. e.g. it must contain `llama` for `Llama2 models`.",
-)
-
-_OUTPUT_DIR = flags.DEFINE_string(
-    "output_dir",
-    None,
-    "The output directory.",
-)
-
-_DATASET_NAME = flags.DEFINE_string(
-    "dataset_name",
-    None,
-    "The dataset name in huggingface.",
-)
-
-_LORA_RANK = flags.DEFINE_integer(
-    "lora_rank",
-    16,
-    "The rank of the update matrices, expressed in int. Lower rank results in"
-    " smaller update matrices with fewer trainable parameters, referring to"
-    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
-)
-
-_LORA_ALPHA = flags.DEFINE_integer(
-    "lora_alpha",
-    32,
-    "LoRA scaling factor, referring to"
-    " https://huggingface.co/docs/peft/conceptual_guides/lora.",
-)
-
-_LORA_DROPOUT = flags.DEFINE_float(
-    "lora_dropout",
-    0.05,
-    "dropout probability of the LoRA layers, referring to"
-    " https://huggingface.co/docs/peft/task_guides/token-classification-lora.",
-)
-
-_NUM_EPOCHS = flags.DEFINE_integer(
-    "num_epochs",
-    None,
-    "The number of training epochs.",
-)
-
-_BATCH_SIZE = flags.DEFINE_integer(
-    "batch_size",
-    32,
-    "The batch size.",
-)
-
-_LEARNING_RATE = flags.DEFINE_float(
-    "learning_rate",
-    2e-4,
-    "The learning rate after the potential warmup period.",
-)
-
-
-def finetune_sequence_classification(
-    pretrained_model_id: str,
-    dataset_name: str,
-    output_dir: str,
-    lora_rank: int = 8,
-    lora_alpha: int = 16,
-    lora_dropout: float = 0.1,
-    num_epochs: int = 20,
-    batch_size: int = 32,
-    learning_rate: float = 3e-4,
-) -> None:
-  """Finetunes sequence classification."""
-  task = "mrpc"
-  device = "cuda"
-
-  peft_config = LoraConfig(
-      task_type="SEQ_CLS",
-      inference_mode=False,
-      r=lora_rank,
-      lora_alpha=lora_alpha,
-      lora_dropout=lora_dropout,
-  )
-  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
-    padding_side = "left"
-  else:
-    padding_side = "right"
-
-  tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_id, padding_side=padding_side
-  )
-  if getattr(tokenizer, "pad_token_id") is None:
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-
-  datasets = load_dataset(dataset_name, task)
-  metric = evaluate.load(dataset_name, task)
-
-  def tokenize_function(examples):
-    # max_length=None => use the model max length (it's actually the default)
-    outputs = tokenizer(
-        examples["sentence1"],
-        examples["sentence2"],
-        truncation=True,
-        max_length=None,
-    )
-    return outputs
-
-  tokenized_datasets = datasets.map(
-      tokenize_function,
-      batched=True,
-      remove_columns=["idx", "sentence1", "sentence2"],
-  )
-
-  # We also rename the 'label' column to 'labels' which is the expected name for
-  # labels by the models of the transformers library.
-  tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
-
-  def collate_fn(examples):
-    return tokenizer.pad(examples, padding="longest", return_tensors="pt")
-
-  # Instantiate dataloaders.
-  train_dataloader = DataLoader(
-      tokenized_datasets["train"],
-      shuffle=True,
-      collate_fn=collate_fn,
-      batch_size=batch_size,
-  )
-  eval_dataloader = DataLoader(
-      tokenized_datasets["validation"],
-      shuffle=False,
-      collate_fn=collate_fn,
-      batch_size=batch_size,
-  )
-
-  model = AutoModelForSequenceClassification.from_pretrained(
-      pretrained_model_id, return_dict=True
-  )
-  model = get_peft_model(model, peft_config)
-  model.print_trainable_parameters()
-
-  optimizer = AdamW(params=model.parameters(), lr=learning_rate)
-
-  # Instantiate scheduler
-  lr_scheduler = get_linear_schedule_with_warmup(
-      optimizer=optimizer,
-      num_warmup_steps=0.06 * (len(train_dataloader) * num_epochs),
-      num_training_steps=(len(train_dataloader) * num_epochs),
-  )
-
-  model.to(device)
-  for epoch in range(num_epochs):
-    model.train()
-    for _, batch in enumerate(tqdm(train_dataloader)):
-      batch.to(device)
-      outputs = model(**batch)
-      loss = outputs.loss
-      loss.backward()
-      optimizer.step()
-      lr_scheduler.step()
-      optimizer.zero_grad()
-
-    model.eval()
-    for _, batch in enumerate(tqdm(eval_dataloader)):
-      batch.to(device)
-      with torch.no_grad():
-        outputs = model(**batch)
-      predictions = outputs.logits.argmax(dim=-1)
-      references = batch["labels"]
-      metric.add_batch(
-          predictions=predictions,
-          references=references,
-      )
-
-    eval_metric = metric.compute()
-    print(f"epoch {epoch}:", eval_metric)
-
-  model.save_pretrained(output_dir)
-
-
-def main(unused_argv: Sequence[str]) -> None:
-  if dataset_validation_util.is_gcs_path(_PRETRAINED_MODEL_ID.value):
-    pretrained_model_id = dataset_validation_util.download_gcs_uri_to_local(
-        _PRETRAINED_MODEL_ID.value
-    )
-  else:
-    pretrained_model_id = _PRETRAINED_MODEL_ID.value
-  pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
-      pretrained_model_id
-  )
-  output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
-
-  finetune_sequence_classification(
-      pretrained_model_id=pretrained_model_path,
-      dataset_name=_DATASET_NAME.value,
-      output_dir=output_dir,
-      lora_rank=_LORA_RANK.value,
-      lora_alpha=_LORA_ALPHA.value,
-      lora_dropout=_LORA_DROPOUT.value,
-      num_epochs=int(_NUM_EPOCHS.value),
-      batch_size=_BATCH_SIZE.value,
-      learning_rate=_LEARNING_RATE.value,
-  )
-
-
-if __name__ == "__main__":
-  app.run(main)

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

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

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

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

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

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

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

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

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

@@ -3,16 +3,18 @@
 # DO NOT MODIFY: this file is auto-generated
 # See go/vmg-oss-peft-tests#command-builder-genpy
 
+
 class InstructLoraCommandBuilder:
 
   def __init__(self):
     self._config_file = None
     self._task = None
-    self._pretrained_model_id = None
-    self._dataset_name = None
-    self._train_split_name = None
-    self._template = None
-    self._instruct_column_in_dataset = None
+    self._gcs_rsync_interval_secs = None
+    self._pretrained_model_name_or_path = None
+    self._train_dataset = None
+    self._train_split = None
+    self._train_template = None
+    self._train_column = None
     self._output_dir = None
     self._merge_base_and_lora_output_dir = None
     self._logging_output_dir = None
@@ -22,14 +24,14 @@ class InstructLoraCommandBuilder:
     self._lora_alpha = None
     self._lora_dropout = None
     self._max_steps = None
-    self._num_epochs = None
+    self._num_train_epochs = None
     self._max_seq_length = None
     self._learning_rate = None
     self._lr_scheduler_type = None
     self._precision_mode = None
     self._train_precision = None
-    self._enable_gradient_checkpointing = None
-    self._use_example_packing = None
+    self._gradient_checkpointing = None
+    self._example_packing = None
     self._attn_implementation = None
     self._optimizer = None
     self._warmup_ratio = None
@@ -37,14 +39,14 @@ class InstructLoraCommandBuilder:
     self._save_steps = None
     self._logging_steps = None
     self._huggingface_access_token = None
-    self._eval_dataset_path = None
+    self._eval_dataset = None
     self._eval_column = None
     self._eval_template = None
     self._eval_split = None
     self._eval_steps = None
-    self._eval_tasks = None
     self._eval_metric_name = None
-    self._completion_only = None
+    self._metric_for_best_model = None
+    self._input_masking = None
     self._max_grad_norm = None
     self._logger_level = None
     self._benchmark_out_file = None
@@ -52,6 +54,7 @@ class InstructLoraCommandBuilder:
     self._enable_peft = None
     self._merge_model_precision_mode = None
     self._target_modules = None
+    self._unnamed_args = None
 
   @property
   def config_file(self):
@@ -70,44 +73,52 @@ class InstructLoraCommandBuilder:
     self._task = val
 
   @property
-  def pretrained_model_id(self):
-    return self._pretrained_model_id
+  def gcs_rsync_interval_secs(self):
+    return self._gcs_rsync_interval_secs
 
-  @pretrained_model_id.setter
-  def pretrained_model_id(self, val: str):
-    self._pretrained_model_id = val
+  @gcs_rsync_interval_secs.setter
+  def gcs_rsync_interval_secs(self, val: str):
+    self._gcs_rsync_interval_secs = val
+
+  @property
+  def pretrained_model_name_or_path(self):
+    return self._pretrained_model_name_or_path
+
+  @pretrained_model_name_or_path.setter
+  def pretrained_model_name_or_path(self, val: str):
+    self._pretrained_model_name_or_path = val
 
   @property
   def train_dataset(self):
-    return self._dataset_name
+    return self._train_dataset
 
   @train_dataset.setter
   def train_dataset(self, val: str):
-    self._dataset_name = val
+    self._train_dataset = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def template(self):
-    return self._template
+  def train_template(self):
+    return self._train_template
 
-  @template.setter
-  def template(self, val: str):
-    self._template = val
+  @train_template.setter
+  def train_template(self, val: str):
+    self._train_template = val
 
   @property
-  def instruct_column(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column.setter
-  def instruct_column(self, val: str):
-    self._instruct_column_in_dataset = val
+  @train_column.setter
+  def train_column(self, val: str):
+    self._train_column = val
 
   @property
   def ckpt_dir(self):
@@ -182,12 +193,12 @@ class InstructLoraCommandBuilder:
     self._max_steps = val
 
   @property
-  def num_epochs(self):
-    return self._num_epochs
+  def num_train_epochs(self):
+    return self._num_train_epochs
 
-  @num_epochs.setter
-  def num_epochs(self, val: float):
-    self._num_epochs = val
+  @num_train_epochs.setter
+  def num_train_epochs(self, val: float):
+    self._num_train_epochs = val
 
   @property
   def max_seq_length(self):
@@ -231,19 +242,19 @@ class InstructLoraCommandBuilder:
 
   @property
   def gradient_checkpointing(self):
-    return self._enable_gradient_checkpointing
+    return self._gradient_checkpointing
 
   @gradient_checkpointing.setter
   def gradient_checkpointing(self, val: bool):
-    self._enable_gradient_checkpointing = val
+    self._gradient_checkpointing = val
 
   @property
   def example_packing(self):
-    return self._use_example_packing
+    return self._example_packing
 
   @example_packing.setter
   def example_packing(self, val: bool):
-    self._use_example_packing = val
+    self._example_packing = val
 
   @property
   def attn_implementation(self):
@@ -303,18 +314,18 @@ class InstructLoraCommandBuilder:
 
   @property
   def eval_dataset(self):
-    return self._eval_dataset_path
+    return self._eval_dataset
 
   @eval_dataset.setter
   def eval_dataset(self, val: str):
-    self._eval_dataset_path = val
+    self._eval_dataset = val
 
   @property
-  def eval_instruct_column(self):
+  def eval_column(self):
     return self._eval_column
 
-  @eval_instruct_column.setter
-  def eval_instruct_column(self, val: str):
+  @eval_column.setter
+  def eval_column(self, val: str):
     self._eval_column = val
 
   @property
@@ -326,11 +337,11 @@ class InstructLoraCommandBuilder:
     self._eval_template = val
 
   @property
-  def eval_split_name(self):
+  def eval_split(self):
     return self._eval_split
 
-  @eval_split_name.setter
-  def eval_split_name(self, val: str):
+  @eval_split.setter
+  def eval_split(self, val: str):
     self._eval_split = val
 
   @property
@@ -341,14 +352,6 @@ class InstructLoraCommandBuilder:
   def eval_steps(self, val: int):
     self._eval_steps = val
 
-  @property
-  def eval_tasks(self):
-    return self._eval_tasks
-
-  @eval_tasks.setter
-  def eval_tasks(self, val: str):
-    self._eval_tasks = val
-
   @property
   def eval_metric_name(self):
     return self._eval_metric_name
@@ -358,12 +361,20 @@ class InstructLoraCommandBuilder:
     self._eval_metric_name = val
 
   @property
-  def completion_only(self):
-    return self._completion_only
+  def metric_for_best_model(self):
+    return self._metric_for_best_model
 
-  @completion_only.setter
-  def completion_only(self, val: bool):
-    self._completion_only = val
+  @metric_for_best_model.setter
+  def metric_for_best_model(self, val: str):
+    self._metric_for_best_model = val
+
+  @property
+  def input_masking(self):
+    return self._input_masking
+
+  @input_masking.setter
+  def input_masking(self, val: bool):
+    self._input_masking = val
 
   @property
   def max_grad_norm(self):
@@ -421,9 +432,22 @@ class InstructLoraCommandBuilder:
   def target_modules(self, val: str):
     self._target_modules = val
 
-  def build_cmd(self) -> str:
+  @property
+  def unnamed_args(self):
+    return self._unnamed_args
+
+  @unnamed_args.setter
+  def unnamed_args(self, val: list):
+    self._unnamed_args = val
+
+  def build_cmd(self) -> list[str]:
     cmd = []
+    args = ''
     for k, v in self.__dict__.items():
+      if k == '_unnamed_args' and v is not None:
+        args += ' '.join(v)
+        continue
       if v is not None:
         cmd.append(f'--{k[1:]}={v}')
+    cmd.append(f'{args}')
     return cmd

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

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

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

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

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

@@ -31,7 +31,7 @@ class AdapterTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -39,9 +39,9 @@ class AdapterTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -87,8 +87,8 @@ class AdapterTest(test_util.TestBase):
   def test_llama_adapters(self, model_name):
     test_function_name = inspect.stack()[0][3]
     self.setup_output_dir(f'{test_function_name}-{model_name}')
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)

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

@@ -4,19 +4,81 @@
 
 import os
 import time
-
+import unittest
 from absl.testing import absltest
 from absl.testing import parameterized
 import instruct_lora_command_builder as task_cmd_builder
 import test_util
 
 
+class EvalConfigTest(test_util.TestBase):
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = 10
+    self.task_cmd_builder.save_steps = 1000
+    self.task_cmd_builder.logging_steps = 1
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.example_packing = True
+    self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
+    self.task_cmd_builder.logging_dir = '/tmp/logs'
+    self.task_cmd_builder.eval_steps = 10
+    self.task_cmd_builder.eval_dataset = 'mlabonne/guanaco-llama2'
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = 'text'
+    self.task_cmd_builder.eval_template = 'openassistant-guanaco'
+
+  @parameterized.named_parameters(
+      ('all_eval_metric', 'loss,perplexity,bleu,google_bleu,rouge1', 0),
+      ('invalid_metric', 'invalid_metric', 1),
+      ('only_loss', 'loss', 0),
+      ('perplexity_without_loss', 'perplexity,bleu', 0),
+      ('unsupported_eval_metric', 'f1', 1),
+  )
+  def test_hf_eval_metrics(self, eval_metric_name, expected_return_code):
+    self.task_cmd_builder.eval_metric_name = eval_metric_name
+    self.assertEqual(self.run_cmd(), expected_return_code)
+
+  @parameterized.named_parameters(
+      ('valid_best_model_metric', 'loss,perplexity', 'perplexity', 0),
+      ('only_loss', None, 'loss', 0),
+      ('invalid_best_model_metric', 'loss', 'invalid_metric', 1),
+  )
+  def test_metric_for_best_model(
+      self, eval_metric_name, metric_for_best_model, expected_return_code
+  ):
+    self.task_cmd_builder.eval_metric_name = eval_metric_name
+    self.task_cmd_builder.metric_for_best_model = metric_for_best_model
+    self.assertEqual(self.run_cmd(), expected_return_code)
+
+
 class GcsUploadDownloadTest(test_util.TestBase):
 
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -24,9 +86,9 @@ class GcsUploadDownloadTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -45,9 +107,11 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_single_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  def test_model_download_single_process(self, pretrained_model_name_or_path):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(
+            pretrained_model_name_or_path
+        )
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
@@ -61,24 +125,26 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ),
       ('llama2_7b_hf', 'NousResearch/Llama-2-7b-hf'),
   )
-  def test_model_download_multi_process(self, pretrained_model_id):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(pretrained_model_id)
+  def test_model_download_multi_process(self, pretrained_model_name_or_path):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(
+            pretrained_model_name_or_path
+        )
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
     self.assertLess(end_time - start_time, 5 * 60.0)
 
-  def test_70b_model_download(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        'gs://vertex-model-garden-public-us/llama3/llama3-70b-hf'
+  def test_8b_model_download(self):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        'gs://vertex-model-garden-public-us/llama3/llama3-8b-hf'
     )
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
@@ -90,8 +156,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge(self, output_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
 
     ckpt_dir = os.path.join(
@@ -107,9 +173,13 @@ class GcsUploadDownloadTest(test_util.TestBase):
     end_time = time.time()
     self.assertLess(end_time - start_time, 5 * 60.0)
 
+  @unittest.skipIf(
+      not test_util.is_gpu_h100(),
+      'Skipping because this test is only for H100',
+  )
   def test_model_fp8_conversion(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
 
     ckpt_dir = f'/tmp/output/output-{test_util.get_timestamp()}'
@@ -128,17 +198,17 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('merged-and-upload-to-gcs', 'gs://vmg-test-ttl-1y/tests/merged'),
   )
   def test_model_merge_and_upload_deepspeed(self, merged_model_dir):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/deepspeed_zero3_8gpu.yaml'
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
     )
     self.task_cmd_builder.merged_model_dir = os.path.join(
         merged_model_dir, f'merged-{test_util.get_timestamp()}'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
@@ -149,8 +219,8 @@ class GcsUploadDownloadTest(test_util.TestBase):
       ('save-multiple-times', 1),
   )
   def test_llama3_8b_save_and_merge_8_gpus_fsdp(self, save_steps):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.save_steps = save_steps
     self.task_cmd_builder.max_steps = 3
@@ -159,7 +229,7 @@ class GcsUploadDownloadTest(test_util.TestBase):
     )
     self.task_cmd_builder.merged_model_dir = '/tmp/merged'
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     start_time = time.time()
     self.assertEqual(self.run_cmd(), 0)
     end_time = time.time()
@@ -171,7 +241,7 @@ class TemplateAndDataStatsTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -187,35 +257,42 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
-  def test_openai_chat_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+  @parameterized.named_parameters(
+      ('multi-chat-string-content', 'openai-multi-chat-example-data.jsonl'),
+      (
+          'multi-chat-array-content',
+          'openai-multi-chat-example-data-array-content.jsonl',
+      ),
+  )
+  def test_openai_chat_template(self, example_dataset):
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
-        'openai-multi-chat-example-data.jsonl'
+        example_dataset
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'openai-chat'
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_openai_completion_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-completion-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'prompt'
+    self.task_cmd_builder.train_template = 'openai-completion'
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_data_stats_chat_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
@@ -223,17 +300,17 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-multi-chat-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'messages'
-    self.task_cmd_builder.template = 'llama3'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'llama3'
     self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
 
     self.assertEqual(self.run_cmd(), 0)
 
   def test_data_stats_completion_template(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.config_file = (
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
@@ -241,12 +318,12 @@ class TemplateAndDataStatsTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'openai-completion-example-data.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'prompt'
-    self.task_cmd_builder.template = 'openai-completion'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'prompt'
+    self.task_cmd_builder.train_template = 'openai-completion'
     self.task_cmd_builder.tuning_data_stats_file = '/tmp/data-stats.json'
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
     self.assertEqual(self.run_cmd(), 0)
 
 
@@ -255,7 +332,7 @@ class TargetModulesTest(test_util.TestBase):
   def setUp(self):
     super().setUp()
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
     self.task_cmd_builder.task = 'instruct-lora'
@@ -263,9 +340,9 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
         'peft_train_sample.jsonl'
     )
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'input_text'
-    self.task_cmd_builder.template = 'llama3-text-bison'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
     self.task_cmd_builder.gradient_accumulation_steps = 1
     self.task_cmd_builder.lora_rank = 16
     self.task_cmd_builder.lora_alpha = 32
@@ -278,8 +355,8 @@ class TargetModulesTest(test_util.TestBase):
     self.task_cmd_builder.ckpt_dir = '/tmp'
 
   def test_target_modules(self):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id('llama3.1-8b-hf')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path('llama3.1-8b-hf')
     )
     self.task_cmd_builder.target_modules = 'q_proj, v_proj, k_proj'
 

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

@@ -43,9 +43,9 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.example_packing = True
     self.task_cmd_builder.train_dataset = 'mlabonne/guanaco-llama2-1k'
-    self.task_cmd_builder.train_split_name = 'train'
-    self.task_cmd_builder.instruct_column = 'text'
-    self.task_cmd_builder.template = 'openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
     self.task_cmd_builder.ckpt_dir = '/tmp/adapter'
     self.task_cmd_builder.logging_dir = '/tmp/logs'
 
@@ -54,23 +54,24 @@ class TrainerThroughputTest(test_util.TestBase):
     if ret != 0:
       with open(self.task_cmd_builder.benchmark_out_file, 'a') as f:
         max_seq_length = self.task_cmd_builder.max_seq_length
-        f.write(f'{max_seq_length/1024.0:.1f}k | failed | n/a\n')
+        f.write(f'{max_seq_length/1024.0:.1f} | failed | n/a\n')
     return ret
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
       ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
   )
   def test_model_single_gpu(self, model_name, precision, max_seq_length):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -78,29 +79,30 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir, f'bm_{model_name}_{precision}.txt'
     )
 
-    self.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
       model_name=[
-          'llama3-70b-hf',
+          'llama3.1-8b-hf',
           'llama3.1-70b-hf',
           'Mistral-7B-v0.1',
           'Mixtral-8x7B-v0.1',
-          'Gemma2-9b-it',
+          'gemma-2-9b-it',
+          'Qwen2.5-32B-Instruct',
       ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
-      config=['deepspeed_zero2', 'deepspeed_zero3'],
+      config=['deepspeed_zero2'],
   )
   def test_model_multi_gpu_deepspeed(
       self, model_name, precision, max_seq_length, num_gpus, config
   ):
     self.assertTrue(num_gpus == 4 or num_gpus == 8)
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -112,14 +114,18 @@ class TrainerThroughputTest(test_util.TestBase):
     self.task_cmd_builder.config_file = (
         f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
     )
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
-      model_name=['llama3.1-70b-hf'],
+      model_name=[
+          'llama3.1-8b-hf',
+          'llama3.1-70b-hf',
+          'Qwen2.5-32B-Instruct',
+      ],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
@@ -127,8 +133,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_lora(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -136,18 +142,108 @@ class TrainerThroughputTest(test_util.TestBase):
         self.test_suite_output_dir,
         f'bm_fsdp_{num_gpus}gpu_{model_name}_{precision}.txt',
     )
-    self.task_cmd_builder.config_file = (
-        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
-    )
+    if 'llama' in model_name.lower():
+      self.task_cmd_builder.config_file = (
+          'vertex_vision_model_garden_peft/llama2_fsdp_8gpu.yaml'
+      )
+    elif 'qwen' in model_name.lower():
+      self.task_cmd_builder.config_file = (
+          'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
+      )
+    else:
+      self.fail(f'Unsupported model: {model_name}')
 
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 
     self.assertEqual(self.run_cmd_and_handle_failure(), 0)
 
   @parameterized.product(
-      model_name=['llama3.1-70b-hf'],
+      model_name=['llama3.1-8b-hf', 'llama3.1-70b-hf'],
+      precision=['4bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      config=['deepspeed_zero2', 'fsdp'],
+  )
+  def test_peft_train_image_automated_test_llama(
+      self, model_name, precision, max_seq_length, config
+  ):
+    num_gpus = 8
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.benchmark_out_file = benchmark_out_file
+    if config == 'fsdp':
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/llama_{config}_{num_gpus}gpu.yaml'
+      )
+    else:
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
+      )
+
+    self.command_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+    self.run_cmd_and_handle_failure()
+    if test_util.is_gpu_h100():
+      self.assertEqual(
+          test_util.check_benchmark_results(
+              benchmark_out_file, 'llama', 10.0, max_seq_length
+          ),
+          True,
+      )
+
+  @parameterized.product(
+      model_name=['gemma-2-2b-it', 'gemma-2-9b-it', 'gemma-2-27b-it'],
+      precision=['4bit', 'bfloat16'],
+      max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
+      config=['deepspeed_zero2', 'deepspeed_zero3', 'fsdp'],
+  )
+  def test_peft_train_image_automated_test_gemma(
+      self, model_name, precision, max_seq_length, config
+  ):
+    num_gpus = 8
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.attn_implementation = 'eager'
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.load_precision = precision
+    benchmark_out_file = os.path.join(
+        self.test_suite_output_dir,
+        f'bm_{config}_{num_gpus}gpu_{model_name}_{precision}.txt',
+    )
+    self.task_cmd_builder.benchmark_out_file = benchmark_out_file
+    if config == 'fsdp':
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/gemma2_{config}_{num_gpus}gpu.yaml'
+      )
+    else:
+      self.task_cmd_builder.config_file = (
+          f'vertex_vision_model_garden_peft/{config}_{num_gpus}gpu.yaml'
+      )
+
+    self.command_builder.add_env_var(
+        'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
+    )
+    self.run_cmd_and_handle_failure()
+    if test_util.is_gpu_h100():
+      self.assertEqual(
+          test_util.check_benchmark_results(
+              benchmark_out_file, 'gemma', 10.0, max_seq_length
+          ),
+          True,
+      )
+
+  @parameterized.product(
+      model_name=['llama3.1-8b-hf', 'llama3.1-70b-hf'],
       precision=['bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
@@ -155,8 +251,8 @@ class TrainerThroughputTest(test_util.TestBase):
   def test_model_multi_gpu_fsdp_full_finetuning(
       self, model_name, precision, max_seq_length, num_gpus
   ):
-    self.task_cmd_builder.pretrained_model_id = (
-        test_util.get_pretrained_model_id(model_name)
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
     )
     self.task_cmd_builder.max_seq_length = max_seq_length
     self.task_cmd_builder.load_precision = precision
@@ -169,7 +265,7 @@ class TrainerThroughputTest(test_util.TestBase):
     )
     self.task_cmd_builder.enable_peft = False
 
-    self.docker_builder.add_env_var(
+    self.command_builder.add_env_var(
         'CUDA_VISIBLE_DEVICES', ','.join([str(x) for x in range(0, num_gpus)])
     )
 

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

@@ -0,0 +1,209 @@
+# pylint: disable=missing-function-docstring
+# pylint: disable=missing-class-docstring
+"""Tests to make sure trained model achieves decent quality.
+
+Right now, the metric is loss decreasing and we'll eyeball the TB graphs.
+"""
+
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import instruct_lora_command_builder as task_cmd_builder
+import test_util
+
+
+class TrainedModelQualityTest(test_util.TestBase):
+
+  _TEST_OUTPUT_DIR = os.path.expanduser('~/output')
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.test_suite_output_dir = os.path.join(
+        cls._TEST_OUTPUT_DIR,
+        os.path.splitext(os.path.basename(__file__))[0],
+        cls.__class__.__name__,
+    )
+
+  def setUp(self):
+    super().setUp()
+
+    self.task_cmd_builder = task_cmd_builder.InstructLoraCommandBuilder()
+    self.task_cmd_builder.task = 'instruct-lora'
+    self.task_cmd_builder.eval_metric_name = 'loss'
+    self.task_cmd_builder.per_device_batch_size = 1
+    self.task_cmd_builder.gradient_accumulation_steps = 8
+    self.task_cmd_builder.lora_rank = 16
+    self.task_cmd_builder.lora_alpha = 32
+    self.task_cmd_builder.lora_dropout = 0.05
+    self.task_cmd_builder.learning_rate = 5e-5
+    self.task_cmd_builder.num_train_epochs = 2.0
+    self.task_cmd_builder.warmup_ratio = 0.01
+    self.task_cmd_builder.max_steps = -1
+    self.task_cmd_builder.save_steps = 10
+    self.task_cmd_builder.eval_steps = 10
+    self.task_cmd_builder.max_seq_length = 4096
+    self.task_cmd_builder.load_precision = '4bit'
+    self.task_cmd_builder.gradient_checkpointing = True
+    self.task_cmd_builder.input_masking = True
+    self.task_cmd_builder.attn_implementation = 'flash_attention_2'
+    self.task_cmd_builder.report_to = 'tensorboard'
+
+  def setup_output_dir(self, testcase_name: str):
+    testcase_output_dir = os.path.join(
+        self.test_suite_output_dir, testcase_name
+    )
+    self.task_cmd_builder.ckpt_dir = os.path.join(
+        testcase_output_dir, 'adapter'
+    )
+    self.task_cmd_builder.logging_dir = os.path.join(
+        testcase_output_dir, 'logs'
+    )
+    self.task_cmd_builder.merged_model_dir = os.path.join(
+        testcase_output_dir, 'merged'
+    )
+
+  @parameterized.named_parameters(
+      ('llama3-8b', 'llama3-8b-hf'),
+      ('llama3.1-8b', 'llama3.1-8b-hf'),
+  )
+  def test_8b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'peft_train_sample.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'input_text'
+    self.task_cmd_builder.train_template = 'llama3-text-bison'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'peft_eval_sample.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('llama3-70b', 'llama3-70b-hf'),
+      ('llama3.1-70b', 'llama3.1-70b-hf'),
+  )
+  def test_70b_model_fsdp(self, model_name):
+    self.setup_output_dir(f'test_fsdp_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/llama_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = 'timdettmers/openassistant-guanaco'
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'text'
+    self.task_cmd_builder.train_template = 'openassistant-guanaco'
+    self.task_cmd_builder.eval_dataset = self.task_cmd_builder.train_dataset
+    self.task_cmd_builder.eval_split = 'test'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
+  )
+  def test_qwen_model_deepspeed(self, model_name):
+    self.setup_output_dir(f'test_deepspeed_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = model_name
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-train.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'qwen2_5'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-eval.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    # Note(lavrai): The following parameters are needed for the opposite-word
+    # dataset to converge properly.
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.num_train_epochs = 10.0
+    self.task_cmd_builder.logging_steps = 1
+
+    self.assertEqual(self.run_cmd(), 0)
+
+  @parameterized.named_parameters(
+      ('Qwen2.5-32B-Instruct', 'Qwen2.5-32B-Instruct'),
+  )
+  def test_qwen_model_fsdp(self, model_name):
+    self.setup_output_dir(f'test_fsdp_{model_name}')
+    self.task_cmd_builder.pretrained_model_name_or_path = (
+        test_util.get_pretrained_model_name_or_path(model_name)
+    )
+    self.task_cmd_builder.config_file = (
+        'vertex_vision_model_garden_peft/qwen2_fsdp_8gpu.yaml'
+    )
+    self.task_cmd_builder.train_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-train.jsonl'
+    )
+    self.task_cmd_builder.train_split = 'train'
+    self.task_cmd_builder.train_column = 'messages'
+    self.task_cmd_builder.train_template = 'qwen2_5'
+    self.task_cmd_builder.eval_dataset = test_util.get_test_data_path(
+        'llama-tuning-test/opposite-examples-eval.jsonl'
+    )
+    self.task_cmd_builder.eval_split = 'train'
+    self.task_cmd_builder.eval_column = self.task_cmd_builder.train_column
+    self.task_cmd_builder.eval_template = self.task_cmd_builder.train_template
+
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0,1,2,3,4,5,6,7')
+    # Note(lavrai): The following parameters are needed for the opposite-word
+    # dataset to converge properly.
+    self.task_cmd_builder.gradient_accumulation_steps = 1
+    self.task_cmd_builder.num_train_epochs = 10.0
+    self.task_cmd_builder.logging_steps = 1
+
+    self.assertEqual(self.run_cmd(), 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

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

@@ -31,7 +31,7 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       dict(
           testcase_name="out_of_range_rows",
           validate_top_k_rows=100000,
-          expected_result=1,
+          expected_result=0,
       ),
   )
   def test_validate_dataset_with_template_top_k_rows(
@@ -40,8 +40,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -79,8 +79,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = (
         "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
     )
@@ -92,6 +92,32 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
     result = self.run_cmd()
     self.assertEqual(result, expected_result)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="small_max_seq_length",
+          max_seq_length=10,
+      ),
+      dict(
+          testcase_name="large_max_seq_length",
+          max_seq_length=1024,
+      ),
+  )
+  def test_validate_dataset_with_template_max_seq_length(
+      self,
+      max_seq_length,
+  ):
+    self.task_cmd_builder.dataset_name = "timdettmers/openassistant-guanaco"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
+    self.task_cmd_builder.template = (
+        "gs://cloud-nas-260507-tmp-20240724/openassistant-guanaco.json"
+    )
+    self.task_cmd_builder.max_seq_length = max_seq_length
+    self.task_cmd_builder.validate_k_rows_of_dataset = None
+    self.task_cmd_builder.use_multiprocessing = True
+    result = self.run_cmd()
+    self.assertEqual(result, 0)
+
   @parameterized.named_parameters(
       dict(
           testcase_name="invalid_default_input_column",
@@ -204,8 +230,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = split
-    self.task_cmd_builder.instruct_column_in_dataset = input_column
+    self.task_cmd_builder.train_split = split
+    self.task_cmd_builder.train_column = input_column
     self.task_cmd_builder.template = template
     self.task_cmd_builder.validate_percentage_of_dataset = (
         validate_percentage_of_dataset
@@ -311,8 +337,8 @@ class ValidateDatasetWithTemplateTest(test_util.TestBase):
       expected_result,
   ):
     self.task_cmd_builder.dataset_name = dataset_name
-    self.task_cmd_builder.train_split_name = "train"
-    self.task_cmd_builder.instruct_column_in_dataset = "text"
+    self.task_cmd_builder.train_split = "train"
+    self.task_cmd_builder.train_column = "text"
     self.task_cmd_builder.template = template
     self.task_cmd_builder.validate_percentage_of_dataset = (
         validate_percentage_of_dataset

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

@@ -10,8 +10,9 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._task = None
     self._template = None
     self._dataset_name = None
-    self._train_split_name = None
-    self._instruct_column_in_dataset = None
+    self._train_split = None
+    self._train_column = None
+    self._max_seq_length = None
     self._use_multiprocessing = None
     self._validate_k_rows_of_dataset = None
     self._validate_percentage_of_dataset = None
@@ -41,20 +42,28 @@ class ValidateDatasetWithTemplateCommandBuilder:
     self._dataset_name = val
 
   @property
-  def train_split_name(self):
-    return self._train_split_name
+  def train_split(self):
+    return self._train_split
 
-  @train_split_name.setter
-  def train_split_name(self, val: str):
-    self._train_split_name = val
+  @train_split.setter
+  def train_split(self, val: str):
+    self._train_split = val
 
   @property
-  def instruct_column_in_dataset(self):
-    return self._instruct_column_in_dataset
+  def train_column(self):
+    return self._train_column
 
-  @instruct_column_in_dataset.setter
-  def instruct_column_in_dataset(self, val: str):
-    self._instruct_column_in_dataset = val
+  @train_column.setter
+  def train_column(self, val: str):
+    self._train_column = val
+
+  @property
+  def max_seq_length(self):
+    return self._max_seq_length
+
+  @max_seq_length.setter
+  def max_seq_length(self, val: int):
+    self._max_seq_length = val
 
   @property
   def use_multiprocessing(self):

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

@@ -8,16 +8,18 @@ environment. Otherwise, `python3` is used.
 """
 
 import argparse
-import json
-import os
+from collections.abc import MutableSequence, Sequence
+import multiprocessing
 import subprocess
-from typing import List, Optional, Sequence
+import sys
 from absl import app
 from absl import flags
 from absl import logging
 from util import dataset_validation_util
+from util import cluster_spec
 from vertex_vision_model_garden_peft.train.vmg import utils
 from util import constants
+from util import gcs_syncer
 from util import hypertune_utils
 
 
@@ -37,20 +39,21 @@ _TASK_TO_SCRIPT = {
     constants.INSTRUCT_LORA: (
         'vertex_vision_model_garden_peft/train/vmg/instruct_lora.py'
     ),
-    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: 'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py',
-    constants.QUANTIZE_MODEL: (
-        'vertex_vision_model_garden_peft/train/vmg/quantize_model.py'
+    constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA: (
+        'vertex_vision_model_garden_peft/train/vmg/merge_causal_language_model_lora.py'
     ),
-    constants.SEQUENCE_CLASSIFICATION_LORA: 'vertex_vision_model_garden_peft/train/vmg/sequence_classification_lora.py',
-    constants.VALIDATE_DATASET_WITH_TEMPLATE: 'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py',
+    constants.VALIDATE_DATASET_WITH_TEMPLATE: (
+        'vertex_vision_model_garden_peft/train/vmg/validate_dataset_with_template.py'
+    ),
+    constants.RUN_TESTS: 'vertex_vision_model_garden_peft/tests/run_tests.py',
 }
 
 
 def launch_script_cmd(
     script: str,
-    config_file: Optional[str],
+    config_file: str | None,
     accelerate_args: argparse.Namespace = argparse.Namespace(),
-) -> List[str]:
+) -> MutableSequence[str]:
   """Returns the command to launch the script."""
   if config_file:
     cmd = [
@@ -69,59 +72,23 @@ def launch_script_cmd(
 
 def _get_accelerate_args() -> argparse.Namespace:
   """Returns the accelerate args."""
-  # For the format of the cluster spec, see
-  # https://cloud.google.com/vertex-ai/docs/training/distributed-training#cluster-spec-format # pylint: disable=line-too-long
-  cluster_spec = os.getenv('CLUSTER_SPEC', default=None)
-  if not cluster_spec:
-    return argparse.Namespace()
-  logging.info('CLUSTER_SPEC: %s', cluster_spec)
-
-  cluster_data = json.loads(cluster_spec)
-  if (
-      'workerpool1' not in cluster_data['cluster']
-      or not cluster_data['cluster']['workerpool1']
-  ):
-    return argparse.Namespace()
-
-  # Get primary node info
-  primary_node = cluster_data['cluster']['workerpool0'][0]
-  logging.info('primary node: %s', primary_node)
-  primary_node_addr, primary_node_port = primary_node.split(':')
-  logging.info('primary node address: %s', primary_node_addr)
-  logging.info('primary node port: %s', primary_node_port)
-
-  # Determine node rank of this machine
-  workerpool = cluster_data['task']['type']
-  if workerpool == 'workerpool0':
-    node_rank = 0
-  elif workerpool == 'workerpool1':
-    # Add 1 for the primary node, since `index` is the index of workerpool1.
-    node_rank = cluster_data['task']['index'] + 1
-  else:
-    raise ValueError(
-        'Only workerpool0 and workerpool1 are supported. Unknown workerpool:'
-        f' {workerpool}'
-    )
-  logging.info('node rank: %s', node_rank)
-
-  # Calculate total nodes
-  num_worker_nodes = len(cluster_data['cluster']['workerpool1'])
-  num_nodes = num_worker_nodes + 1  # Add 1 for the primary node
-  logging.info('num nodes: %s', num_nodes)
-
+  primary_node_addr, primary_node_port, node_rank, num_nodes = (
+      cluster_spec.get_cluster_spec()
+  )
   accelerate_args = argparse.Namespace()
-  accelerate_args.machine_rank = node_rank
-  accelerate_args.num_machines = num_nodes
-  accelerate_args.main_process_ip = primary_node_addr
-  accelerate_args.main_process_port = primary_node_port
-  accelerate_args.max_restarts = 0
-  accelerate_args.monitor_interval = 120
+  if num_nodes > 1:
+    accelerate_args.machine_rank = node_rank
+    accelerate_args.num_machines = num_nodes
+    accelerate_args.main_process_ip = primary_node_addr
+    accelerate_args.main_process_port = primary_node_port
+    accelerate_args.max_restarts = 0
+    accelerate_args.monitor_interval = 120
 
   return accelerate_args
 
 
 def _append_args_to_command_in_place(
-    args: argparse.Namespace, command: List[str]
+    args: argparse.Namespace, command: MutableSequence[str]
 ):
   for key, value in vars(args).items():
     # If not specified, skip.
@@ -129,15 +96,25 @@ def _append_args_to_command_in_place(
       command.append(f'--{key}={value}')
 
 
-def _get_train_cmd_and_maybe_merge_cmd(
-    task: str, config_file: str, unknown: Sequence[str]
+def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
+    task_type: str, config_file: str, unknown: Sequence[str]
 ) -> Sequence[Sequence[str]]:
-  """Returns the training command and maybe the merge command if applicable."""
+  """Returns the training and merge command(if applicable) and dirs to sync.
 
+  Args:
+    task_type: The task type.
+    config_file: The accelerate config file path.
+    unknown: The unknown args which are not recognised by the parser.
+
+  Returns:
+    The bash commands to execute and the directories to sync.
+  """
+  dirs_to_sync = []
   # Only populated when multi-node is used.
   accelerate_args = _get_accelerate_args()
+  node_rank = getattr(accelerate_args, 'machine_rank', 0)
   training_cmd = launch_script_cmd(
-      _TASK_TO_SCRIPT[task],
+      _TASK_TO_SCRIPT[task_type],
       config_file,
       accelerate_args=accelerate_args,
   )
@@ -152,17 +129,30 @@ def _get_train_cmd_and_maybe_merge_cmd(
       dataset_validation_util.force_gcs_fuse_path(training_args.output_dir)
   )
 
+  local_output_dir, gcs_output_dir = gcs_syncer.manage_sync_path(
+      training_args.output_dir, node_rank
+  )
+  training_args.output_dir = local_output_dir
+  if gcs_syncer.is_gcs_or_gcsfuse_path(gcs_output_dir):
+    dirs_to_sync.append((local_output_dir, gcs_output_dir))
+
   # Merge only flags.
   merge_parser = argparse.ArgumentParser()
   merge_parser.add_argument('--merge_model_precision_mode')
-  merge_parser.add_argument('--executor_input')
-  merge_parser.add_argument('--restrict_model_upload_docker_uri')
   merge_parser.add_argument('--merge_base_and_lora_output_dir')
   merge_args, unknown = merge_parser.parse_known_args(unknown)
 
+  if merge_args.merge_base_and_lora_output_dir:
+    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
+        merge_args.merge_base_and_lora_output_dir, None
+    )
+    merge_args.merge_base_and_lora_output_dir = merge_local_dir
+    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
+      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
+
   # Common flags shared by merging and training.
   common_parser = argparse.ArgumentParser()
-  common_parser.add_argument('--pretrained_model_id', required=True)
+  common_parser.add_argument('--pretrained_model_name_or_path', required=True)
   common_parser.add_argument('--huggingface_access_token')
   common_args, remaining = common_parser.parse_known_args(unknown)
 
@@ -173,17 +163,19 @@ def _get_train_cmd_and_maybe_merge_cmd(
   commands = [training_cmd]
 
   # Only the main node runs merging.
-  if (
-      merge_args.merge_base_and_lora_output_dir
-      and getattr(accelerate_args, 'machine_rank', 0) == 0
-  ):
+  if merge_args.merge_base_and_lora_output_dir and node_rank == 0:
     lora_dir = utils.get_final_checkpoint_path(training_args.output_dir)
+    lora_local_dir, lora_gcs_dir = gcs_syncer.manage_sync_path(
+        lora_dir, node_rank
+    )
+    if gcs_syncer.is_gcs_or_gcsfuse_path(lora_gcs_dir):
+      dirs_to_sync.append((lora_local_dir, lora_gcs_dir))
 
     merge_cmd = [
         'WORLD_SIZE=1',  # To ignore other nodes in multi-node setting.
         'python3',
         _TASK_TO_SCRIPT[constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA],
-        f'--finetuned_lora_model_dir={lora_dir}',
+        f'--finetuned_lora_model_dir={lora_local_dir}',
     ]
     _append_args_to_command_in_place(merge_args, merge_cmd)
     _append_args_to_command_in_place(common_args, merge_cmd)
@@ -196,16 +188,51 @@ def _get_train_cmd_and_maybe_merge_cmd(
     ]
     commands.append(conda_run_cmd)
 
-  return commands
+  return commands, dirs_to_sync
+
+
+def _get_merge_cmd_and_dirs_to_sync(
+    task_type: str, config_file: str, unknown: Sequence[str]
+) -> Sequence[Sequence[str]]:
+  """Returns the merge command and dirs to sync.
+
+  Args:
+    task_type: The task type.
+    config_file: The accelerate config file path.
+    unknown: The unknown args which are not recognised by the parser.
+
+  Returns:
+    The bash commands to execute and the directories to sync.
+  """
+  # Merge only flags.
+  merge_parser = argparse.ArgumentParser()
+  merge_parser.add_argument('--merge_base_and_lora_output_dir')
+  merge_args, unknown = merge_parser.parse_known_args(unknown)
+
+  dirs_to_sync = []
+  if merge_args.merge_base_and_lora_output_dir:
+    merge_local_dir, merge_gcs_dir = gcs_syncer.manage_sync_path(
+        merge_args.merge_base_and_lora_output_dir, None
+    )
+    merge_args.merge_base_and_lora_output_dir = merge_local_dir
+    if gcs_syncer.is_gcs_or_gcsfuse_path(merge_gcs_dir):
+      dirs_to_sync.append((merge_local_dir, merge_gcs_dir))
+
+  cmd = launch_script_cmd(_TASK_TO_SCRIPT[task_type], config_file)
+  _append_args_to_command_in_place(merge_args, cmd)
+  cmd.extend(unknown)
+  return [cmd], dirs_to_sync
 
 
 def main(unused_argv: Sequence[str]) -> None:
   parser = argparse.ArgumentParser()
   parser.add_argument('--config_file')
   parser.add_argument('--task')
+  parser.add_argument('--gcs_rsync_interval_secs', type=int, default=60)
   args, unknown = parser.parse_known_args()
 
   task = args.task
+  dirs_to_sync = None
 
   if task in _TEXT_TO_IMAGE_TASKS_SCRIPTS:
     # Setup accelerate config before running trainer.
@@ -224,8 +251,12 @@ def main(unused_argv: Sequence[str]) -> None:
     ] + list(map(dataset_validation_util.force_gcs_fuse_path, unknown))
     commands = [config_gen_cmd, task_cmd]
   elif task in [constants.INSTRUCT_LORA]:
-    commands = _get_train_cmd_and_maybe_merge_cmd(
-        task=task, config_file=args.config_file, unknown=unknown
+    commands, dirs_to_sync = _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
+        task_type=task, config_file=args.config_file, unknown=unknown
+    )
+  elif task in [constants.MERGE_CAUSAL_LANGUAGE_MODEL_LORA]:
+    commands, dirs_to_sync = _get_merge_cmd_and_dirs_to_sync(
+        task_type=task, config_file=args.config_file, unknown=unknown
     )
   else:
     assert task in _TASK_TO_SCRIPT
@@ -233,10 +264,29 @@ def main(unused_argv: Sequence[str]) -> None:
     cmd.extend(unknown)
     commands = [cmd]
 
+  rsync_process = None
+  mp_queue = multiprocessing.Queue(maxsize=1)
+  if dirs_to_sync:
+    rsync_process = gcs_syncer.setup_gcs_rsync(
+        dirs_to_sync, mp_queue, args.gcs_rsync_interval_secs
+    )
+
   for cmd in commands:
     logging.info('launching task=%s with cmd: \n%s', task, ' \\\n'.join(cmd))
-    subprocess.run(cmd, check=True)
+    # Both absl logging and python's logging module writes to stderr by default.
+    # Redirect output to stdout on purpose, such that log entries do not get
+    # marked as `Error` in Cloud's Log Explorer.
+    try:
+      subprocess.run(cmd, stdout=sys.stdout, stderr=sys.stdout, check=True)
+    except subprocess.CalledProcessError as e:
+      if rsync_process is not None and rsync_process.is_alive():
+        logging.info('Terminating GCS rsync process.')
+        rsync_process.terminate()
+      raise e
+  if rsync_process is not None:
+    gcs_syncer.cleanup_gcs_rsync(rsync_process, mp_queue)
 
 
 if __name__ == '__main__':
+  logging.get_absl_handler().python_handler.stream = sys.stdout
   app.run(main, flags_parser=lambda _args: flags.FLAGS(_args, known_only=True))

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

@@ -1,232 +1,44 @@
 """Common libraries for PEFT."""
 
-import dataclasses
+from collections.abc import Mapping, Sequence
 import datetime
 import gc
-import multiprocessing as mp
 import os
-import subprocess
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 
 from absl import logging
 import accelerate
 from accelerate import DistributedType
 from accelerate import PartialState
-from google.protobuf import json_format
-from kfp.pipeline_spec import pipeline_spec_pb2
-import numpy as np
 import peft
 from peft import PeftModel
 from peft import prepare_model_for_kbit_training
-import pynvml
 import torch
 import transformers
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 from transformers import BitsAndBytesConfig
 from transformers import FbgemmFp8Config
-from transformers.integrations import is_deepspeed_zero3_enabled
 import trl
 
 from util import dataset_validation_util
 from util import constants
-from util import fileutils
 
-
-_MODELS_REQUIRING_PAD_TOKEN = ("llama", "falcon", "mistral", "mixtral")
-_MODELS_REQUIRING_EOS_TOEKN = ("gemma-2b", "gemma-7b")
 _LLAMA_3_1_405B_MODEL_ID = "Meta-Llama-3.1-405B"
 _LOCAL_MERGED_MODEL_DIR = "/tmp/merged_model"
-
-
-
-class GcsOrLocalDirectory(os.PathLike):
-  """A class to represent a directory with upload support if GCS path is given.
-
-  This class is used to represent a directory. It can be used for a temporary
-  local directory and for uploading files to the GCS directory later if the
-  given path is a GCS directory. If the given path is a local directory, a call
-  to gcs_dir attribute will raise an error. This class has multi-node and
-  multi-process support with accelerate.
-
-  Attributes:
-    local_dir: The local directory to store the files.
-    gcs_dir: The path to the GCS directory.
-  """
-
-  def __init__(
-      self,
-      path: str,
-      check_empty: bool = False,
-      upload_from_all_nodes: bool = False,
-  ):
-    """Initializes the GcsOrLocalDirectory.
-
-    Args:
-      path: The path to the directory.
-      check_empty: If True, check if the GCS directory is empty. No-op for local
-        directory.
-      upload_from_all_nodes: If True, upload the local directory to GCS from all
-        nodes.
-    """
-    if len(path) > 1:
-      path = path.rstrip("/")
-
-    self._upload_from_all_nodes = upload_from_all_nodes
-
-    if path.startswith(constants.GCS_URI_PREFIX) or path.startswith(
-        constants.GCSFUSE_URI_PREFIX
-    ):
-      self._is_gcs_path = True
-      self._local_dir = _get_local_dir_from_gcs_dir(path)
-      self._gcs_dir = fileutils.force_gcs_path(path)
-      os.makedirs(self.local_dir, exist_ok=True)
-
-      with PartialState().main_process_first():
-        if (
-            check_empty
-            and PartialState().is_main_process
-            and not _is_gcs_dir_empty(self._gcs_dir)
-        ):
-          raise ValueError(f"{self._gcs_dir} needs to be empty.")
-    else:
-      self._is_gcs_path = False
-      self._local_dir = path
-      self._gcs_dir = path
-
-  def __fspath__(self) -> str:
-    return self.local_dir
-
-  @property
-  def local_dir(self) -> str:
-    return self._local_dir
-
-  @property
-  def gcs_dir(self) -> str:
-    """Returns the GCS directory path.
-
-    Returns:
-      The GCS directory path.
-
-    Raises:
-      ValueError: If the path is not a GCS path.
-    """
-    if not self._is_gcs_path:
-      raise ValueError(f"{self._gcs_dir} is not a GCS path.")
-    return self._gcs_dir
-
-  def upload_to_gcs(
-      self,
-      skip_if_exists: bool = True,
-      force_upload: bool = False,
-  ):
-    """Uploads the local directory to GCS."""
-    if not self._is_gcs_path:
-      logging.info(
-          "Not uploading to GCS since %s is not a GCS path.", self.local_dir
-      )
-      return
-
-    if not os.listdir(self.local_dir):
-      logging.info("Not uploading to GCS since %s is empty.", self.local_dir)
-      return
-
-    target = os.path.dirname(self.gcs_dir) + "/"
-    # Avoid race condition uploading the same file from multiple processes.
-    with PartialState().main_process_first():
-      if not PartialState().is_local_main_process:
-        # Non local main processes don't upload.
-        pass
-      elif self._upload_from_all_nodes or PartialState().is_main_process:
-        logging.info("Uploading %s to %s...", self.local_dir, target)
-        cmd = [
-            "gsutil",
-            "-m",
-            "cp",
-            "-r",
-        ]
-        if skip_if_exists:
-          cmd.append("-n")
-        if force_upload:
-          cmd.append("-f")
-        cmd.extend([self.local_dir, target])
-        subprocess.check_output(cmd)
-        logging.info("%s uploaded.", self.local_dir)
-
-
-def _get_local_dir_from_gcs_dir(path: str) -> str:
-  return os.path.join(
-      constants.LOCAL_OUTPUT_DIR,
-      dataset_validation_util.force_gcs_fuse_path(path)[1:],
-  )
-
-
-def _is_gcs_dir_empty(path: str) -> bool:
-  """Checks if a GCS directory is empty.
-
-  Args:
-    path: The GCS directory path.
-
-  Returns:
-    True if the directory is empty.
-
-  Raises:
-    subprocess.CalledProcessError: If the gsutil command failure reason is not
-      because the dir is empty.
-  """
-  path = path.rstrip("/") + "/"
-  try:
-    subprocess.check_output(["gsutil", "ls", path], stderr=subprocess.STDOUT)
-  except subprocess.CalledProcessError as e:
-    if (
-        str(e.output, encoding="utf-8")
-        == "CommandException: One or more URLs matched no objects.\n"
-    ):
-      return True
-    else:
-      logging.info(str(e.output, encoding="utf-8"))
-      raise
-  else:
-    return False
-
-
-def load_tokenizer(
-    pretrained_model_id: str,
-    padding_side: Optional[str] = None,
-    access_token: Optional[str] = None,
-) -> AutoTokenizer:
-  """Loads tokenizer based on `pretrained_model_id`."""
-  tokenizer_kwargs = {}
-  if should_add_eos_token(pretrained_model_id):
-    tokenizer_kwargs["add_eos_token"] = True
-  if padding_side:
-    tokenizer_kwargs["padding_side"] = padding_side
-
-  with PartialState().local_main_process_first():
-    tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_id,
-        trust_remote_code=False,
-        use_fast=True,
-        token=access_token,
-        **tokenizer_kwargs,
-    )
-
-  if should_add_pad_token(pretrained_model_id):
-    tokenizer.add_special_tokens({"pad_token": "[PAD]"})
-
-  return tokenizer
+_GEMMA2_MODEL = "gemma-2"
 
 
 def load_model(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     tokenizer: AutoTokenizer,
     precision_mode: str = None,
-    enable_gradient_checkpointing: bool = False,
-    gradient_checkpointing_kwargs: Optional[Dict[str, Any]] = None,
-    access_token: Optional[str] = None,
-    attn_implementation: Optional[str] = None,
-    train_precision: Optional[str] = None,
-    device_map: Optional[str] = None,
+    gradient_checkpointing: bool = False,
+    gradient_checkpointing_kwargs: Mapping[str, Any] | None = None,
+    access_token: str | None = None,
+    attn_implementation: str | None = None,
+    train_precision: str | None = None,
+    device_map: str | None = None,
     is_training: bool = True,
 ) -> AutoModelForCausalLM:
   """Loads models from the local dir if specified or from huggingface."""
@@ -304,25 +116,37 @@ def load_model(
     raise ValueError(f"Invalid precision mode: {precision_mode}")
   logging.info("using torch_type=%s", torch_dtype)
 
+  model_kwargs = {
+      "use_cache": not gradient_checkpointing,
+      "device_map": device_map,
+      "torch_dtype": torch_dtype,
+      "quantization_config": quantization_config,
+      "trust_remote_code": False,
+      "token": access_token,
+      "attn_implementation": attn_implementation,
+  }
+  if _GEMMA2_MODEL in pretrained_model_name_or_path:
+    # The cache_implementation for Gemma 2 is set to hybrid by default. This
+    # param is only supported by Gemma 2. The default 'hybrid' value causes an
+    # issue when use_cache is set to False. So we have to use 'None' in such
+    # cases.
+    # https://github.com/huggingface/transformers/commit/238b13478df209ab534f2195a397dc64a3930883
+    model_kwargs["cache_implementation"] = (
+        None if gradient_checkpointing else "hybrid"
+    )
+
   model = AutoModelForCausalLM.from_pretrained(
-      pretrained_model_id,
-      use_cache=not enable_gradient_checkpointing,
-      device_map=device_map,
-      torch_dtype=torch_dtype,
-      quantization_config=quantization_config,
-      trust_remote_code=True,
-      token=access_token,
-      attn_implementation=attn_implementation,
+      pretrained_model_name_or_path, **model_kwargs
   )
 
   if precision_mode in (constants.PRECISION_MODE_4, constants.PRECISION_MODE_8):
     model = prepare_model_for_kbit_training(
         model,
-        use_gradient_checkpointing=enable_gradient_checkpointing,
+        use_gradient_checkpointing=gradient_checkpointing,
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs,
     )
 
-  if enable_gradient_checkpointing:
+  if gradient_checkpointing:
     model.gradient_checkpointing_enable(
         gradient_checkpointing_kwargs=gradient_checkpointing_kwargs
     )
@@ -345,8 +169,10 @@ def load_model(
     # https://stackoverflow.com/a/77408076
     model.config.use_cache = False
 
-  if should_add_pad_token(pretrained_model_id):
-    model.resize_token_embeddings(len(tokenizer))
+  if dataset_validation_util.should_add_pad_token(
+      pretrained_model_name_or_path
+  ):
+    model.resize_token_embeddings(len(tokenizer), mean_resizing=False)
     if is_training:
       # The following is needed since we added a new token that needs to be
       # learned.
@@ -357,22 +183,24 @@ def load_model(
 
 
 def _merge_causal_language_model_with_lora_internal(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     merge_precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Internal function to merges the base model with the lora adapter."""
   logging.info("loading tokenizer...")
-  tokenizer = load_tokenizer(pretrained_model_id)
+  tokenizer = dataset_validation_util.load_tokenizer(
+      pretrained_model_name_or_path
+  )
 
   # Note: merging peft adapter requires loading model in 16 bits, so merging
   # is done on CPU on purpose in case one GPU cannot hold the base model.
-  logging.info("loading model %s...", pretrained_model_id)
+  logging.info("loading model %s...", pretrained_model_name_or_path)
   device_map = "cpu"
   model = load_model(
-      pretrained_model_id=pretrained_model_id,
+      pretrained_model_name_or_path=pretrained_model_name_or_path,
       tokenizer=tokenizer,
       precision_mode=merge_precision_mode,
       access_token=access_token,
@@ -402,42 +230,12 @@ def _merge_causal_language_model_with_lora_internal(
   )
 
 
-def merge_causal_language_model_with_lora_fsdp(
-    pretrained_model_id: str,
-    merge_precision_mode: str,
-    finetuned_lora_model_dir: str,
-    merged_model_output_dir: str,
-    access_token: Optional[str] = None,
-) -> None:
-  """Merges the base model with the lora adapter for FSDP.
-
-  Only the main process should call this function.
-
-  Args:
-    pretrained_model_id: Predefined base model name or path to directory
-      containing model checkpoints.
-    merge_precision_mode: Precision mode for saving model weights.
-    finetuned_lora_model_dir: Path to directory containing PEFT-finetuned model
-      weights.
-    merged_model_output_dir: Path to directory to save the merged model.
-    access_token: Access token for accessing the model.
-  """
-  assert PartialState().is_main_process
-  _merge_causal_language_model_with_lora_internal(
-      pretrained_model_id=pretrained_model_id,
-      merge_precision_mode=merge_precision_mode,
-      finetuned_lora_model_dir=finetuned_lora_model_dir,
-      merged_model_output_dir=merged_model_output_dir,
-      access_token=access_token,
-  )
-
-
 def merge_causal_language_model_with_lora(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     precision_mode: str,
     finetuned_lora_model_dir: str,
     merged_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Merges the base model with the lora adapter."""
 
@@ -452,39 +250,13 @@ def merge_causal_language_model_with_lora(
 
   if PartialState().is_main_process:
     logging.info("Starting merging job...")
-    # When deepspeed Zero3 is enabled, users are not allowed to specify
-    # `device_map` when loading the model (even on CPU).
-    #
-    # To work-around this, we kick off another process (from the
-    # is_main_process) and set up the environment to avoid using Deepspeed when
-    # doing the merging.
-    if is_deepspeed_zero3_enabled():
-      ctx = mp.get_context("spawn")
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "false"
-      merge_job = ctx.Process(
-          target=_merge_causal_language_model_with_lora_internal,
-          args=(
-              pretrained_model_id,
-              merge_precision_mode,
-              finetuned_lora_model_dir,
-              local_merged_model_dir,
-          ),
-          kwargs={
-              "access_token": access_token,
-          },
-      )
-      merge_job.start()
-      merge_job.join()
-      os.environ["ACCELERATE_USE_DEEPSPEED"] = "true"
-    else:
-      _merge_causal_language_model_with_lora_internal(
-          pretrained_model_id=pretrained_model_id,
-          merge_precision_mode=merge_precision_mode,
-          finetuned_lora_model_dir=finetuned_lora_model_dir,
-          merged_model_output_dir=local_merged_model_dir,
-          access_token=access_token,
-      )
-
+    _merge_causal_language_model_with_lora_internal(
+        pretrained_model_name_or_path=pretrained_model_name_or_path,
+        merge_precision_mode=merge_precision_mode,
+        finetuned_lora_model_dir=finetuned_lora_model_dir,
+        merged_model_output_dir=local_merged_model_dir,
+        access_token=access_token,
+    )
     logging.info("merging job is done")
 
   # Wait for all processes to sync here.
@@ -492,7 +264,7 @@ def merge_causal_language_model_with_lora(
 
   if precision_mode == constants.PRECISION_MODE_FP8:
     convert_model_to_fp8(
-        pretrained_model_name_or_path=pretrained_model_id,
+        pretrained_model_name_or_path=pretrained_model_name_or_path,
         merged_model_output_dir=local_merged_model_dir,
         quantized_model_output_dir=merged_model_output_dir,
         access_token=access_token,
@@ -503,7 +275,7 @@ def convert_model_to_fp8(
     pretrained_model_name_or_path: str,
     merged_model_output_dir: str,
     quantized_model_output_dir: str,
-    access_token: Optional[str] = None,
+    access_token: str | None = None,
 ) -> None:
   """Converts the model to fp8.
 
@@ -533,166 +305,12 @@ def convert_model_to_fp8(
   PartialState().wait_for_everyone()
 
 
-@dataclasses.dataclass
-class TuningDataStats:
-  tuning_dataset_example_count: int
-  total_billable_token_count: int
-  tuning_step_count: int
-
-
-def get_dataset_stats(
-    dataset: Any,
-    tokenizer: transformers.PreTrainedTokenizer,
-    column: str,
-    effective_batch_size: int,
-) -> TuningDataStats:
-  """Calculates dataset statistics, e.g., total number of tokens."""
-  tokenized_dataset = dataset.map(lambda x: tokenizer(x[column]))
-  inputs = tokenized_dataset["input_ids"]
-  tuning_dataset_example_count = int(len(inputs))
-  total_billable_token_count = int(np.sum([len(ex) for ex in inputs]))
-  tuning_step_count = (
-      tuning_dataset_example_count + effective_batch_size - 1
-  ) // effective_batch_size
-  return TuningDataStats(
-      tuning_dataset_example_count,
-      total_billable_token_count,
-      tuning_step_count,
-  )
-
-
 def force_gc():
   """Collects garbage immediately to release unused CPU/GPU resources."""
   gc.collect()
   torch.cuda.empty_cache()
 
 
-def should_add_pad_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special pad token."""
-  return any(s.lower() in model_id.lower() for s in _MODELS_REQUIRING_PAD_TOKEN)
-
-
-def should_add_eos_token(model_id: str) -> bool:
-  """Returns whether the model requires adding a special eos token."""
-  return any(m in model_id for m in _MODELS_REQUIRING_EOS_TOEKN)
-
-
-def write_kfp_outputs(
-    executor_input: str, output_artifacts: Dict[str, str]
-) -> None:
-  """Writes KFP outputs given a dict of output artifact names and URIs."""
-  # Only the main process writes to avoid race condition.
-  if PartialState().is_main_process:
-    executor_input = json_format.Parse(
-        executor_input, pipeline_spec_pb2.ExecutorInput()
-    )
-    outputs = executor_input.outputs
-    # set all artifacts
-    for name, uri in output_artifacts.items():
-      artifact_list = outputs.artifacts.get(name)
-      if not artifact_list or not artifact_list.artifacts:
-        raise ValueError(f"Artifact name={name} does not exist.")
-      artifact_list.artifacts[0].uri = uri
-
-    # write output file
-    executor_output = pipeline_spec_pb2.ExecutorOutput(
-        artifacts=outputs.artifacts
-    )
-    os.makedirs(os.path.dirname(outputs.output_file), exist_ok=True)
-    with open(outputs.output_file, "w") as f:
-      f.write(json_format.MessageToJson(executor_output, indent=None))
-
-  # Wait for the main process to finish before moving on to the next task.
-  PartialState().wait_for_everyone()
-
-
-def upload_local_dir_to_gcs(local_dir: str, gcs_path: str):
-  """Uploads local dir to GCS."""
-
-  if PartialState().is_main_process:
-    logging.info("uploading %s to %s...", local_dir, gcs_path)
-    subprocess.check_output([
-        "gsutil",
-        "-m",
-        "cp",
-        "-r",
-        local_dir,
-        gcs_path,
-    ])
-    logging.info("%s uploaded.", local_dir)
-
-  PartialState().wait_for_everyone()
-
-
-def write_first_party_model_metadata(output_dir: str, docker_uri: str) -> None:
-  """Multi-process friendly version of fileutils.write_first_party_model_metadata."""
-  if PartialState().is_main_process:
-    fileutils.write_first_party_model_metadata(output_dir, docker_uri)
-  PartialState().wait_for_everyone()
-
-
-@dataclasses.dataclass
-class GpuStats:
-  """Holds information about GPU usage stats.
-
-  For memory related, see
-  https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-management
-  """
-
-  # total memory
-  total_mem: float
-  # memory occupied.
-  occupied: float
-  # memory reserved, but not used.
-  unused: float
-  # nvidia-smi usually reports more memory usages than pytorch (for driver,
-  # kernel and etc). `smi_diff` tracks this difference.
-  smi_diff: float
-  # Gpu utilization.
-  util: float
-
-  # Allows unpacking operation like
-  # total_mem, occupied, unused, smi_diff, util = GpuStats(...)
-  # See https://stackoverflow.com/a/70753113
-  def __iter__(self):
-    return iter(dataclasses.astuple(self))
-
-
-def gpu_stats() -> GpuStats:
-  """Reports GPU memory usage and utilization."""
-  # See https://pytorch.org/docs/stable/notes/cuda.html#memory-management
-  bytes_per_gb = 1024.0**3
-  device = torch.cuda.current_device()
-  occupied = torch.cuda.memory_allocated(device) / bytes_per_gb
-  reserved = torch.cuda.memory_reserved(device) / bytes_per_gb
-  unused = reserved - occupied
-
-  def smi_mem(device):
-    try:
-      pynvml.nvmlInit()
-      handle = pynvml.nvmlDeviceGetHandleByIndex(device)
-      info = pynvml.nvmlDeviceGetMemoryInfo(handle)
-      return info.used / bytes_per_gb
-    except pynvml.NVMLError:
-      return 0.0
-
-  mem_used_smi = smi_mem(device)
-  smi_diff = mem_used_smi - reserved
-
-  util = torch.cuda.utilization(device)
-  return GpuStats(mem_used_smi, occupied, unused, smi_diff, util)
-
-
-def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
-  if stats is None:
-    stats = gpu_stats()
-  total, occupied, unused, smi_diff, util = stats
-  return (
-      f"GPU memory: {total:.2f}({occupied=:.2f}, {unused=:.2f},"
-      f" {smi_diff=:.2f}) GB. Utilization: {util:.2f}%"
-  )
-
-
 def init_partial_state(
     timeout: datetime.timedelta = datetime.timedelta(seconds=600),
 ) -> None:
@@ -722,7 +340,7 @@ def get_final_checkpoint_path(output_dir: str) -> str:
 
 def _maybe_get_modules_to_not_convert_by_model_id(
     pretrained_model_name_or_path: str,
-) -> Optional[Sequence[str]]:
+) -> Sequence[str] | None:
   """Returns the modules to not convert for the model."""
   if _LLAMA_3_1_405B_MODEL_ID in pretrained_model_name_or_path:
     return _get_llama_3_1_405b_modules_to_not_convert()

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

@@ -17,15 +17,15 @@ _DATASET_NAME = flags.DEFINE_string(
     required=True,
 )
 
-_TRAIN_SPLIT_NAME = flags.DEFINE_string(
-    'train_split_name',
+_TRAIN_SPLIT = flags.DEFINE_string(
+    'train_split',
     'train',
     'The train split name.',
 )
 
-_INSTRUCT_COLUMN_IN_DATASET = flags.DEFINE_string(
-    'instruct_column_in_dataset',
-    constants.DEFAULT_INSTRUCT_COLUMN_IN_DATASET,
+_TRAIN_COLUMN = flags.DEFINE_string(
+    'train_column',
+    constants.DEFAULT_TRAIN_COLUMN,
     'The instruct column in dataset.',
 )
 
@@ -38,6 +38,12 @@ _TEMPLATE = flags.DEFINE_string(
     required=True,
 )
 
+_MAX_SEQ_LENGTH = flags.DEFINE_integer(
+    'max_seq_length',
+    None,
+    'The maximum sequence length.',
+)
+
 _VALIDATE_PERCENTAGE_OF_DATASET = flags.DEFINE_integer(
     'validate_percentage_of_dataset',
     None,
@@ -64,9 +70,10 @@ def main(unused_argv: Sequence[str]) -> None:
 
   dataset_validation_util.validate_dataset_with_template(
       dataset_name=_DATASET_NAME.value,
-      split=_TRAIN_SPLIT_NAME.value,
-      input_column=_INSTRUCT_COLUMN_IN_DATASET.value,
+      split=_TRAIN_SPLIT.value,
+      input_column=_TRAIN_COLUMN.value,
       template=_TEMPLATE.value,
+      max_seq_length=_MAX_SEQ_LENGTH.value,
       use_multiprocessing=_USE_MULTIPROCESSING.value,
       validate_percentage_of_dataset=_VALIDATE_PERCENTAGE_OF_DATASET.value,
       validate_k_rows_of_dataset=_VALIDATE_K_ROWS_OF_DATASET.value,

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

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

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

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

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

@@ -1,10 +1,13 @@
 """Fileutil lib to copy files between gcs and local."""
 
-import glob
+import filecmp
+import fnmatch
 import os
 import pathlib
 import shutil
-from typing import Tuple
+import subprocess
+import time
+from typing import List, Optional, Tuple
 import uuid
 
 from absl import logging
@@ -13,6 +16,17 @@ from google.cloud import storage
 from util import constants
 
 
+_GCS_CLIENT = None
+
+
+def _get_gcs_client() -> storage.Client:
+  """Gets the default GCS client."""
+  global _GCS_CLIENT
+  if _GCS_CLIENT is None:
+    _GCS_CLIENT = storage.Client()
+  return _GCS_CLIENT
+
+
 def generate_tmp_path(extension: str = '') -> str:
   """Generates a temporary file path with UUID.
 
@@ -36,6 +50,123 @@ def force_gcs_fuse_path(gcs_uri: str) -> str:
     return gcs_uri
 
 
+def force_gcs_path(uri: str) -> str:
+  """Converts /gcs/ uris to their gs:// equivalents. No-op for other uris."""
+  if uri.startswith(constants.GCSFUSE_URI_PREFIX):
+    return uri.replace(
+        constants.GCSFUSE_URI_PREFIX, constants.GCS_URI_PREFIX, 1
+    )
+  else:
+    return uri
+
+
+def is_file_available(
+    file_path: str, retry_interval_secs: int = 60, timeout_secs: int = 3600
+) -> bool:
+  """Checks and waits for a file to be available in GCS.
+
+  Args:
+    file_path: The file path to check.
+    retry_interval_secs: The interval in seconds to check the file.
+    timeout_secs: The timeout in seconds to wait for the file.
+
+  Returns:
+    True if the file is available, False otherwise.
+  """
+  start_time = time.time()
+  while True:
+    try:
+      file_check_cmd = ['gcloud', 'storage', 'ls', file_path]
+      result = subprocess.run(
+          file_check_cmd, capture_output=True, text=True, check=True
+      )
+      if file_path in result.stdout:
+        logging.info('File %s exists.', file_path)
+        return True
+    except subprocess.CalledProcessError as e:
+      elapsed_time = time.time() - start_time
+      if elapsed_time > timeout_secs:
+        logging.info(
+            "Timeout: File '%s' not found after %d seconds. Error: %s",
+            file_path,
+            elapsed_time,
+            e,
+        )
+        return False
+
+      logging.info(
+          "File '%s' not found yet. Checking again in %d seconds. Error: %s",
+          file_path,
+          retry_interval_secs,
+          e,
+      )
+      time.sleep(retry_interval_secs)
+
+
+def compare_dirs(
+    local_dir: str,
+    gcsfuse_dir: str,
+    retry_interval_secs: int = 30,
+    timeout_secs: int = 3600,
+) -> bool:
+  """Compares two directories and returns True if they are the same.
+
+  Args:
+    local_dir: The local directory.
+    gcsfuse_dir: The gcsfuse directory.
+    retry_interval_secs: The interval in seconds to check the directories.
+    timeout_secs: The timeout in seconds to wait for the directories.
+
+  Returns:
+    True if the directories are the same, False otherwise.
+  """
+  start_time = time.time()
+  while True:
+    if os.path.exists(local_dir) and os.path.exists(gcsfuse_dir):
+      comparison = filecmp.dircmp(local_dir, gcsfuse_dir)
+      if (
+          not comparison.left_only
+          and not comparison.right_only
+          and not comparison.diff_files
+      ):
+        return True
+    elapsed_time = time.time() - start_time
+    if elapsed_time > timeout_secs:
+      logging.info(
+          "Timeout: Directories '%s' and '%s' do not match after %d seconds.",
+          local_dir,
+          gcsfuse_dir,
+          elapsed_time,
+      )
+      return False
+
+    logging.info(
+        "Directories '%s' and '%s' do not match yet. Checking again in %d"
+        ' seconds.',
+        local_dir,
+        gcsfuse_dir,
+        retry_interval_secs,
+    )
+    time.sleep(retry_interval_secs)
+
+
+def download_gcs_file_to_memory(gcs_uri: str) -> bytes:
+  """Downloads a gcs file to in memory.
+
+  Args:
+    gcs_uri: A string of GCS uri.
+
+  Returns:
+    The content of the gcs file in byte format.
+  """
+  bucket = gcs_uri.split('/')[2]
+  file_path = gcs_uri[len(constants.GCS_URI_PREFIX + bucket + '/') :]
+  client = _get_gcs_client()
+  bucket = client.bucket(bucket)
+  blob = bucket.blob(file_path)
+  return blob.download_as_bytes()
+
+
 def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str):
   """Download a gcs file to a local dir.
 
@@ -62,15 +193,47 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str):
     raise ValueError(
         f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.'
     )
-  client = storage.Client()
+  client = _get_gcs_client()
   os.makedirs(os.path.dirname(local_path), exist_ok=True)
   with open(local_path, 'wb') as f:
     client.download_blob_to_file(gcs_uri, f)
 
 
+def download_gcs_file_list_to_local(
+    gcs_uri_list: List[str], local_dir: str
+) -> List[str]:
+  """Downloads a list of GCS files to a local directory.
+
+  Args:
+    gcs_uri_list: A list of GCS file paths.
+    local_dir: Local directory in which the GCS files are saved.
+
+  Returns:
+    The local file paths corresponding to the input GCS file paths.
+
+  Raises:
+    ValueError: An input file path is not a GCS path.
+  """
+  local_paths = []
+  for gcs_uri in gcs_uri_list:
+    if not is_gcs_path(gcs_uri):
+      raise ValueError(
+          f'{gcs_uri} is not a GCS path starting with'
+          f' {constants.GCS_URI_PREFIX}.'
+      )
+    local_path = os.path.join(local_dir, gcs_uri.replace('gs://', ''))
+    download_gcs_file_to_local(gcs_uri, local_path)
+    local_paths.append(local_path)
+  return local_paths
+
+
 def download_gcs_dir_to_local(
-    gcs_dir: str, local_dir: str, skip_hf_model_bin: bool = False
-):
+    gcs_dir: str,
+    local_dir: str,
+    skip_hf_model_bin: bool = False,
+    allow_patterns: Optional[List[str]] = None,
+    log: bool = True,
+) -> None:
   """Downloads files in a GCS directory to a local directory.
 
   For example:
@@ -78,16 +241,21 @@ def download_gcs_dir_to_local(
     gs://bucket/foo/a -> /tmp/bar/a
     gs://bucket/foo/b/c -> /tmp/bar/b/c
 
-  Arguments:
+  Args:
     gcs_dir: A string of directory path on GCS.
     local_dir: A string of local directory path.
     skip_hf_model_bin: True to skip downloading HF model bin files.
+    allow_patterns: A list of allowed patterns. If provided, only files matching
+      one or more patterns are downloaded.
+    log: True to log each downloaded file.
   """
   if not is_gcs_path(gcs_dir):
     raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.')
   bucket_name = gcs_dir.split('/')[2]
-  prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/')
-  client = storage.Client()
+  prefix = (
+      gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/') + '/'
+  )
+  client = _get_gcs_client()
   blobs = client.list_blobs(bucket_name, prefix=prefix)
   for blob in blobs:
     if blob.name[-1] == '/':
@@ -95,43 +263,63 @@ def download_gcs_dir_to_local(
     file_path = blob.name[len(prefix) :].strip('/')
     local_file_path = os.path.join(local_dir, file_path)
     os.makedirs(os.path.dirname(local_file_path), exist_ok=True)
+    if allow_patterns and all(
+        [not fnmatch.fnmatch(file_path, p) for p in allow_patterns]
+    ):
+      continue
     if (
         file_path.endswith(constants.HF_MODEL_WEIGHTS_SUFFIX)
         and skip_hf_model_bin
     ):
-      logging.info('Skip downloading model bin %s', file_path)
+      if log:
+        logging.info('Skip downloading model bin %s', file_path)
       with open(local_file_path, 'w') as f:
-        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}/{file_path}')
+        f.write(f'{constants.GCS_URI_PREFIX}{bucket_name}/{prefix}{file_path}')
     else:
-      logging.info('Downloading %s to %s', file_path, local_file_path)
+      if log:
+        logging.info('Downloading %s to %s', file_path, local_file_path)
       blob.download_to_filename(local_file_path)
 
 
+def _get_relative_paths(base_dir: str) -> List[str]:
+  """Gets relative paths of all files in a local base directory."""
+  path = pathlib.Path(base_dir)
+  relative_paths = []
+  for local_file in path.rglob('*'):
+    if os.path.isfile(local_file):
+      relative_path = os.path.relpath(local_file, base_dir)
+      relative_paths.append(relative_path)
+  return relative_paths
+
+
+def _upload_local_files_to_gcs(
+    relative_paths: List[str], local_dir: str, gcs_dir: str
+):
+  """Uploads local files to gcs."""
+  bucket_name = gcs_dir.split('/')[2]
+  blob_dir = '/'.join(gcs_dir.split('/')[3:])
+  client = _get_gcs_client()
+  bucket = client.bucket(bucket_name)
+  for relative_path in relative_paths:
+    blob = bucket.blob(os.path.join(blob_dir, relative_path))
+    blob.upload_from_filename(os.path.join(local_dir, relative_path))
+
+
 def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str):
   """Uploads local dir to gcs.
 
   For example:
     upload_local_dir_to_gcs(/tmp/bar, gs://bucket/foo)
-    gs://bucket/foo/a -> /tmp/bar/a
-    gs://bucket/foo/b/c -> /tmp/bar/b/c
+    /tmp/bar/a -> gs://bucket/foo/a
+    /tmp/bar/b/c -> gs://bucket/foo/b/c
 
   Arguments:
     local_dir: A string of local directory path.
     gcs_dir: A string of directory path on GCS.
   """
-  bucket_name = gcs_dir.split('/')[2]
-  blob_dir = '/'.join(gcs_dir.split('/')[3:])
-  client = storage.Client()
-  bucket = client.bucket(bucket_name)
-  for local_file in glob.glob(local_dir + '/**'):
-    if os.path.isfile(local_file):
-      logging.info(
-          'Uploading %s to %s',
-          local_file,
-          os.path.join(constants.GCS_URI_PREFIX, bucket_name, blob_dir),
-      )
-      blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file)))
-      blob.upload_from_filename(local_file)
+  # Relative paths of all files in local_dir.
+  relative_paths = _get_relative_paths(local_dir)
+  _upload_local_files_to_gcs(relative_paths, local_dir, gcs_dir)
 
 
 def upload_file_to_gcs_path(
@@ -155,7 +343,7 @@ def upload_file_to_gcs_path(
   if not source_path_obj.exists():
     raise RuntimeError(f'Source path does not exist: {source_path}')
 
-  storage_client = storage.Client()
+  storage_client = _get_gcs_client()
   source_file_path = source_path
   destination_file_uri = destination_uri
   logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri)
@@ -174,7 +362,9 @@ def is_gcs_path(input_path: str) -> bool:
   Returns:
       True if the input path is a GCS path, False otherwise.
   """
-  return input_path.startswith(constants.GCS_URI_PREFIX)
+  return input_path is not None and input_path.startswith(
+      constants.GCS_URI_PREFIX
+  )
 
 
 def release_text_assets(
@@ -232,13 +422,10 @@ def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]:
   """
   _, local_video_file_name = os.path.split(video_file_path)
   file_extension = os.path.splitext(video_file_path)[1]
-  if file_extension:
-    remote_video_file_name = local_video_file_name.replace(
-        file_extension, '_overlay.mp4'
-    )
-  else:
-    remote_video_file_name = local_video_file_name + '_overlay.mp4'
-  local_file_path = generate_tmp_path(file_extension)
+  remote_video_file_name = local_video_file_name.replace(
+      file_extension, '_overlay.mp4'
+  )
+  local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1])
   logging.info('Downloading %s to %s...', video_file_path, local_file_path)
   download_gcs_file_to_local(video_file_path, local_file_path)
   return local_file_path, remote_video_file_name
@@ -254,10 +441,19 @@ def get_output_video_file(video_output_file_path: str) -> str:
     str: Local video output file path.
   """
   file_extension = os.path.splitext(video_output_file_path)[1]
-  if file_extension:
-    out_local_video_file_name = video_output_file_path.replace(
-        file_extension, '_overlay' + file_extension
-    )
-  else:
-    out_local_video_file_name = video_output_file_path + '_overlay'
+  out_local_video_file_name = video_output_file_path.replace(
+      file_extension, '_overlay' + file_extension
+  )
   return out_local_video_file_name
+
+
+def delete_local_file(local_file_path: str) -> None:
+  """Deletes a local file."""
+  if os.path.exists(local_file_path):
+    os.remove(local_file_path)
+
+
+def delete_local_dir(local_dir: str) -> None:
+  """Deletes a local directory recursively."""
+  if os.path.exists(local_dir):
+    shutil.rmtree(local_dir)

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

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

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

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

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

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

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

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

community-content/vertex_pipeline_kfpv2/customjob.py

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

community-content/vertex_pipeline_kfpv2/pipelinejob.py

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

notebooks/community/CODEOWNERS

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

notebooks/community/bigquery_ml/Anomaly_detection_in_Cloud_Audit_logs_with_BQML.ipynb

@@ -37,18 +37,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/bigquery_ml/Anomaly_detection_in_Cloud_Audit_logs_with_BQML.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/bigquery_ml/Anomaly_detection_in_Cloud_Audit_logs_with_BQML.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/bigquery_ml/Anomaly_detection_in_Cloud_Audit_logs_with_BQML.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>                                                                                               \n",

notebooks/community/bigquery_ml/bq_ml_with_vision_translation_nlp.ipynb

@@ -47,18 +47,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/bigquery_ml/bq_ml_with_vision_translation_nlp.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/bigquery_ml/bq_ml_with_vision_translation_nlp.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/bigquery_ml/bq_ml_with_vision_translation_nlp.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/bigquery_ml/bqml-online-prediction.ipynb

@@ -35,18 +35,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/bigquery_ml/bqml-online-prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/official/bigquery_ml/bqml-online-prediction.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/official/bigquery_ml/bqml-online-prediction.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>                                                                                               \n",

notebooks/community/cohere/cohere_embedding_with_matching_engine.ipynb

@@ -38,7 +38,7 @@
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/cohere/cohere_embedding_with_matching_engine.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/experiments/vertex_ai_model_experimentation.ipynb

@@ -36,18 +36,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/vertex_ai_experiments_classification.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/vertex_ai_experiments_classification.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/vertex_ai_experiments_classification.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>                                                                                               \n",

notebooks/community/explainable_ai/SDK_Custom_Container_XAI.ipynb

@@ -33,18 +33,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/explainable_ai/SDK_Custom_Container_XAI.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/explainable_ai/SDK_Custom_Container_XAI.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/explainable_ai/SDK_Custom_Container_XAI.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>                                                                                               \n",

notebooks/community/feature_store/gapic-feature-store.ipynb

@@ -33,12 +33,12 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/feature_store/gapic-feature-store.ipynb\"\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/official/feature_store/gapic-feature-store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/feature_store/get_started_vertex_feature_store.ipynb

@@ -34,20 +34,20 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage6/get_started_vertex_feature_store_serving.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "    \n",
         "  <td>\n",
         "        <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/ml_ops/stage6/get_started_vertex_feature_store_serving.ipynb\">\n",
-        "        <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "        <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "        </a>\n",
         "  </td>\n",
         "    \n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/notebooks/community/ml_ops/stage6/get_started_vertex_feature_store_serving.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/feature_store/mobile_gaming/mobile_gaming_feature_store.ipynb

@@ -33,18 +33,18 @@
         "\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/vertex-ai-samples/blob/main/notebooks/community/feature_store/mobile_gaming/mobile_gaming_feature_store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/mobile_gaming/mobile_gaming_feature_store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/mobile_gaming/mobile_gaming_feature_store.ipynb\">\n",
-        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\">\n",
+        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\">\n",
         "      Open in Vertex AI Workbench\n",
         "    </a>\n",
         "  </td>                                                                                               \n",

notebooks/community/feature_store/sdk-feature-store.ipynb

@@ -32,12 +32,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/sdk-feature-store.ipynb\"\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/feature_store/sdk-feature-store.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_classification_batch.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_classification_export_edge.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_export_edge.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_export_edge.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_classification_online.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_classification_online_proxy.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_online_proxy.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_classification_online_proxy.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_object_detection_batch.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_object_detection_export_edge.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_export_edge.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_export_edge.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_object_detection_online.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_object_detection_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_segmentation_batch.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_segmentation_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_segmentation_batch.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",

notebooks/community/gapic/automl/showcase_automl_image_segmentation_online.ipynb

@@ -34,12 +34,12 @@
         "<table align=\"left\">\n",
         "  <td>\n",
         "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_segmentation_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Colab logo\"> Run in Colab\n",
+        "      <img src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Colab logo\"> Run in Colab\n",
         "    </a>\n",
         "  </td>\n",
         "  <td>\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/notebooks/community/gapic/automl/showcase_automl_image_segmentation_online.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+        "      <img width=\"32px\" src=\"https://www.svgrepo.com/download/217753/github.svg\" alt=\"GitHub logo\">\n",
         "      View on GitHub\n",
         "    </a>\n",
         "  </td>\n",