Update CODEOWNERS

PiperOrigin-RevId: 744731980

.github/workflows/linter/requirements.txt

@@ -4,7 +4,7 @@ jupyter
 nbconvert
 black==25.1.0
 pyupgrade==3.19.1
-isort==6.0.0
-flake8==7.1.1
+isort==6.0.1
+flake8==7.2.0
 nbqa==1.9.1
 

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

@@ -6,8 +6,10 @@ import io
 import json
 import os
 import subprocess
+import time
 from typing import Any, Dict, Sequence
 
+from google import auth
 from google.cloud import storage
 import matplotlib.pyplot as plt
 import numpy as np
@@ -533,6 +535,7 @@ def get_resource_id(
       "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_TESLA_T4": "nvidia_t4_gpus",
       "TPU_V5e": "tpu_v5e",
       "TPU_V3": "tpu_v3",
@@ -620,3 +623,75 @@ def check_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

@@ -541,4 +541,3 @@ def validate_dataset_with_template(
           os.path.basename(dataset_name), os.path.basename(template)
       )
   )
-

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

@@ -0,0 +1,440 @@
+"""Test util class."""
+
+import copy
+import dataclasses
+import datetime
+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 frozendict
+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://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.',
+)
+
+_GCS_TESTDATA_DIR = 'peft-train-image-test'
+
+_THROUGHPUT_TEST_EXCEPTIONS = frozendict.frozendict({
+    ('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:
+    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 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):
+    download_from_gcs(
+        os.path.join(_GCS_INPUT_DIR.value, name), _LOCAL_INPUT_DIR.value
+    )
+
+  return local_data
+
+
+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_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/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
@@ -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,27 @@ 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)
       logging.info(
-          'on_step_end: %s, throughput: %.2f s/it',
+          'on_step_end: Throughput: %.2f token/s. %s, %s',
+          throughput,
           utils.gpu_stats_str(gpu_stats),
-          delta_t,
+          utils.cpu_stats_str(),
       )
 
   def on_train_begin(
@@ -61,7 +93,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',
+          utils.gpu_stats_str(),
+          utils.cpu_stats_str(),
+      )
 
   def on_train_end(
       self,
@@ -72,15 +108,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/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,27 @@
 --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
 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,12 @@ 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/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/
 
 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,286 @@
 """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.
+    num_fewshot: The number of few-shot examples to use 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
+  num_fewshot: int | None
+  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/gcs_syncer.py

@@ -0,0 +1,86 @@
+"""Sync local directory to GCS directory using rsync."""
+
+from collections.abc import Sequence
+import multiprocessing
+import os
+import subprocess
+import time
+
+from absl import logging
+
+_GCS_COMMAND_RETRIES = 3
+_RSYNC_RETRY_INTERVAL_SECS = 30
+
+
+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/peft/train/vmg/instruct_lora.py

@@ -1,10 +1,11 @@
 """Instruct/Chat with LoRA models."""
 
+from collections.abc import Callable, Mapping, Sequence
 import dataclasses
 import datetime
 import json
 import os
-from typing import Any, Dict, Optional, Sequence
+from typing import Any
 import warnings
 
 from absl import app
@@ -13,14 +14,12 @@ 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
@@ -31,14 +30,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 +48,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 +119,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 +130,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 +146,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 +170,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 +181,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,23 +219,14 @@ _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.',
 )
@@ -253,17 +244,19 @@ _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.',
+    ' this can be any Hugging Face dataset name or path.',
 )
 
 # We set the default eval split as `test`, based on observation from
@@ -288,6 +281,13 @@ _EVAL_COLUMN = flags.DEFINE_string(
     'Eval column name in the eval dataset for `builtin_eval`.',
 )
 
+_METRIC_FOR_BEST_MODEL = flags.DEFINE_string(
+    'metric_for_best_model',
+    None,
+    'If set, the best model is saved at the end of training based on the'
+    ' metric',
+)
+
 _TRAIN_PRECISION = flags.DEFINE_enum(
     'train_precision',
     constants.PRECISION_MODE_16B,
@@ -299,15 +299,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 +356,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 +410,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 +479,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 +534,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',
+      utils.gpu_stats_str(),
+      utils.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 +588,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_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,
   )
 
@@ -509,9 +622,9 @@ def finetune_instruct(
           effective_batch_size,
       )
       train_dataset_stats = utils.get_dataset_stats(
-          train_dataset,
+          train_dataset_with_template,
           tokenizer,
-          instruct_column_in_dataset,
+          train_column,
           effective_batch_size,
       )
       logging.info('stats: %s', train_dataset_stats)
@@ -522,10 +635,10 @@ def finetune_instruct(
         json.dump(dataclasses.asdict(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,
@@ -568,14 +681,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 +695,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 +723,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 +733,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 +779,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 +789,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 +803,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 +816,42 @@ 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,
+        per_device_batch_size=_PER_DEVICE_EVAL_BATCH_SIZE.value,
         num_fewshot=_EVAL_NUM_FEWSHOT.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 +860,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 +902,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,
 )
 
@@ -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,24 @@ 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.value
-  )
-
-  merge_base_and_lora_output_dir = utils.GcsOrLocalDirectory(
-      _MERGE_BASE_AND_LORA_OUTPUT_DIR.value
-  )
-
   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.value,
+      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

@@ -19,14 +19,15 @@ from transformers import get_linear_schedule_with_warmup
 from util import dataset_validation_util
 
 
-_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`.",
 )
 
 _OUTPUT_DIR = flags.DEFINE_string(
@@ -63,8 +64,8 @@ _LORA_DROPOUT = flags.DEFINE_float(
     " https://huggingface.co/docs/peft/task_guides/token-classification-lora.",
 )
 
-_NUM_EPOCHS = flags.DEFINE_integer(
-    "num_epochs",
+_NUM_TRAIN_EPOCHS = flags.DEFINE_integer(
+    "num_train_epochs",
     None,
     "The number of training epochs.",
 )
@@ -83,13 +84,13 @@ _LEARNING_RATE = flags.DEFINE_float(
 
 
 def finetune_sequence_classification(
-    pretrained_model_id: str,
+    pretrained_model_name_or_path: str,
     dataset_name: str,
     output_dir: str,
     lora_rank: int = 8,
     lora_alpha: int = 16,
     lora_dropout: float = 0.1,
-    num_epochs: int = 20,
+    num_train_epochs: int = 20,
     batch_size: int = 32,
     learning_rate: float = 3e-4,
 ) -> None:
@@ -104,13 +105,13 @@ def finetune_sequence_classification(
       lora_alpha=lora_alpha,
       lora_dropout=lora_dropout,
   )
-  if any(k in pretrained_model_id for k in ("gpt", "opt", "bloom")):
+  if any(k in pretrained_model_name_or_path for k in ("gpt", "opt", "bloom")):
     padding_side = "left"
   else:
     padding_side = "right"
 
   tokenizer = AutoTokenizer.from_pretrained(
-      pretrained_model_id, padding_side=padding_side
+      pretrained_model_name_or_path, padding_side=padding_side
   )
   if getattr(tokenizer, "pad_token_id") is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
@@ -156,7 +157,7 @@ def finetune_sequence_classification(
   )
 
   model = AutoModelForSequenceClassification.from_pretrained(
-      pretrained_model_id, return_dict=True
+      pretrained_model_name_or_path, return_dict=True
   )
   model = get_peft_model(model, peft_config)
   model.print_trainable_parameters()
@@ -166,12 +167,12 @@ def finetune_sequence_classification(
   # 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),
+      num_warmup_steps=0.06 * (len(train_dataloader) * num_train_epochs),
+      num_training_steps=(len(train_dataloader) * num_train_epochs),
   )
 
   model.to(device)
-  for epoch in range(num_epochs):
+  for epoch in range(num_train_epochs):
     model.train()
     for _, batch in enumerate(tqdm(train_dataloader)):
       batch.to(device)
@@ -201,25 +202,27 @@ def finetune_sequence_classification(
 
 
 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
+  if dataset_validation_util.is_gcs_path(_PRETRAINED_MODEL_NAME_OR_PATH.value):
+    pretrained_model_name_or_path = (
+        dataset_validation_util.download_gcs_uri_to_local(
+            _PRETRAINED_MODEL_NAME_OR_PATH.value
+        )
     )
   else:
-    pretrained_model_id = _PRETRAINED_MODEL_ID.value
+    pretrained_model_name_or_path = _PRETRAINED_MODEL_NAME_OR_PATH.value
   pretrained_model_path = dataset_validation_util.force_gcs_fuse_path(
-      pretrained_model_id
+      pretrained_model_name_or_path
   )
   output_dir = dataset_validation_util.force_gcs_fuse_path(_OUTPUT_DIR.value)
 
   finetune_sequence_classification(
-      pretrained_model_id=pretrained_model_path,
+      pretrained_model_name_or_path=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),
+      num_train_epochs=int(_NUM_TRAIN_EPOCHS.value),
       batch_size=_BATCH_SIZE.value,
       learning_rate=_LEARNING_RATE.value,
   )

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/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/command_builder.py

@@ -0,0 +1,87 @@
+"""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._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 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
+    )
+
+
+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

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,23 @@ 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',
       ],
       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 +78,29 @@ 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',
       ],
       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 +112,14 @@ 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'],
       precision=['4bit', '8bit', 'bfloat16'],
       max_seq_length=list(range(4 * 1024, 24 * 1024 + 1, 4 * 1024)),
       num_gpus=[8],
@@ -127,8 +127,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
@@ -140,14 +140,97 @@ class TrainerThroughputTest(test_util.TestBase):
         'vertex_vision_model_garden_peft/llama_fsdp_8gpu.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'],
+      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 +238,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 +252,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

@@ -31,7 +31,6 @@ class TrainedModelQualityTest(test_util.TestBase):
 
     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
@@ -39,7 +38,7 @@ class TrainedModelQualityTest(test_util.TestBase):
     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.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
@@ -47,7 +46,7 @@ class TrainedModelQualityTest(test_util.TestBase):
     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.input_masking = True
     self.task_cmd_builder.attn_implementation = 'flash_attention_2'
     self.task_cmd_builder.report_to = 'tensorboard'
 
@@ -71,25 +70,23 @@ class TrainedModelQualityTest(test_util.TestBase):
   )
   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.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_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.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.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.docker_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
+    self.command_builder.add_env_var('CUDA_VISIBLE_DEVICES', '0')
 
     self.assertEqual(self.run_cmd(), 0)
 
@@ -99,24 +96,22 @@ class TrainedModelQualityTest(test_util.TestBase):
   )
   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.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_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.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.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.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)
 
@@ -126,24 +121,22 @@ class TrainedModelQualityTest(test_util.TestBase):
   )
   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.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_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.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.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.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)
 

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,20 @@ environment. Otherwise, `python3` is used.
 """
 
 import argparse
+from collections.abc import MutableSequence, Sequence
 import json
+import multiprocessing
 import os
 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 vertex_vision_model_garden_peft.train.vmg import gcs_syncer
 from vertex_vision_model_garden_peft.train.vmg import utils
 from util import constants
+from util import fileutils
 from util import hypertune_utils
 
 
@@ -38,19 +42,17 @@ _TASK_TO_SCRIPT = {
         '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.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.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 = [
@@ -121,7 +123,7 @@ def _get_accelerate_args() -> argparse.Namespace:
 
 
 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 +131,64 @@ 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]
-) -> Sequence[Sequence[str]]:
-  """Returns the training command and maybe the merge command if applicable."""
+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_training_path(path: str, node_rank: int) -> 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. Othereise, 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,
+        dataset_validation_util.force_gcs_fuse_path(path)[1:],
+    )
+    gcs_dir = fileutils.force_gcs_path(path)
+    os.makedirs(local_dir, exist_ok=True)
+  return local_dir, os.path.join(gcs_dir, f'node-{node_rank}')
+
+
+def _get_train_and_maybe_merge_cmd_and_dirs_to_sync(
+    task_type: str, config_file: str, unknown: Sequence[str]
+) -> Sequence[Sequence[str]]:
+  """Returns the training and merge command(if applicable) and dirs to sync.
+
+  Args:
+    task_type: The task type.
+    config_file: The accelerate config file path.
+    unknown: The unknown args which are not recognised by the parser.
+
+  Returns:
+    The bash commands to execute and the directories to sync.
+  """
+  dirs_to_sync = []
   # Only populated when multi-node is used.
   accelerate_args = _get_accelerate_args()
+  node_rank = getattr(accelerate_args, 'machine_rank', 0)
   training_cmd = launch_script_cmd(
-      _TASK_TO_SCRIPT[task],
+      _TASK_TO_SCRIPT[task_type],
       config_file,
       accelerate_args=accelerate_args,
   )
@@ -152,17 +203,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 = _manage_training_path(
+      training_args.output_dir, node_rank
+  )
+  training_args.output_dir = local_output_dir
+  if _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 = _manage_training_path(
+        merge_args.merge_base_and_lora_output_dir, node_rank
+    )
+    merge_args.merge_base_and_lora_output_dir = merge_local_dir
+    if _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 +237,17 @@ 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 = _manage_training_path(lora_dir, node_rank)
+    if _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 +260,60 @@ def _get_train_cmd_and_maybe_merge_cmd(
     ]
     commands.append(conda_run_cmd)
 
-  return commands
+  return commands, dirs_to_sync
+
+
+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=gcs_syncer.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('training finished')
+  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 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 +332,8 @@ 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
     )
   else:
     assert task in _TASK_TO_SCRIPT
@@ -233,10 +341,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 = _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:
+    _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,23 +1,21 @@
 """Common libraries for PEFT."""
 
+from collections.abc import Mapping, Sequence
 import dataclasses
 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 psutil
 import pynvml
 import torch
 import transformers
@@ -25,208 +23,27 @@ 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 +121,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": True,
+      "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 +174,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 +188,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 +235,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 +255,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 +269,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 +280,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.
 
@@ -567,70 +344,6 @@ def force_gc():
   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.
@@ -683,7 +396,7 @@ def gpu_stats() -> GpuStats:
   return GpuStats(mem_used_smi, occupied, unused, smi_diff, util)
 
 
-def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
+def gpu_stats_str(stats: GpuStats | None = None) -> str:
   if stats is None:
     stats = gpu_stats()
   total, occupied, unused, smi_diff, util = stats
@@ -693,6 +406,75 @@ def gpu_stats_str(stats: Optional[GpuStats] = None) -> str:
   )
 
 
+@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}%"
+
+
 def init_partial_state(
     timeout: datetime.timedelta = datetime.timedelta(seconds=600),
 ) -> None:
@@ -722,7 +504,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/constants.py

@@ -90,12 +90,11 @@ TEXT_TO_IMAGE_DREAMBOOTH_LORA = 'text-to-image-dreambooth-lora'
 TEXT_TO_IMAGE_DREAMBOOTH_LORA_SDXL = 'text-to-image-dreambooth-lora-sdxl'
 SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora'
 MERGE_CAUSAL_LANGUAGE_MODEL_LORA = 'merge-causal-language-model-lora'
-QUANTIZE_MODEL = 'quantize-model'
 INSTRUCT_LORA = 'instruct-lora'
 VALIDATE_DATASET_WITH_TEMPLATE = 'validate-dataset-with-template'
+RUN_TESTS = 'test'
 DEFAULT_TEXT_COLUMN_IN_DATASET = 'quote'
-DEFAULT_TEXT_COLUMN_IN_QUANTIZATION_DATASET = 'text'
-DEFAULT_INSTRUCT_COLUMN_IN_DATASET = 'text'
+DEFAULT_TRAIN_COLUMN = 'text'
 
 FINAL_CHECKPOINT_DIRNAME = 'checkpoint-final'
 
@@ -113,13 +112,17 @@ PRECISION_MODE_16 = 'float16'
 PRECISION_MODE_16B = 'bfloat16'
 PRECISION_MODE_32 = 'float32'
 
-# Quantization modes.
-GPTQ = 'gptq'
-AWQ = 'awq'
+ROUGE_VARIANTS = ('rouge1', 'rouge2', 'rougeL', 'rougeLsum')
 
-# AWQ versions.
-GEMM = 'GEMM'
-GEMV = 'GEMV'
+# 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'

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

@@ -57,6 +57,23 @@ def force_gcs_path(uri: str) -> str:
     return uri
 
 
+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.
 
@@ -335,24 +352,3 @@ def get_output_video_file(video_output_file_path: str) -> str:
       file_extension, '_overlay' + file_extension
   )
   return out_local_video_file_name
-
-
-def write_first_party_model_metadata(
-    output_path: str, required_container_uri: str
-) -> None:
-  """Write Vertex internal model metadata for first party artifacts."""
-  model_metadata_fname = 'model_metadata.jsonl'
-  if len(required_container_uri) > 126:
-    raise ValueError(f'Docker URI exceeds 126 chars: {required_container_uri}')
-  payload = '\n{}{}'.format(  # serialized proto
-      chr(len(required_container_uri)),
-      required_container_uri,
-  )
-  os.makedirs(output_path, exist_ok=True)
-  output_dirs = [output_path]
-  if output_path.startswith('/gcs'):
-    # include all parent dirs, except "/", "/gcs"
-    output_dirs.extend([str(p) for p in pathlib.Path(output_path).parents][:-2])
-  for output_dir in output_dirs:
-    with open(os.path.join(output_dir, model_metadata_fname), 'w') as f:
-      f.write(payload)

notebooks/community/model_garden/gke_model_ui_deployment_notebook.ipynb

@@ -0,0 +1,422 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "Pr9TgOcV9vAXeqGiyTaTI5kS",
+      "metadata": {
+        "cellView": "form",
+        "id": "Pr9TgOcV9vAXeqGiyTaTI5kS"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "M1CpgYundFwz",
+      "metadata": {
+        "id": "M1CpgYundFwz"
+      },
+      "source": [
+        "# Get started with your deployed model on GKE\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fgke_model_ui_deployment_notebook.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/gke_model_ui_deployment_notebook.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "t2jj2XOgkS4F",
+      "metadata": {
+        "id": "t2jj2XOgkS4F"
+      },
+      "source": [
+        "# Overview\n",
+        "\n",
+        "This notebook will guide you through the initial step of testing your recently deployed model with text prompts. Depending on your deployed model's inference setup, the notebook utilizes either Text Generation Inference [TGI](https://huggingface.co/docs/text-generation-inference/en/index) or [vLLM](https://developers.googleblog.com/en/inference-with-gemma-using-dataflow-and-vllm/#:~:text=model%20frameworks%20simple.-,What%20is%20vLLM%3F,-vLLM%20is%20an), two efficient serving frameworks that enhance the performance of your GPU model. Ready to see your deployed model respond? Run the cells below and start experimenting with different prompts!\n",
+        "\n",
+        "### Prerequisites\n",
+        "\n",
+        "Before proceeding with this notebook, ensure you have already deployed a model using the Google Cloud Console. You can find an overview of AI and Machine Learning services on [GKE AI/ML](https://console.cloud.google.com/kubernetes/aiml/overview).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "Enable prompt-based testing of the AI model deployed on GKE\n",
+        "\n",
+        "### GPUs\n",
+        "\n",
+        "GPUs let you accelerate specific workloads running on your nodes, such as machine learning and data processing. GKE provides a range of machine type options for node configuration, including machine types with NVIDIA H100, L4, and A100 GPUs.\n",
+        "\n",
+        "### Understanding the Inference Frameworks\n",
+        "\n",
+        "Your model is running on one of two popular and efficient serving frameworks: vLLM or Text Generation Inference (TGI). The following sections provide a brief overview of each to give you context on the underlying technology powering your model.\n",
+        "\n",
+        "\n",
+        "#### TGI\n",
+        "\n",
+        "TGI is a highly optimized open-source LLM serving framework that can increase serving throughput on GPUs. TGI includes features such as:\n",
+        "\n",
+        "* Optimized transformer implementation with PagedAttention\n",
+        "* Continuous batching to improve the overall serving throughput\n",
+        "* Tensor parallelism and distributed serving on multiple GPUs\n",
+        "\n",
+        "To learn more, refer to the [TGI documentation](https://github.com/huggingface/text-generation-inference/blob/main/README.md)\n",
+        "\n",
+        "#### vLLM\n",
+        "\n",
+        "vLLM is another fast and easy-to-use library for LLM inference and serving. It's known for its high throughput and efficiency, and it leverages PagedAttention. Key features include:\n",
+        "\n",
+        "* PagedAttention: Efficient memory management for handling long sequences and dynamic workloads.\n",
+        "* Continuous batching: Maximizes GPU utilization by batching incoming requests.\n",
+        "* High-throughput serving: Designed for production-level serving with low latency.\n",
+        "* Optimized CUDA kernels.\n",
+        "\n",
+        "To learn more, refer to the [vLLM documentation](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/vllm/use-vllm)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "XMf-T58TkDy1",
+      "metadata": {
+        "cellView": "form",
+        "id": "XMf-T58TkDy1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Connect to Google Cloud Project\n",
+        "# @markdown #### Run this cell to configure your Google Cloud environment for Kubernetes (GKE) operations.\n",
+        "\n",
+        "# @markdown #### Actions:\n",
+        "# @markdown 1.  **Connects to Project & Region:** Retrieves and sets your Google Cloud project ID and region.\n",
+        "# @markdown 3.  **Installs `kubectl`:** Installs the Kubernetes command-line tool.\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Set up gcloud.\n",
+        "! gcloud config set project \"$PROJECT_ID\"\n",
+        "! gcloud services enable container.googleapis.com\n",
+        "\n",
+        "# Add kubectl to the set of available tools.\n",
+        "! mkdir -p /tools/google-cloud-sdk/.install\n",
+        "! gcloud components install kubectl --quiet"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1oG8ymQenHyD",
+      "metadata": {
+        "cellView": "form",
+        "id": "1oG8ymQenHyD"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Select Cluster and Deployment { vertical-output: true }\n",
+        "\n",
+        "# @markdown ## Instruction:\n",
+        "\n",
+        "# @markdown This cell provides interactive dropdown menus to select a Google Kubernetes Engine (GKE) cluster and a deployment within that cluster.\n",
+        "\n",
+        "# @markdown ***Please select a cluster and deployment before proceeding.***\n",
+        "\n",
+        "import json\n",
+        "import subprocess\n",
+        "\n",
+        "import ipywidgets as widgets\n",
+        "from IPython.display import display\n",
+        "\n",
+        "SELECTED_DEPLOYMENT = None\n",
+        "\n",
+        "\n",
+        "def get_clusters(p, r):\n",
+        "    try:\n",
+        "        return (\n",
+        "            subprocess.run(\n",
+        "                [\n",
+        "                    \"gcloud\",\n",
+        "                    \"container\",\n",
+        "                    \"clusters\",\n",
+        "                    \"list\",\n",
+        "                    \"--project\",\n",
+        "                    p,\n",
+        "                    \"--region\",\n",
+        "                    r,\n",
+        "                    \"--format=value(name)\",\n",
+        "                ],\n",
+        "                capture_output=True,\n",
+        "                text=True,\n",
+        "                check=True,\n",
+        "            )\n",
+        "            .stdout.strip()\n",
+        "            .split(\"\\n\")\n",
+        "        )\n",
+        "    except subprocess.CalledProcessError as e:\n",
+        "        print(f\"Error: {e}\")\n",
+        "        return []\n",
+        "\n",
+        "\n",
+        "def get_deployments(c, r):\n",
+        "    try:\n",
+        "        subprocess.run(\n",
+        "            [\n",
+        "                \"gcloud\",\n",
+        "                \"container\",\n",
+        "                \"clusters\",\n",
+        "                \"get-credentials\",\n",
+        "                c,\n",
+        "                \"--location\",\n",
+        "                r,\n",
+        "            ],\n",
+        "            capture_output=True,\n",
+        "            text=True,\n",
+        "            check=True,\n",
+        "        )\n",
+        "        deployments = json.loads(\n",
+        "            subprocess.run(\n",
+        "                [\"kubectl\", \"get\", \"deployments\", \"-o\", \"json\"],\n",
+        "                capture_output=True,\n",
+        "                text=True,\n",
+        "                check=True,\n",
+        "            ).stdout\n",
+        "        )\n",
+        "        return [i[\"metadata\"][\"name\"] for i in deployments[\"items\"]]\n",
+        "    except subprocess.CalledProcessError as e:\n",
+        "        print(f\"Error: {e}\")\n",
+        "        return []\n",
+        "\n",
+        "\n",
+        "def create_deployment_dropdown(cluster_name, region, on_select_deployment):\n",
+        "    deployments = get_deployments(cluster_name, region)\n",
+        "    deployments_with_prompt = [\"Select Deployment\"] + deployments\n",
+        "    deployment_dropdown = widgets.Dropdown(\n",
+        "        options=deployments_with_prompt,\n",
+        "        description=\"Deployments\",\n",
+        "        disabled=False,\n",
+        "        width=\"4000px\",\n",
+        "    )\n",
+        "    deployment_dropdown.observe(\n",
+        "        lambda c: on_select_deployment(c[\"new\"])\n",
+        "        if c[\"type\"] == \"change\" and c[\"name\"] == \"value\"\n",
+        "        else None,\n",
+        "        names=\"value\",\n",
+        "    )\n",
+        "    return deployment_dropdown\n",
+        "\n",
+        "\n",
+        "def on_deployment_select(deployment_name):\n",
+        "    global SELECTED_DEPLOYMENT\n",
+        "    SELECTED_DEPLOYMENT = deployment_name\n",
+        "    print(f\"Selected deployment: {SELECTED_DEPLOYMENT}\")\n",
+        "\n",
+        "\n",
+        "def on_cluster_change(change):\n",
+        "    if change[\"type\"] == \"change\" and change[\"name\"] == \"value\":\n",
+        "        if change[\"new\"] == \"Select Cluster\":\n",
+        "            return\n",
+        "        deployment_dropdown = create_deployment_dropdown(\n",
+        "            change[\"new\"], REGION, on_deployment_select\n",
+        "        )\n",
+        "        display(deployment_dropdown)\n",
+        "\n",
+        "\n",
+        "clusters = get_clusters(PROJECT_ID, REGION)\n",
+        "if clusters:\n",
+        "    # @markdown Run this cell to display the Cluster dropdown menu:\n",
+        "    clusters_with_prompt = [\"Select Cluster\"] + clusters\n",
+        "    cluster_dropdown = widgets.Dropdown(\n",
+        "        options=clusters_with_prompt, description=\"Clusters\", disabled=False\n",
+        "    )\n",
+        "    cluster_dropdown.observe(on_cluster_change, names=\"value\")\n",
+        "    display(cluster_dropdown)\n",
+        "else:\n",
+        "    print(f\"No clusters found in {PROJECT_ID}/{REGION}.\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "IKGTaN84p8rX",
+      "metadata": {
+        "cellView": "form",
+        "id": "IKGTaN84p8rX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Chat completion for text-only models {run:\"auto\", vertical-output: true}\n",
+        "\n",
+        "# @markdown You may send prompts to the model server for prediction.\n",
+        "# @markdown\n",
+        "# @markdown * **user_prompt (string):** This is the text prompt you provide to the language model. It's the question or instruction e (e.g., \"Explain neural networks\").\n",
+        "\n",
+        "# @markdown * **temperature (number):** This  parameter controls the randomness of the model's output. It influences how the model selects the next token in the sequence it generates. Typical values range from 0.2 to 1.0.\n",
+        "\n",
+        "# @markdown * **max_tokens (number):** This parameter refers to the maximum number of tokens (words or sub-word units) that the model is allowed to generate in its response.\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "\n",
+        "\n",
+        "def get_deployment_pod_name(deployment):\n",
+        "    try:\n",
+        "        label = deployment + \"-app\"\n",
+        "        pods = json.loads(\n",
+        "            subprocess.run(\n",
+        "                [\"kubectl\", \"get\", \"pods\", \"-o\", \"json\", \"-l\", f\"app={label}\"],\n",
+        "                capture_output=True,\n",
+        "                check=True,\n",
+        "            ).stdout\n",
+        "        )\n",
+        "        return pods[\"items\"][0][\"metadata\"][\"name\"] if pods[\"items\"] else None\n",
+        "    except (\n",
+        "        subprocess.CalledProcessError,\n",
+        "        json.JSONDecodeError,\n",
+        "        KeyError,\n",
+        "        IndexError,\n",
+        "    ):\n",
+        "        return None\n",
+        "\n",
+        "\n",
+        "def check_vllm_label(pod_name):\n",
+        "    \"\"\"Checks if the pod has the 'ai.gke.io/inference-server=vllm' label.\"\"\"\n",
+        "    try:\n",
+        "        result = subprocess.run(\n",
+        "            [\"kubectl\", \"get\", \"pod\", pod_name, \"-o\", \"json\"],\n",
+        "            capture_output=True,\n",
+        "            check=True,\n",
+        "        )\n",
+        "        labels = json.loads(result.stdout)[\"metadata\"][\"labels\"]\n",
+        "        return labels.get(\"ai.gke.io/inference-server\") == \"vllm\"\n",
+        "    except (subprocess.CalledProcessError, KeyError, json.JSONDecodeError):\n",
+        "        return False\n",
+        "\n",
+        "\n",
+        "def process_response(request, pod_name, pod_endpoint, is_vllm):\n",
+        "    response = !kubectl exec -t {pod_name} -- curl -X POST http://{pod_endpoint}/generate -H \"Content-Type: application/json\" -d '{json.dumps(request)}' 2> /dev/null\n",
+        "    try:\n",
+        "        data = json.loads(response[0])\n",
+        "        if is_vllm:\n",
+        "            return data[\"predictions\"][0]\n",
+        "        else:\n",
+        "            return data[\"generated_text\"]\n",
+        "    except (json.JSONDecodeError, KeyError, IndexError) as e:\n",
+        "        return f\"Error: {e}, Raw: {response}\"\n",
+        "\n",
+        "\n",
+        "deployment_pod = get_deployment_pod_name(SELECTED_DEPLOYMENT)\n",
+        "is_vllm_inference = check_vllm_label(deployment_pod)\n",
+        "\n",
+        "user_prompt = \"What is AI?\"  # @param {type: \"string\"}\n",
+        "temperature = 0.50  # @param {type: \"number\"}\n",
+        "max_tokens = 250  # @param {type: \"number\"}\n",
+        "\n",
+        "request = {\n",
+        "    \"max_tokens\": 250 if max_tokens is None else max_tokens,\n",
+        "    \"temperature\": 0.5 if temperature is None else temperature,\n",
+        "}\n",
+        "\n",
+        "if is_vllm_inference:\n",
+        "    request[\"prompt\"] = user_prompt\n",
+        "else:\n",
+        "    request[\"inputs\"] = user_prompt\n",
+        "\n",
+        "model_service = SELECTED_DEPLOYMENT + \"-service\"\n",
+        "output = !kubectl get endpoints {model_service}\n",
+        "pod_endpoint = output[1].split()[1]\n",
+        "\n",
+        "# @markdown ### Response:\n",
+        "response = process_response(request, deployment_pod, pod_endpoint, is_vllm_inference)\n",
+        "HTML(\n",
+        "    '<div style=\"overflow-x: auto; font-size: 16px; line-height:'\n",
+        "    f' 1.8;\">{response}</div>'\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "id": "5b6ZM2K3fux0",
+      "metadata": {
+        "id": "5b6ZM2K3fux0"
+      },
+      "source": [
+        "# Next Steps: Integrating the GKE Service Endpoint\n",
+        "\n",
+        "After successfully deploying a model on Google Kubernetes Engine (GKE) and verifying it via a notebook, the next step is to integrate it into various applications. This involves making HTTP requests to the service's endpoint from your application code.\n",
+        "\n",
+        "### Exposing the Service\n",
+        "\n",
+        "To make your deployed model accessible to applications, you'll need to expose its service endpoint. Google Kubernetes Engine offers several ways to do this:\n",
+        "\n",
+        "1.  **Ingress:** Configure an Ingress resource to route external HTTP(S) traffic to your service. Set up Ingress for either an internal Load Balancer (accessible only within your VPC) or an external Load Balancer (accessible from the internet). [Learn more about GKE Ingress](https://cloud.google.com/kubernetes-engine/docs/concepts/ingress).\n",
+        "2.  **Gateway API:** A more modern and feature-rich API for managing traffic routing in Kubernetes. Similar to Ingress, Gateway API allows you to define how external and internal traffic should be directed to your services. [Explore GKE Gateway API](https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api).\n",
+        "\n",
+        "### Setting Up Autoscaling\n",
+        "\n",
+        "Ensure your model serving can handle varying traffic by configuring the Horizontal Pod Autoscaler (HPA). HPA automatically scales the number of Pods based on resource utilization or custom metrics, optimizing performance and cost. [See how to configure HPA](https://cloud.google.com/kubernetes-engine/docs/how-to/horizontal-pod-autoscaling).\n",
+        "\n",
+        "### Setting Up Monitoring\n",
+        "\n",
+        "Monitor the health and performance of your deployed model using Google Cloud Managed Service for Prometheus. Configure your model serving to expose Prometheus metrics for comprehensive insights. [Get started with Google Cloud Managed Prometheus](https://cloud.google.com/kubernetes-engine/docs/how-to/configure-automatic-application-monitoring).\n",
+        "\n",
+        "### Additional Resources:\n",
+        "\n",
+        "* #### Kubernetes Documentation:\n",
+        "   * Services: https://kubernetes.io/docs/concepts/services-networking/service/\n",
+        "\n",
+        "* #### Google Cloud Documentation:\n",
+        "   * Google Kubernetes Engine (GKE): https://cloud.google.com/kubernetes-engine\n",
+        "   * Cloud Load Balancing: https://cloud.google.com/load-balancing/docs/ingress\n",
+        "   * Gateway API on GKE: https://cloud.google.com/kubernetes-engine/docs/concepts/gateway-api\n",
+        "   * Learn about GPUs in GKE: https://cloud.google.com/kubernetes-engine/docs/concepts/gpus\n",
+        "\n",
+        "* #### Python requests Library:\n",
+        "   * https://requests.readthedocs.io/en/latest/\n",
+        "\n",
+        "* #### LangChain with Google Integrations:\n",
+        "   * The Langchain documentation is very useful: https://python.langchain.com/docs/integrations/providers/google/"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "gke_model_ui_deployment_notebook.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_advanced_features.ipynb

@@ -114,26 +114,19 @@
         "The quota for A100_80GB deployment `Custom model serving Nvidia A100 80GB GPUs per region` is 0. You need to request at least 4 for 70B model and 1 for 8B model following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota)."
       ]
     },
-    {
-      "cell_type": "markdown",
-      "id": "bdbcf9e7",
-      "metadata": {
-        "id": "jAELn17_1GfV"
-      },
-      "source": [
-        "### Setup Google Cloud project"
-      ]
-    },
     {
       "cell_type": "code",
       "execution_count": null,
       "id": "L3dqbxovo5t6",
+      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "L3dqbxovo5t6"
       },
       "outputs": [],
       "source": [
+        "# @title Setup Google Cloud Project\n",
+        "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
         "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
@@ -144,14 +137,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Install and import the necessary packages\n",
         "! pip install -q openai google-auth requests\n",

notebooks/community/model_garden/model_garden_codegemma_deployment_on_vertex.ipynb

@@ -131,6 +131,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -322,7 +328,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -416,6 +422,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -586,6 +593,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -638,6 +646,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -676,6 +688,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_codegemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",

notebooks/community/model_garden/model_garden_deployment_tutorial.ipynb

@@ -0,0 +1,424 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GyJPhQ2_Om3X"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "L1BRnURwUb87"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Deployment Tutorial\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_deployment_tutorial.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_deployment_tutorial.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "K4W8U0MQUb87"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "You can deploy open models (including Hugging Face models) by using [Google Gen AI SDK or Google Cloud CLI](https://cloud.google.com/vertex-ai/generative-ai/docs/model-garden/use-models#deploy_a_model).\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "x7Z25I3xUb87"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "WnEYdVEKUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"my-endpoint\"\n",
+        "\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2NRInIKZ3nSJ"
+      },
+      "source": [
+        "## Deploy"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xpIv3skP3nSK"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Choose the model to deploy\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "# @markdown List all deployable models and then get the ID of the model to deploy.\n",
+        "\n",
+        "# @markdown You can also use the Hugging Face model ID.\n",
+        "list_hf_models = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown You can also filter by model name.\n",
+        "model_filter = \"gemma\"  # @param {type:\"string\"}\n",
+        "\n",
+        "model_garden.list_deployable_models(\n",
+        "    list_hf_models=list_hf_models, model_filter=model_filter\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "uaLkN2mta5T2"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Select a model from the list above.\n",
+        "MODEL_ID = \"google/gemma3@gemma-3-1b-it\"  # @param [\"google/gemma3@gemma-3-1b-it\", \"google/gemma-3-1b-it\"] {isTemplate:true}\n",
+        "\n",
+        "# @markdown Follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "fR9wl0-Na5T2"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Fw9W4TrOUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rKjAOh7FUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"my-endpoint\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wxe8GIwpUb87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\n",
+        "        \"my-endpoint\"\n",
+        "    ].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"my-endpoint\"].name\n",
+        ")\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cqHoUEEMXNAH"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "5qKldg25Ub87"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_deployment_tutorial.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma2_deployment_on_vertex.ipynb

@@ -144,6 +144,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -285,7 +291,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -379,6 +385,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -604,6 +611,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_gemma2_finetuning_on_vertex.ipynb

@@ -10,7 +10,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -101,17 +101,11 @@
       "source": [
         "# @title Install Python Packages for Finetuning\n",
         "\n",
-        "# @markdown 1. Install google-cloud-aiplatform package and restart the session if instructed.\n",
-        "! pip install --upgrade --quiet 'google-cloud-aiplatform>=1.66.0'\n",
-        "\n",
-        "# @markdown 2. Install packages to validate dataset with template.\n",
+        "# @markdown 1. Install packages to validate dataset with template.\n",
         "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
-        "! pip install --upgrade --quiet accelerate==0.31.0\n",
-        "! pip install --upgrade --quiet transformers==4.43.1\n",
-        "! pip install --upgrade --quiet datasets==2.19.2\n",
-        "\n",
-        "# Load local tensorboard.\n",
-        "%load_ext tensorboard"
+        "! pip install --upgrade --quiet accelerate==0.34.2\n",
+        "! pip install --upgrade --quiet transformers==4.47.1\n",
+        "! pip install --upgrade --quiet datasets==2.20.0"
       ]
     },
     {
@@ -135,7 +129,9 @@
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-west1, europe-west4, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -147,7 +143,7 @@
         "\n",
         "# Import the necessary packages\n",
         "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "! cd vertex-ai-samples && git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
+        "! cd vertex-ai-samples && git reset --hard 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
         "\n",
         "import datetime\n",
         "import importlib\n",
@@ -170,6 +166,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -252,7 +254,7 @@
         "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
         "\n",
         "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
-        "# @markdown You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
+        "# @markdown You can set `train_dataset` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `train_column` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `train_column` to `text` in this notebook.\n",
         "\n",
         "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
         "\n",
@@ -261,13 +263,13 @@
         "# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
         "# @markdown ```\n",
         "\n",
-        "# @markdown The JSON object has a key `text`, which should match `instruct_column_in_dataset`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
+        "# @markdown The JSON object has a key `text`, which should match `train_column`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
         "\n",
         "# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
         "\n",
         "# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
         "\n",
-        "# @markdown Optionally update the `instruct_column_in_dataset` field below if your JSON objects use a key other than the default `text`.\n",
+        "# @markdown Optionally update the `train_column` field below if your JSON objects use a key other than the default `text`.\n",
         "\n",
         "# @markdown ### (Optional) Format your data with custom JSON template\n",
         "\n",
@@ -293,23 +295,25 @@
         "# @markdown\n",
         "# @markdown To try such custom dataset, you can make the following changes:\n",
         "# @markdown 1. Set `template` to `llama3-text-bison`\n",
-        "# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
-        "# @markdown 1. Set `train_split_name` to `train`\n",
-        "# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
-        "# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
-        "# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
+        "# @markdown 1. Set `train_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split` to `train`\n",
+        "# @markdown 1. Set `eval_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `train_column` as `input_text`.\n",
         "\n",
         "# Template name or gs:// URI to a custom template.\n",
         "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
         "\n",
         "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
-        "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
-        "train_split_name = \"train\"  # @param {type:\"string\"}\n",
-        "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
-        "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
+        "train_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split = \"test\"  # @param {type:\"string\"}\n",
         "\n",
         "# Name of the dataset column containing training text input.\n",
-        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
+        "train_column = \"text\"  # @param {type:\"string\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}"
       ]
     },
     {
@@ -341,8 +345,6 @@
         "\n",
         "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
         "\n",
-        "import transformers\n",
-        "\n",
         "dataset_validation_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
         ")\n",
@@ -358,32 +360,30 @@
         "    tokenizer_path = pretrained_model_id\n",
         "    access_token = HF_TOKEN\n",
         "\n",
-        "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
-        "    tokenizer_path,\n",
-        "    trust_remote_code=False,\n",
-        "    use_fast=True,\n",
-        "    token=access_token,\n",
-        ")\n",
+        "tokenizer = dataset_validation_util.load_tokenizer(tokenizer_path, None, access_token)\n",
         "\n",
         "# Validate the train dataset.\n",
         "dataset_validation_util.validate_dataset_with_template(\n",
-        "    dataset_name=train_dataset_name,\n",
-        "    split=train_split_name,\n",
-        "    input_column=instruct_column_in_dataset,\n",
+        "    dataset_name=train_dataset,\n",
+        "    split=train_split,\n",
+        "    input_column=train_column,\n",
         "    template=template,\n",
+        "    max_seq_length=max_seq_length,\n",
         "    use_multiprocessing=False,\n",
         "    tokenizer=tokenizer,\n",
         ")\n",
         "\n",
-        "# Validate the eval dataset.\n",
-        "dataset_validation_util.validate_dataset_with_template(\n",
-        "    dataset_name=eval_dataset_name,\n",
-        "    split=eval_split_name,\n",
-        "    input_column=instruct_column_in_dataset,\n",
-        "    template=template,\n",
-        "    use_multiprocessing=False,\n",
-        "    tokenizer=tokenizer,\n",
-        ")"
+        "# Validate the eval dataset if it exists.\n",
+        "if eval_dataset:\n",
+        "    dataset_validation_util.validate_dataset_with_template(\n",
+        "        dataset_name=eval_dataset,\n",
+        "        split=eval_split,\n",
+        "        input_column=train_column,\n",
+        "        template=template,\n",
+        "        max_seq_length=max_seq_length,\n",
+        "        use_multiprocessing=False,\n",
+        "        tokenizer=tokenizer,\n",
+        "    )"
       ]
     },
     {
@@ -407,10 +407,12 @@
         "# @markdown 1. If `max_steps > 0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
         "\n",
         "# @markdown Accelerator type to use for training.\n",
-        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "# fmt: off\n",
+        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "# fmt: on\n",
         "\n",
         "# The pre-built training docker image.\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
         "    repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
         "    is_restricted_image = True\n",
         "    is_dynamic_workload_scheduler = False\n",
@@ -425,30 +427,28 @@
         "    }\n",
         "\n",
         "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20240909\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250213\"\n",
         ")\n",
         "\n",
         "# Worker pool spec.\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
         "    per_node_accelerator_count = 8\n",
-        "    machine_type = \"a2-ultragpu-8g\"\n",
-        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    training_machine_type = \"a2-ultragpu-8g\"\n",
+        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
         "    per_node_accelerator_count = 8\n",
-        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    training_machine_type = \"a3-highgpu-8g\"\n",
         "else:\n",
         "    raise ValueError(\n",
-        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `per_node_accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
+        "        f\"Recommended machine settings not found for: {training_accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` by clicking `Show Code` and then modifying the code.\"\n",
         "    )\n",
         "\n",
         "# @markdown Batch size for finetuning.\n",
         "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
         "# @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
         "gradient_accumulation_steps = 4  # @param{type:\"integer\"}\n",
-        "# @markdown Maximum sequence length.\n",
-        "max_seq_length = 4096  # @param{type:\"integer\"}\n",
-        "# @markdown Setting a positive `max_steps` here will override `num_epochs`.\n",
+        "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
         "max_steps = -1  # @param{type:\"integer\"}\n",
-        "num_epochs = 1.0  # @param{type:\"number\"}\n",
+        "num_train_epochs = 1.0  # @param{type:\"number\"}\n",
         "# @markdown Precision mode for finetuning.\n",
         "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
         "# @markdown Learning rate.\n",
@@ -460,7 +460,7 @@
         "lora_alpha = 32  # @param{type:\"integer\"}\n",
         "lora_dropout = 0.05  # @param{type:\"number\"}\n",
         "# Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
-        "enable_gradient_checkpointing = True\n",
+        "gradient_checkpointing = True\n",
         "# Attention implementation to use in the model.\n",
         "attn_implementation = \"eager\"\n",
         "# The optimizer for which to schedule the learning rate.\n",
@@ -476,12 +476,17 @@
         "# Train precision of the model.\n",
         "train_precision = \"bfloat16\"\n",
         "\n",
+        "# @markdown Evaluation metrics to compute. Supported eval metrics: loss, perplexity, bleu, google_bleu, rouge1, rouge2, rougeL, rougeLsum.\n",
+        "eval_metric_name = \"loss,perplexity,bleu\"  # @param{type:\"string\"}\n",
+        "# @markdown Metric to use for best model selection. This will save the best checkpoint based on the eval metric.\n",
+        "metric_for_best_model = \"perplexity\"  # @param{type:\"string\"}\n",
+        "\n",
         "replica_count = 1\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count * replica_count,\n",
         "    is_for_training=True,\n",
         "    is_restricted_image=is_restricted_image,\n",
@@ -508,23 +513,23 @@
         "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
         "\n",
         "eval_args = [\n",
-        "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
-        "    f\"--eval_column={instruct_column_in_dataset}\",\n",
+        "    f\"--eval_dataset={eval_dataset}\",\n",
+        "    f\"--eval_column={train_column}\",\n",
         "    f\"--eval_template={template}\",\n",
-        "    f\"--eval_split={eval_split_name}\",\n",
+        "    f\"--eval_split={eval_split}\",\n",
         "    f\"--eval_steps={save_steps}\",\n",
-        "    \"--eval_tasks=builtin_eval\",\n",
-        "    \"--eval_metric_name=loss\",\n",
+        "    f\"--eval_metric_name={eval_metric_name}\",\n",
+        "    f\"--metric_for_best_model={metric_for_best_model}\",\n",
         "]\n",
         "\n",
         "train_job_args = [\n",
         "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml\",\n",
         "    \"--task=instruct-lora\",\n",
-        "    \"--completion_only=True\",\n",
-        "    f\"--pretrained_model_id={pretrained_model_id}\",\n",
-        "    f\"--dataset_name={train_dataset_name}\",\n",
-        "    f\"--train_split_name={train_split_name}\",\n",
-        "    f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
+        "    \"--input_masking=True\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
+        "    f\"--train_dataset={train_dataset}\",\n",
+        "    f\"--train_split={train_split}\",\n",
+        "    f\"--train_column={train_column}\",\n",
         "    f\"--output_dir={lora_output_dir}\",\n",
         "    f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
         "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
@@ -538,8 +543,9 @@
         "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
         "    f\"--precision_mode={finetuning_precision_mode}\",\n",
         "    f\"--train_precision={train_precision}\",\n",
-        "    f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
-        "    f\"--num_epochs={num_epochs}\",\n",
+        "    f\"--merge_model_precision_mode={train_precision}\",\n",
+        "    f\"--gradient_checkpointing={gradient_checkpointing}\",\n",
+        "    f\"--num_train_epochs={num_train_epochs}\",\n",
         "    f\"--attn_implementation={attn_implementation}\",\n",
         "    f\"--optimizer={optimizer}\",\n",
         "    f\"--warmup_ratio={warmup_ratio}\",\n",
@@ -547,7 +553,7 @@
         "    f\"--logging_output_dir={base_output_dir}\",\n",
         "    f\"--save_steps={save_steps}\",\n",
         "    f\"--logging_steps={logging_steps}\",\n",
-        "    f\"--template={template}\",\n",
+        "    f\"--train_template={template}\",\n",
         "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
         "] + eval_args\n",
         "\n",
@@ -563,8 +569,8 @@
         "train_job.run(\n",
         "    args=train_job_args,\n",
         "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    machine_type=training_machine_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count,\n",
         "    boot_disk_size_gb=500,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
@@ -576,6 +582,8 @@
         "# Wait until resource has been created.\n",
         "train_job.wait_for_resource_creation()\n",
         "\n",
+        "merged_model_output_dir = os.path.join(merged_model_output_dir, \"node-0\")\n",
+        "\n",
         "print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
         "print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
         "\n",
@@ -623,41 +631,44 @@
         "print(\"Deploying models in:\", merged_model_output_dir)\n",
         "\n",
         "# The pre-built serving docker image for vLLM.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240815_1634_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
         "# @markdown Accelerator type to use for serving.\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "serving_accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
         "\n",
         "if \"2b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 1 L4 (24G) to deploy Gemma 2 2B models.\n",
-        "        machine_type = \"g2-standard-12\"\n",
-        "        accelerator_count = 1\n",
+        "        serving_machine_type = \"g2-standard-12\"\n",
+        "        serving_accelerator_count = 1\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "elif \"9b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 2 L4 (24G) to deploy Gemma 2 9B models.\n",
-        "        machine_type = \"g2-standard-24\"\n",
-        "        accelerator_count = 2\n",
+        "        serving_machine_type = \"g2-standard-24\"\n",
+        "        serving_accelerator_count = 2\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "elif \"27b\" in base_model_id:\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "    if serving_accelerator_type == \"NVIDIA_L4\":\n",
         "        # Sets 4 L4 (24G) to deploy Gemma 2 27B models.\n",
-        "        machine_type = \"g2-standard-48\"\n",
-        "        accelerator_count = 4\n",
+        "        serving_machine_type = \"g2-standard-48\"\n",
+        "        serving_accelerator_count = 4\n",
         "    else:\n",
         "        raise ValueError(\n",
         "            \"Recommended machine settings not found for accelerator type: %s\"\n",
-        "            % accelerator_type\n",
+        "            % serving_accelerator_type\n",
         "        )\n",
         "else:\n",
         "    raise ValueError(\n",
@@ -667,21 +678,32 @@
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
         "    is_for_training=False,\n",
         ")\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to True if the endpoint is [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) enabled.\n",
-        "use_dedicated_endpoint = False  # @param {type:\"boolean\"}\n",
-        "\n",
         "gpu_memory_utilization = 0.85\n",
         "max_model_len = 4096  # Maximum context length.\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
         "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
@@ -693,11 +715,15 @@
         "    enable_trust_remote_code: bool = False,\n",
         "    enforce_eager: bool = False,\n",
         "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
         "    max_loras: int = 1,\n",
         "    max_cpu_loras: int = 8,\n",
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -736,9 +762,24 @@
         "    if enable_lora:\n",
         "        vllm_args.append(\"--enable-lora\")\n",
         "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -761,6 +802,9 @@
         "        serving_container_environment_variables=env_vars,\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
         "    )\n",
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
@@ -772,6 +816,10 @@
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma2_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
+        "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
@@ -781,10 +829,12 @@
         "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma2-vllm-serve\"),\n",
         "    model_id=merged_model_output_dir,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma2\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
+        "    machine_type=serving_machine_type,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
         "    gpu_memory_utilization=gpu_memory_utilization,\n",
         "    max_model_len=max_model_len,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",

notebooks/community/model_garden/model_garden_gemma3_deployment_on_vertex.ipynb

@@ -0,0 +1,963 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Gemma 3 (Deployment)\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma3_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma3_deployment_on_vertex.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying Gemma 3 models on GPU using [vLLM](https://github.com/vllm-project/vllm).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy Gemma 3 with vLLM on GPU\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2707b02ef5df"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"vllm_gpu\"\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "GFfqpQm8BNwZ"
+      },
+      "source": [
+        "## Deploy Gemma 3 1B models with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "W1PYsnz2JXHm"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"gemma-3-1b-it\"  # @param [\"gemma-3-1b-pt\", \"gemma-3-1b-it\"] {isTemplate:true}\n",
+        "hf_model_id = \"google/\" + base_model_name\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/gemma3@{base_model_name}\"\n",
+        "model_id = f\"gs://vertex-model-garden-restricted-us/gemma3/{base_model_name}\"\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "machine_type = \"g2-standard-12\"\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "d2SO5S5fMKE2"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TBNJYZMlBNwZ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Gemma 3 1B models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes 15 minutes to 30 minutes to finish.\n",
+        "\n",
+        "gpu_memory_utilization = 0.95\n",
+        "max_model_len = 32768\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma3_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma3\",\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rDHsCOqvFYBi"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LSG9ITWTbTb7"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        ")\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "_0IYHNsYJO55"
+      },
+      "source": [
+        "## Deploy Gemma 3 4B, 12B and 27B multimodal models with vLLM on GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zpKbaFK_Jeny"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"gemma-3-4b-it\"  # @param [\"gemma-3-4b-pt\", \"gemma-3-4b-it\", \"gemma-3-12b-pt\", \"gemma-3-12b-it\", \"gemma-3-27b-pt\", \"gemma-3-27b-it\"] {isTemplate:true}\n",
+        "hf_model_id = \"google/\" + base_model_name\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/google/models/gemma3@{base_model_name}\"\n",
+        "model_id = f\"gs://vertex-model-garden-restricted-us/gemma3/{base_model_name}\"\n",
+        "\n",
+        "# The pre-built serving docker image.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "if \"4b\" in model_id or \"12b\" in model_id:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "    accelerator_count = 2\n",
+        "elif \"27b\" in model_id:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-4g\"\n",
+        "    accelerator_count = 4\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for model: {base_model_name}.\"\n",
+        "    )\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "CA_dsYyjLy1b"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "plirtyNxJO55"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Gemma 3 multimodal models to Model Registry and deploys them to a Vertex Prediction Endpoint. It takes 15 minutes to 1 hour to finish.\n",
+        "\n",
+        "gpu_memory_utilization = 0.95\n",
+        "max_model_len = 131072\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma3_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma3-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma3\",\n",
+        "    base_model_id=hf_model_id,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "eMBKLoUPJO55"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion with text-only requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        ")\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "b3c2f6559a67"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Chat completion with multimodal requests\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        ")\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_image = \"https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Blue_Marble_%28remastered%29.jpg/580px-The_Blue_Marble_%28remastered%29.jpg\"  # @param {type: \"string\"}\n",
+        "user_message = \"What is in the image?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[\n",
+        "        {\n",
+        "            \"role\": \"user\",\n",
+        "            \"content\": [\n",
+        "                {\"type\": \"image_url\", \"image_url\": {\"url\": user_image}},\n",
+        "                {\"type\": \"text\", \"text\": user_message},\n",
+        "            ],\n",
+        "        }\n",
+        "    ],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        ")\n",
+        "print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma3_deployment_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma3_finetuning_on_vertex.ipynb

@@ -0,0 +1,925 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "7d9bbf86da5e"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "language": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Gemma 3 Finetuning\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_gemma3_finetuning_on_vertex.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_gemma3_finetuning_on_vertex.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates finetuning and deploying Gemma 3 models with [Vertex AI Custom Training Job](https://cloud.google.com/vertex-ai/docs/training/create-custom-job). All of the examples in this notebook use parameter efficient finetuning methods [PEFT (LoRA)](https://github.com/huggingface/peft) to reduce training and storage costs. LoRA (Low-Rank Adaptation) is one approach of Parameter Efficient FineTuning (PEFT), where pretrained model weights are frozen and rank decomposition matrices representing the change in model weights are trained during finetuning. Read more about LoRA in the following publication: [Hu, E.J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang, S., Wang, L. and Chen, W., 2021. Lora: Low-rank adaptation of large language models. *arXiv preprint arXiv:2106.09685*](https://arxiv.org/abs/2106.09685).\n",
+        "\n",
+        "\n",
+        "After tuning, we can deploy models on Vertex with GPU.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Finetune and deploy Gemma 3 models with Vertex AI Custom Training Jobs.\n",
+        "- Send prediction requests to your finetuned Gemma 3 model.\n",
+        "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "aS52SK74gDoB"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Install Python Packages for Finetuning\n",
+        "\n",
+        "# @markdown 1. Install packages to validate dataset with template.\n",
+        "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
+        "! pip install --upgrade --quiet accelerate==0.34.2\n",
+        "! pip install --upgrade --quiet transformers==4.47.1\n",
+        "! pip install --upgrade --quiet datasets==2.20.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "855d6b96f291"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. For finetuning, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Frestricted_image_training_nvidia_a100_80gb_gpus)** to check if your project already has the required 8 Nvidia A100 80 GB GPUs in the us-central1 region. If yes, then run this notebook in the us-central1 region. If you do not have 8 Nvidia A100 80 GPUs or have more GPU requirements than this, then schedule your job with Nvidia H100 GPUs via Dynamic Workload Scheduler using [these instructions](https://cloud.google.com/vertex-ai/docs/training/schedule-jobs-dws). For Dynamic Workload Scheduler, check the [us-central1](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) or [europe-west4](https://console.cloud.google.com/iam-admin/quotas?location=europe-west4&metric=aiplatform.googleapis.com%2Fcustom_model_training_preemptible_nvidia_h100_gpus) quota for Nvidia H100 GPUs. If you do not have enough GPUs, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request quota.\n",
+        "\n",
+        "# @markdown 3. For serving, **[click here](https://console.cloud.google.com/iam-admin/quotas?location=us-central1&metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_l4_gpus)** to check if your project already has the required 1 L4 GPU in the us-central1 region.  If yes, then run this notebook in the us-central1 region. If you need more L4 GPUs for your project, then you can follow [these instructions](https://cloud.google.com/docs/quotas/view-manage#viewing_your_quota_console) to request more. Alternatively, if you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "\n",
+        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 5. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! cd vertex-ai-samples && git reset --hard 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
+        "\n",
+        "import datetime\n",
+        "import importlib\n",
+        "import os\n",
+        "import uuid\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "from google.cloud.aiplatform.compat.types import \\\n",
+        "    custom_job as gca_custom_job_compat\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Cloud Storage bucket for storing the experiment artifacts.\n",
+        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
+        "# prefer using your own GCS bucket, change the value yourself below.\n",
+        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
+        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "\n",
+        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
+        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
+        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
+        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
+        "else:\n",
+        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
+        "    bucket_region = shell_output[0].strip().lower()\n",
+        "    if bucket_region != REGION:\n",
+        "        raise ValueError(\n",
+        "            \"Bucket region %s is different from notebook region %s\"\n",
+        "            % (bucket_region, REGION)\n",
+        "        )\n",
+        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
+        "\n",
+        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
+        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"gemma3\")\n",
+        "\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "\n",
+        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
+        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
+        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "\n",
+        "# @markdown ## Access Gemma 3 Models\n",
+        "\n",
+        "# @markdown You must provide a Hugging Face User Access Token (read) to access the Gemma 3 models. You can follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate:true}\n",
+        "assert HF_TOKEN, \"Provide a read HF_TOKEN to load models from Hugging Face.\"\n",
+        "\n",
+        "model_path_prefix = \"google/\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cb56d402e84a"
+      },
+      "source": [
+        "## Finetune with HuggingFace PEFT and Deploy with vLLM on GPUs"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "KwAW99YZHTdy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set dataset\n",
+        "\n",
+        "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
+        "# @markdown You can set `train_dataset` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `train_column` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `train_column` to `text` in this notebook.\n",
+        "\n",
+        "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
+        "\n",
+        "# @markdown You can prepare a JSONL file where each line is a valid JSON string as your custom training dataset. For example, here is one line from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset:\n",
+        "# @markdown ```\n",
+        "# @markdown {\"text\": \"### Human: Hola### Assistant: \\u00a1Hola! \\u00bfEn qu\\u00e9 puedo ayudarte hoy?\"}\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown The JSON object has a key `text`, which should match `train_column`; The value should be one training data point, i.e. a string. After you prepared your JSONL file, you can either upload it to [Hugging Face datasets](https://huggingface.co/datasets) or [Google Cloud Storage](https://cloud.google.com/storage).\n",
+        "\n",
+        "# @markdown - To upload a JSONL dataset to [Hugging Face datasets](https://huggingface.co/datasets), follow the instructions on [Uploading Datasets](https://huggingface.co/docs/hub/en/datasets-adding). Then, set `dataset_name` to the name of your newly created dataset on Hugging Face.\n",
+        "\n",
+        "# @markdown - To upload a JSONL dataset to [Google Cloud Storage](https://cloud.google.com/storage), follow the instructions on [Upload objects from a filesystem](https://cloud.google.com/storage/docs/uploading-objects). Then, set `dataset_name` to the `gs://` URI to your JSONL file. For example: `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`.\n",
+        "\n",
+        "# @markdown Optionally update the `train_column` field below if your JSON objects use a key other than the default `text`.\n",
+        "\n",
+        "# @markdown ### (Optional) Format your data with custom JSON template\n",
+        "\n",
+        "# @markdown Sometimes, your dataset might have multiple text columns and you want to construct the training data with a template. You can prepare a JSON template in the following format:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown {\n",
+        "# @markdown   \"description\": \"Template that accepts text-bison format.\",\n",
+        "# @markdown   \"source\": \"https://cloud.google.com/vertex-ai/generative-ai/docs/models/tune-text-models-supervised#dataset-format\",\n",
+        "# @markdown   \"prompt_input\": \"\\n\\n<|start_header_id|>user<|end_header_id|>\\n\\n{input_text}<|eot_id|>\\n\\n<|start_header_id|>assistant<|end_header_id|>\\n\\n{output_text}<|eot_id|>\",\n",
+        "# @markdown   \"instruction_separator\": \"<|start_header_id|>user<|end_header_id|>\\n\\n\",\n",
+        "# @markdown   \"response_separator\": \"<|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
+        "# @markdown }\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown As an example, the template above can be used to format the following training data (this line comes from `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`):\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown {\"input_text\":\"TRANSCRIPT: \\nREASON FOR EVALUATION:,\\n\\n LABEL:\",\"output_text\":\"Chiropractic\"}\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown This example template simply concatenates `input_text` with `output_text` with some special tokens in between.\n",
+        "# @markdown\n",
+        "# @markdown To try such custom dataset, you can make the following changes:\n",
+        "# @markdown 1. Set `template` to `llama3-text-bison`\n",
+        "# @markdown 1. Set `train_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split` to `train`\n",
+        "# @markdown 1. Set `eval_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `train_column` as `input_text`.\n",
+        "\n",
+        "# Template name or gs:// URI to a custom template.\n",
+        "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
+        "train_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split = \"test\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# Name of the dataset column containing training text input.\n",
+        "train_column = \"text\"  # @param {type:\"string\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "NoueJWi72OSo"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set model\n",
+        "\n",
+        "# @markdown Select a model variant of Gemma 3.\n",
+        "base_model_id = \"gemma-3-1b-pt\"  # @param [\"gemma-3-1b-pt\", \"gemma-3-1b-it\"] {isTemplate: true}\n",
+        "pretrained_model_id = os.path.join(model_path_prefix, base_model_id)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-NlLSiCOvru1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Validate Dataset with Template\n",
+        "\n",
+        "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
+        "\n",
+        "dataset_validation_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
+        ")\n",
+        "\n",
+        "if dataset_validation_util.is_gcs_path(pretrained_model_id):\n",
+        "    # Download tokenizer.\n",
+        "    ! mkdir tokenizer\n",
+        "    ! gsutil cp {pretrained_model_id}/tokenizer.json ./tokenizer\n",
+        "    ! gsutil cp {pretrained_model_id}/config.json ./tokenizer\n",
+        "    tokenizer_path = \"./tokenizer\"\n",
+        "    access_token = \"\"\n",
+        "else:\n",
+        "    tokenizer_path = pretrained_model_id\n",
+        "    access_token = HF_TOKEN\n",
+        "\n",
+        "tokenizer = dataset_validation_util.load_tokenizer(tokenizer_path, None, access_token)\n",
+        "\n",
+        "# Validate the train dataset.\n",
+        "dataset_validation_util.validate_dataset_with_template(\n",
+        "    dataset_name=train_dataset,\n",
+        "    split=train_split,\n",
+        "    input_column=train_column,\n",
+        "    template=template,\n",
+        "    max_seq_length=max_seq_length,\n",
+        "    use_multiprocessing=False,\n",
+        "    tokenizer=tokenizer,\n",
+        ")\n",
+        "\n",
+        "# Validate the eval dataset if it exists.\n",
+        "if eval_dataset:\n",
+        "    dataset_validation_util.validate_dataset_with_template(\n",
+        "        dataset_name=eval_dataset,\n",
+        "        split=eval_split,\n",
+        "        input_column=train_column,\n",
+        "        template=template,\n",
+        "        max_seq_length=max_seq_length,\n",
+        "        use_multiprocessing=False,\n",
+        "        tokenizer=tokenizer,\n",
+        "    )"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "ivVGS9dHXPOz"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Finetune\n",
+        "\n",
+        "# @markdown This section demonstrates how to finetune the Gemma 3 text only model and merge the finetuned LoRA adapter with the base model on Vertex AI. It uses the Vertex AI SDK to create and run the custom training jobs.\n",
+        "\n",
+        "# @markdown The training job takes approximately between 10 to 20 mins to set-up. Once done, the training job is expected to take around 20 mins with the default configuration. To find the training time, throughput, and memory usage of your training job, you can go to the training logs and check the log line of the last training epoch.\n",
+        "\n",
+        "# @markdown **Note**:\n",
+        "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
+        "# @markdown 1. If `max_steps > 0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
+        "\n",
+        "# @markdown Accelerator type to use for training.\n",
+        "# fmt: off\n",
+        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "# fmt: on\n",
+        "\n",
+        "# The pre-built training docker image.\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
+        "    is_restricted_image = True\n",
+        "    is_dynamic_workload_scheduler = False\n",
+        "    dws_kwargs = {}\n",
+        "else:\n",
+        "    repo = \"us-docker.pkg.dev/vertex-ai\"\n",
+        "    is_restricted_image = False\n",
+        "    is_dynamic_workload_scheduler = True\n",
+        "    dws_kwargs = {\n",
+        "        \"max_wait_duration\": 1800,  # 30 minutes\n",
+        "        \"scheduling_strategy\": gca_custom_job_compat.Scheduling.Strategy.FLEX_START,\n",
+        "    }\n",
+        "\n",
+        "TRAIN_DOCKER_URI = (\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:gemma3_20250312\"\n",
+        ")\n",
+        "\n",
+        "# Worker pool spec.\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    per_node_accelerator_count = 8\n",
+        "    training_machine_type = \"a2-ultragpu-8g\"\n",
+        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    per_node_accelerator_count = 8\n",
+        "    training_machine_type = \"a3-highgpu-8g\"\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended machine settings not found for: {training_accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` by clicking `Show Code` and then modifying the code.\"\n",
+        "    )\n",
+        "\n",
+        "# @markdown Batch size for finetuning.\n",
+        "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
+        "# @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
+        "gradient_accumulation_steps = 4  # @param{type:\"integer\"}\n",
+        "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
+        "max_steps = -1  # @param{type:\"integer\"}\n",
+        "num_train_epochs = 1.0  # @param{type:\"number\"}\n",
+        "# @markdown Precision mode for finetuning.\n",
+        "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
+        "# @markdown Learning rate.\n",
+        "learning_rate = 5e-5  # @param{type:\"number\"}\n",
+        "# @markdown The scheduler type to use.\n",
+        "lr_scheduler_type = \"cosine\"  # @param{type:\"string\"}\n",
+        "# @markdown LoRA parameters.\n",
+        "lora_rank = 16  # @param{type:\"integer\"}\n",
+        "lora_alpha = 32  # @param{type:\"integer\"}\n",
+        "lora_dropout = 0.05  # @param{type:\"number\"}\n",
+        "# Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
+        "gradient_checkpointing = True\n",
+        "# Attention implementation to use in the model.\n",
+        "attn_implementation = \"eager\"\n",
+        "# The optimizer for which to schedule the learning rate.\n",
+        "optimizer = \"paged_adamw_32bit\"\n",
+        "# Define the proportion of training to be dedicated to a linear warmup where learning rate gradually increases.\n",
+        "warmup_ratio = \"0.01\"\n",
+        "# The list or string of integrations to report the results and logs to.\n",
+        "report_to = \"tensorboard\"\n",
+        "# Number of updates steps before two checkpoint saves.\n",
+        "save_steps = 10\n",
+        "# Number of update steps between two logs.\n",
+        "logging_steps = save_steps\n",
+        "# Train precision of the model.\n",
+        "train_precision = \"bfloat16\"\n",
+        "\n",
+        "# @markdown Evaluation metrics to compute. Supported eval metrics: loss, perplexity, bleu, google_bleu, rouge1, rouge2, rougeL, rougeLsum.\n",
+        "eval_metric_name = \"loss,perplexity,bleu\"  # @param{type:\"string\"}\n",
+        "# @markdown Metric to use for best model selection. This will save the best checkpoint based on the eval metric.\n",
+        "metric_for_best_model = \"perplexity\"  # @param{type:\"string\"}\n",
+        "\n",
+        "replica_count = 1\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=training_accelerator_type,\n",
+        "    accelerator_count=per_node_accelerator_count * replica_count,\n",
+        "    is_for_training=True,\n",
+        "    is_restricted_image=is_restricted_image,\n",
+        "    is_dynamic_workload_scheduler=is_dynamic_workload_scheduler,\n",
+        ")\n",
+        "\n",
+        "job_name = common_util.get_job_name_with_datetime(\"gemma3-lora-train\")\n",
+        "\n",
+        "base_output_dir = os.path.join(STAGING_BUCKET, job_name)\n",
+        "# Create a GCS folder to store the LORA adapter.\n",
+        "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
+        "# Create a GCS folder to store the merged model with the base model and the\n",
+        "# finetuned LORA adapter.\n",
+        "merged_model_output_dir = os.path.join(base_output_dir, \"merged-model\")\n",
+        "\n",
+        "# Add labels for the finetuning job.\n",
+        "labels = {\n",
+        "    \"mg-source\": \"notebook\",\n",
+        "    \"mg-notebook-name\": \"model_garden_gemma3_finetuning_on_vertex.ipynb\".split(\".\")[0],\n",
+        "}\n",
+        "\n",
+        "labels[\"mg-tune\"] = \"publishers-google-models-gemma-3\"\n",
+        "versioned_model_id = base_model_id.lower().replace(\".\", \"-\")\n",
+        "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
+        "\n",
+        "eval_args = [\n",
+        "    f\"--eval_dataset={eval_dataset}\",\n",
+        "    f\"--eval_column={train_column}\",\n",
+        "    f\"--eval_template={template}\",\n",
+        "    f\"--eval_split={eval_split}\",\n",
+        "    f\"--eval_steps={save_steps}\",\n",
+        "    f\"--eval_metric_name={eval_metric_name}\",\n",
+        "    f\"--metric_for_best_model={metric_for_best_model}\",\n",
+        "]\n",
+        "\n",
+        "train_job_args = [\n",
+        "    \"--config_file=vertex_vision_model_garden_peft/deepspeed_zero2_8gpu.yaml\",\n",
+        "    \"--task=instruct-lora\",\n",
+        "    \"--input_masking=True\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
+        "    f\"--train_dataset={train_dataset}\",\n",
+        "    f\"--train_split={train_split}\",\n",
+        "    f\"--train_column={train_column}\",\n",
+        "    f\"--output_dir={lora_output_dir}\",\n",
+        "    f\"--merge_base_and_lora_output_dir={merged_model_output_dir}\",\n",
+        "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
+        "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
+        "    f\"--lora_rank={lora_rank}\",\n",
+        "    f\"--lora_alpha={lora_alpha}\",\n",
+        "    f\"--lora_dropout={lora_dropout}\",\n",
+        "    f\"--max_steps={max_steps}\",\n",
+        "    f\"--max_seq_length={max_seq_length}\",\n",
+        "    f\"--learning_rate={learning_rate}\",\n",
+        "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
+        "    f\"--precision_mode={finetuning_precision_mode}\",\n",
+        "    f\"--train_precision={train_precision}\",\n",
+        "    f\"--gradient_checkpointing={gradient_checkpointing}\",\n",
+        "    f\"--num_train_epochs={num_train_epochs}\",\n",
+        "    f\"--attn_implementation={attn_implementation}\",\n",
+        "    f\"--optimizer={optimizer}\",\n",
+        "    f\"--warmup_ratio={warmup_ratio}\",\n",
+        "    f\"--report_to={report_to}\",\n",
+        "    f\"--logging_output_dir={base_output_dir}\",\n",
+        "    f\"--save_steps={save_steps}\",\n",
+        "    f\"--logging_steps={logging_steps}\",\n",
+        "    f\"--train_template={template}\",\n",
+        "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
+        "] + eval_args\n",
+        "\n",
+        "# Pass training arguments and launch job.\n",
+        "train_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=job_name,\n",
+        "    container_uri=TRAIN_DOCKER_URI,\n",
+        "    labels=labels,\n",
+        ")\n",
+        "\n",
+        "print(\"Running training job with args:\")\n",
+        "print(\" \\\\\\n\".join(train_job_args))\n",
+        "train_job.run(\n",
+        "    args=train_job_args,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=training_machine_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
+        "    accelerator_count=per_node_accelerator_count,\n",
+        "    boot_disk_size_gb=500,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    base_output_dir=base_output_dir,\n",
+        "    sync=False,  # Non-blocking call to run.\n",
+        "    **dws_kwargs,\n",
+        ")\n",
+        "\n",
+        "# Wait until resource has been created.\n",
+        "train_job.wait_for_resource_creation()\n",
+        "\n",
+        "print(\"LoRA adapter will be saved in:\", lora_output_dir)\n",
+        "print(\"Trained and merged models will be saved in:\", merged_model_output_dir)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lu-uxrFBmZ0s"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Run TensorBoard\n",
+        "# @markdown This section shows how to launch TensorBoard in a [Cloud Shell](https://cloud.google.com/shell/docs).\n",
+        "# @markdown 1. Click the Cloud Shell icon(![terminal](https://github.com/google/material-design-icons/blob/master/png/action/terminal/materialicons/24dp/1x/baseline_terminal_black_24dp.png?raw=true)) on the top right to open the Cloud Shell.\n",
+        "# @markdown 2. Copy the `tensorboard` command shown below by running this cell.\n",
+        "# @markdown 3. Paste and run the command in the Cloud Shell to launch TensorBoard.\n",
+        "# @markdown 4. Once the command runs (You may have to click `Authorize` if prompted), click the link starting with `http://localhost`.\n",
+        "\n",
+        "# @markdown Note: You may need to wait around 10 minutes after the job starts in order for the TensorBoard logs to be written to the GCS bucket.\n",
+        "print(f\"Command to copy: tensorboard --logdir {base_output_dir}/logs\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "qmHW6m8xG_4U"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes 15 minutes to 1 hour to finish.\n",
+        "\n",
+        "if train_job.end_time is None:\n",
+        "    print(\"Waiting for the training job to finish...\")\n",
+        "    train_job.wait()\n",
+        "    print(\"The training job has finished.\")\n",
+        "\n",
+        "print(\"Deploying models in:\", merged_model_output_dir)\n",
+        "\n",
+        "# The pre-built serving docker image for vLLM.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250312_0916_RC01\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
+        "# @markdown Accelerator type to use for serving.\n",
+        "serving_accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "serving_machine_type = \"g2-standard-12\"\n",
+        "serving_accelerator_count = 1\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "gpu_memory_utilization = 0.85\n",
+        "max_model_len = 4096  # Maximum context length.\n",
+        "\n",
+        "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_gemma3_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"gemma3-vllm-serve\"),\n",
+        "    model_id=merged_model_output_dir,\n",
+        "    publisher=\"google\",\n",
+        "    publisher_model_id=\"gemma3\",\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        "    machine_type=serving_machine_type,\n",
+        "    accelerator_type=serving_accelerator_type,\n",
+        "    accelerator_count=serving_accelerator_count,\n",
+        "    gpu_memory_utilization=gpu_memory_utilization,\n",
+        "    max_model_len=max_model_len,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2UYUNn60G_4U"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "af21a3cff1e0"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# Delete the train job.\n",
+        "train_job.delete()\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_gemma3_finetuning_on_vertex.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_gemma_deployment_on_vertex.ipynb

@@ -131,6 +131,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -310,7 +316,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -404,6 +410,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -524,6 +531,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -550,8 +558,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "REGION = _region\n",
         "\n",
@@ -682,6 +703,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -734,6 +756,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -772,6 +798,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_deployment_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -900,6 +927,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -926,8 +954,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]

notebooks/community/model_garden/model_garden_gemma_evaluation.ipynb

@@ -133,6 +133,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_gemma_finetuning_on_vertex.ipynb

@@ -133,7 +133,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -168,6 +168,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -680,6 +686,18 @@
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
@@ -704,6 +722,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -756,6 +775,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -794,6 +817,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -1225,7 +1249,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -1319,6 +1343,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gemma_finetuning_on_vertex.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_gradio_streaming_chat_completions.ipynb

@@ -208,6 +208,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -260,6 +261,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -298,6 +303,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gradio_streaming_chat_completions.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -361,6 +367,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_gradio_streaming_chat_completions.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_hf_paligemma2_deployment.ipynb

@@ -109,14 +109,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -153,6 +153,7 @@
         "\n",
         "\n",
         "def deploy_model(\n",
+        "    model_name: str = None,\n",
         "    model_id: str = None,\n",
         "    task: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
@@ -166,6 +167,7 @@
         "    \"\"\"Deploys a model to a real-time prediction endpoint.\n",
         "\n",
         "    Args:\n",
+        "        model_name: The base name of the model.\n",
         "        model_id: The model ID.\n",
         "        task: The task to perform.\n",
         "        machine_type: The machine type.\n",
@@ -181,7 +183,7 @@
         "    \"\"\"\n",
         "\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
-        "        display_name=common_util.get_job_name_with_datetime(prefix=\"paligemma-2\")\n",
+        "        display_name=common_util.get_job_name_with_datetime(prefix=model_name)\n",
         "    )\n",
         "    serving_env = {\n",
         "        \"MODEL_ID\": model_id,\n",
@@ -204,9 +206,7 @@
         "        accelerator_count=accelerator_count,\n",
         "        service_account=service_account,\n",
         "        sync=False,\n",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": (\"model_garden_hf_paligemma2_deployment_templated.ipynb\")\n",
-        "        },\n",
+        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_hf_paligemma2_deployment.ipynb\"},\n",
         "    )\n",
         "    return endpoint, model\n",
         "\n",
@@ -276,20 +276,28 @@
         "\n",
         "# @markdown Select the desired resolution and precision of prebuilt model to deploy, leaving the optional `custom_paligemma_model_uri` as is. Higher resolution and precision_type can result in better inference results, but may require additional GPU.\n",
         "\n",
+        "MODEL_NAME = \"paligemma2-3b-pt-224\"  # @param [\"paligemma2-3b-pt-224\", \"paligemma2-3b-mix-224\", \"paligemma2-3b-ft-docci-448\", \"paligemma2-3b-mix-448\", \"paligemma2-3b-pt-448\", \"paligemma2-3b-pt-896\", \"paligemma2-10b-mix-224\", \"paligemma2-10b-pt-224\", \"paligemma2-10b-ft-docci-448\", \"paligemma2-10b-mix-448\", \"paligemma2-10b-pt-448\", \"paligemma2-10b-pt-896\", \"paligemma2-28b-mix-224\", \"paligemma2-28b-pt-224\", \"paligemma2-28b-mix-448\", \"paligemma2-28b-pt-448\", \"paligemma2-28b-pt-896\"]\n",
+        "GCS_PREFIX = \"gs://vertex-model-garden-restricted-us/paligemma2\"\n",
         "\n",
-        "MODEL_ID = \"paligemma2-3b-pt-224\"  # @param [\"paligemma2-3b-pt-224\", \"paligemma2-3b-pt-448\", \"paligemma2-10b-ft-docci-448\"]\n",
-        "\n",
-        "GSC_PREFIX = \"gs://vertex-model-garden-public-us/paligemma2/\"\n",
-        "MODEL_ID = GSC_PREFIX + MODEL_ID\n",
-        "\n",
-        "TASK = \"paligemma_VQA\"  # @param [\"paligemma_VQA\"]\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"]\n",
-        "accelerator_count = 1  # @param [1]\n",
-        "machine_type = \"g2-standard-8\"  # @param  [\"g2-standard-8\"]\n",
+        "MODEL_ID = os.path.join(GCS_PREFIX, MODEL_NAME)\n",
         "\n",
         "\n",
         "# @markdown If you want to use other accelerator types not listed above, then check other Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute. You may need to manually set the `machine_type`, `accelerator_type`, and `accelerator_count` in the code by clicking `Show code` first.\n",
         "\n",
+        "if \"3b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_count = 1\n",
+        "elif \"10b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "    accelerator_count = 1\n",
+        "elif \"28b\" in MODEL_NAME:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    accelerator_count = 8\n",
+        "else:\n",
+        "    raise ValueError(f\"Recommended GPU setting not found for: {MODEL_NAME}.\")\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
@@ -299,7 +307,10 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "TASK = \"paligemma_VQA\"\n",
+        "\n",
         "endpoints[\"paligemma2\"], models[\"paligemma2\"] = deploy_model(\n",
+        "    model_name=MODEL_NAME,\n",
         "    model_id=MODEL_ID,\n",
         "    task=TASK,\n",
         "    machine_type=machine_type,\n",
@@ -364,7 +375,7 @@
         "image_url = \"https://images.pexels.com/photos/1006293/pexels-photo-1006293.jpeg\"  # @param {type:\"string\"}\n",
         "\n",
         "# @markdown You may leave question prompts empty and they will be ignored.\n",
-        "question_prompt = \"What animal is shown in the picture?\"  # @param {type: \"string\"}\n",
+        "question_prompt = \"What is shown in the picture?\"  # @param {type: \"string\"}\n",
         "\n",
         "# @markdown The question prompt can be non-English languages.\n",
         "\n",

notebooks/community/model_garden/model_garden_huggingface_pytorch_inference_deployment.ipynb

@@ -62,6 +62,10 @@
         "- Download and deploy the `distilbert/distilbert-base-uncased-finetuned-sst-2-english` model with Hugging Face Pytorch Inference\n",
         "- Send prediction request to the deployed endpoint\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -97,6 +101,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "# Upgrade Vertex AI SDK.\n",
         "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
@@ -111,6 +123,7 @@
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -118,16 +131,11 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
         "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
-        "! gcloud config set project $PROJECT_ID\n",
         "\n",
-        "models, endpoints = {}, {}\n",
+        "! gcloud config set project $PROJECT_ID\n",
         "\n",
         "HF_TOKEN = \"\"\n",
         "# Dedicated endpoint not supported yet\n",
@@ -172,6 +180,7 @@
         "        is_for_training=False,\n",
         "    )\n",
         "\n",
+        "models, endpoints = {}, {}\n",
         "\n",
         "GCS_PREFIX = \"gs://\"\n",
         "\n",

notebooks/community/model_garden/model_garden_huggingface_tei_deployment.ipynb

@@ -116,6 +116,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_huggingface_tgi_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_tgi_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -97,13 +102,22 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 3. You must agree to the license on the [model card](https://huggingface.co/google/gemma-2-2b-it) before accessing the Gemma 2 models.\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
-        "# @markdown 4. Follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "\n",
+        "# @markdown 5. You must agree to the license on the [model card](https://huggingface.co/google/gemma-2-2b-it) before accessing the Gemma 2 models.\n",
+        "\n",
+        "# @markdown 6. Follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
         "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
         "import importlib\n",
         "import os\n",
@@ -111,6 +125,10 @@
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
@@ -121,6 +139,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -134,7 +158,62 @@
         "\n",
         "models, endpoints = {}, {}\n",
         "\n",
-        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}"
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"google/gemma-2-2b-it\"  # @param {type: \"string\", isTemplate: true}\n",
+        "\n",
+        "machine_type = \"g2-standard-12\"  # @param {type: \"string\", isTemplate: true}\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
+        "accelerator_count = 1  # @param {type: \"integer\", isTemplate: true}\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# The pre-built serving docker image for TGI.\n",
+        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-4.ubuntu2204.py311\"\n",
+        "SERVING_CONTAINER_IMAGE_URI = TGI_DOCKER_URI\n",
+        "LABEL = \"tgi\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LSkIcKP_FrkO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(HUGGING_FACE_MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
       ]
     },
     {
@@ -146,31 +225,11 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy with TGI from Hugging Face\n",
+        "# @title [Option 2] Deploy with customized configs\n",
         "\n",
         "# @markdown This section downloads the `Gemma-2-2b-it` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
         "# @markdown It takes ~20 minutes to complete the deployment.\n",
         "\n",
-        "MODEL_ID = \"google/gemma-2-2b-it\"  # @param {type: \"string\", isTemplate: true}\n",
-        "\n",
-        "# The pre-built serving docker image for TGI.\n",
-        "TGI_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/gcr.io/huggingface-text-generation-inference-cu124.2-4.ubuntu2204.py311\"\n",
-        "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
-        "machine_type = \"g2-standard-8\"  # @param {type: \"string\", isTemplate: true}\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\"] {isTemplate: true}\n",
-        "accelerator_count = 1  # @param {type: \"integer\", isTemplate: true}\n",
-        "\n",
-        "common_util.check_quota(\n",
-        "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    is_for_training=False,\n",
-        ")\n",
-        "\n",
         "\n",
         "def deploy_model_tgi(\n",
         "    model_name: str,\n",
@@ -228,14 +287,15 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_tgi_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
         "models[\"tgi\"], endpoints[\"tgi\"] = deploy_model_tgi(\n",
-        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
-        "    model_id=MODEL_ID,\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=HUGGING_FACE_MODEL_ID),\n",
+        "    model_id=HUGGING_FACE_MODEL_ID,\n",
         "    publisher=\"google\",\n",
         "    publisher_model_id=\"gemma2\",\n",
         "    service_account=\"\",\n",

notebooks/community/model_garden/model_garden_huggingface_tgi_vllm_deployment.ipynb

@@ -0,0 +1,604 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qZCBJ69YYS0S"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "08f4AuF5eXzO"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Hugging Face Text Generation with vLLM Container Deployment\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_huggingface_tgi_vllm_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_huggingface_tgi_vllm_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "j0yLwcdReXzO"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying [qwen/qwq-32b](https://huggingface.co/Qwen/QwQ-32B) model with vLLM container from Hugging Face. In additional to `qwen/qwq-32b`, You can view and change the code to deploy a different Hugging Face `text-generation` model with appropriate machine specs. **Note that some models might fail to deploy, even if they have `text-generation` tags on the Hugging Face model card page.**\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Download and deploy the `qwen/qwq-32b` model with TGI\n",
+        "- Send prediction request to the deployed endpoint\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wkhp_C5leXzO"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "D8wjhOsKeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# @markdown 4. Follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create a **read** access token and put it in the `HF_TOKEN` field below.\n",
+        "\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "\n",
+        "# The pre-built serving docker image for TGI with vLLM.\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/vllm-inference.cu121.0-6.ubuntu2204.py310\"\n",
+        "SERVING_CONTAINER_IMAGE_URI = VLLM_DOCKER_URI\n",
+        "LABEL = \"vllm\"\n",
+        "\n",
+        "MACHINE_SPEC_MAP = {\n",
+        "    \"2 NVIDIA_H100_80GB a3-highgpu-2g\": {\n",
+        "        \"machine_type\": \"a3-highgpu-2g\",\n",
+        "        \"accelerator_type\": \"NVIDIA_H100_80GB\",\n",
+        "        \"accelerator_count\": 2,\n",
+        "    },\n",
+        "    \"4 NVIDIA_L4 g2-standard-48\": {\n",
+        "        \"machine_type\": \"g2-standard-48\",\n",
+        "        \"accelerator_type\": \"NVIDIA_L4\",\n",
+        "        \"accelerator_count\": 4,\n",
+        "    },\n",
+        "}\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9j0VbjSNwGJM"
+      },
+      "source": [
+        "## Deploy and predict"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "e8-PBYQ7wGJM"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"qwen/qwq-32b\"  # @param {type: \"string\", isTemplate: true}\n",
+        "DEPLOY_CONFIG = \"2 NVIDIA_H100_80GB a3-highgpu-2g\"  # @param [\"2 NVIDIA_H100_80GB a3-highgpu-2g\", \"4 NVIDIA_L4 g2-standard-48\",] {isTemplate: true}\n",
+        "\n",
+        "machine_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"machine_type\"]\n",
+        "accelerator_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_type\"]\n",
+        "accelerator_count = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_count\"]\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(HUGGING_FACE_MODEL_ID)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    hugging_face_access_token=HF_TOKEN,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "qbNI0WkVeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section downloads the `qwen/qwq-32b` model from Hugging Face and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "\n",
+        "HUGGING_FACE_MODEL_ID = \"qwen/qwq-32b\"  # @param {type: \"string\", isTemplate: true}\n",
+        "DEPLOY_CONFIG = \"2 NVIDIA_H100_80GB a3-highgpu-2g\"  # @param [\"2 NVIDIA_H100_80GB a3-highgpu-2g\", \"4 NVIDIA_L4 g2-standard-48\"] {isTemplate: true}\n",
+        "\n",
+        "machine_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"machine_type\"]\n",
+        "accelerator_type = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_type\"]\n",
+        "accelerator_count = MACHINE_SPEC_MAP[DEPLOY_CONFIG][\"accelerator_count\"]\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    service_account: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_huggingface_tgi_vllm_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"vllm\"], endpoints[\"vllm\"] = deploy_model_vllm(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=HUGGING_FACE_MODEL_ID),\n",
+        "    model_id=HUGGING_FACE_MODEL_ID,\n",
+        "    publisher=HUGGING_FACE_MODEL_ID.split(\"/\")[0].lower(),\n",
+        "    publisher_model_id=HUGGING_FACE_MODEL_ID.split(\"/\")[1].lower(),\n",
+        "    service_account=\"\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "K0B98U88eXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Sampling parameters supported by vLLM can be found [here](https://docs.vllm.ai/en/latest/dev/sampling_params.html).\n",
+        "\n",
+        "# @markdown Example:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown Human: What is a car?\n",
+        "# @markdown Assistant:  A car, or a motor car, is a road-connected human-transportation system used to move people or goods from one place to another. The term also encompasses a wide range of vehicles, including motorboats, trains, and aircrafts. Cars typically have four wheels, a cabin for passengers, and an engine or motor. They have been around since the early 19th century and are now one of the most popular forms of transportation, used for daily commuting, shopping, and other purposes.\n",
+        "# @markdown ```\n",
+        "# @markdown Additionally, you can moderate the generated text with Vertex AI. See [Moderate text documentation](https://cloud.google.com/natural-language/docs/moderating-text) for more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"What is a car?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter an issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, by lowering `max_tokens`.\n",
+        "max_tokens = 50  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "top_k = 1  # @param {type:\"integer\"}\n",
+        "# @markdown Set `raw_response` to `True` to obtain the raw model output. Set `raw_response` to `False` to apply additional formatting in the structure of `\"Prompt:\\n{prompt.strip()}\\nOutput:\\n{output}\"`.\n",
+        "raw_response = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# Overrides parameters for inferences.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"max_tokens\": max_tokens,\n",
+        "        \"temperature\": temperature,\n",
+        "        \"top_p\": top_p,\n",
+        "        \"top_k\": top_k,\n",
+        "        \"raw_response\": raw_response,\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"vllm\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "iK3kcE3CeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict: Streaming Chat Completions\n",
+        "\n",
+        "# @markdown You can also send chat completions requests in streaming mode, if supported by the model.\n",
+        "\n",
+        "# @markdown Here we use an example from the [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) to show the chat completions output:\n",
+        "\n",
+        "# @markdown ```\n",
+        "# @markdown User: How would the Future of AI in 10 Years look?\n",
+        "# @markdown Assistant: Predicting the future is always a challenging task, but here are some possible ways that AI could evolve over the next 10 years: Continued advancements in deep learning: Deep learning has been one of the main drivers of recent AI breakthroughs, and we can expect continued advancements in this area. This may include improvements to existing algorithms, as well as the development of new architectures that are better suited to specific types of data and tasks. Increased use of AI in healthcare: AI has the potential to revolutionize healthcare, by improving the accuracy of diagnoses, developing new treatments, and personalizing patient care. We can expect to see continued investment in this area, with more healthcare providers and researchers using AI to improve patient outcomes. Greater automation in the workplace: Automation is already transforming many industries, and AI is likely to play an increasingly important role in this process. We can expect to see more jobs being automated, as well as the development of new types of jobs that require a combination of human and machine skills. More natural and intuitive interactions with technology: As AI becomes more advanced, we can expect to see more natural and intuitive ways of interacting with technology. This may include voice and gesture recognition, as well as more sophisticated chatbots and virtual assistants. Increased focus on ethical considerations: As AI becomes more powerful, there will be a growing need to consider its ethical implications. This may include issues such as bias in AI algorithms, the impact of automation on employment, and the use of AI in surveillance and policing. Overall, the future of AI in 10 years is likely to be shaped by a combination of technological advancements, societal changes, and ethical considerations. While there are many exciting possibilities for AI in the future, it will be important to carefully consider its potential impact on society and to work towards ensuring that its benefits are shared fairly and equitably.\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "import json\n",
+        "\n",
+        "import requests\n",
+        "\n",
+        "# Initialize request argument(s).\n",
+        "prompt = \"How would the Future of AI in 10 Years look?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 128  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 0.9  # @param {type:\"number\"}\n",
+        "\n",
+        "# Send streaming chat completions request.\n",
+        "endpoint = endpoints[\"vllm\"]\n",
+        "url = f\"https://{endpoint.location}-aiplatform.googleapis.com/v1beta1/{endpoint.resource_name}/chat/completions\"\n",
+        "if use_dedicated_endpoint:\n",
+        "    url = f\"https://{endpoint.gca_resource.dedicated_endpoint_dns}/v1beta1/{endpoint.resource_name}/chat/completions\"\n",
+        "\n",
+        "payload = {\n",
+        "    \"messages\": [{\"role\": \"user\", \"content\": prompt}],\n",
+        "    \"max_tokens\": max_tokens,\n",
+        "    \"temperature\": temperature,\n",
+        "    \"top_p\": top_p,\n",
+        "    \"stream\": True,\n",
+        "}\n",
+        "access_token = ! gcloud auth print-access-token\n",
+        "access_token = access_token[0]\n",
+        "\n",
+        "response = requests.post(\n",
+        "    url,\n",
+        "    headers={\"Authorization\": f\"Bearer {access_token}\"},\n",
+        "    json=payload,\n",
+        "    stream=True,\n",
+        ")\n",
+        "\n",
+        "if not response.ok:\n",
+        "    raise ValueError(response.text)\n",
+        "\n",
+        "for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False):\n",
+        "    if chunk:\n",
+        "        chunk = chunk.decode(\"utf-8\").removeprefix(\"data:\").strip()\n",
+        "        if chunk == \"[DONE]\":\n",
+        "            break\n",
+        "        data = json.loads(chunk)\n",
+        "        if type(data) is not dict or \"error\" in data:\n",
+        "            raise ValueError(data)\n",
+        "        delta = data[\"choices\"][0][\"delta\"].get(\"content\")\n",
+        "        if delta:\n",
+        "            print(delta, end=\"\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UowpdPJOeXzO"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "JgvATV8aeXzO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_huggingface_tgi_vllm_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_jax_owl_vit_v2.ipynb

@@ -118,14 +118,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "import base64\n",
         "import datetime\n",
@@ -155,6 +155,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_jax_paligemma_deployment.ipynb

@@ -139,6 +139,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_jax_paligemma_finetuning.ipynb

@@ -108,7 +108,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -142,6 +142,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_keras_yolov8.ipynb

@@ -108,14 +108,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Configs for all notebooks.\n",
         "! pip3 install --quiet keras-cv==0.9.0\n",
@@ -154,6 +154,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_llama3_2_evaluation.ipynb

@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_llama_guard_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_llama_guard_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -62,6 +67,10 @@
         "- Use the Llama Guard models to safeguard LLM inputs and outputs with the Vertex Llama 3.1 API service\n",
         "- Use the Llama Guard models to safeguard LLM vision inputs and outputs with the Vertex Llama 3.2 API service\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -110,44 +119,44 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
         "import re\n",
-        "import uuid\n",
         "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "LABEL = \"vllm_gpu\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -155,52 +164,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama_guard\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown # Access Llama Guard models on Vertex AI\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -221,15 +195,7 @@
         "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD = parsed_gcs_url.group()\n",
         "assert VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD.startswith(\n",
         "    \"gs://\"\n",
-        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\"\n",
-        "print(\n",
-        "    \"Copying Llama Guard model artifacts from\",\n",
-        "    VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD,\n",
-        "    \"to \",\n",
-        "    MODEL_BUCKET,\n",
-        ")\n",
-        "\n",
-        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD/* $MODEL_BUCKET"
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD is expected to be a GCS URI and must start with `gs://`.\""
       ]
     },
     {
@@ -238,7 +204,7 @@
         "id": "z-XybZjtgF9M"
       },
       "source": [
-        "## Deploy Llama Guard with vLLM on GPU"
+        "## Deploy Llama Guard"
       ]
     },
     {
@@ -246,27 +212,30 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "E8OiHHNNE_wj"
+        "id": "kRiRTAMxxUoq"
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
         "\n",
-        "# @markdown This section uploads prebuilt Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "# @markdown Select one of the three model variations.\n",
         "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"Llama-Guard-3-8B\"  # @param [\"Llama-Guard-3-8B\", \"Llama-Guard-3-1B\", \"Llama-Guard-3-11B-Vision\"] {allow-input: true, isTemplate: true}\n",
-        "model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
+        "base_model_name = \"Llama-Guard-3-11B-Vision\"  # @param [\"Llama-Guard-3-8B\", \"Llama-Guard-3-1B\", \"Llama-Guard-3-11B-Vision\"] {allow-input: true, isTemplate: true}\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_GUARD, base_model_name)\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "version_id = base_model_name.lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama-guard@{version_id}\"\n",
         "\n",
         "# The pre-built serving docker images.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240925_0916_RC01\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00_maas\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "if \"3-1B\" in base_model_name or \"3-8B\" in base_model_name:\n",
         "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    machine_type = \"g2-standard-8\"\n",
+        "    machine_type = \"g2-standard-12\"\n",
         "    accelerator_count = 1\n",
         "    max_num_seqs = 256\n",
         "elif \"3-11B\" in base_model_name:\n",
@@ -283,7 +252,45 @@
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    is_for_training=False,\n",
-        ")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2DiRl36FzauJ"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configurations\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama Guard models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
         "gpu_memory_utilization = 0.9\n",
         "max_model_len = 4096\n",
@@ -294,7 +301,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -313,6 +319,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -365,6 +372,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -400,9 +411,9 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_llama_guard_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -410,15 +421,11 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-guard\"),\n",
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
-        "    publisher=\"llama-guard\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
+        "    publisher_model_id=\"llama-guard\",\n",
         "    base_model_id=hf_model_id,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
@@ -428,6 +435,7 @@
         "    enforce_eager=True,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         "    max_num_seqs=max_num_seqs,\n",
+        "    enable_llama_tool_parser=True,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -549,8 +557,8 @@
         "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
         ")\n",
         "\n",
-        "prediction = response.predictions[0]\n",
-        "print(\"Llama Guard prediction:\", prediction[0][\"message\"][\"content\"])"
+        "prediction = response.predictions[\"choices\"][0][\"message\"][\"content\"]\n",
+        "print(\"Llama Guard prediction:\", prediction)"
       ]
     },
     {
@@ -616,29 +624,6 @@
       "source": [
         "# @markdown Define input message in conversation and get output message from model.\n",
         "\n",
-        "\n",
-        "from urllib.request import urlopen\n",
-        "\n",
-        "from google.cloud import storage\n",
-        "\n",
-        "IMAGE_GCS_PATH = \"input_image.jpg\"\n",
-        "\n",
-        "\n",
-        "def image_url_to_gcs(url):\n",
-        "    # Uploads image to GCS and returns GCS URL.\n",
-        "    storage_client = storage.Client()\n",
-        "    bucket_name = BUCKET_URI.replace(\"gs://\", \"\")\n",
-        "    bucket = storage_client.get_bucket(bucket_name)\n",
-        "    blob = bucket.blob(IMAGE_GCS_PATH)\n",
-        "\n",
-        "    with urlopen(url) as response:\n",
-        "        image_data = response.read()\n",
-        "        image_type = response.info().get_content_type()\n",
-        "        blob.upload_from_string(image_data, content_type=image_type)\n",
-        "\n",
-        "    return os.path.join(BUCKET_URI, IMAGE_GCS_PATH)\n",
-        "\n",
-        "\n",
         "user_image = \"https://upload.wikimedia.org/wikipedia/commons/thumb/c/cb/The_Blue_Marble_%28remastered%29.jpg/580px-The_Blue_Marble_%28remastered%29.jpg\"  # @param {type: \"string\"}\n",
         "user_message = \"What is in the image?\"  # @param {type: \"string\"}\n",
         "\n",
@@ -654,23 +639,12 @@
         "        ],\n",
         "    }\n",
         "]\n",
-        "messages_gcs = [\n",
-        "    {\n",
-        "        \"role\": \"user\",\n",
-        "        \"content\": [\n",
-        "            {\n",
-        "                \"type\": \"image_url\",\n",
-        "                \"image_url\": {\"url\": image_url_to_gcs(user_image)},\n",
-        "            },\n",
-        "            {\"type\": \"text\", \"text\": user_message},\n",
-        "        ],\n",
-        "    }\n",
-        "]\n",
-        "print(\"Conversation [turn 1]:\", messages_gcs)\n",
+        "\n",
+        "print(\"Conversation [turn 1]:\", messages)\n",
         "\n",
         "response = client.chat.completions.create(\n",
         "    model=LLAMA3_90B_VISION_INSTRUCT,\n",
-        "    messages=messages_gcs,\n",
+        "    messages=messages,\n",
         ")\n",
         "print(\"Response:\", response)\n",
         "\n",
@@ -680,13 +654,8 @@
         "        \"content\": response.choices[0].message.content,\n",
         "    }\n",
         ")\n",
-        "messages_gcs.append(\n",
-        "    {\n",
-        "        \"role\": response.choices[0].message.role,\n",
-        "        \"content\": response.choices[0].message.content,\n",
-        "    }\n",
-        ")\n",
-        "print(\"Conversation [turn 2]:\", messages_gcs)"
+        "\n",
+        "print(\"Conversation [turn 2]:\", messages)"
       ]
     },
     {
@@ -713,8 +682,8 @@
         "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
         ")\n",
         "\n",
-        "prediction = response.predictions[0]\n",
-        "print(\"Llama Guard prediction:\", prediction[0][\"message\"][\"content\"])"
+        "prediction = response.predictions[\"choices\"][0][\"message\"][\"content\"]\n",
+        "print(\"Llama Guard prediction:\", prediction)"
       ]
     },
     {
@@ -745,11 +714,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_mediapipe_face_stylizer.ipynb

@@ -110,7 +110,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -140,6 +140,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_mediapipe_object_detection.ipynb

@@ -104,7 +104,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_mediapipe_text_classification.ipynb

@@ -106,7 +106,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -410,7 +416,7 @@
         "    ],\n",
         "}\n",
         "\n",
-        "labels[\"mg-tune\"] = \"publishers-google-models-mdeiapipe\"\n",
+        "labels[\"mg-tune\"] = \"publishers-google-models-mediapipe\"\n",
         "versioned_model_id = model_architecture.lower().replace(\"_\", \"-\")\n",
         "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
         "\n",

notebooks/community/model_garden/model_garden_movinet_action_recognition.ipynb

@@ -110,7 +110,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -145,6 +145,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -363,7 +369,7 @@
         "    staging_bucket=STAGING_BUCKET,\n",
         ")\n",
         "\n",
-        "data_converter_custom_job.run()\n",
+        "data_converter_custom_job.run(service_account=SERVICE_ACCOUNT)\n",
         "\n",
         "input_train_data_path = os.path.join(data_converter_output_dir, \"train.tfrecord*\")\n",
         "input_validation_data_path = os.path.join(data_converter_output_dir, \"val.tfrecord*\")\n",
@@ -580,7 +586,7 @@
         "    search_algorithm=None,\n",
         ")\n",
         "\n",
-        "train_hpt_job.run()\n",
+        "train_hpt_job.run(service_account=SERVICE_ACCOUNT)\n",
         "\n",
         "print(\"model_dir is:\", model_dir)"
       ]
@@ -691,7 +697,7 @@
         "    staging_bucket=STAGING_BUCKET,\n",
         ")\n",
         "\n",
-        "model_export_custom_job.run()\n",
+        "model_export_custom_job.run(service_account=SERVICE_ACCOUNT)\n",
         "\n",
         "print(\"best model is saved to: \", container_args[\"export_path\"])"
       ]
@@ -760,7 +766,6 @@
         "    serving_container_predict_route=\"/predict\",\n",
         "    serving_container_health_route=\"/ping\",\n",
         "    serving_container_environment_variables=serving_env,\n",
-        "    model_garden_source_model_name=\"publishers/google/models/tfvision-movinet-var\",\n",
         ")\n",
         "\n",
         "models[\"model_var\"].wait()\n",
@@ -801,8 +806,6 @@
         "\n",
         "# Path to the prediction input JSONL file.\n",
         "test_jsonl_path = \"\"  # @param {type:\"string\"}\n",
-        "# Full service account name with the suffix `gserviceaccount.com`.\n",
-        "batch_predict_service_account = \"\"  # @param {type:\"string\"}\n",
         "\n",
         "predict_job_name = common_util.get_job_name_with_datetime(\n",
         "    f\"{PREDICTION_JOB_PREFIX}_{model_name}\"\n",
@@ -817,7 +820,7 @@
         "    accelerator_count=PREDICTION_ACCELERATOR_COUNT,\n",
         "    accelerator_type=PREDICTION_ACCELERATOR_TYPE,\n",
         "    max_replica_count=1,\n",
-        "    service_account=batch_predict_service_account,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
         ")\n",
         "\n",
         "batch_prediction_job.wait()\n",
@@ -836,7 +839,7 @@
       },
       "outputs": [],
       "source": [
-        "# @title Read the prediction response.\n",
+        "# @title Read the prediction response\n",
         "\n",
         "\n",
         "def print_response_instance(json_str: str, label_map: dict[int, str]):\n",

notebooks/community/model_garden/model_garden_movinet_clip_classification.ipynb

@@ -384,9 +384,9 @@
         "2. Create hyperparameter tuning jobs to train new models\n",
         "3. Find and export best models\n",
         "\n",
-        "If you already trained models, please go to the section `Test Trained models`.\n",
+        "If you already trained models, go to the section `Test Trained models`.\n",
         "\n",
-        "Please select a model:\n",
+        "Select a model:\n",
         "* `model_id`: MoViNet model variant ID, one of `a0`, `a1`, `a2`, `a3`, `a4`, `a5`. The model with a larger number requires more resources to train, and is expected to have a higher accuracy and latency. Here, we use `a0` for demonstration purpose.\n",
         "* `model_mode`: MoViNet model type, either `base` or `stream`. The base model has a slightly higher accuracy, while the streaming model is optimized for streaming and faster CPU inference. See [official MoViNet docs](https://github.com/tensorflow/models/tree/master/official/projects/movinet) for more information.\n",
         "\n",
@@ -552,7 +552,7 @@
         "config_file = f\"https://raw.githubusercontent.com/tensorflow/models/master/official/projects/movinet/configs/yaml/movinet_{config_file}_gpu.yaml\"\n",
         "config_file = upload_config_to_gcs(config_file)\n",
         "\n",
-        "# The parameters here are mainly for demonstration purpose. Please update them\n",
+        "# The parameters here are mainly for demonstration purpose. Update them\n",
         "# for better performance.\n",
         "trainer_args = {\n",
         "    \"experiment\": \"movinet_kinetics600\",\n",
@@ -770,7 +770,6 @@
         "    serving_container_predict_route=\"/predict\",\n",
         "    serving_container_health_route=\"/ping\",\n",
         "    serving_container_environment_variables=serving_env,\n",
-        "    model_garden_source_model_name=\"publishers/google/models/tfvision-movinet-vcn\"\n",
         ")\n",
         "\n",
         "model.wait()\n",
@@ -808,7 +807,7 @@
         "\n",
         "We will now run batch predictions with the trained MoViNet clip classification model with [Vertex AI Batch Prediction](https://cloud.google.com/vertex-ai/docs/predictions/get-batch-predictions).\n",
         "\n",
-        "Please prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/classification/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information."
+        "Prepare an input JSONL file where each line follows [this format](https://cloud.google.com/vertex-ai/docs/video-data/classification/get-predictions?hl=en#input_data_requirements) and store it in a Cloud Storage bucket. The service account should have read access to the buckets containing the trained model and the input data. See [Service accounts overview](https://cloud.google.com/iam/docs/service-account-overview) for more information."
       ]
     },
     {

notebooks/community/model_garden/model_garden_nvidia_cosmos_deployment.ipynb

@@ -0,0 +1,639 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "1gcBBbBCW_CV"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "wKzYxAA1W_CV"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Nvidia Cosmos 1.0\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_nvidia_cosmos_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_nvidia_cosmos_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2WwEeH8BW_CV"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates deploying Nvidia Cosmos world foundation models (WFM) on Vertex AI for online prediction.\n",
+        "  - [nvidia/Cosmos-1.0-Diffusion-7B-Text2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Text2World)\n",
+        "  - [nvidia/Cosmos-1.0-Diffusion-14B-Text2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Text2World)\n",
+        "  - [nvidia/Cosmos-1.0-Diffusion-7B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-7B-Video2World)\n",
+        "  - [nvidia/Cosmos-1.0-Diffusion-14B-Video2World](https://huggingface.co/nvidia/Cosmos-1.0-Diffusion-14B-Video2World)\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Upload the model to [Model Registry](https://cloud.google.com/vertex-ai/docs/model-registry/introduction).\n",
+        "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
+        "- Run online predictions for `text-to-world` and `video-to-world`.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "TAKAyLQvW_CV"
+      },
+      "source": [
+        "## Run the notebook"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "sGzHHcL3W_CV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "from IPython.display import HTML\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "# Gets the default SERVICE_ACCOUNT.\n",
+        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
+        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
+        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
+        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "q36QziORW_CV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy the [Text2World] model to Vertex for online predictions\n",
+        "\n",
+        "# @markdown This section uploads the [Text2World] model to Model Registry and deploys it on the Endpoint with the specified accelerator.\n",
+        "\n",
+        "# @markdown The deployment process takes approximately 15-30 minutes to complete.\n",
+        "# @markdown A valid HF_TOKEN is required for model deployment.\n",
+        "# @markdown Follow the instructions at [Hugging Face Token Guide](https://huggingface.co/docs/hub/en/security-tokens) to obtain your HF_TOKEN.\n",
+        "# @markdown Additionally, ensure you have access to the model by following the instructions on its Hugging Face model card page.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "if not HF_TOKEN:\n",
+        "    print(\"Error: HF_TOKEN is required to deploy the model.\")\n",
+        "\n",
+        "# @markdown The inference timeout is set to 30 minutes, as the video generation process can take a long time.\n",
+        "INFERENCE_TIMEOUT_SECS = 1800\n",
+        "model_id = \"nvidia/Cosmos-1.0-Diffusion-7B-Text2World\"  # @param [\"nvidia/Cosmos-1.0-Diffusion-7B-Text2World\", \"nvidia/Cosmos-1.0-Diffusion-14B-Text2World\"]\n",
+        "task = \"text-to-world\"\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_H100_80GB\", \"NVIDIA_A100_80GB\"]\n",
+        "\n",
+        "machine_type_map = {\n",
+        "    \"NVIDIA_A100_80GB\": \"a2-ultragpu-1g\",\n",
+        "    \"NVIDIA_H100_80GB\": \"a3-highgpu-2g\",\n",
+        "}\n",
+        "\n",
+        "machine_type = machine_type_map.get(accelerator_type)\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "    accelerator_count = 2\n",
+        "\n",
+        "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cosmos:20250314\"\n",
+        "\n",
+        "\n",
+        "def deploy_model(model_id, task, machine_type, accelerator_type, accelerator_count):\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=False,\n",
+        "    )\n",
+        "\n",
+        "    model_name = model_id\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=True,\n",
+        "        sync=True,\n",
+        "        inference_timeout=INFERENCE_TIMEOUT_SECS,\n",
+        "    )\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        \"HUGGING_FACE_HUB_TOKEN\": HF_TOKEN,\n",
+        "        \"OFFLOAD_NETWORK\": \"false\",\n",
+        "        \"OFFLOAD_TOKENIZER\": \"false\",\n",
+        "        \"OFFLOAD_TEXT_ENCODER_MODEL\": \"false\",\n",
+        "        \"OFFLOAD_GUARDRAIL_MODELS\": \"true\",\n",
+        "        \"OFFLOAD_PROMPT_UPSAMPLER\": \"true\",\n",
+        "    }\n",
+        "\n",
+        "    # Also offload the text encoder model for 14B models, to avoid CUDA OOM issue.\n",
+        "    if model_id.lower().includes(\"14b\"):\n",
+        "        serving_env[\"OFFLOAD_TEXT_ENCODER_MODEL\"] = \"true\"\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=\"publishers/nvidia/models/cosmos\",\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_nvidia_cosmos_deployment.ipynb\"},\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
+        "    model_id=model_id,\n",
+        "    task=task,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "\n",
+        "print(\"endpoint_name:\", endpoints[\"endpoint\"].name)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TKJsEJoeW_CV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Text2World] Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The inference takes:\n",
+        "\n",
+        "# @markdown - ~800s with 1 A100 80GB GPU.\n",
+        "# @markdown\n",
+        "# @markdown - ~420s with 2 H100 80GB GPU\n",
+        "\n",
+        "# @markdown Example:\n",
+        "# @markdown ```json\n",
+        "# @markdown {\n",
+        "# @markdown   \"instances\":[\n",
+        "# @markdown     {\n",
+        "# @markdown       \"text\":\"A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves.\",\n",
+        "# @markdown     }\n",
+        "# @markdown   ],\n",
+        "# @markdown    \"parameters\": {\n",
+        "# @markdown      \"negative_prompt\": \"\",\n",
+        "# @markdown      \"guidance\": 7.0,\n",
+        "# @markdown      \"num_steps\": 30,\n",
+        "# @markdown      \"height\": 704,\n",
+        "# @markdown      \"width\": 1280,\n",
+        "# @markdown      \"fps\": 24,\n",
+        "# @markdown      \"num_video_frames\": 121,\n",
+        "# @markdown      \"seed\": 42\n",
+        "# @markdown    }\n",
+        "# @markdown  }\n",
+        "# @markdown }\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown You can adjust the parameters below to use your own text prompt.\n",
+        "# @markdown The `negative_prompt` parameter is optional. If not specified, a default value will be used.\n",
+        "# @markdown You can find the default value here: [Inference Utils (Line 104)](https://github.com/NVIDIA/Cosmos/blob/main/cosmos1/models/diffusion/inference/inference_utils.py#L104).\n",
+        "# @markdown\n",
+        "# @markdown For inference tasks exceeding 10 minutes, we recommend using CURL for predictions. Refer to the following sections for detailed instructions.\n",
+        "\n",
+        "text = \"A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves.\"  # @param {type: \"string\"}\n",
+        "\n",
+        "instances = [{\"text\": text}]\n",
+        "parameters = {\n",
+        "    \"negative_prompt\": \"\",\n",
+        "    \"guidance\": 7.0,\n",
+        "    \"num_steps\": 30,\n",
+        "    \"height\": 704,\n",
+        "    \"width\": 1280,\n",
+        "    \"fps\": 24,\n",
+        "    \"num_video_frames\": 121,\n",
+        "    \"seed\": 42,\n",
+        "}\n",
+        "\n",
+        "\n",
+        "response = endpoints[\"endpoint\"].predict(\n",
+        "    instances=instances, parameters=parameters, use_dedicated_endpoint=True\n",
+        ")\n",
+        "\n",
+        "video_bytes = response.predictions[0][\"output\"]\n",
+        "\n",
+        "video_html = f\"\"\"\n",
+        "<video width=\"1280\" height=\"704\" controls>\n",
+        "<source src=\"data:video/mp4;base64,{video_bytes}\" type=\"video/mp4\">\n",
+        "Your browser does not support the video tag.\n",
+        "</video>\n",
+        "\"\"\"  # Assumes MP4. Change type if needed (e.g., video/webm)\n",
+        "\n",
+        "display(HTML(video_html))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3TurWPvt8wSf"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy the [Video2World] model to Vertex for online predictions\n",
+        "\n",
+        "# @markdown This section uploads the [Video2World] model to Model Registry and deploys it on the Endpoint with the specified accelerator.\n",
+        "\n",
+        "# @markdown The deployment process takes approximately 15-30 minutes to complete.\n",
+        "# @markdown A valid HF_TOKEN is required for model deployment.\n",
+        "# @markdown Follow the instructions at [Hugging Face Token Guide](https://huggingface.co/docs/hub/en/security-tokens) to obtain your HF_TOKEN.\n",
+        "# @markdown Additionally, ensure you have access to the model by following the instructions on its Hugging Face model card page.\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "if not HF_TOKEN:\n",
+        "    print(\"Error: HF_TOKEN is required to deploy the model.\")\n",
+        "# @markdown The inference timeout is set to 30 minutes, as the video generation process can take a long time.\n",
+        "INFERENCE_TIMEOUT_SECS = 1800\n",
+        "\n",
+        "model_id = \"nvidia/Cosmos-1.0-Diffusion-7B-Video2World\"  # @param [\"nvidia/Cosmos-1.0-Diffusion-7B-Video2World\", \"nvidia/Cosmos-1.0-Diffusion-14B-Video2World\"]\n",
+        "task = \"video-to-world\"\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_H100_80GB\", \"NVIDIA_A100_80GB\"]\n",
+        "\n",
+        "machine_type_map = {\n",
+        "    \"NVIDIA_A100_80GB\": \"a2-ultragpu-1g\",\n",
+        "    \"NVIDIA_H100_80GB\": \"a3-highgpu-2g\",\n",
+        "}\n",
+        "\n",
+        "machine_type = machine_type_map.get(accelerator_type)\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "    accelerator_count = 2\n",
+        "\n",
+        "\n",
+        "# The pre-built serving docker image. It contains serving scripts and models.\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-cosmos:20250314\"\n",
+        "\n",
+        "\n",
+        "def deploy_model(model_id, task, machine_type, accelerator_type, accelerator_count):\n",
+        "    \"\"\"Create a Vertex AI Endpoint and deploy the specified model to the endpoint.\"\"\"\n",
+        "    common_util.check_quota(\n",
+        "        project_id=PROJECT_ID,\n",
+        "        region=REGION,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        is_for_training=False,\n",
+        "    )\n",
+        "\n",
+        "    model_name = model_id\n",
+        "\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=True,\n",
+        "        sync=True,\n",
+        "        inference_timeout=INFERENCE_TIMEOUT_SECS,\n",
+        "    )\n",
+        "    serving_env = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"TASK\": task,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        \"HUGGING_FACE_HUB_TOKEN\": HF_TOKEN,\n",
+        "        \"OFFLOAD_NETWORK\": \"false\",\n",
+        "        \"OFFLOAD_TOKENIZER\": \"false\",\n",
+        "        \"OFFLOAD_TEXT_ENCODER_MODEL\": \"false\",\n",
+        "        \"OFFLOAD_GUARDRAIL_MODELS\": \"true\",\n",
+        "        \"OFFLOAD_PROMPT_UPSAMPLER\": \"true\",\n",
+        "    }\n",
+        "\n",
+        "    # Also offload the text encoder model for 14B models, to avoid CUDA OOM issue.\n",
+        "    if model_id.lower().includes(\"14b\"):\n",
+        "        serving_env[\"OFFLOAD_TEXT_ENCODER_MODEL\"] = \"true\"\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
+        "        serving_container_ports=[7080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables=serving_env,\n",
+        "        model_garden_source_model_name=\"publishers/nvidia/models/cosmos\",\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        service_account=SERVICE_ACCOUNT,\n",
+        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_nvidia_cosmos_deployment.ipynb\"},\n",
+        "    )\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
+        "    model_id=model_id,\n",
+        "    task=task,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        ")\n",
+        "\n",
+        "print(\"endpoint_name:\", endpoints[\"endpoint\"].name)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ajOOy8Qv8wSf"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Video2World] Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. The inference takes:\n",
+        "\n",
+        "# @markdown - ~400s with 1 A100 GPU.\n",
+        "# @markdown\n",
+        "# @markdown - ~400s with 2 H100 GPU\n",
+        "\n",
+        "# @markdown Example:\n",
+        "# @markdown ```json\n",
+        "# @markdown  {\n",
+        "# @markdown    \"instances\": [\n",
+        "# @markdown      {\n",
+        "# @markdown        \"gcs_uri\": \"gs://vertex-model-garden-public-us/cosmos/video2world_input0.jpg\",\n",
+        "# @markdown        \"num_input_frames\": 1\n",
+        "# @markdown      }\n",
+        "# @markdown    ],\n",
+        "# @markdown    \"parameters\": {\n",
+        "# @markdown      \"negative_prompt\": \"\",\n",
+        "# @markdown      \"guidance\": 7.0,\n",
+        "# @markdown      \"num_steps\": 25,\n",
+        "# @markdown      \"height\": 704,\n",
+        "# @markdown      \"width\": 1280,\n",
+        "# @markdown      \"fps\": 24,\n",
+        "# @markdown      \"num_video_frames\": 121,\n",
+        "# @markdown      \"seed\": 42\n",
+        "# @markdown    }\n",
+        "# @markdown  }\n",
+        "# @markdown ```\n",
+        "\n",
+        "# @markdown You can adjust the parameters below to use your own video.\n",
+        "# @markdown The model also supports single-image input by setting `num_input_frames = 1`.\n",
+        "# @markdown Note that `num_input_frames` should match the actual number of frames in your video.\n",
+        "# @markdown The `negative_prompt` parameter is optional. If not specified, a default value will be used.\n",
+        "# @markdown You can find the default value here: [Inference Utils (Line 104)](https://github.com/NVIDIA/Cosmos/blob/main/cosmos1/models/diffusion/inference/inference_utils.py#L104).\n",
+        "\n",
+        "# @markdown\n",
+        "# @markdown For inference tasks exceeding 10 minutes, we recommend using CURL for predictions. Refer to the following sections for detailed instructions.\n",
+        "\n",
+        "gcs_uri = \"gs://vertex-model-garden-public-us/cosmos/video2world_input0.jpg\"  # @param {type: \"string\"}\n",
+        "num_input_frames = 1  # @param {type: \"integer\"}\n",
+        "negative_prompt = \"\"  # @param {type: \"string\"}\n",
+        "\n",
+        "instances = [{\"gcs_uri\": gcs_uri, \"num_input_frames\": num_input_frames}]\n",
+        "parameters = {\n",
+        "    \"negative_prompt\": negative_prompt,\n",
+        "    \"guidance\": 7.0,\n",
+        "    \"num_steps\": 25,\n",
+        "    \"height\": 704,\n",
+        "    \"width\": 1280,\n",
+        "    \"fps\": 24,\n",
+        "    \"num_video_frames\": 121,\n",
+        "    \"seed\": 42,\n",
+        "}\n",
+        "\n",
+        "\n",
+        "response = endpoints[\"endpoint\"].predict(\n",
+        "    instances=instances, parameters=parameters, use_dedicated_endpoint=True\n",
+        ")\n",
+        "\n",
+        "video_bytes = response.predictions[0][\"output\"]\n",
+        "\n",
+        "video_html = f\"\"\"\n",
+        "<video width=\"1280\" height=\"704\" controls>\n",
+        "<source src=\"data:video/mp4;base64,{video_bytes}\" type=\"video/mp4\">\n",
+        "Your browser does not support the video tag.\n",
+        "</video>\n",
+        "\"\"\"  # Assumes MP4. Change type if needed (e.g., video/webm)\n",
+        "\n",
+        "display(HTML(video_html))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "pzZo5t_mqNDy"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict with CURL for long-running prediction tasks\n",
+        "\n",
+        "# @markdown For inference tasks exceeding 10 minutes, we recommend using CURL for predictions.\n",
+        "\n",
+        "os.environ[\"ENDPOINT_ID\"] = endpoints[\"endpoint\"].name\n",
+        "os.environ[\"PROJECT_ID\"] = project_number\n",
+        "os.environ[\"REGION\"] = REGION"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "2FOfZRLbqNDy"
+      },
+      "outputs": [],
+      "source": [
+        "%%bash\n",
+        "\n",
+        "# Leverage CURL in shell for predictions, especially for long-running tasks (exceeding 10 minutes). \n",
+        "ENDPOINT_URL=\"https://${ENDPOINT_ID}.${REGION}-${PROJECT_ID}.prediction.vertexai.goog/v1/projects/${PROJECT_ID}/locations/${REGION}/endpoints/${ENDPOINT_ID}:predict\"\n",
+        "TEXT=\"A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves.\"\n",
+        "DATA='{\"instances\": [{\"text\":\"'${TEXT}'\"}], \"parameters\": {\"negative_prompt\":\"\", \"guidance\":7.0,\"num_steps\":35,\"height\":704,\"width\":1280,\"fps\":24,\"num_video_frames\":121,\"seed\":42}}'\n",
+        "\n",
+        "curl \\\n",
+        "  -X POST \\\n",
+        "  -H \"Authorization: Bearer $(gcloud auth print-access-token)\" \\\n",
+        "  -H \"Content-Type: application/json\" \\\n",
+        "  \"${ENDPOINT_URL}\" \\\n",
+        "  -d \"${DATA}\" > /content/t2w_response.json"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "HAPiswf9qNDy"
+      },
+      "outputs": [],
+      "source": [
+        "import json\n",
+        "\n",
+        "with open(\"/content/t2w_response.json\", \"r\") as f:\n",
+        "    response_data = json.load(f)\n",
+        "\n",
+        "video_bytes = response_data[\"predictions\"][0][\"output\"]\n",
+        "print(video_bytes)\n",
+        "\n",
+        "video_html = f\"\"\"\n",
+        "<video width=\"1280\" height=\"704\" controls>\n",
+        "<source src=\"data:video/mp4;base64,{video_bytes}\" type=\"video/mp4\">\n",
+        "Your browser does not support the video tag.\n",
+        "</video>\n",
+        "\"\"\"  # Assumes MP4. Change type if needed (e.g., video/webm)\n",
+        "\n",
+        "display(HTML(video_html))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "42leJGJFW_CV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Clean up resources\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_nvidia_cosmos_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_ollama_deployment.ipynb

@@ -0,0 +1,451 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "20qcPG1PmFUM"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "QXYOa1odnikj"
+      },
+      "source": [
+        "# Vertex AI Model Garden - Ollama Deployment\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_ollama_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_ollama_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cbDI9ag4oR4C"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to deploy GPT-Generated Unified Format (GGUF) models with Vertex Model Garden released Ollama serving dockers, which are mainly based on [Ollama](https://github.com/ollama/ollama/tree/main).\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy the deepseek-r1 1.5b and 671b GGUF models with Ollama\n",
+        "- Send prediction request to the deployed endpoint\n",
+        "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "hQJWRopioSKT"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "J_jmxcIZoSxU"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "WWNxEb-vlosS"
+      },
+      "source": [
+        "## Deploy with Ollama from Hugging Face"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "USB7dvYqvNdu"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section downloads the `deepseek-r1:1.5b` or `deepseek-r1:671b` model from Ollama and deploys it to a Vertex AI Endpoint.\n",
+        "# @markdown It takes ~20 minutes to complete the deployment.\n",
+        "\n",
+        "MODEL_ID = \"deepseek-r1:1.5b\"  # @param [\"deepseek-r1:1.5b\", \"deepseek-r1:671b\"]\n",
+        "\n",
+        "# The pre-built serving docker image for Ollama.\n",
+        "OLLAMA_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/ollama-serve.cu125.0-5.ubuntu2204.py310\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if \"1.5b\" in MODEL_ID:\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_type = \"NVIDIA_L4\"\n",
+        "    accelerator_count = 1\n",
+        "elif \"671b\" in MODEL_ID:\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    machine_type = \"a3-highgpu-8g\"\n",
+        "    accelerator_count = 8\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended GPU setting not found for: {accelerator_type} and {MODEL_ID}.\"\n",
+        "    )\n",
+        "\n",
+        "context_length = 131072 if \"1.5b\" in MODEL_ID else 16384\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def deploy_model_ollama(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    context_length: int,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Ollama on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"CONTEXT_LENGTH\": context_length,\n",
+        "    }\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=OLLAMA_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=3600,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_ollama_deployment.ipynb\",\n",
+        "            \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"ollama\"], endpoints[\"ollama\"] = deploy_model_ollama(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=MODEL_ID),\n",
+        "    model_id=MODEL_ID,\n",
+        "    publisher=\"deepseek-ai\",\n",
+        "    publisher_model_id=\"deepseek-r1\",\n",
+        "    context_length=context_length,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cFB1B17ylx0a"
+      },
+      "source": [
+        "## Predict"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Aa4e1-6FvRAP"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Raw Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details.\n",
+        "\n",
+        "# Loads an existing endpoint instance using the endpoint name:\n",
+        "# - Using `endpoint_name = endpoint.name` allows us to get the\n",
+        "#   endpoint name of the endpoint `endpoint` created in the cell\n",
+        "#   above.\n",
+        "# - Alternatively, you can set `endpoint_name = \"1234567890123456789\"` to load\n",
+        "#   an existing endpoint with the ID 1234567890123456789.\n",
+        "# You may uncomment the code below to load an existing endpoint.\n",
+        "\n",
+        "# endpoint_name = \"\"  # @param {type:\"string\"}\n",
+        "# aip_endpoint_name = (\n",
+        "#     f\"projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint_name}\"\n",
+        "# )\n",
+        "# endpoint = aiplatform.Endpoint(aip_endpoint_name)\n",
+        "\n",
+        "prompt = \"Why is the sky blue?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 128  # @param {type:\"integer\"}\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 0.7  # @param {type:\"number\"}\n",
+        "top_k = -1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# Overrides max_tokens and top_k parameters during inferences.\n",
+        "instances = [\n",
+        "    {\n",
+        "        \"prompt\": prompt,\n",
+        "        \"options\": {\n",
+        "            \"num_predict\": max_tokens,\n",
+        "            \"temperature\": temperature,\n",
+        "            \"top_p\": top_p,\n",
+        "            \"top_k\": top_k,\n",
+        "        },\n",
+        "    },\n",
+        "]\n",
+        "response = endpoints[\"ollama\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    print(prediction)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "KRWGdMl3WEO5"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict: Streaming Chat Completions\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ollama\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"ollama\"].name\n",
+        ")\n",
+        "\n",
+        "# @title Chat Completions Inference\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint using the OpenAI SDK.\n",
+        "\n",
+        "# @markdown First you will need to install the SDK and some auth-related dependencies.\n",
+        "\n",
+        "! pip install -qU openai google-auth requests\n",
+        "\n",
+        "# @markdown Next fill out some request parameters:\n",
+        "\n",
+        "user_message = \"How is your day going?\"  # @param {type: \"string\"}\n",
+        "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
+        "max_tokens = 50  # @param {type: \"integer\"}\n",
+        "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
+        "\n",
+        "# @markdown Now we can send a request.\n",
+        "\n",
+        "import google.auth\n",
+        "import openai\n",
+        "\n",
+        "creds, project = google.auth.default()\n",
+        "auth_req = google.auth.transport.requests.Request()\n",
+        "creds.refresh(auth_req)\n",
+        "\n",
+        "BASE_URL = (\n",
+        "    f\"https://{REGION}-aiplatform.googleapis.com/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        ")\n",
+        "try:\n",
+        "    if use_dedicated_endpoint:\n",
+        "        BASE_URL = f\"https://{DEDICATED_ENDPOINT_DNS}/v1beta1/{ENDPOINT_RESOURCE_NAME}\"\n",
+        "except NameError:\n",
+        "    pass\n",
+        "\n",
+        "client = openai.OpenAI(base_url=BASE_URL, api_key=creds.token)\n",
+        "\n",
+        "model_response = client.chat.completions.create(\n",
+        "    model=\"\",\n",
+        "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
+        "    temperature=temperature,\n",
+        "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
+        ")\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tAelDidov5AW"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "8SeZCFo5v7z-"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_ollama_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_phi3_deployment.ipynb

@@ -102,14 +102,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
         "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
@@ -136,6 +136,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -378,6 +384,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -430,6 +437,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -468,6 +479,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_phi3_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -653,7 +665,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -747,6 +759,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_phi3_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_phi4_deployment.ipynb

@@ -103,14 +103,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "import datetime\n",
@@ -134,6 +134,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -279,6 +285,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -331,6 +338,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -369,6 +380,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_phi4_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -540,7 +552,7 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
+        "    service_account: str = None,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -634,6 +646,7 @@
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_phi4_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",

notebooks/community/model_garden/model_garden_pytorch_autogluon.ipynb

@@ -133,6 +133,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -209,9 +215,9 @@
         "job_name = common_util.get_job_name_with_datetime(prefix=\"pytorch-autogluon\")\n",
         "\n",
         "finetuning_workdir = os.path.join(BUCKET_URI, job_name)\n",
-        "train_data_path = (\n",
-        "    \"https://raw.githubusercontent.com/mli/ag-docs/main/knot_theory/train.csv\"\n",
-        ")\n",
+        "\n",
+        "# @markdown `train_data_path` - The path to the training data. It can also be a GCS path.\n",
+        "train_data_path = \"https://raw.githubusercontent.com/mli/ag-docs/main/knot_theory/train.csv\"  # @param {type:\"string\"}\n",
         "\n",
         "# @markdown `label` - The column id to predict.\n",
         "label = \"signature\"  # @param {type:\"string\"}\n",
@@ -221,13 +227,14 @@
         "\n",
         "docker_args_list = [\n",
         "    \"--train_data_path\",\n",
-        "    train_data_path,\n",
+        "    common_util.gcs_fuse_path(train_data_path),\n",
         "    \"--label\",\n",
         "    label,\n",
         "    \"--model_save_path\",\n",
-        "    f\"{common_util.gcs_fuse_path(finetuning_workdir)}\",\n",
+        "    common_util.gcs_fuse_path(finetuning_workdir),\n",
         "]\n",
-        "print(docker_args_list)\n",
+        "print(\"Running training job with args:\")\n",
+        "print(\" \\\\\\n\".join(docker_args_list))\n",
         "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",

notebooks/community/model_garden/model_garden_pytorch_bart_large_cnn.ipynb

@@ -113,6 +113,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_pytorch_biogpt_serve.ipynb

@@ -70,6 +70,9 @@
         "| :- | :- |\n",
         "| [microsoft/biogpt](https://huggingface.co/docs/transformers/model_doc/biogpt) | N |\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
         "\n",
         "### Costs\n",
         "\n",
@@ -103,29 +106,24 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
+        "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
@@ -142,52 +140,138 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"biogpt_serve\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_biogpt_serve.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint"
       ]
     },
     {
@@ -227,7 +311,6 @@
         "def deploy_model(\n",
         "    model_name,\n",
         "    model_id,\n",
-        "    service_account,\n",
         "    task,\n",
         "    max_length,\n",
         "    num_return_sequences,\n",
@@ -264,7 +347,6 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_biogpt_serve.ipynb\"},\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -281,7 +363,6 @@
         "models[\"biogpt_model\"], endpoints[\"biogpt_endpoint\"] = deploy_model(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=TASK),\n",
         "    model_id=MODEL_ID,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    task=TASK,\n",
         "    max_length=200,\n",
         "    num_return_sequences=10,\n",
@@ -337,11 +418,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_biomedclip.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_biomedclip.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -62,6 +67,10 @@
         "- Serve BiomedCLIP using Vertex AI\n",
         "- Run zero-shot image classification with BiomedCLIP\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -96,42 +105,41 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
         "from typing import Tuple\n",
         "\n",
         "import numpy as np\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "LABEL = \"biomedclip_serve\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -141,52 +149,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"biomedclip\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\""
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
       ]
     },
     {
@@ -198,35 +171,86 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Set the model parameters\n",
+        "\n",
+        "base_model_name = \"biomedclip\"\n",
+        "PUBLISHER_MODEL_NAME = (\n",
+        "    f\"publishers/microsoft/models/microsoft-biomedclip@{base_model_name}\"\n",
+        ")\n",
+        "model_id = \"hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224\"\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
+        "TASK = \"zero-shot-image-classification\"\n",
+        "PRECISION = \"amp\"\n",
+        "SERVE_PORT = 7080\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_L4\":\n",
+        "    machine_type = \"g2-standard-8\"\n",
+        "    accelerator_count = 1\n",
+        "elif accelerator_type == \"NVIDIA_TESLA_V100\":\n",
+        "    machine_type = \"n1-standard-8\"\n",
+        "    accelerator_count = 4\n",
+        "else:\n",
+        "    raise ValueError(f\"Recommended GPU setting not found for: {accelerator_type}\")\n",
+        "\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=1,\n",
+        "    is_for_training=False,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6RXJEDT-dLlL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "OYNaiywCdNYS"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
         "\n",
         "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint.\n",
         "\n",
         "# @markdown The model deployment step will take ~10 minutes to complete.\n",
         "\n",
-        "model_id = \"hf-hub:microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224\"\n",
-        "serving_accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_V100\"]\n",
-        "TASK = \"zero-shot-image-classification\"\n",
-        "PRECISION = \"amp\"\n",
-        "SERVE_PORT = 7080\n",
-        "\n",
         "# The pre-built serving docker image.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-serve:20250128_1414_RC00\"\n",
-        "\n",
-        "if serving_accelerator_type == \"NVIDIA_L4\":\n",
-        "    serving_machine_type = \"g2-standard-8\"\n",
-        "elif serving_accelerator_type == \"NVIDIA_TESLA_V100\":\n",
-        "    serving_machine_type = \"n1-standard-8\"\n",
-        "else:\n",
-        "    raise ValueError(\n",
-        "        f\"Recommended GPU setting not found for: {serving_accelerator_type}\"\n",
-        "    )\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-open-clip-serve:20250313_0922_RC00\"\n",
         "\n",
         "\n",
         "def deploy_model(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
-        "    service_account: str,\n",
         "    task: str,\n",
         "    precision: str,\n",
         "    machine_type: str,\n",
@@ -259,28 +283,21 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_biomedclip.ipynb\"},\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_biomedclip.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "common_util.check_quota(\n",
-        "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
-        "    accelerator_type=serving_accelerator_type,\n",
-        "    accelerator_count=1,\n",
-        "    is_for_training=False,\n",
-        ")\n",
-        "\n",
-        "models[\"biomedclip_model\"], endpoints[\"biomedclip_endpoint\"] = deploy_model(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"biomedclip-serve\"),\n",
         "    model_id=model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    task=TASK,\n",
         "    precision=PRECISION,\n",
-        "    machine_type=serving_machine_type,\n",
-        "    accelerator_type=serving_accelerator_type,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=1,\n",
         ")"
       ]
@@ -345,7 +362,7 @@
         "    {\"text\": \"This is a photo of hematoxylin and eosin histopathology\"},\n",
         "    {\"text\": \"This is a photo of pie chart\"},\n",
         "]\n",
-        "response = endpoints[\"biomedclip_endpoint\"].predict(instances=instances)\n",
+        "response = endpoints[LABEL].predict(instances=instances)\n",
         "\n",
         "print(response.predictions)\n",
         "\n",
@@ -372,11 +389,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_blip_vqa.ipynb

@@ -114,6 +114,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_pytorch_codellama.ipynb

@@ -100,14 +100,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
@@ -134,6 +134,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -265,6 +271,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -317,6 +324,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -355,6 +366,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_codellama.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",

notebooks/community/model_garden/model_garden_pytorch_csm_deployment.ipynb

@@ -0,0 +1,353 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "IayCWUliSqzV"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "tpUJmjb5SqzV"
+      },
+      "source": [
+        "# Vertex AI Model Garden - CSM-1B Deployment\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_csm_deployment.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_csm_deployment.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "-oUrHK7JSqzV"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates serving the [sesame/csm-1b](https://huggingface.co/sesame/csm-1b) model on Vertex. CSM (Conversational Speech Model) is a speech generation model from Sesame that generates RVQ audio codes from text and audio inputs. The model architecture employs a Llama backbone and a smaller audio decoder that produces Mimi audio codes.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "\n",
+        "- Deploy [sesame/csm-1b](https://huggingface.co/sesame/csm-1b) on Vertex AI.\n",
+        "- Generate conversation audios with the deployed endpoint.\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2s8qdX4USqzV"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "GpuzpvUVSqzV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# @markdown 4. To use the model, you need to accept the agreement of [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B) and [CSM-1B](https://huggingface.co/sesame/CSM-1B) on Hugging Face.\n",
+        "\n",
+        "# @markdown 5. Set Hugging Face access token in `HF_TOKEN` field. If you don't already have a \"read\" access token, follow the [Hugging Face documentation](https://huggingface.co/docs/hub/en/security-tokens) to create an access token with \"read\" permission. You can find your existing access tokens in the Hugging Face [Access Token](https://huggingface.co/settings/tokens) page.\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import base64\n",
+        "import importlib\n",
+        "import os\n",
+        "from typing import Tuple\n",
+        "\n",
+        "from google.cloud import aiplatform\n",
+        "from IPython.core.display import display\n",
+        "from IPython.display import Audio\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
+        "\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "\n",
+        "HF_TOKEN = \"\"  # @param {type:\"string\", isTemplate: true}\n",
+        "assert HF_TOKEN, \"Set Hugging Face access token in `HF_TOKEN`.\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dlozukBDSqzV"
+      },
+      "source": [
+        "## Deploy and predict with the CSM-1B model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "W5D8X2LFSqzV"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads the CSM-1B model to Model Registry and deploys it to a Vertex Prediction Endpoint with 1 `NVIDIA_L4` GPU and `g2-standard-8` machine type.\n",
+        "# @markdown It takes ~10 minutes to finish.\n",
+        "\n",
+        "# @markdown It's recommended to use the region selected by the deployment button on the model card. If the deployment button is not available, it's recommended to stay with the default region of the notebook.\n",
+        "\n",
+        "model_id = \"sesame/csm-1b\"\n",
+        "publisher, publisher_model_id = model_id.split(\"/\")\n",
+        "\n",
+        "PYTORCH_DOCKER_URI = (\n",
+        "    \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-csm-serve\"\n",
+        ")\n",
+        "\n",
+        "# @markdown Use a [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint) for the deployment.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "accelerator_count = 1\n",
+        "machine_type = \"g2-standard-8\"\n",
+        "\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "\n",
+        "def deploy_model_pytorch(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    task: str,\n",
+        "    handler: str = \"\",\n",
+        "    service_account: str | None = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys models with Model Garden Pytorch Inference on GPU in Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": model_id,\n",
+        "        \"TASK\": task,\n",
+        "    }\n",
+        "\n",
+        "    if handler:\n",
+        "        env_vars[\"HANDLER\"] = handler\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=PYTORCH_DOCKER_URI,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/predict\",\n",
+        "        serving_container_health_route=\"/health\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=3600,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_csm_deployment.ipynb\",\n",
+        "            \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "        },\n",
+        "        service_account=service_account,\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
+        "\n",
+        "models[\"pytorch_gpu\"], endpoints[\"pytorch_gpu\"] = deploy_model_pytorch(\n",
+        "    model_name=common_util.get_job_name_with_datetime(prefix=\"csm-1b-serve\"),\n",
+        "    model_id=model_id,\n",
+        "    publisher=publisher,\n",
+        "    publisher_model_id=publisher_model_id,\n",
+        "    task=\"text-to-speech\",\n",
+        "    service_account=None,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        ")\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "mRDvRmOJX2IX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Predict\n",
+        "\n",
+        "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts. Note that the first few prompts will take longer to execute.\n",
+        "\n",
+        "# @markdown Here, we use the following conversation example:\n",
+        "# @markdown - Speaker 0: I just won a million dollar lottery.\n",
+        "# @markdown - Speaker 1: You're kidding me!\n",
+        "\n",
+        "instances = [\n",
+        "    {\"speaker\": 0, \"text\": \"I just won a million dollar lottery.\"},\n",
+        "    {\"speaker\": 1, \"text\": \"You're kidding me!\"},\n",
+        "]\n",
+        "\n",
+        "response = endpoints[\"pytorch_gpu\"].predict(\n",
+        "    instances=instances, use_dedicated_endpoint=use_dedicated_endpoint\n",
+        ")\n",
+        "\n",
+        "for prediction in response.predictions:\n",
+        "    display(Audio(base64.b64decode(prediction[\"audio\"])))\n",
+        "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "EsfASp4WX2IX"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "cak7S2PnX2IX"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the models and endpoints\n",
+        "\n",
+        "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
+        "\n",
+        "# Undeploy model and delete endpoint.\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
+        "\n",
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_csm_deployment.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_deployed_model_agent_engine.ipynb

@@ -0,0 +1,612 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "SgQ6t5bqZVlH"
+      },
+      "outputs": [],
+      "source": [
+        "# Copyright 2025 Google LLC\n",
+        "#\n",
+        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+        "# you may not use this file except in compliance with the License.\n",
+        "# You may obtain a copy of the License at\n",
+        "#\n",
+        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+        "#\n",
+        "# Unless required by applicable law or agreed to in writing, software\n",
+        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+        "# See the License for the specific language governing permissions and\n",
+        "# limitations under the License."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99c1c3fc2ca5"
+      },
+      "source": [
+        "# Vertex AI Model Garden + Agent Engine - Build, Deploy and Test Agents using a Self-deployed Endpoint\n",
+        "\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_deployed_model_agent_engine.ipynb\">\n",
+        "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_deployed_model_agent_engine.ipynb\">\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
+        "    </a>\n",
+        "  </td>\n",
+        "</tr></tbody></table>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3de7470326a2"
+      },
+      "source": [
+        "## Overview\n",
+        "\n",
+        "This notebook demonstrates how to build, deploy and test three types of agents using [Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview) with self-deployed model in Vertex AI.\n",
+        "\n",
+        "[Agent Engine](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview) (LangChain on Vertex AI) is a managed service in Vertex AI that helps you build and deploy model-based agents. It gives you the flexibility to choose how much reasoning you want to delegate to the LLM and how much you want to handle with custom code.\n",
+        "\n",
+        "A previous [notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb) demonstrates how to use Llama 3.1 models as Model-as-a-service (MaaS) to build `chatbot` and `translator` agents.\n",
+        "\n",
+        "\n",
+        "### Objective\n",
+        "    \n",
+        "- Integrate with Agent Engine: Use the Vertex AI SDK to build three simple agents with the deployed endpoint:\n",
+        "    - A Chatbot Agent\n",
+        "    - A Translator Agent\n",
+        "    - An Agent that uses [an Exchange Rate Tool](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/develop#define-function)\n",
+        "- Test your agent locally.\n",
+        "- Deploy and test your agent on the Agent Engine.\n",
+        "\n",
+        "\n",
+        "### Costs\n",
+        "\n",
+        "This tutorial uses billable components of Google Cloud:\n",
+        "\n",
+        "* Vertex AI\n",
+        "* Cloud Storage\n",
+        "\n",
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "264c07757582"
+      },
+      "source": [
+        "## Before you begin"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "language": "python",
+      "metadata": {
+        "cellView": "form",
+        "id": "YXFGIp1l-qtT"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Setup Google Cloud project\n",
+        "\n",
+        "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
+        "\n",
+        "# @markdown 2. Create a bucket created for agent engine.\n",
+        "\n",
+        "BUCKET_NAME = \"\"  # @param {type:\"string\", placeholder: \"[your-bucket-name]\"}\n",
+        "STAGING_BUCKET = f\"gs://{BUCKET_NAME}\"\n",
+        "\n",
+        "# @markdown 3. You can find the deployed model endpoint in the [Vertex AI console](https://console.cloud.google.com/vertex-ai/endpoints).\n",
+        "DEPLOYED_MODEL_ENDPOINT = \"\"  # @param {type:\"string\", placeholder: \"[your-deployed-model-endpoint]\"}\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet \\\n",
+        "    \"google-cloud-aiplatform>=1.64.0\" \\\n",
+        "    cloudpickle==3.0.0 \\\n",
+        "    pydantic==2.10.6 \\\n",
+        "    requests \\\n",
+        "    langchain-openai\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "import os\n",
+        "import requests\n",
+        "from typing import Tuple\n",
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "# Get the default cloud project id.\n",
+        "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
+        "\n",
+        "# Get the default region for launching jobs.\n",
+        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "    \n",
+        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
+        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
+        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "blbSoA-3gSyN"
+      },
+      "source": [
+        "### Initialization"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "KsF63jInoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Authenticate your notebook environment (Colab only)\n",
+        "import sys\n",
+        "\n",
+        "if \"google.colab\" in sys.modules:\n",
+        "\n",
+        "    from google.colab import auth\n",
+        "\n",
+        "    auth.authenticate_user()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "SLq7whlDoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Initialize Vertex AI SDK for Python\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "TFkSqK91ki8E"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Import libraries\n",
+        "\n",
+        "# @markdown Import libraries to use in this tutorial.\n",
+        "\n",
+        "import google.auth\n",
+        "from langchain_core.output_parsers import StrOutputParser\n",
+        "from langchain_core.prompts import PromptTemplate\n",
+        "from langchain_openai import ChatOpenAI\n",
+        "from vertexai import agent_engines\n",
+        "from vertexai.preview import reasoning_engines"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "71WcCJ57gJZM"
+      },
+      "source": [
+        "### Chat with `Agent Engine`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "sgG-NyAvoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title `Agent Engine` use self-deployed API endpoint with different configuration\n",
+        "\n",
+        "# @markdown To use the self-deployed API endpoint with Agent Engine capabilities, you need to request the access token and configure the langchain ChatOpenAI to point to the API endpoint.\n",
+        "\n",
+        "# @markdown In previous [notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb), we demonstrated how to `Ask Llama 3.1 using different model configuration`.\n",
+        "\n",
+        "# @markdown In this colab, we will show you how to use the `Agent Engine` to send a request to the self-deployed API endpoint with different model configuration.\n",
+        "\n",
+        "\n",
+        "def model_builder(\n",
+        "    *,\n",
+        "    model_name: str,\n",
+        "    model_kwargs=None,\n",
+        "    project: str,  # Specified via vertexai.init\n",
+        "    location: str,  # Specified via vertexai.init\n",
+        "    **kwargs,\n",
+        "):\n",
+        "\n",
+        "    # Note: the credential lives for 1 hour by default.\n",
+        "    # After expiration, it must be refreshed.\n",
+        "    creds, _ = google.auth.default(\n",
+        "        scopes=[\"https://www.googleapis.com/auth/cloud-platform\"]\n",
+        "    )\n",
+        "    auth_req = google.auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    if model_kwargs is None:\n",
+        "        model_kwargs = {}\n",
+        "\n",
+        "    return ChatOpenAI(\n",
+        "        model=\"\",\n",
+        "        base_url=DEPLOYED_MODEL_ENDPOINT,\n",
+        "        api_key=creds.token,\n",
+        "        **model_kwargs,\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "# @markdown Use the following parameters to generate different answers:\n",
+        "# @markdown *   `temperature` to control the randomness of the response\n",
+        "# @markdown *   `top_p` to control the quality of the response\n",
+        "\n",
+        "temperature = 1.0  # @param {type:\"number\"}\n",
+        "top_p = 1.0  # @param {type:\"number\"}\n",
+        "\n",
+        "agent = reasoning_engines.LangchainAgent(\n",
+        "    model=\"\",  # Required.\n",
+        "    model_builder=model_builder,  # Required.\n",
+        "    model_kwargs={\n",
+        "        \"temperature\": temperature,  # Optional.\n",
+        "        \"top_p\": top_p,  # Optional.\n",
+        "        \"extra_body\": {},\n",
+        "    },\n",
+        ")\n",
+        "\n",
+        "# @markdown Now we can test the model and agent behavior to ensure that it's working as expected before we deploy it:\n",
+        "\n",
+        "response = agent.query(input=\"Hello, how are you!\")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rn-cHx7xoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy your agent on Vertex AI\n",
+        "\n",
+        "# @markdown Now that you've specified a model, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!\n",
+        "\n",
+        "remote_agent = agent_engines.create(\n",
+        "    agent,\n",
+        "    requirements=[\n",
+        "        \"google-cloud-aiplatform[langchain,agent_engines]\",\n",
+        "        \"cloudpickle==3.0.0\",\n",
+        "        \"pydantic==2.10.6\",\n",
+        "        \"requests\",\n",
+        "        \"langchain-openai\",\n",
+        "    ],\n",
+        ")\n",
+        "\n",
+        "response = remote_agent.query(input=\"Hello, how are you!\")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "MjwPLr0LoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Reusing your deployed agent from other applications or SDKs\n",
+        "\n",
+        "# @markdown The remotely deployed `Agent Engine` is now available for import and use. You can access it within your current notebook session, a different notebook, or a Python script.\n",
+        "\n",
+        "AGENT_ENGINE_RESOURCE_NAME = remote_agent.resource_name\n",
+        "print(AGENT_ENGINE_RESOURCE_NAME)\n",
+        "\n",
+        "# Afterwards, you can use the below code:\n",
+        "\n",
+        "# from vertexai.preview import agent_engines`\n",
+        "\n",
+        "# remote_agent = agent_engines.get(AGENT_ENGINE_RESOURCE_NAME)`\n",
+        "# response = remote_agent.query(input=query)`\n",
+        "\n",
+        "# @markdown Alternatively, you can query your agent from other programming languages using any of the [available client libraries in Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries), including C#, Java, Node.js, Python, Go, or REST API."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ato0F6XFoByL"
+      },
+      "source": [
+        "### Simple Translator Agent"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "bI3JqaeMoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Use Agent Engine to build a simple translator agent\n",
+        "\n",
+        "# @markdown In previous [notebook](https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_openai_api_llama3_1.ipynb), we demonstrates how to use `LangChain Expression Language` (LCEL) to build a simple chain which translates some `text_to_translate` to the specified `target_language`.\n",
+        "\n",
+        "# @markdown In this colab, we will show you how to use the `Agent Engine` to build and deploy the agent.\n",
+        "\n",
+        "\n",
+        "def lcel_builder(*, model, **kwargs):\n",
+        "\n",
+        "    template = \"\"\"Translate the following {text} to {target_language}:\"\"\"\n",
+        "    prompt = PromptTemplate(\n",
+        "        input_variables=[\"text\", \"target_language\"], template=template\n",
+        "    )\n",
+        "\n",
+        "    return prompt | model | StrOutputParser()\n",
+        "\n",
+        "\n",
+        "agent = reasoning_engines.LangchainAgent(\n",
+        "    model=\"\",\n",
+        "    model_builder=model_builder,\n",
+        "    runnable_builder=lcel_builder,\n",
+        ")\n",
+        "\n",
+        "text_to_translate = \"\"  # @param {type:\"string\", placeholder:\"Hello, how are you!\"}\n",
+        "target_language = \"\"  # @param {type:\"string\", placeholder:\"Italian\"}\n",
+        "\n",
+        "response = agent.query(\n",
+        "    input={\"text\": text_to_translate, \"target_language\": target_language}\n",
+        ")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "XNF9slsSEHLz"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy your agent on Vertex AI\n",
+        "\n",
+        "# @markdown Now that you've specified a model, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!\n",
+        "\n",
+        "remote_agent = agent_engines.create(\n",
+        "    agent,\n",
+        "    requirements=[\n",
+        "        \"google-cloud-aiplatform[langchain,agent_engines]\",\n",
+        "        \"cloudpickle==3.0.0\",\n",
+        "        \"pydantic==2.10.6\",\n",
+        "        \"requests\",\n",
+        "        \"langchain-openai\",\n",
+        "    ],\n",
+        ")\n",
+        "\n",
+        "response = remote_agent.query(\n",
+        "    input={\"text\": text_to_translate, \"target_language\": target_language}\n",
+        ")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "tnbhCYwbERSx"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Reusing your deployed agent from other applications or SDKs\n",
+        "\n",
+        "# @markdown The remotely deployed `Agent Engine` is now available for import and use. You can access it within your current notebook session, a different notebook, or a Python script.\n",
+        "\n",
+        "AGENT_ENGINE_RESOURCE_NAME = remote_agent.resource_name\n",
+        "print(AGENT_ENGINE_RESOURCE_NAME)\n",
+        "\n",
+        "# Afterwards, you can use the below code:\n",
+        "\n",
+        "# from vertexai.preview import agent_engines`\n",
+        "\n",
+        "# remote_agent = agent_engines.get(AGENT_ENGINE_RESOURCE_NAME)`\n",
+        "# response = remote_agent.query(input=query)`\n",
+        "\n",
+        "# @markdown Alternatively, you can query your agent from other programming languages using any of the [available client libraries in Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries), including C#, Java, Node.js, Python, Go, or REST API."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0lI86Eu7oByL"
+      },
+      "source": [
+        "### Exchange Rate Tool\n",
+        "\n",
+        "[Function calling](https://cloud.google.com/vertex-ai/docs/generative-ai/multimodal/function-calling) lets developers create a description of a function in their code, then pass that description to a language model in a request. The response from the model includes the name of a function that matches the description and the arguments to call it with.\n",
+        "\n",
+        "In this example, we will use an Exchange Rate tool in the Agent Engine."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Dmv42OOOoByL"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Agent that uses an Exchange Rate Tool\n",
+        "\n",
+        "# @markdown Tools and functions enable the generative model to interact with external systems, databases, document stores, and other APIs so that the model can get the most up-to-date information or take action with those systems.\n",
+        "\n",
+        "# @markdown In this example, you'll define a function called get_exchange_rate that uses the requests library to retrieve real-time currency exchange information from an API:\n",
+        "\n",
+        "\n",
+        "def get_exchange_rate(\n",
+        "    currency_from: str = \"USD\",\n",
+        "    currency_to: str = \"EUR\",\n",
+        "    currency_date: str = \"latest\",\n",
+        "):\n",
+        "    \"\"\"Retrieves the exchange rate between two currencies on a specified date.\n",
+        "    Args:\n",
+        "        currency_from: The source currency code.\n",
+        "        currency_to: The target currency code.\n",
+        "        currency_date: The date to retrieve the exchange rate.\n",
+        "    Returns:\n",
+        "        Exchange rate between two currencies on a specified date.\n",
+        "    \"\"\"\n",
+        "    response = requests.get(\n",
+        "        f\"https://api.frankfurter.app/{currency_date}\",\n",
+        "        params={\"from\": currency_from, \"to\": currency_to},\n",
+        "    )\n",
+        "    return response.json()\n",
+        "\n",
+        "\n",
+        "get_exchange_rate(currency_from=\"USD\", currency_to=\"SEK\")\n",
+        "\n",
+        "\n",
+        "agent = reasoning_engines.LangchainAgent(\n",
+        "    model=\"\",  # Required.\n",
+        "    model_builder=model_builder,  # Required.\n",
+        "    tools=[get_exchange_rate],  # Optional.\n",
+        "    agent_executor_kwargs={\n",
+        "        \"return_intermediate_steps\": True,\n",
+        "        \"stream_runnable\": False,\n",
+        "    },  # Optional.\n",
+        ")\n",
+        "\n",
+        "# @markdown Test the function with sample inputs to ensure that it's working as expected:\n",
+        "response = agent.query(\n",
+        "    input=\"What's the exchange rate from US dollars to Swedish currency at 2024-07-26?\"\n",
+        ")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "6y2G_bjbDam_"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy your agent on Vertex AI\n",
+        "\n",
+        "# @markdown Now that you've specified a model, and reasoning for your agent and tested it out, you're ready to deploy your agent as a remote service in Vertex AI!\n",
+        "\n",
+        "remote_agent = agent_engines.create(\n",
+        "    agent,\n",
+        "    requirements=[\n",
+        "        \"google-cloud-aiplatform[langchain,agent_engines]\",\n",
+        "        \"cloudpickle==3.0.0\",\n",
+        "        \"pydantic==2.10.6\",\n",
+        "        \"requests\",\n",
+        "        \"langchain-openai\",\n",
+        "    ],\n",
+        ")\n",
+        "\n",
+        "response = remote_agent.query(\n",
+        "    input=\"What's the exchange rate from US dollars to Swedish currency at 2024-07-26?\"\n",
+        ")\n",
+        "print(response)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "oEbI1hm1KoQE"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Reusing your deployed agent from other applications or SDKs\n",
+        "\n",
+        "# @markdown The remotely deployed `Agent Engine` is now available for import and use. You can access it within your current notebook session, a different notebook, or a Python script.\n",
+        "\n",
+        "AGENT_ENGINE_RESOURCE_NAME = remote_agent.resource_name\n",
+        "print(AGENT_ENGINE_RESOURCE_NAME)\n",
+        "\n",
+        "# Afterwards, you can use the below code:\n",
+        "\n",
+        "# from vertexai.preview import agent_engines`\n",
+        "\n",
+        "# remote_agent = agent_engines.get(AGENT_ENGINE_RESOURCE_NAME)`\n",
+        "# response = remote_agent.query(input=query)`\n",
+        "\n",
+        "# @markdown Alternatively, you can query your agent from other programming languages using any of the [available client libraries in Vertex AI](https://cloud.google.com/vertex-ai/docs/start/client-libraries), including C#, Java, Node.js, Python, Go, or REST API."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "JETd33jIDcjm"
+      },
+      "source": [
+        "## Clean up resources"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "911406c1561e"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Delete the buckets and agent engines\n",
+        "\n",
+        "delete_bucket = False  # @param {type:\"boolean\"}\n",
+        "if delete_bucket:\n",
+        "    ! gsutil -m rm -r $BUCKET_NAME\n",
+        "\n",
+        "delete_agent_engine = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "if delete_agent_engine:\n",
+        "    remote_agent.delete()"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "model_garden_pytorch_deployed_model_agent_engine.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

notebooks/community/model_garden/model_garden_pytorch_deployed_model_reasoning_engine.ipynb

@@ -179,9 +179,7 @@
         "# @title Initialize Vertex AI SDK for Python\n",
         "import vertexai\n",
         "\n",
-        "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "MODEL_ID = \"meta/llama-3.1-405b-instruct-maas\"  # this is a placeholder only, and it is not being used by the agent."
+        "vertexai.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)"
       ]
     },
     {
@@ -252,7 +250,7 @@
         "        model_kwargs = {}\n",
         "\n",
         "    return ChatOpenAI(\n",
-        "        model=model_name,\n",
+        "        model=\"\",\n",
         "        base_url=DEPLOYED_MODEL_ENDPOINT,\n",
         "        api_key=creds.token,\n",
         "        **model_kwargs,\n",
@@ -267,7 +265,7 @@
         "top_p = 1.0  # @param {type:\"number\"}\n",
         "\n",
         "agent = reasoning_engines.LangchainAgent(\n",
-        "    model=MODEL_ID,  # Required.\n",
+        "    model=\"\",  # Required.\n",
         "    model_builder=model_builder,  # Required.\n",
         "    model_kwargs={\n",
         "        \"temperature\": temperature,  # Optional.\n",
@@ -372,7 +370,7 @@
         "\n",
         "\n",
         "agent = reasoning_engines.LangchainAgent(\n",
-        "    model=MODEL_ID,\n",
+        "    model=\"\",\n",
         "    model_builder=model_builder,\n",
         "    runnable_builder=lcel_builder,\n",
         ")\n",
@@ -495,7 +493,7 @@
         "\n",
         "\n",
         "agent = reasoning_engines.LangchainAgent(\n",
-        "    model=MODEL_ID,  # Required.\n",
+        "    model=\"\",  # Required.\n",
         "    model_builder=model_builder,  # Required.\n",
         "    tools=[get_exchange_rate],  # Optional.\n",
         "    agent_executor_kwargs={\n",

notebooks/community/model_garden/model_garden_pytorch_flux.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_flux.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -62,6 +67,10 @@
         "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-to-image.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -94,32 +103,38 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "\n",
-        "import base64\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
-        "from io import BytesIO\n",
         "\n",
         "from google.cloud import aiplatform\n",
-        "from PIL import Image\n",
+        "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "models, endpoints = {}, {}\n",
+        "LABEL = \"xdit_gpu\"\n",
+        "\n",
         "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -128,76 +143,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"flux\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
         "\n",
-        "models, endpoints = {}, {}\n",
-        "\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "\n",
-        "common_util = importlib.import_module(\n",
-        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
-        ")\n",
-        "\n",
-        "\n",
-        "def base64_to_image(image_str):\n",
-        "    \"\"\"Convert base64 encoded string to an image.\"\"\"\n",
-        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
-        "    return image\n",
-        "\n",
-        "\n",
-        "def image_grid(imgs, rows=2, cols=2):\n",
-        "    w, h = imgs[0].size\n",
-        "    grid = Image.new(\n",
-        "        mode=\"RGB\", size=(cols * w + 10 * cols, rows * h), color=(255, 255, 255)\n",
-        "    )\n",
-        "    for i, img in enumerate(imgs):\n",
-        "        grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))\n",
-        "    return grid"
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
       ]
     },
     {
@@ -209,27 +165,73 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy the model to Vertex for online predictions\n",
+        "# @title Set the model parameters\n",
+        "\n",
+        "base_model_name = \"flux.1-schnell\"\n",
+        "PUBLISHER_MODEL_NAME = (\n",
+        "    f\"publishers/black-forest-labs/models/flux1-schnell@{base_model_name}\"\n",
+        ")\n",
+        "\n",
+        "MODEL_ID = \"black-forest-labs/FLUX.1-schnell\"\n",
+        "TASK = \"text-to-image\"\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_TESLA_A100\"  # @param [\"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "accelerator_count = 1\n",
+        "\n",
+        "if accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "    machine_type = \"a2-highgpu-1g\"\n",
+        "elif accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "    machine_type = \"a2-ultragpu-1g\"\n",
+        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    machine_type = \"a3-highgpu-2g\"\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported accelerator type: {accelerator_type}\")\n",
+        "\n",
+        "# Dedicated endpoint is not supported.\n",
+        "use_dedicated_endpoint = False"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "No2EqZPiYamO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "lSD2g1pYYamO"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
         "\n",
         "# @markdown This section uploads the [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) model to Model Registry and deploys it on the Endpoint with 1 A100 80G GPU.\n",
         "\n",
         "# @markdown The deployment takes ~15 minutes to finish.\n",
         "\n",
-        "model_id = \"black-forest-labs/FLUX.1-schnell\"\n",
-        "task = \"text-to-image\"\n",
-        "\n",
-        "accelerator_type = \"NVIDIA_TESLA_A100\"  # @param [\"NVIDIA_TESLA_A100\", \"NVIDIA_A100_80GB\"]\n",
-        "\n",
-        "machine_type_map = {\n",
-        "    \"NVIDIA_TESLA_A100\": \"a2-highgpu-1g\",\n",
-        "    \"NVIDIA_A100_80GB\": \"a2-ultragpu-1g\",\n",
-        "}\n",
-        "\n",
-        "machine_type = machine_type_map.get(accelerator_type)\n",
-        "accelerator_count = 1\n",
-        "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
-        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/pytorch-inference.cu125.0-1.ubuntu2204.py310\"\n",
+        "SERVE_DOCKER_URI = \"us-docker.pkg.dev/deeplearning-platform-release/vertex-model-garden/xdit-serve.cu125.0-1.ubuntu2204.py310\"\n",
         "\n",
         "\n",
         "def deploy_model(model_id, task, machine_type, accelerator_type, accelerator_count):\n",
@@ -251,6 +253,12 @@
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
         "    }\n",
         "\n",
+        "    # xDiT serving parameters\n",
+        "    serving_env[\"USE_TORCH_COMPILE\"] = \"true\"\n",
+        "    serving_env[\"N_GPUS\"] = accelerator_count\n",
+        "    if accelerator_count == 2:\n",
+        "        serving_env[\"RING_DEGREE\"] = \"2\"\n",
+        "\n",
         "    model = aiplatform.Model.upload(\n",
         "        display_name=model_name,\n",
         "        serving_container_image_uri=SERVE_DOCKER_URI,\n",
@@ -267,21 +275,23 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=SERVICE_ACCOUNT,\n",
-        "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_flux.ipynb\"},\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_flux.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
         "    )\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "models[\"model\"], endpoints[\"endpoint\"] = deploy_model(\n",
-        "    model_id=model_id,\n",
-        "    task=task,\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
+        "    model_id=MODEL_ID,\n",
+        "    task=TASK,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         ")\n",
         "\n",
-        "print(\"endpoint_name:\", endpoints[\"endpoint\"].name)"
+        "print(\"endpoint_name:\", endpoints[LABEL].name)"
       ]
     },
     {
@@ -297,8 +307,6 @@
         "\n",
         "# @markdown Once deployment succeeds, you can send requests to the endpoint with text prompts.\n",
         "\n",
-        "# @markdown The inference takes ~3s with 1 A100 GPU.\n",
-        "\n",
         "# @markdown Example:\n",
         "\n",
         "# @markdown ```\n",
@@ -321,11 +329,12 @@
         "\n",
         "# The default num inference steps is set to 4 in the serving container, but\n",
         "# you can change it to your own preference for image quality in the request.\n",
-        "response = endpoints[\"endpoint\"].predict(instances=instances, parameters=parameters)\n",
+        "response = endpoints[LABEL].predict(instances=instances, parameters=parameters)\n",
         "images = [\n",
-        "    base64_to_image(prediction.get(\"output\")) for prediction in response.predictions\n",
+        "    common_util.base64_to_image(prediction.get(\"output\"))\n",
+        "    for prediction in response.predictions\n",
         "]\n",
-        "image_grid(images, rows=1)"
+        "common_util.image_grid(images, rows=1)"
       ]
     },
     {
@@ -347,11 +356,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_flux_gradio.ipynb

@@ -97,14 +97,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "! pip3 install --upgrade gradio==4.29.0 opencv-python\n",
         "# Uninstall nest-asyncio and uvloop as a workaround to https://github.com/gradio-app/gradio/issues/8238#issuecomment-2101066984\n",

notebooks/community/model_garden/model_garden_pytorch_gemma_peft_finetuning_hf.ipynb

@@ -130,7 +130,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -165,6 +165,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -665,6 +671,18 @@
         "max_model_len = 2048\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
@@ -689,6 +707,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -741,6 +760,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -779,6 +802,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_gemma_peft_finetuning_hf.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",

notebooks/community/model_garden/model_garden_pytorch_instant_id.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_instant_id.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -61,6 +66,10 @@
         "- Deploy the model to a [Vertex AI Endpoint resource](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-to-image.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -93,28 +102,39 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
         "# Import the necessary packages\n",
-        "import datetime\n",
+        "\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
+        "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
         "models, endpoints = {}, {}\n",
+        "LABEL = \"vllm_gpu\"\n",
         "\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -124,52 +144,19 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"instant_id\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# @markdown Click \"Show code\" to see more details."
       ]
@@ -192,20 +179,17 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy the InstantID model to Vertex for online predictions\n",
-        "\n",
-        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes ~15 minutes to finish.\n",
-        "# @markdown Click \"Show Code\" to see more details.\n",
-        "\n",
+        "base_model_name = \"instantid\"\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/instantx/models/instant-id@{base_model_name}\"\n",
         "model_id = \"instantx/instantid\"\n",
         "task = \"instant-id\"\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
         "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
         "\n",
-        "ACCELERATOR_TYPE = \"NVIDIA_L4\"\n",
-        "ACCELERATOR_COUNT = 1\n",
-        "MACHINE_TYPE = \"g2-standard-8\"\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "accelerator_count = 1\n",
+        "machine_type = \"g2-standard-8\"\n",
         "\n",
         "\n",
         "def deploy_model(model_id, task, accelerator_type, machine_type, accelerator_count=1):\n",
@@ -234,7 +218,6 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=SERVICE_ACCOUNT,\n",
         "        system_labels={\"NOTEBOOK_NAME\": \"model_garden_pytorch_instant_id.ipynb\"},\n",
         "    )\n",
         "    print(\"To load this existing endpoint from a different session:\")\n",
@@ -252,12 +235,151 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=False,  # Change to True if you have accepted the EULA on the model card.\n",
+        ")\n",
+        "\n",
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads the model to Model Registry and deploys it on the Endpoint. It takes ~15 minutes to finish.\n",
+        "# @markdown Click \"Show Code\" to see more details.\n",
+        "\n",
+        "\n",
+        "def deploy_model_vllm(\n",
+        "    model_name: str,\n",
+        "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
+        "    base_model_id: str = None,\n",
+        "    machine_type: str = \"g2-standard-8\",\n",
+        "    accelerator_type: str = \"NVIDIA_L4\",\n",
+        "    accelerator_count: int = 1,\n",
+        "    gpu_memory_utilization: float = 0.9,\n",
+        "    max_model_len: int = 4096,\n",
+        "    dtype: str = \"auto\",\n",
+        "    enable_trust_remote_code: bool = False,\n",
+        "    enforce_eager: bool = False,\n",
+        "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
+        "    max_loras: int = 1,\n",
+        "    max_cpu_loras: int = 8,\n",
+        "    use_dedicated_endpoint: bool = False,\n",
+        "    max_num_seqs: int = 256,\n",
+        "    model_type: str = None,\n",
+        ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
+        "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
+        "    endpoint = aiplatform.Endpoint.create(\n",
+        "        display_name=f\"{model_name}-endpoint\",\n",
+        "        dedicated_endpoint_enabled=use_dedicated_endpoint,\n",
+        "    )\n",
+        "\n",
+        "    if not base_model_id:\n",
+        "        base_model_id = model_id\n",
+        "\n",
+        "    # See https://docs.vllm.ai/en/latest/models/engine_args.html for a list of possible arguments with descriptions.\n",
+        "    vllm_args = [\n",
+        "        \"python\",\n",
+        "        \"-m\",\n",
+        "        \"vllm.entrypoints.api_server\",\n",
+        "        \"--host=0.0.0.0\",\n",
+        "        \"--port=8080\",\n",
+        "        f\"--model={model_id}\",\n",
+        "        f\"--tensor-parallel-size={accelerator_count}\",\n",
+        "        \"--swap-space=16\",\n",
+        "        f\"--gpu-memory-utilization={gpu_memory_utilization}\",\n",
+        "        f\"--max-model-len={max_model_len}\",\n",
+        "        f\"--dtype={dtype}\",\n",
+        "        f\"--max-loras={max_loras}\",\n",
+        "        f\"--max-cpu-loras={max_cpu_loras}\",\n",
+        "        f\"--max-num-seqs={max_num_seqs}\",\n",
+        "        \"--disable-log-stats\",\n",
+        "    ]\n",
+        "\n",
+        "    if enable_trust_remote_code:\n",
+        "        vllm_args.append(\"--trust-remote-code\")\n",
+        "\n",
+        "    if enforce_eager:\n",
+        "        vllm_args.append(\"--enforce-eager\")\n",
+        "\n",
+        "    if enable_lora:\n",
+        "        vllm_args.append(\"--enable-lora\")\n",
+        "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
+        "    if model_type:\n",
+        "        vllm_args.append(f\"--model-type={model_type}\")\n",
+        "\n",
+        "    env_vars = {\n",
+        "        \"MODEL_ID\": base_model_id,\n",
+        "        \"DEPLOY_SOURCE\": \"notebook\",\n",
+        "    }\n",
+        "\n",
+        "    # HF_TOKEN is not a compulsory field and may not be defined.\n",
+        "    try:\n",
+        "        if HF_TOKEN:\n",
+        "            env_vars[\"HF_TOKEN\"] = HF_TOKEN\n",
+        "    except NameError:\n",
+        "        pass\n",
+        "\n",
+        "    model = aiplatform.Model.upload(\n",
+        "        display_name=model_name,\n",
+        "        serving_container_image_uri=VLLM_DOCKER_URI,\n",
+        "        serving_container_args=vllm_args,\n",
+        "        serving_container_ports=[8080],\n",
+        "        serving_container_predict_route=\"/generate\",\n",
+        "        serving_container_health_route=\"/ping\",\n",
+        "        serving_container_environment_variables=env_vars,\n",
+        "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
+        "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
+        "    )\n",
+        "    print(\n",
+        "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
+        "    )\n",
+        "    model.deploy(\n",
+        "        endpoint=endpoint,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
+        "        deploy_request_timeout=1800,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_instant_id.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "        },\n",
+        "    )\n",
+        "    print(\"endpoint_name:\", endpoint.name)\n",
+        "\n",
+        "    return model, endpoint\n",
+        "\n",
         "\n",
         "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model(\n",
         "    model_id=model_id,\n",
         "    task=task,\n",
-        "    accelerator_type=ACCELERATOR_TYPE,\n",
-        "    machine_type=MACHINE_TYPE,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    machine_type=machine_type,\n",
         ")\n",
         "print(\"endpoint_name:\", endpoints[\"vllm_gpu\"].name)\n",
         "\n",
@@ -405,11 +527,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb

@@ -4,11 +4,12 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
+        "cellView": "form",
         "id": "7d9bbf86da5e"
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -31,19 +32,23 @@
       "source": [
         "# Vertex AI Model Garden - InstructPix2Pix\n",
         "\n",
-        "<table align=\"left\">\n",
-        "  <td>\n",
+        "<table><tbody><tr>\n",
+        "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_instructpix2pix.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
         "  </td>\n",
-        "  <td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/main/notebooks/community/model_garden/model_garden_pytorch_instructpix2pix.ipynb\">\n",
-        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br>\n",
-        "      View on GitHub\n",
+        "      <img alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"><br> View on GitHub\n",
         "    </a>\n",
         "  </td>\n",
-        "</table>"
+        "</tr></tbody></table>"
       ]
     },
     {
@@ -62,6 +67,10 @@
         "- Deploy the model on [Endpoint](https://cloud.google.com/vertex-ai/docs/predictions/using-private-endpoints).\n",
         "- Run online predictions for text-guided image editing.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -69,7 +78,7 @@
         "* Vertex AI\n",
         "* Cloud Storage\n",
         "\n",
-        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing) and [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
+        "Learn about [Vertex AI pricing](https://cloud.google.com/vertex-ai/pricing), [Cloud Storage pricing](https://cloud.google.com/storage/pricing), and use the [Pricing Calculator](https://cloud.google.com/products/calculator/) to generate a cost estimate based on your projected usage."
       ]
     },
     {
@@ -94,105 +103,130 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. [Optional] [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
-        "import base64\n",
+        "REGION = \"\"  # @param {type:\"string\"}\n",
+        "\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
+        "\n",
+        "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
+        "# @markdown | ----------- | ----------- | ----------- |\n",
+        "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
+        "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "# Import the necessary packages\n",
+        "import importlib\n",
         "import os\n",
-        "import sys\n",
-        "import uuid\n",
-        "from datetime import datetime\n",
-        "from io import BytesIO\n",
         "\n",
-        "import requests\n",
         "from google.cloud import aiplatform\n",
-        "from PIL import Image\n",
+        "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
+        "common_util = importlib.import_module(\n",
+        "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
+        ")\n",
+        "\n",
+        "LABEL = \"diffusers_gpu\"\n",
+        "models, endpoints = {}, {}\n",
+        "\n",
         "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
         "# Get the default region for launching jobs.\n",
-        "REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
+        "if not REGION:\n",
+        "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
+        "# Initialize Vertex AI API.\n",
+        "print(\"Initializing Vertex AI API.\")\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, please change the value yourself below.\n",
-        "now = datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_URI = \"gs://\"  # @param {type: \"string\"}\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
+        "! gcloud config set project $PROJECT_ID\n",
+        "import vertexai\n",
         "\n",
-        "# Create a unique GCS bucket for this notebook, if not specified by the user.\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "# Set up the default SERVICE_ACCOUNT.\n",
-        "SERVICE_ACCOUNT = None\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
-        "\n",
-        "if \"google.colab\" in sys.modules:\n",
-        "    from google.colab import auth\n",
-        "\n",
-        "    auth.authenticate_user(project_id=PROJECT_ID)\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xqDwCV8Dwvg1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Set the model parameters\n",
         "\n",
         "# The pre-built serving docker image. It contains serving scripts and models.\n",
         "SERVE_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-diffusers-serve-opt:20240605_1400_RC00\"\n",
         "\n",
+        "VERSION_ID = \"instruct-pix2pix\"\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/timbrooks/models/instruct-pix2pix@{VERSION_ID}\"\n",
         "\n",
-        "# Define common functions.\n",
-        "def create_job_name(prefix):\n",
-        "    now = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
-        "    job_name = f\"{prefix}-{now}\"\n",
-        "    return job_name\n",
+        "# The machine and accelerator specs for model deployment.\n",
+        "accelerator_type = \"NVIDIA_L4\"\n",
+        "machine_type = \"g2-standard-12\"\n",
+        "accelerator_count = 1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "MyjO08oPwvg1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "J7WFlBNXwvg1"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section deploys the InstructPix2Pix model for the text-guided image-to-image task.\n",
+        "\n",
+        "# @markdown The model deployment step will take ~15 minutes to complete.\n",
         "\n",
         "\n",
-        "def download_image(url):\n",
-        "    response = requests.get(url)\n",
-        "    return Image.open(BytesIO(response.content))\n",
-        "\n",
-        "\n",
-        "def image_to_base64(image, format=\"JPEG\"):\n",
-        "    buffer = BytesIO()\n",
-        "    image.save(buffer, format=format)\n",
-        "    image_str = base64.b64encode(buffer.getvalue()).decode(\"utf-8\")\n",
-        "    return image_str\n",
-        "\n",
-        "\n",
-        "def base64_to_image(image_str):\n",
-        "    image = Image.open(BytesIO(base64.b64decode(image_str)))\n",
-        "    return image\n",
-        "\n",
-        "\n",
-        "def image_grid(imgs, rows=2, cols=2):\n",
-        "    w, h = imgs[0].size\n",
-        "    grid = Image.new(\"RGB\", size=(cols * w, rows * h), color=(255, 255, 255))\n",
-        "    for i, img in enumerate(imgs):\n",
-        "        grid.paste(img, box=(i % cols * w + 10 * i, i // cols * h))\n",
-        "    return grid\n",
-        "\n",
-        "\n",
-        "def deploy_model(model_id, task):\n",
+        "def deploy_model(\n",
+        "    model_id: str,\n",
+        "    task: str,\n",
+        "    machine_type: str,\n",
+        "    accelerator_type: str,\n",
+        "    accelerator_count: int,\n",
+        "):\n",
         "    model_name = \"instruct-pix2pix\"\n",
         "    endpoint = aiplatform.Endpoint.create(display_name=f\"{model_name}-endpoint\")\n",
         "    serving_env = {\n",
@@ -208,38 +242,36 @@
         "        serving_container_predict_route=\"/predictions/diffusers_serving\",\n",
         "        serving_container_health_route=\"/ping\",\n",
         "        serving_container_environment_variables=serving_env,\n",
-        "        model_garden_source_model_name=\"publishers/timbrooks/models/instruct-pix2pix\"\n",
+        "        model_garden_source_model_name=\"publishers/timbrooks/models/instruct-pix2pix\",\n",
         "    )\n",
         "    model.deploy(\n",
         "        endpoint=endpoint,\n",
-        "        machine_type=\"g2-standard-8\",\n",
-        "        accelerator_type=\"NVIDIA_L4\",\n",
-        "        accelerator_count=1,\n",
+        "        machine_type=machine_type,\n",
+        "        accelerator_type=accelerator_type,\n",
+        "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_instructpix2pix.ipynb\"\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_instructpix2pix.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
-        "    return model, endpoint"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "b4b46c28d8b1"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Upload and deploy model\n",
+        "    return model, endpoint\n",
         "\n",
-        "# @markdown This section deploys the InstructPix2Pix model for the text-guided image-to-image task.\n",
         "\n",
-        "# @markdown The model deployment step will take ~15 minutes to complete.\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
         "\n",
-        "model, endpoint = deploy_model(\n",
-        "    model_id=\"timbrooks/instruct-pix2pix\", task=\"instruct-pix2pix\"\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model(\n",
+        "    model_id=\"timbrooks/instruct-pix2pix\",\n",
+        "    task=\"instruct-pix2pix\",\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
         ")"
       ]
     },
@@ -262,17 +294,17 @@
         "image = \"https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png\"  # @param {type: \"string\"}\n",
         "num_inference_steps = 25  # @param {type:\"number\"}\n",
         "\n",
-        "init_image = download_image(image)\n",
+        "init_image = common_util.download_image(image)\n",
         "instances = [\n",
         "    {\n",
         "        \"prompt\": prompt,\n",
-        "        \"image\": image_to_base64(init_image),\n",
+        "        \"image\": common_util.image_to_base64(init_image),\n",
         "        \"num_inference_steps\": num_inference_steps,\n",
         "    },\n",
         "]\n",
-        "response = endpoint.predict(instances=instances)\n",
-        "images = [base64_to_image(image) for image in response.predictions]\n",
-        "image_grid([init_image, images[0]], rows=1, cols=2)"
+        "response = endpoints[LABEL].predict(instances=instances)\n",
+        "images = [common_util.base64_to_image(image) for image in response.predictions]\n",
+        "common_util.image_grid([init_image, images[0]], rows=1, cols=2)"
       ]
     },
     {
@@ -287,18 +319,15 @@
         "# @title Clean up resources\n",
         "\n",
         "# @markdown  Delete the experiment models and endpoints to recycle the resources\n",
-        "# @markdown  and avoid unnecessary continouous charges that may incur.\n",
+        "# @markdown  and avoid unnecessary continuous charges that may incur.\n",
         "\n",
         "# Undeploy model and delete endpoint.\n",
-        "endpoint.delete(force=True)\n",
+        "for endpoint in endpoints.values():\n",
+        "    endpoint.delete(force=True)\n",
         "\n",
-        "# Delete model.\n",
-        "model.delete()\n",
-        "\n",
-        "# Delete bucket.\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "# Delete models.\n",
+        "for model in models.values():\n",
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_llama3_1_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_1_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -64,6 +69,10 @@
         "- Deploy Llama 3.1 8B, 70B and 405B with standard vLLM on GPU, optionally with dynamic LoRA adapters.\n",
         "- Deploy Llama 3.1 8B and 70B with optimized vLLM on GPU.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -102,6 +111,7 @@
     {
       "cell_type": "code",
       "execution_count": null,
+      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "YXFGIp1l-qtT"
@@ -112,44 +122,46 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
+        "import re\n",
         "from typing import Tuple\n",
         "\n",
         "import requests\n",
+        "from google import auth\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "LABEL = \"vllm_gpu\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
+        "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
         "\n",
@@ -157,52 +169,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3_1\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown # Access Llama 3.1 models on Vertex AI for serving\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -216,7 +193,13 @@
         "VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "assert (\n",
         "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_1\n",
-        "), \"Click the agreement of Llama 3.1 in Vertex AI Model Garden, and get the GCS path of Llama 3.1 model artifacts.\""
+        "), \"Click the agreement of Llama 3.1 in Vertex AI Model Garden, and get the GCS path of Llama 3.1 model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_3_1)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA_3_1.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA_3_1 is expected to be a GCS URI and must start with `gs://`.\""
       ]
     },
     {
@@ -246,14 +229,14 @@
         "# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
         "# @markdown Select one of the four model variations. More model variants will be supported by Hex-LLM in the future.\n",
-        "MODEL_ID = \"Meta-Llama-3.1-8B\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\"] {allow-input: true, isTemplate: true}\n",
+        "MODEL_ID = \"Meta-Llama-3.1-8B-Instruct\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\"] {allow-input: true, isTemplate: true}\n",
         "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
         "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_1, MODEL_ID)\n",
         "hf_model_id = \"meta-llama/\" + MODEL_ID\n",
         "\n",
         "# The pre-built serving docker images.\n",
         "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# @markdown Find Vertex AI prediction TPUv5e machine types in\n",
@@ -315,7 +298,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -404,11 +386,11 @@
         "        machine_type=machine_type,\n",
         "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        min_replica_count=min_replica_count,\n",
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -419,7 +401,6 @@
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3_1\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    base_model_id=hf_model_id,\n",
         "    tensor_parallel_size=tensor_parallel_size,\n",
         "    machine_type=machine_type,\n",
@@ -508,7 +489,7 @@
       "source": [
         "# @title Chat completion\n",
         "\n",
-        "_region = REGION\n",
+        "temp_region = REGION\n",
         "REGION = TPU_DEPLOYMENT_REGION\n",
         "\n",
         "if use_dedicated_endpoint:\n",
@@ -531,6 +512,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -557,10 +539,23 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
         "\n",
-        "REGION = _region\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "REGION = temp_region\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -585,13 +580,18 @@
       "source": [
         "# @title Fast Deployment\n",
         "\n",
+        "# @markdown The Llama 3.1 8B Instruct model will be deployed to a dedicated endpoint on an `a2-ultragpu-1g` machine with Fast Deployment.\n",
+        "# @markdown **Currently, the Fast Deployment is only supported in the `us-central1` region.**\n",
+        "\n",
         "# @markdown This section demonstrates how to use the Fast Deployment feature.\n",
         "\n",
         "# @markdown The Fast Deployment feature prioritizes speed for model exploration, making it ideal for initial testing and experimentation. For sensitive data or production workloads, use the Standard environment for enhanced security and stability.\n",
         "\n",
         "# @markdown Note that only a subset of the models support the Fast Deployment feature.\n",
         "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\"\n",
+        "FAST_DEPLOYMENT_REGION = \"us-central1\"  # @param [\"us-central1\"] {isTemplate:true}\n",
+        "\n",
+        "API_ENDPOINT = f\"{FAST_DEPLOYMENT_REGION}-aiplatform.googleapis.com\"\n",
         "\n",
         "\n",
         "def fast_deploy(\n",
@@ -619,7 +619,7 @@
         "        == 0\n",
         "    ):\n",
         "        raise ValueError(\n",
-        "            f\"No supportedActions.multiDeployVertex found in {REGION}. You can skip\"\n",
+        "            f\"No supportedActions.multiDeployVertex found in {FAST_DEPLOYMENT_REGION}. You can skip\"\n",
         "            \" this section or try a different region.\"\n",
         "        )\n",
         "    deploy_configs = response[\"supportedActions\"][\"multiDeployVertex\"][\n",
@@ -637,7 +637,7 @@
         "        fast_deploy_config = fast_deploy_config[0]\n",
         "    else:\n",
         "        raise ValueError(\n",
-        "            f\"No Fast Deployment config found in {REGION}. You can skip this\"\n",
+        "            f\"No Fast Deployment config found in {FAST_DEPLOYMENT_REGION}. You can skip this\"\n",
         "            \" section or try a different region.\"\n",
         "        )\n",
         "\n",
@@ -661,12 +661,14 @@
         "        serving_container_health_route=container_spec.get(\"healthRoute\"),\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        model_garden_source_model_name=(\n",
         "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
         "        ),\n",
         "    )\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
         "        display_name=model.name + \"-endpoint\",\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        dedicated_endpoint_enabled=True,\n",
         "    )\n",
         "    print(\n",
@@ -690,9 +692,6 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "# @markdown The Llama 3.1 8B Instruct model will be deployed to a dedicated endpoint on an `a2-ultragpu-1g` machine with Fast Deployment.\n",
-        "# @markdown **Currently, the Fast Deployment is only supported in the `us-central1` region.**\n",
-        "\n",
         "use_dedicated_endpoint = True  # Fast Deployment only supports dedicated endpoints.\n",
         "models[\"vllm_fast\"], endpoints[\"vllm_fast\"] = fast_deploy(\n",
         "    \"meta\", \"llama3_1\", \"llama-3.1-8b-instruct\"\n",
@@ -775,6 +774,10 @@
       "source": [
         "# @title Chat completion\n",
         "\n",
+        "temp_region = REGION\n",
+        "REGION = FAST_DEPLOYMENT_REGION\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
         "DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
         "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
@@ -794,6 +797,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -820,8 +824,23 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "REGION = temp_region\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -844,7 +863,7 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Set the model variants\n",
         "\n",
         "# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
@@ -854,15 +873,17 @@
         "\n",
         "# @markdown Set the model to deploy.\n",
         "\n",
-        "base_model_name = \"Meta-Llama-3.1-8B\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\", \"Meta-Llama-3.1-405B-FP8\", \"Meta-Llama-3.1-405B-Instruct-FP8\"] {isTemplate:true}\n",
+        "base_model_name = \"Meta-Llama-3.1-8B-Instruct\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\", \"Meta-Llama-3.1-405B-FP8\", \"Meta-Llama-3.1-405B-Instruct-FP8\"] {isTemplate:true}\n",
         "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_1, base_model_name)\n",
         "ENABLE_DYNAMIC_LORA = True  # @param {type:\"boolean\", isTemplate:true}\n",
+        "version_id = base_model_name.lower()[5:]\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama3_1@{version_id}\"\n",
         "\n",
         "# The pre-built serving docker images.\n",
         "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241210_0916_RC00\"\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "if \"8b\" in base_model_name.lower():\n",
@@ -893,10 +914,55 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "QpOh4c5SLi_c"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "DsRuuOjjLzq-"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama 3.1 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of H100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have H100 quota.\n",
+        "\n",
+        "# @markdown H100 is hard to get for now. It's recommended to use the deployment button in the model card. You can still try to deploy H100 endpoint through the notebook, but there is a chance that resource is not available.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
         "gpu_memory_utilization = 0.95\n",
         "max_model_len = 8192  # Maximum context length.\n",
         "\n",
-        "\n",
         "# Enable automatic prefix caching using GPU HBM\n",
         "enable_prefix_cache = True\n",
         "# Setting this value >0 will use the idle host memory for a second-tier prefix kv\n",
@@ -905,13 +971,27 @@
         "# Setting host_prefix_kv_cache_utilization_target to 0 will disable the host memory prefix kv cache.\n",
         "host_prefix_kv_cache_utilization_target = 0.7\n",
         "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown To enable the auto-scaling in deployment, you can set the following options:\n",
+        "\n",
+        "min_replica_count = 1  # @param {type:\"integer\"}\n",
+        "max_replica_count = 1  # @param {type:\"integer\"}\n",
+        "required_replica_count = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Set the target of GPU duty cycle or CPU usage between 1 and 100 for auto-scaling.\n",
+        "autoscale_by_gpu_duty_cycle_target = 0  # @param {type:\"integer\"}\n",
+        "autoscale_by_cpu_usage_target = 0  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Note: GPU duty cycle is not the most accurate metric for scaling workloads. More advanced auto-scaling metrics are coming soon. See [the public doc](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#autoscaling) for more details.\n",
+        "\n",
         "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -930,6 +1010,13 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        "    required_replica_count: int = 1,\n",
+        "    autoscale_by_gpu_duty_cycle_target: int = 0,\n",
+        "    autoscale_by_cpu_usage_target: int = 0,\n",
+        "    is_spot: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -982,6 +1069,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -1011,29 +1102,74 @@
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
         "    )\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=accelerator_count,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_deployment.ipynb\",\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": min_replica_count,\n",
+        "                \"requiredReplicaCount\": required_replica_count,\n",
+        "                \"maxReplicaCount\": max_replica_count,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
         "        },\n",
-        "    )\n",
+        "    }\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    if autoscale_by_gpu_duty_cycle_target > 0 or autoscale_by_cpu_usage_target > 0:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"autoscalingMetricSpecs\"] = []\n",
+        "        if autoscale_by_gpu_duty_cycle_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle\",\n",
+        "                    \"target\": autoscale_by_gpu_duty_cycle_target,\n",
+        "                }\n",
+        "            )\n",
+        "        if autoscale_by_cpu_usage_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/cpu/utilization\",\n",
+        "                    \"target\": autoscale_by_cpu_usage_target,\n",
+        "                }\n",
+        "            )\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    common_util.poll_and_wait(response.json()[\"name\"], REGION, 7200)\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_1-serve\"),\n",
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3_1\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -1046,7 +1182,15 @@
         "    enable_prefix_cache=enable_prefix_cache,\n",
         "    host_prefix_kv_cache_utilization_target=host_prefix_kv_cache_utilization_target,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    enable_llama_tool_parser=True,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        "    required_replica_count=required_replica_count,\n",
+        "    autoscale_by_gpu_duty_cycle_target=autoscale_by_gpu_duty_cycle_target,\n",
+        "    autoscale_by_cpu_usage_target=autoscale_by_cpu_usage_target,\n",
+        "    is_spot=is_spot,\n",
         ")\n",
+        "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
     },
@@ -1118,9 +1262,9 @@
         "# @title Chat completion\n",
         "\n",
         "if use_dedicated_endpoint:\n",
-        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
         ")\n",
         "\n",
         "# @title Chat Completions Inference\n",
@@ -1137,6 +1281,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -1163,8 +1308,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -1199,7 +1357,7 @@
         "\n",
         "# @markdown Set the model to deploy.\n",
         "\n",
-        "base_model_name = \"Meta-Llama-3.1-8B\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\"] {isTemplate:true}\n",
+        "base_model_name = \"Meta-Llama-3.1-8B-Instruct\"  # @param [\"Meta-Llama-3.1-8B\", \"Meta-Llama-3.1-8B-Instruct\", \"Meta-Llama-3.1-70B\", \"Meta-Llama-3.1-70B-Instruct\"] {isTemplate:true}\n",
         "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_1, base_model_name)\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
         "\n",
@@ -1236,7 +1394,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-12\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -1305,9 +1462,9 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -1324,7 +1481,6 @@
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3_1\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -1426,6 +1582,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -1452,8 +1609,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -1505,11 +1675,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_llama3_1_finetuning.ipynb

@@ -9,7 +9,7 @@
       },
       "outputs": [],
       "source": [
-        "# Copyright 2024 Google LLC\n",
+        "# Copyright 2025 Google LLC\n",
         "#\n",
         "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
         "# you may not use this file except in compliance with the License.\n",
@@ -99,17 +99,11 @@
       "source": [
         "# @title Install Python Packages for Finetuning\n",
         "\n",
-        "# @markdown 1. Install google-cloud-aiplatform package and restart the session if instructed.\n",
-        "! pip install --upgrade --quiet 'google-cloud-aiplatform>=1.66.0'\n",
-        "\n",
-        "# @markdown 2. Install packages to validate dataset with template.\n",
+        "# @markdown 1. Install packages to validate dataset with template.\n",
         "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
-        "! pip install --upgrade --quiet accelerate==0.31.0\n",
-        "! pip install --upgrade --quiet transformers==4.43.1\n",
-        "! pip install --upgrade --quiet datasets==2.19.2\n",
-        "\n",
-        "# Load local tensorboard.\n",
-        "%load_ext tensorboard"
+        "! pip install --upgrade --quiet accelerate==0.34.2\n",
+        "! pip install --upgrade --quiet transformers==4.47.1\n",
+        "! pip install --upgrade --quiet datasets==2.20.0"
       ]
     },
     {
@@ -132,7 +126,9 @@
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-west1, europe-west4, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -144,7 +140,7 @@
         "\n",
         "# Import the necessary packages\n",
         "! rm -rf vertex-ai-samples && git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
-        "! cd vertex-ai-samples && git reset --hard 0727e19520cf7957bceb701c248221bd3dbe4f1f\n",
+        "! cd vertex-ai-samples && git reset --hard 80320a9a1b818534ca785444e704f6953f2a9dd9\n",
         "\n",
         "import datetime\n",
         "import importlib\n",
@@ -167,6 +163,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -293,7 +295,7 @@
         "# @markdown Use the Vertex AI SDK to create and run the custom training jobs.\n",
         "\n",
         "# @markdown This notebook uses [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) dataset as an example.\n",
-        "# @markdown You can set `dataset_name` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `instruct_column_in_dataset` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `instruct_column_in_dataset` to `text` in this notebook.\n",
+        "# @markdown You can set `train_dataset` to any existing [Hugging Face dataset](https://huggingface.co/datasets) name, and set `train_column` to the name of the dataset column containing training data. The [timdettmers/openassistant-guanaco](https://huggingface.co/datasets/timdettmers/openassistant-guanaco) has only one column `text`, and therefore we set `train_column` to `text` in this notebook.\n",
         "\n",
         "# @markdown ### (Optional) Prepare a custom JSONL dataset for finetuning\n",
         "\n",
@@ -334,23 +336,25 @@
         "# @markdown\n",
         "# @markdown To try such custom dataset, you can make the following changes:\n",
         "# @markdown 1. Set `template` to `llama3-text-bison`\n",
-        "# @markdown 1. Set `train_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
-        "# @markdown 1. Set `train_split_name` to `train`\n",
-        "# @markdown 1. Set `eval_dataset_name` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
-        "# @markdown 1. Set `eval_split_name` to `train` (**NOT** `test`)\n",
-        "# @markdown 1. Set `instruct_column_in_dataset` as `input_text`.\n",
+        "# @markdown 1. Set `train_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_train_sample.jsonl`\n",
+        "# @markdown 1. Set `train_split` to `train`\n",
+        "# @markdown 1. Set `eval_dataset` to `gs://cloud-samples-data/vertex-ai/model-evaluation/peft_eval_sample.jsonl`\n",
+        "# @markdown 1. Set `eval_split` to `train` (**NOT** `test`)\n",
+        "# @markdown 1. Set `train_column` as `input_text`.\n",
         "\n",
         "# Template name or gs:// URI to a custom template.\n",
         "template = \"openassistant-guanaco\"  # @param {type:\"string\"}\n",
         "\n",
         "# Hugging Face dataset name or gs:// URI to a custom JSONL dataset.\n",
-        "train_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
-        "train_split_name = \"train\"  # @param {type:\"string\"}\n",
-        "eval_dataset_name = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
-        "eval_split_name = \"test\"  # @param {type:\"string\"}\n",
+        "train_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "train_split = \"train\"  # @param {type:\"string\"}\n",
+        "eval_dataset = \"timdettmers/openassistant-guanaco\"  # @param {type:\"string\"}\n",
+        "eval_split = \"test\"  # @param {type:\"string\"}\n",
         "\n",
         "# Name of the dataset column containing training text input.\n",
-        "instruct_column_in_dataset = \"text\"  # @param {type:\"string\"}"
+        "train_column = \"text\"  # @param {type:\"string\"}\n",
+        "# Maximum sequence length.\n",
+        "max_seq_length = 4096  # @param{type:\"integer\"}"
       ]
     },
     {
@@ -365,7 +369,7 @@
         "# @title Set model\n",
         "\n",
         "# @markdown Select a model variant of Llama 3.1.\n",
-        "base_model_id = \"meta-llama/Meta-Llama-3.1-8B-Instruct\"  # @param [\"meta-llama/Meta-Llama-3.1-8B\", \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"meta-llama/Meta-Llama-3.1-70B\", \"meta-llama/Meta-Llama-3.1-70B-Instruct\", \"meta-llama/Meta-Llama-3.1-405B\", \"meta-llama/Meta-Llama-3.1-405B-Instruct\"] {isTemplate:true}\n",
+        "base_model_id = \"meta-llama/Meta-Llama-3.1-8B-Instruct\"  # @param [\"meta-llama/Meta-Llama-3.1-8B\", \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"meta-llama/Meta-Llama-3.1-70B\", \"meta-llama/Meta-Llama-3.1-70B-Instruct\", \"meta-llama/Meta-Llama-3.1-405B\", \"meta-llama/Meta-Llama-3.1-405B-Instruct\", \"deepseek-ai/DeepSeek-R1-Distill-Llama-70B\"] {isTemplate:true}\n",
         "if LOAD_MODEL_FROM == \"Google Cloud\":\n",
         "    pretrained_model_id = os.path.join(MODEL_BUCKET, base_model_id.split(\"/\")[-1])\n",
         "else:\n",
@@ -385,8 +389,6 @@
         "\n",
         "# @markdown This section validates the train and eval datasets with the template before starting the fine tuning process.\n",
         "\n",
-        "import transformers\n",
-        "\n",
         "dataset_validation_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.dataset_validation_util\"\n",
         ")\n",
@@ -402,38 +404,35 @@
         "    tokenizer_path = pretrained_model_id\n",
         "    access_token = HF_TOKEN\n",
         "\n",
-        "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
-        "    tokenizer_path,\n",
-        "    trust_remote_code=False,\n",
-        "    use_fast=True,\n",
-        "    token=access_token,\n",
-        ")\n",
+        "tokenizer = dataset_validation_util.load_tokenizer(tokenizer_path, None, access_token)\n",
         "\n",
         "# Validate the train dataset.\n",
         "dataset_validation_util.validate_dataset_with_template(\n",
-        "    dataset_name=train_dataset_name,\n",
-        "    split=train_split_name,\n",
-        "    input_column=instruct_column_in_dataset,\n",
+        "    dataset_name=train_dataset,\n",
+        "    split=train_split,\n",
+        "    input_column=train_column,\n",
         "    template=template,\n",
+        "    max_seq_length=max_seq_length,\n",
         "    use_multiprocessing=False,\n",
         "    tokenizer=tokenizer,\n",
         ")\n",
         "\n",
-        "# Validate the eval dataset.\n",
-        "dataset_validation_util.validate_dataset_with_template(\n",
-        "    dataset_name=eval_dataset_name,\n",
-        "    split=eval_split_name,\n",
-        "    input_column=instruct_column_in_dataset,\n",
-        "    template=template,\n",
-        "    use_multiprocessing=False,\n",
-        "    tokenizer=tokenizer,\n",
-        ")"
+        "# Validate the eval dataset if it exists.\n",
+        "if eval_dataset:\n",
+        "    dataset_validation_util.validate_dataset_with_template(\n",
+        "        dataset_name=eval_dataset,\n",
+        "        split=eval_split,\n",
+        "        input_column=train_column,\n",
+        "        template=template,\n",
+        "        max_seq_length=max_seq_length,\n",
+        "        use_multiprocessing=False,\n",
+        "        tokenizer=tokenizer,\n",
+        "    )"
       ]
     },
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "ivVGS9dHXPOz"
@@ -448,14 +447,15 @@
         "\n",
         "# @markdown **Note**:\n",
         "# @markdown 1. We recommend setting `finetuning_precision_mode` to `4bit` because it enables using fewer hardware resources for finetuning.\n",
-        "# @markdown 1. We recommend using NVIDIA_A100_80GB for 8B and 70B models, and NVIDIA_H100_80GB for 405B models.\n",
-        "# @markdown 1. If `max_steps>0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
+        "# @markdown 1. If `max_steps > 0`, it takes precedence over `epochs`. One can set a small `max_steps` value to quickly check the pipeline.\n",
         "\n",
         "# @markdown Acceletor type to use for training.\n",
-        "accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "# fmt: off\n",
+        "training_accelerator_type = \"NVIDIA_A100_80GB\"  # @param [\"NVIDIA_A100_80GB\", \"NVIDIA_H100_80GB\"]\n",
+        "# fmt: on\n",
         "\n",
         "# The pre-built training docker image.\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
         "    repo = \"us-docker.pkg.dev/vertex-ai-restricted\"\n",
         "    is_restricted_image = True\n",
         "    is_dynamic_workload_scheduler = False\n",
@@ -474,21 +474,22 @@
         "    }\n",
         "\n",
         "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20240909\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250213\"\n",
         ")\n",
         "\n",
         "# Worker pool spec.\n",
-        "if accelerator_type == \"NVIDIA_A100_80GB\":\n",
+        "boot_disk_size_gb = 500\n",
+        "if training_accelerator_type == \"NVIDIA_A100_80GB\":\n",
         "    per_node_accelerator_count = 8\n",
-        "    machine_type = \"a2-ultragpu-8g\"\n",
-        "    boot_disk_size_gb = 500\n",
-        "elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "    training_machine_type = \"a2-ultragpu-8g\"\n",
+        "elif training_accelerator_type == \"NVIDIA_H100_80GB\":\n",
         "    per_node_accelerator_count = 8\n",
-        "    machine_type = \"a3-highgpu-8g\"\n",
-        "    boot_disk_size_gb = 2000\n",
+        "    training_machine_type = \"a3-highgpu-8g\"\n",
+        "    if \"405b\" in base_model_id.lower():\n",
+        "        boot_disk_size_gb = 2000\n",
         "else:\n",
         "    raise ValueError(\n",
-        "        f\"Recommended machine settings not found for: {accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `machine_type`, `accelerator_type`, and `per_node_accelerator_count` to the deploy_model_vllm function by clicking `Show Code` and then modifying the code.\"\n",
+        "        f\"Recommended machine settings not found for: {training_accelerator_type}. To use another accelerator type, edit this code block to pass in an appropriate `training_machine_type`, `training_accelerator_type`, and `per_node_accelerator_count` by clicking `Show Code` and then modifying the code.\"\n",
         "    )\n",
         "\n",
         "# @markdown The number of nodes to use for this worker pool in distributed training.\n",
@@ -511,11 +512,9 @@
         "per_device_train_batch_size = 1  # @param{type:\"integer\"}\n",
         "# @markdown Number of updates steps to accumulate the gradients for, before performing a backward/update pass.\n",
         "gradient_accumulation_steps = 4  # @param{type:\"integer\"}\n",
-        "# @markdown Maximum sequence length.\n",
-        "max_seq_length = 4096  # @param{type:\"integer\"}\n",
-        "# @markdown Setting a positive `max_steps` here will override `num_epochs`.\n",
+        "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
         "max_steps = -1  # @param{type:\"integer\"}\n",
-        "num_epochs = 1.0  # @param{type:\"number\"}\n",
+        "num_train_epochs = 1.0  # @param{type:\"number\"}\n",
         "# @markdown Precision mode for finetuning.\n",
         "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
         "# @markdown Learning rate.\n",
@@ -527,7 +526,7 @@
         "lora_alpha = 32  # @param{type:\"integer\"}\n",
         "lora_dropout = 0.05  # @param{type:\"number\"}\n",
         "# Activates gradient checkpointing for the current model (may be referred to as activation checkpointing or checkpoint activations in other frameworks).\n",
-        "enable_gradient_checkpointing = True\n",
+        "gradient_checkpointing = True\n",
         "# Attention implementation to use in the model.\n",
         "attn_implementation = \"flash_attention_2\"\n",
         "# The optimizer for which to schedule the learning rate.\n",
@@ -541,10 +540,15 @@
         "# Number of update steps between two logs.\n",
         "logging_steps = save_steps\n",
         "\n",
+        "# @markdown Evaluation metrics to compute. Supported eval metrics: loss, perplexity, bleu, google_bleu, rouge1, rouge2, rougeL, rougeLsum.\n",
+        "eval_metric_name = \"loss,perplexity,bleu\"  # @param{type:\"string\"}\n",
+        "# @markdown Metric to use for best model selection. This will save the best checkpoint based on the eval metric.\n",
+        "metric_for_best_model = \"perplexity\"  # @param{type:\"string\"}\n",
+        "\n",
         "common_util.check_quota(\n",
         "    project_id=PROJECT_ID,\n",
         "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count * replica_count,\n",
         "    is_for_training=True,\n",
         "    is_restricted_image=is_restricted_image,\n",
@@ -557,7 +561,7 @@
         "# Create a GCS folder to store the LORA adapter.\n",
         "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
         "# Create a GCS folder to store the finetuned LORA adapter.\n",
-        "final_checkpoint = os.path.join(lora_output_dir, \"checkpoint-final\")\n",
+        "final_checkpoint = os.path.join(lora_output_dir, \"node-0\", \"checkpoint-final\")\n",
         "\n",
         "# Add labels for the finetuning job.\n",
         "labels = {\n",
@@ -570,23 +574,23 @@
         "labels[\"versioned-mg-tune\"] = f\"{labels['mg-tune']}-{versioned_model_id}\"\n",
         "\n",
         "eval_args = [\n",
-        "    f\"--eval_dataset_path={eval_dataset_name}\",\n",
-        "    f\"--eval_column={instruct_column_in_dataset}\",\n",
+        "    f\"--eval_dataset={eval_dataset}\",\n",
+        "    f\"--eval_column={train_column}\",\n",
         "    f\"--eval_template={template}\",\n",
-        "    f\"--eval_split={eval_split_name}\",\n",
+        "    f\"--eval_split={eval_split}\",\n",
         "    f\"--eval_steps={save_steps}\",\n",
-        "    \"--eval_tasks=builtin_eval\",\n",
-        "    \"--eval_metric_name=loss\",\n",
+        "    f\"--eval_metric_name={eval_metric_name}\",\n",
+        "    f\"--metric_for_best_model={metric_for_best_model}\",\n",
         "]\n",
         "\n",
         "train_job_args = [\n",
         "    f\"--config_file={config_file}\",\n",
         "    \"--task=instruct-lora\",\n",
-        "    \"--completion_only=True\",\n",
-        "    f\"--pretrained_model_id={pretrained_model_id}\",\n",
-        "    f\"--dataset_name={train_dataset_name}\",\n",
-        "    f\"--train_split_name={train_split_name}\",\n",
-        "    f\"--instruct_column_in_dataset={instruct_column_in_dataset}\",\n",
+        "    \"--input_masking=True\",\n",
+        "    f\"--pretrained_model_name_or_path={pretrained_model_id}\",\n",
+        "    f\"--train_dataset={train_dataset}\",\n",
+        "    f\"--train_split={train_split}\",\n",
+        "    f\"--train_column={train_column}\",\n",
         "    f\"--output_dir={lora_output_dir}\",\n",
         "    f\"--per_device_train_batch_size={per_device_train_batch_size}\",\n",
         "    f\"--gradient_accumulation_steps={gradient_accumulation_steps}\",\n",
@@ -598,8 +602,8 @@
         "    f\"--learning_rate={learning_rate}\",\n",
         "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
         "    f\"--precision_mode={finetuning_precision_mode}\",\n",
-        "    f\"--enable_gradient_checkpointing={enable_gradient_checkpointing}\",\n",
-        "    f\"--num_epochs={num_epochs}\",\n",
+        "    f\"--gradient_checkpointing={gradient_checkpointing}\",\n",
+        "    f\"--num_train_epochs={num_train_epochs}\",\n",
         "    f\"--attn_implementation={attn_implementation}\",\n",
         "    f\"--optimizer={optimizer}\",\n",
         "    f\"--warmup_ratio={warmup_ratio}\",\n",
@@ -607,7 +611,7 @@
         "    f\"--logging_output_dir={base_output_dir}\",\n",
         "    f\"--save_steps={save_steps}\",\n",
         "    f\"--logging_steps={logging_steps}\",\n",
-        "    f\"--template={template}\",\n",
+        "    f\"--train_template={template}\",\n",
         "    f\"--huggingface_access_token={HF_TOKEN}\",\n",
         "] + eval_args\n",
         "\n",
@@ -623,8 +627,8 @@
         "train_job.run(\n",
         "    args=train_job_args,\n",
         "    replica_count=replica_count,\n",
-        "    machine_type=machine_type,\n",
-        "    accelerator_type=accelerator_type,\n",
+        "    machine_type=training_machine_type,\n",
+        "    accelerator_type=training_accelerator_type,\n",
         "    accelerator_count=per_node_accelerator_count,\n",
         "    boot_disk_size_gb=boot_disk_size_gb,\n",
         "    service_account=SERVICE_ACCOUNT,\n",
@@ -665,7 +669,6 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "qmHW6m8xG_4U"
@@ -681,7 +684,7 @@
         "    print(\"The training job has finished.\")\n",
         "\n",
         "# The pre-built serving docker image for vLLM.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240819_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
         "\n",
         "# Find Vertex AI prediction supported accelerators and regions [here](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute).\n",
         "if \"8b\" in base_model_id.lower():\n",
@@ -689,9 +692,9 @@
         "    accelerator_type = \"NVIDIA_L4\"\n",
         "    per_node_accelerator_count = 1\n",
         "elif \"70b\" in base_model_id.lower():\n",
-        "    machine_type = \"g2-standard-96\"\n",
-        "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    per_node_accelerator_count = 8\n",
+        "    machine_type = \"a3-highgpu-4g\"\n",
+        "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
+        "    per_node_accelerator_count = 4\n",
         "elif \"405b\" in base_model_id.lower():\n",
         "    machine_type = \"a3-highgpu-8g\"\n",
         "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
@@ -718,9 +721,23 @@
         "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
+        "    publisher: str,\n",
+        "    publisher_model_id: str,\n",
         "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
@@ -732,11 +749,15 @@
         "    enable_trust_remote_code: bool = False,\n",
         "    enforce_eager: bool = False,\n",
         "    enable_lora: bool = False,\n",
+        "    enable_chunked_prefill: bool = False,\n",
+        "    enable_prefix_cache: bool = False,\n",
+        "    host_prefix_kv_cache_utilization_target: float = 0.0,\n",
         "    max_loras: int = 1,\n",
         "    max_cpu_loras: int = 8,\n",
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -775,9 +796,24 @@
         "    if enable_lora:\n",
         "        vllm_args.append(\"--enable-lora\")\n",
         "\n",
+        "    if enable_chunked_prefill:\n",
+        "        vllm_args.append(\"--enable-chunked-prefill\")\n",
+        "\n",
+        "    if enable_prefix_cache:\n",
+        "        vllm_args.append(\"--enable-prefix-caching\")\n",
+        "\n",
+        "    if 0 < host_prefix_kv_cache_utilization_target < 1:\n",
+        "        vllm_args.append(\n",
+        "            f\"--host-prefix-kv-cache-utilization-target={host_prefix_kv_cache_utilization_target}\"\n",
+        "        )\n",
+        "\n",
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -800,6 +836,9 @@
         "        serving_container_environment_variables=env_vars,\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        model_garden_source_model_name=(\n",
+        "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
+        "        ),\n",
         "    )\n",
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
@@ -811,6 +850,10 @@
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
         "        service_account=service_account,\n",
+        "        system_labels={\n",
+        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_finetuning.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
+        "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
@@ -855,6 +898,8 @@
         "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_1-vllm-serve\"),\n",
         "    model_id=deploy_pretrained_model_id,\n",
+        "    publisher=\"meta\",\n",
+        "    publisher_model_id=\"llama3_1\",\n",
         "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
@@ -871,7 +916,6 @@
     {
       "cell_type": "code",
       "execution_count": null,
-      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "2UYUNn60G_4U"
@@ -899,7 +943,6 @@
         "top_k = 1  # @param {type:\"integer\"}\n",
         "raw_response = False  # @param {type:\"boolean\"}\n",
         "\n",
-        "\n",
         "predict_vllm(\n",
         "    prompt=prompt,\n",
         "    max_tokens=max_tokens,\n",

notebooks/community/model_garden/model_garden_pytorch_llama3_1_qwen2_5_deployment_tpu.ipynb

@@ -138,6 +138,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",

notebooks/community/model_garden/model_garden_pytorch_llama3_1_reasoning_engine.ipynb

@@ -134,14 +134,14 @@
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
@@ -174,6 +174,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -264,7 +270,9 @@
         "\n",
         "# @markdown Note that only a subset of the models support the Fast Deployment feature.\n",
         "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\"\n",
+        "FAST_DEPLOYMENT_REGION = \"us-central1\"  # @param [\"us-central1\"] {isTemplate:true}\n",
+        "\n",
+        "API_ENDPOINT = f\"{FAST_DEPLOYMENT_REGION}-aiplatform.googleapis.com\"\n",
         "\n",
         "\n",
         "def fast_deploy(\n",
@@ -292,7 +300,7 @@
         "        == 0\n",
         "    ):\n",
         "        raise ValueError(\n",
-        "            f\"No supportedActions.multiDeployVertex found in {REGION}. You can skip\"\n",
+        "            f\"No supportedActions.multiDeployVertex found in {FAST_DEPLOYMENT_REGION}. You can skip\"\n",
         "            \" this section or try a different region.\"\n",
         "        )\n",
         "    deploy_configs = response[\"supportedActions\"][\"multiDeployVertex\"][\n",
@@ -310,7 +318,7 @@
         "        fast_deploy_config = fast_deploy_config[0]\n",
         "    else:\n",
         "        raise ValueError(\n",
-        "            f\"No Fast Deployment config found in {REGION}. You can skip this\"\n",
+        "            f\"No Fast Deployment config found in {FAST_DEPLOYMENT_REGION}. You can skip this\"\n",
         "            \" section or try a different region.\"\n",
         "        )\n",
         "\n",
@@ -334,12 +342,14 @@
         "        serving_container_health_route=container_spec.get(\"healthRoute\"),\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        model_garden_source_model_name=(\n",
         "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
         "        ),\n",
         "    )\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
         "        display_name=model.name + \"-endpoint\",\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        dedicated_endpoint_enabled=True,\n",
         "    )\n",
         "    print(\n",
@@ -481,6 +491,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -533,6 +544,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -571,6 +586,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_1_reasoning_engine.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",

notebooks/community/model_garden/model_garden_pytorch_llama3_2_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_2_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -56,7 +61,6 @@
         "\n",
         "This notebook demonstrates downloading, deploying, and serving prebuilt Llama 3.2 models on GPU and TPU.\n",
         "\n",
-        "\n",
         "### Objective\n",
         "\n",
         "- Deploy Llama 3.2 1B and 3B with [Hex-LLM](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm) on TPU.\n",
@@ -64,6 +68,10 @@
         "- Deploy Llama 3.2 11B-vision and 90B-vision with [vLLM](https://github.com/vllm-project/vllm) on GPU with limited functions supported.\n",
         "- Deploy Llama 3.2 11B-vision and 90B-vision with reference server on GPU with full functions supported.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -104,6 +112,7 @@
     {
       "cell_type": "code",
       "execution_count": null,
+      "language": "python",
       "metadata": {
         "cellView": "form",
         "id": "YXFGIp1l-qtT"
@@ -114,46 +123,45 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
         "import re\n",
-        "import uuid\n",
         "from typing import Tuple\n",
         "\n",
         "import requests\n",
+        "from google import auth\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "LABEL = \"vllm_gpu\"\n",
         "models, endpoints = {}, {}\n",
-        "# Dedicated endpoint not supported yet\n",
-        "use_dedicated_endpoint = False\n",
+        "\n",
         "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -162,52 +170,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3_2\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown # Access Llama 3.2 models on Vertex AI for serving\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -252,15 +225,16 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
-        "\n",
-        "# @markdown This section uploads prebuilt Llama 3.2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "# @title Select the model variants\n",
         "\n",
         "# @markdown Select one of the four model variations. More model variants will be supported by Hex-LLM in the future.\n",
-        "MODEL_ID = \"Llama-3.2-1B\"  # @param [\"Llama-3.2-1B\", \"Llama-3.2-1B-Instruct\", \"Llama-3.2-3B\", \"Llama-3.2-3B-Instruct\"] {allow-input: true, isTemplate: true}\n",
+        "\n",
+        "MODEL_ID = \"Llama-3.2-1B-Instruct\"  # @param [\"Llama-3.2-1B\", \"Llama-3.2-1B-Instruct\", \"Llama-3.2-3B\", \"Llama-3.2-3B-Instruct\"] {allow-input: true, isTemplate: true}\n",
         "TPU_DEPLOYMENT_REGION = \"us-west1\"  # @param [\"us-west1\"] {isTemplate:true}\n",
-        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, MODEL_ID)\n",
         "hf_model_id = \"meta-llama/\" + MODEL_ID\n",
+        "version_id = MODEL_ID.lower()\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, MODEL_ID)\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama3-2@{version_id}\"\n",
         "\n",
         "# The pre-built serving docker images.\n",
         "HEXLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai-restricted/vertex-vision-model-garden-dockers/hex-llm-serve:20241210_2323_RC00\"\n",
@@ -281,7 +255,21 @@
         "    accelerator_type=tpu_type,\n",
         "    accelerator_count=tpu_count,\n",
         "    is_for_training=False,\n",
-        ")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "OsafKD_VxgiK"
+      },
+      "outputs": [],
+      "source": [
+        "# @title Deploy\n",
+        "\n",
+        "# @markdown This section uploads prebuilt Llama 3.2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
         "# @markdown Set enable_prefix_cache_hbm to False if you don't want to use [prefix caching](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-hex-llm#prefix-caching).\n",
         "enable_prefix_cache_hbm = True  # @param {type:\"boolean\"}\n",
@@ -300,13 +288,15 @@
         "min_replica_count = 1\n",
         "max_replica_count = 1\n",
         "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "\n",
         "def deploy_model_hexllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    data_parallel_size: int = 1,\n",
         "    tensor_parallel_size: int = 1,\n",
@@ -395,11 +385,11 @@
         "        machine_type=machine_type,\n",
         "        tpu_topology=tpu_topology if num_hosts > 1 else None,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        min_replica_count=min_replica_count,\n",
         "        max_replica_count=max_replica_count,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_2_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    return model, endpoint\n",
@@ -410,7 +400,6 @@
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3-2\",\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    base_model_id=hf_model_id,\n",
         "    tensor_parallel_size=tensor_parallel_size,\n",
         "    machine_type=machine_type,\n",
@@ -498,9 +487,11 @@
       "source": [
         "# @title Chat completion\n",
         "\n",
-        "_region = REGION\n",
+        "temp_region = REGION\n",
         "REGION = TPU_DEPLOYMENT_REGION\n",
         "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"hexllm_tpu\"].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
         "    PROJECT_ID, REGION, endpoints[\"hexllm_tpu\"].name\n",
         ")\n",
@@ -519,6 +510,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -545,10 +537,23 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
         "\n",
-        "REGION = _region\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "REGION = temp_region\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -573,13 +578,18 @@
       "source": [
         "# @title Fast Deployment\n",
         "\n",
+        "# @markdown The Llama 3.2 1B/3B Instruct model will be deployed to a dedicated endpoint on an `a2-ultragpu-1g` machine with Fast Deployment.\n",
+        "# @markdown **Currently, the Fast Deployment is only supported in the `us-central1` region.**\n",
+        "\n",
         "# @markdown This section demonstrates how to use the Fast Deployment feature.\n",
         "\n",
         "# @markdown The Fast Deployment feature prioritizes speed for model exploration, making it ideal for initial testing and experimentation. For sensitive data or production workloads, use the Standard environment for enhanced security and stability.\n",
         "\n",
         "# @markdown Note that only a subset of the models support the Fast Deployment feature.\n",
         "\n",
-        "API_ENDPOINT = f\"{REGION}-aiplatform.googleapis.com\"\n",
+        "FAST_DEPLOYMENT_REGION = \"us-central1\"  # @param [\"us-central1\"] {isTemplate:true}\n",
+        "\n",
+        "API_ENDPOINT = f\"{FAST_DEPLOYMENT_REGION}-aiplatform.googleapis.com\"\n",
         "\n",
         "\n",
         "def fast_deploy(\n",
@@ -607,7 +617,7 @@
         "        == 0\n",
         "    ):\n",
         "        raise ValueError(\n",
-        "            f\"No supportedActions.multiDeployVertex found in {REGION}. You can skip\"\n",
+        "            f\"No supportedActions.multiDeployVertex found in {FAST_DEPLOYMENT_REGION}. You can skip\"\n",
         "            \" this section or try a different region.\"\n",
         "        )\n",
         "    deploy_configs = response[\"supportedActions\"][\"multiDeployVertex\"][\n",
@@ -625,7 +635,7 @@
         "        fast_deploy_config = fast_deploy_config[0]\n",
         "    else:\n",
         "        raise ValueError(\n",
-        "            f\"No Fast Deployment config found in {REGION}. You can skip this\"\n",
+        "            f\"No Fast Deployment config found in {FAST_DEPLOYMENT_REGION}. You can skip this\"\n",
         "            \" section or try a different region.\"\n",
         "        )\n",
         "\n",
@@ -649,12 +659,14 @@
         "        serving_container_health_route=container_spec.get(\"healthRoute\"),\n",
         "        serving_container_shared_memory_size_mb=(16 * 1024),  # 16 GB\n",
         "        serving_container_deployment_timeout=7200,\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        model_garden_source_model_name=(\n",
         "            f\"publishers/{publisher}/models/{publisher_model_id}\"\n",
         "        ),\n",
         "    )\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
         "        display_name=model.name + \"-endpoint\",\n",
+        "        location=FAST_DEPLOYMENT_REGION,\n",
         "        dedicated_endpoint_enabled=True,\n",
         "    )\n",
         "    print(\n",
@@ -678,10 +690,9 @@
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "# @markdown The Llama 3.2 1B/3B Instruct model will be deployed to a dedicated endpoint on an `a2-ultragpu-1g` machine with Fast Deployment.\n",
-        "# @markdown **Currently, the Fast Deployment is only supported in the `us-central1` region.**\n",
+        "# Fast Deployment only supports dedicated endpoints.\n",
+        "use_dedicated_endpoint = True\n",
         "\n",
-        "use_dedicated_endpoint = True  # Fast Deployment only supports dedicated endpoints.\n",
         "base_model_name = \"Llama-3.2-1B-Instruct\"  # @param [\"Llama-3.2-1B-Instruct\", \"Llama-3.2-3B-Instruct\"] {isTemplate:true}\n",
         "models[\"vllm_fast\"], endpoints[\"vllm_fast\"] = fast_deploy(\n",
         "    \"meta\", \"llama3-2\", base_model_name.lower()\n",
@@ -764,7 +775,11 @@
       "source": [
         "# @title Chat completion\n",
         "\n",
-        "DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
+        "temp_region = REGION\n",
+        "REGION = FAST_DEPLOYMENT_REGION\n",
+        "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"vllm_fast\"].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
         "    PROJECT_ID, REGION, endpoints[\"vllm_fast\"].name\n",
         ")\n",
@@ -783,6 +798,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -809,8 +825,23 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
+        "\n",
+        "REGION = temp_region\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -821,7 +852,7 @@
         "id": "lBeZoHJCrHew"
       },
       "source": [
-        "## Deploy prebuilt Llama 3.2 11B-vision and 90B-vision with vLLM"
+        "## Deploy prebuilt Llama 3.2 11B-Vision and 90B-Vision with vLLM"
       ]
     },
     {
@@ -833,20 +864,19 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
+        "# @title Select the model variants\n",
         "\n",
-        "# @markdown This section uploads prebuilt Llama 3.2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
-        "\n",
-        "# @markdown Currently vLLM can be used with limited inputs for multi-modality models, which are \"text only\" format and \"single leading image + text\" format. More input formats will be supported later.\n",
-        "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"Llama-3.2-11B-Vision\"  # @param [\"Llama-3.2-1B\", \"Llama-3.2-1B-Instruct\", \"Llama-3.2-3B\", \"Llama-3.2-3B-Instruct\", \"Llama-3.2-11B-Vision\", \"Llama-3.2-11B-Vision-Instruct\", \"Llama-3.2-90B-Vision\", \"Llama-3.2-90B-Vision-Instruct\"] {isTemplate:true}\n",
-        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, base_model_name)\n",
+        "base_model_name = \"Llama-3.2-11B-Vision-Instruct\"  # @param [\"Llama-3.2-11B-Vision\", \"Llama-3.2-11B-Vision-Instruct\", \"Llama-3.2-90B-Vision\", \"Llama-3.2-90B-Vision-Instruct\"] {isTemplate:true}\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, base_model_name)\n",
+        "version_id = base_model_name.replace(\"3.2\", \"3-2\").lower()\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama3-2@{version_id}\"\n",
         "\n",
         "# The pre-built serving docker images.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241001_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241202_0916_RC00_maas\"\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
         "\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "if \"3.2-1B\" in base_model_name or \"3.2-3B\" in base_model_name:\n",
@@ -872,17 +902,72 @@
         "    is_for_training=False,\n",
         ")\n",
         "\n",
+        "# @markdown Click \"Show Code\" to see more details."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "zgxn-i945wCK"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Z8IY5x7j5zZm"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "# @markdown This section uploads prebuilt Llama 3.2 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
+        "\n",
+        "# @markdown Currently vLLM can be used with limited inputs for multi-modality models, which are \"text only\" format and \"single leading image + text\" format. More input formats will be supported later.\n",
+        "\n",
         "gpu_memory_utilization = 0.9\n",
         "max_model_len = 4096\n",
         "max_num_seqs = 12\n",
         "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown To enable the auto-scaling in deployment, you can set the following options:\n",
+        "\n",
+        "min_replica_count = 1  # @param {type:\"integer\"}\n",
+        "max_replica_count = 1  # @param {type:\"integer\"}\n",
+        "required_replica_count = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Set the target of GPU duty cycle or CPU usage between 1 and 100 for auto-scaling.\n",
+        "autoscale_by_gpu_duty_cycle_target = 0  # @param {type:\"integer\"}\n",
+        "autoscale_by_cpu_usage_target = 0  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Note: GPU duty cycle is not the most accurate metric for scaling workloads. More advanced auto-scaling metrics are coming soon. See [the public doc](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#autoscaling) for more details.\n",
+        "\n",
         "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -901,6 +986,13 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        "    required_replica_count: int = 1,\n",
+        "    autoscale_by_gpu_duty_cycle_target: int = 0,\n",
+        "    autoscale_by_cpu_usage_target: int = 0,\n",
+        "    is_spot: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -953,6 +1045,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -982,29 +1078,74 @@
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
         "    )\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=accelerator_count,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_2_deployment.ipynb\",\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": min_replica_count,\n",
+        "                \"requiredReplicaCount\": required_replica_count,\n",
+        "                \"maxReplicaCount\": max_replica_count,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_2_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
         "        },\n",
-        "    )\n",
+        "    }\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    if autoscale_by_gpu_duty_cycle_target > 0 or autoscale_by_cpu_usage_target > 0:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"autoscalingMetricSpecs\"] = []\n",
+        "        if autoscale_by_gpu_duty_cycle_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle\",\n",
+        "                    \"target\": autoscale_by_gpu_duty_cycle_target,\n",
+        "                }\n",
+        "            )\n",
+        "        if autoscale_by_cpu_usage_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/cpu/utilization\",\n",
+        "                    \"target\": autoscale_by_cpu_usage_target,\n",
+        "                }\n",
+        "            )\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    common_util.poll_and_wait(response.json()[\"name\"], REGION, 7200)\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3_2-serve-vllm\"),\n",
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3-2\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -1013,7 +1154,13 @@
         "    enforce_eager=True,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
         "    max_num_seqs=max_num_seqs,\n",
-        "    model_type=\"llama3.1\",\n",
+        "    enable_llama_tool_parser=True,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        "    required_replica_count=required_replica_count,\n",
+        "    autoscale_by_gpu_duty_cycle_target=autoscale_by_gpu_duty_cycle_target,\n",
+        "    autoscale_by_cpu_usage_target=autoscale_by_cpu_usage_target,\n",
+        "    is_spot=is_spot,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -1027,10 +1174,12 @@
       },
       "outputs": [],
       "source": [
-        "# @title Chat completion for text-only models\n",
+        "# @title Chat completion with text-only requests\n",
         "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
         ")\n",
         "\n",
         "# @title Chat Completions Inference\n",
@@ -1047,6 +1196,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -1073,8 +1223,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -1090,8 +1253,10 @@
       "source": [
         "# @title Chat completion for vision models\n",
         "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[LABEL].gca_resource.dedicated_endpoint_dns\n",
         "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
-        "    PROJECT_ID, REGION, endpoints[\"vllm_gpu\"].name\n",
+        "    PROJECT_ID, REGION, endpoints[LABEL].name\n",
         ")\n",
         "\n",
         "# @title Chat Completions Inference\n",
@@ -1155,7 +1320,7 @@
         "id": "4AoiFsQWhYtM"
       },
       "source": [
-        "## Deploy prebuilt Llama 3.2 11B-vision and 90B-vision with reference server"
+        "## Deploy prebuilt Llama 3.2 11B-Vision and 90B-Vision with reference server"
       ]
     },
     {
@@ -1175,7 +1340,7 @@
         "\n",
         "# @markdown Set the model to deploy.\n",
         "\n",
-        "base_model_name = \"Llama-3.2-11B-Vision-meta\"  # @param [\"Llama-3.2-1B-meta\", \"Llama-3.2-1B-Instruct-meta\", \"Llama-3.2-3B-meta\", \"Llama-3.2-3B-Instruct-meta\", \"Llama-3.2-11B-Vision-meta\", \"Llama-3.2-11B-Vision-Instruct-meta\", \"Llama-3.2-90B-Vision-meta\", \"Llama-3.2-90B-Vision-Instruct-meta\"] {isTemplate:true}\n",
+        "base_model_name = \"Llama-3.2-11B-Vision-Instruct-meta\"  # @param [\"Llama-3.2-11B-Vision-meta\", \"Llama-3.2-11B-Vision-Instruct-meta\", \"Llama-3.2-90B-Vision-meta\", \"Llama-3.2-90B-Vision-Instruct-meta\"] {isTemplate:true}\n",
         "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_2, base_model_name)\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
         "\n",
@@ -1183,15 +1348,11 @@
         "REF_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-peft-serve:llama_ref_impl\"\n",
         "\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "if \"3.2-1B\" in base_model_name or \"3.2-3B\" in base_model_name:\n",
-        "    accelerator_type = \"NVIDIA_L4\"\n",
-        "    machine_type = \"g2-standard-8\"\n",
-        "    accelerator_count = 1\n",
-        "elif \"3.2-11B\" in base_model_name:\n",
+        "if \"11B\" in base_model_name:\n",
         "    accelerator_type = \"NVIDIA_TESLA_A100\"\n",
         "    machine_type = \"a2-highgpu-1g\"\n",
         "    accelerator_count = 1\n",
-        "elif \"3.2-90B\" in base_model_name:\n",
+        "elif \"90B\" in base_model_name:\n",
         "    accelerator_type = \"NVIDIA_H100_80GB\"\n",
         "    machine_type = \"a3-highgpu-8g\"\n",
         "    accelerator_count = 8\n",
@@ -1212,7 +1373,6 @@
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -1273,9 +1433,9 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_2_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -1291,7 +1451,6 @@
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3-2\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -1311,6 +1470,12 @@
       "source": [
         "# @title Chat completion for text-only models\n",
         "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ref_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"ref_gpu\"].name\n",
+        ")\n",
+        "\n",
         "# @markdown Once deployment succeeds, you can send requests to the endpoint using the Vertex SDK.\n",
         "\n",
         "# @markdown Next fill out some request parameters:\n",
@@ -1350,6 +1515,12 @@
       "source": [
         "# @title Chat completion for vision models\n",
         "\n",
+        "if use_dedicated_endpoint:\n",
+        "    DEDICATED_ENDPOINT_DNS = endpoints[\"ref_gpu\"].gca_resource.dedicated_endpoint_dns\n",
+        "ENDPOINT_RESOURCE_NAME = \"projects/{}/locations/{}/endpoints/{}\".format(\n",
+        "    PROJECT_ID, REGION, endpoints[\"ref_gpu\"].name\n",
+        ")\n",
+        "\n",
         "# @markdown Once deployment succeeds, you can send requests to the endpoint using the Vertex SDK.\n",
         "\n",
         "# @markdown Next fill out some request parameters:\n",
@@ -1436,11 +1607,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_llama3_3_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_3_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -61,6 +66,10 @@
         "\n",
         "- Deploy Llama 3.3 70B Instruct with vLLM on GPU, optionally with dynamic LoRA adapters.\n",
         "\n",
+        "### File a bug\n",
+        "\n",
+        "File a bug on [GitHub](https://github.com/GoogleCloudPlatform/vertex-ai-samples/issues/new) if you encounter any issue with the notebook.\n",
+        "\n",
         "### Costs\n",
         "\n",
         "This tutorial uses billable components of Google Cloud:\n",
@@ -107,41 +116,43 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
-        "import uuid\n",
+        "import re\n",
         "from typing import Tuple\n",
         "\n",
+        "import requests\n",
+        "from google import auth\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
+        "LABEL = \"vllm_gpu\"\n",
         "models, endpoints = {}, {}\n",
         "\n",
         "# Get the default cloud project id.\n",
@@ -151,52 +162,17 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3-3\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")\n",
         "\n",
         "# @markdown # Access Llama 3.3 models on Vertex AI for serving\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -210,7 +186,13 @@
         "VERTEX_AI_MODEL_GARDEN_LLAMA_3_3 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "assert (\n",
         "    VERTEX_AI_MODEL_GARDEN_LLAMA_3_3\n",
-        "), \"Click the agreement of Llama 3.3 in Vertex AI Model Garden, and get the GCS path of Llama 3.3 model artifacts.\""
+        "), \"Click the agreement of Llama 3.3 in Vertex AI Model Garden, and get the GCS path of Llama 3.3 model artifacts.\"\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(/)?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA_3_3)\n",
+        "if parsed_gcs_url:\n",
+        "    VERTEX_AI_MODEL_GARDEN_LLAMA3 = parsed_gcs_url.group()\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA3.startswith(\n",
+        "    \"gs://\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA3 is expected to be a GCS URI and must start with `gs://`.\""
       ]
     },
     {
@@ -231,28 +213,28 @@
       },
       "outputs": [],
       "source": [
-        "# @title Deploy\n",
-        "\n",
-        "# @markdown This section uploads Llama 3.3 to Model Registry and deploys it to a Vertex AI Endpoint. It takes ~30 minutes.\n",
-        "\n",
-        "# @markdown The serving efficiency of L4 GPUs is inferior to that of H100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have H100 quota.\n",
-        "\n",
-        "# @markdown H100 is hard to get for now. It's recommended to use the deployment button in the model card. You can still try to deploy H100 endpoint through the notebook, but there is a chance that resource is not available.\n",
+        "# @title Select the model variants\n",
         "\n",
         "# @markdown Set the model to deploy.\n",
         "\n",
+        "# fmt: off\n",
         "base_model_name = \"Llama-3.3-70B-Instruct\"  # @param [\"Llama-3.3-70B-Instruct\"] {isTemplate:true}\n",
+        "# fmt: on\n",
+        "\n",
         "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA_3_3, base_model_name)\n",
         "ENABLE_DYNAMIC_LORA = True  # @param {type:\"boolean\", isTemplate:true}\n",
         "hf_model_id = \"meta-llama/\" + base_model_name\n",
+        "version_id = \"llama-3.3-70b-instruct\"\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama3-3@{version_id}\"\n",
         "\n",
         "accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_H100_80GB\", \"NVIDIA_L4\"]\n",
         "\n",
         "# The pre-built serving docker images.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241001_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250114_0916_RC00_maas\"\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
         "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
+        "\n",
         "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
         "if accelerator_type == \"NVIDIA_L4\":\n",
         "    machine_type = \"g2-standard-96\"\n",
@@ -273,18 +255,74 @@
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
         "    is_for_training=False,\n",
-        ")\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "EYEJBfsbNNmR"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "Hl4l047_NSb3"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads Llama 3.3 to Model Registry and deploys it to a Vertex AI Endpoint. It takes ~30 minutes.\n",
+        "\n",
+        "# @markdown The serving efficiency of L4 GPUs is inferior to that of H100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have H100 quota.\n",
+        "\n",
+        "# @markdown H100 is hard to get for now. It's recommended to use the deployment button in the model card. You can still try to deploy H100 endpoint through the notebook, but there is a chance that resource is not available.\n",
         "\n",
         "gpu_memory_utilization = 0.95\n",
         "max_model_len = 8192  # Maximum context length.\n",
         "\n",
+        "# @markdown Choose whether to use a [Spot VM](https://cloud.google.com/compute/docs/instances/spot) for the deployment.\n",
+        "is_spot = False  # @param {type:\"boolean\"}\n",
+        "\n",
+        "# @markdown To enable the auto-scaling in deployment, you can set the following options:\n",
+        "\n",
+        "min_replica_count = 1  # @param {type:\"integer\"}\n",
+        "max_replica_count = 1  # @param {type:\"integer\"}\n",
+        "required_replica_count = 1  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Set the target of GPU duty cycle or CPU usage between 1 and 100 for auto-scaling.\n",
+        "autoscale_by_gpu_duty_cycle_target = 0  # @param {type:\"integer\"}\n",
+        "autoscale_by_cpu_usage_target = 0  # @param {type:\"integer\"}\n",
+        "\n",
+        "# @markdown Note: GPU duty cycle is not the most accurate metric for scaling workloads. More advanced auto-scaling metrics are coming soon. See [the public doc](https://cloud.google.com/vertex-ai/docs/predictions/configure-compute#autoscaling) for more details.\n",
+        "\n",
         "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -303,6 +341,13 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
+        "    min_replica_count: int = 1,\n",
+        "    max_replica_count: int = 1,\n",
+        "    required_replica_count: int = 1,\n",
+        "    autoscale_by_gpu_duty_cycle_target: int = 0,\n",
+        "    autoscale_by_cpu_usage_target: int = 0,\n",
+        "    is_spot: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -355,6 +400,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -384,29 +433,74 @@
         "    print(\n",
         "        f\"Deploying {model_name} on {machine_type} with {accelerator_count} {accelerator_type} GPU(s).\"\n",
         "    )\n",
-        "    model.deploy(\n",
-        "        endpoint=endpoint,\n",
-        "        machine_type=machine_type,\n",
-        "        accelerator_type=accelerator_type,\n",
-        "        accelerator_count=accelerator_count,\n",
-        "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
-        "        system_labels={\n",
-        "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_3_deployment.ipynb\",\n",
+        "\n",
+        "    creds, _ = auth.default()\n",
+        "    auth_req = auth.transport.requests.Request()\n",
+        "    creds.refresh(auth_req)\n",
+        "\n",
+        "    url = f\"https://{REGION}-aiplatform.googleapis.com/ui/projects/{PROJECT_ID}/locations/{REGION}/endpoints/{endpoint.name}:deployModel\"\n",
+        "    headers = {\n",
+        "        \"Content-Type\": \"application/json\",\n",
+        "        \"Authorization\": f\"Bearer {creds.token}\",\n",
+        "    }\n",
+        "    data = {\n",
+        "        \"deployedModel\": {\n",
+        "            \"model\": model.resource_name,\n",
+        "            \"displayName\": model_name,\n",
+        "            \"dedicatedResources\": {\n",
+        "                \"machineSpec\": {\n",
+        "                    \"machineType\": machine_type,\n",
+        "                    \"acceleratorType\": accelerator_type,\n",
+        "                    \"acceleratorCount\": accelerator_count,\n",
+        "                },\n",
+        "                \"minReplicaCount\": min_replica_count,\n",
+        "                \"requiredReplicaCount\": required_replica_count,\n",
+        "                \"maxReplicaCount\": max_replica_count,\n",
+        "            },\n",
+        "            \"system_labels\": {\n",
+        "                \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_3_deployment.ipynb\",\n",
+        "                \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
+        "            },\n",
         "        },\n",
-        "    )\n",
+        "    }\n",
+        "    if is_spot:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"spot\"] = True\n",
+        "    if autoscale_by_gpu_duty_cycle_target > 0 or autoscale_by_cpu_usage_target > 0:\n",
+        "        data[\"deployedModel\"][\"dedicatedResources\"][\"autoscalingMetricSpecs\"] = []\n",
+        "        if autoscale_by_gpu_duty_cycle_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/accelerator/duty_cycle\",\n",
+        "                    \"target\": autoscale_by_gpu_duty_cycle_target,\n",
+        "                }\n",
+        "            )\n",
+        "        if autoscale_by_cpu_usage_target > 0:\n",
+        "            data[\"deployedModel\"][\"dedicatedResources\"][\n",
+        "                \"autoscalingMetricSpecs\"\n",
+        "            ].append(\n",
+        "                {\n",
+        "                    \"metricName\": \"aiplatform.googleapis.com/prediction/online/cpu/utilization\",\n",
+        "                    \"target\": autoscale_by_cpu_usage_target,\n",
+        "                }\n",
+        "            )\n",
+        "    response = requests.post(url, headers=headers, json=data)\n",
+        "    print(f\"Deploy Model response: {response.json()}\")\n",
+        "    if response.status_code != 200 or \"name\" not in response.json():\n",
+        "        raise ValueError(f\"Failed to deploy model: {response.text}\")\n",
+        "    common_util.poll_and_wait(response.json()[\"name\"], REGION, 7200)\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
         "    return model, endpoint\n",
         "\n",
         "\n",
-        "models[\"vllm_gpu\"], endpoints[\"vllm_gpu\"] = deploy_model_vllm(\n",
+        "models[LABEL], endpoints[LABEL] = deploy_model_vllm(\n",
         "    model_name=common_util.get_job_name_with_datetime(prefix=\"llama3-3-serve\"),\n",
         "    model_id=model_id,\n",
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3-3\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -417,7 +511,13 @@
         "    enable_lora=ENABLE_DYNAMIC_LORA,\n",
         "    enable_chunked_prefill=not ENABLE_DYNAMIC_LORA,\n",
         "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
-        "    model_type=\"llama3.1\",\n",
+        "    enable_llama_tool_parser=True,\n",
+        "    min_replica_count=min_replica_count,\n",
+        "    max_replica_count=max_replica_count,\n",
+        "    required_replica_count=required_replica_count,\n",
+        "    autoscale_by_gpu_duty_cycle_target=autoscale_by_gpu_duty_cycle_target,\n",
+        "    autoscale_by_cpu_usage_target=autoscale_by_cpu_usage_target,\n",
+        "    is_spot=is_spot,\n",
         ")\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -509,6 +609,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -535,8 +636,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -570,11 +684,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],

notebooks/community/model_garden/model_garden_pytorch_llama3_3_finetuning.ipynb

@@ -104,9 +104,9 @@
         "\n",
         "# @markdown 2. Install packages to validate dataset with template.\n",
         "! pip install --upgrade --quiet gcsfs==2024.3.1\n",
-        "! pip install --upgrade --quiet accelerate==0.31.0\n",
-        "! pip install --upgrade --quiet transformers==4.45.2\n",
-        "! pip install --upgrade --quiet datasets==2.19.2\n",
+        "! pip install --upgrade --quiet accelerate==0.34.2\n",
+        "! pip install --upgrade --quiet transformers==4.47.1\n",
+        "! pip install --upgrade --quiet datasets==2.20.0\n",
         "\n",
         "# Load local tensorboard.\n",
         "%load_ext tensorboard"
@@ -134,7 +134,7 @@
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# @markdown 4. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
         "\n",
@@ -169,6 +169,12 @@
         "\n",
         "# Get the default region for launching jobs.\n",
         "if not REGION:\n",
+        "    if not os.environ.get(\"GOOGLE_CLOUD_REGION\"):\n",
+        "        raise ValueError(\n",
+        "            \"REGION must be set. See\"\n",
+        "            \" https://cloud.google.com/vertex-ai/docs/general/locations for\"\n",
+        "            \" available cloud locations.\"\n",
+        "        )\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
         "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
@@ -470,7 +476,7 @@
         "    }\n",
         "\n",
         "TRAIN_DOCKER_URI = (\n",
-        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20241205\"\n",
+        "    f\"{repo}/vertex-vision-model-garden-dockers/pytorch-peft-train:stable_20250320\"\n",
         ")\n",
         "\n",
         "# Worker pool spec.\n",
@@ -509,9 +515,9 @@
         "gradient_accumulation_steps = 4  # @param{type:\"integer\"}\n",
         "# @markdown Maximum sequence length.\n",
         "max_seq_length = 4096  # @param{type:\"integer\"}\n",
-        "# @markdown Setting a positive `max_steps` here will override `num_epochs`.\n",
+        "# @markdown Setting a positive `max_steps` here will override `num_train_epochs`.\n",
         "max_steps = -1  # @param{type:\"integer\"}\n",
-        "num_epochs = 1.0  # @param{type:\"number\"}\n",
+        "num_train_epochs = 1.0  # @param{type:\"number\"}\n",
         "# @markdown Precision mode for finetuning.\n",
         "finetuning_precision_mode = \"4bit\"  # @param [\"4bit\", \"8bit\", \"float16\"]\n",
         "# @markdown Learning rate.\n",
@@ -553,7 +559,7 @@
         "# Create a GCS folder to store the LORA adapter.\n",
         "lora_output_dir = os.path.join(base_output_dir, \"adapter\")\n",
         "# Create a GCS folder to store the finetuned LORA adapter.\n",
-        "final_checkpoint = os.path.join(lora_output_dir, \"checkpoint-final\")\n",
+        "final_checkpoint = os.path.join(lora_output_dir, \"node-0\", \"checkpoint-final\")\n",
         "\n",
         "# Add labels for the finetuning job.\n",
         "labels = {\n",
@@ -571,7 +577,6 @@
         "    f\"--eval_template={template}\",\n",
         "    f\"--eval_split={eval_split_name}\",\n",
         "    f\"--eval_steps={save_steps}\",\n",
-        "    \"--eval_tasks=builtin_eval\",\n",
         "    \"--eval_metric_name=loss\",\n",
         "]\n",
         "\n",
@@ -595,7 +600,7 @@
         "    f\"--lr_scheduler_type={lr_scheduler_type}\",\n",
         "    f\"--precision_mode={finetuning_precision_mode}\",\n",
         "    f\"--gradient_checkpointing={enable_gradient_checkpointing}\",\n",
-        "    f\"--num_epochs={num_epochs}\",\n",
+        "    f\"--num_train_epochs={num_train_epochs}\",\n",
         "    f\"--attn_implementation={attn_implementation}\",\n",
         "    f\"--optimizer={optimizer}\",\n",
         "    f\"--warmup_ratio={warmup_ratio}\",\n",
@@ -677,7 +682,7 @@
         "    print(\"The training job has finished.\")\n",
         "\n",
         "# The pre-built serving docker image for vLLM.\n",
-        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20241001_0916_RC00\"\n",
+        "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20250116_0916_RC00\"\n",
         "\n",
         "serve_accelerator_type = \"NVIDIA_H100_80GB\"  # @param [\"NVIDIA_H100_80GB\", \"NVIDIA_L4\"]\n",
         "\n",
@@ -712,6 +717,18 @@
         "    raise ValueError(\"max_model_len cannot exceed 8192\")\n",
         "\n",
         "\n",
+        "def get_deploy_source() -> str:\n",
+        "    \"\"\"Gets deploy_source string based on running environment.\"\"\"\n",
+        "    vertex_product = os.environ.get(\"VERTEX_PRODUCT\", \"\")\n",
+        "    if vertex_product == \"COLAB_ENTERPRISE\":\n",
+        "        return \"notebook_colab_enterprise\"\n",
+        "    elif vertex_product == \"WORKBENCH_INSTANCE\":\n",
+        "        return \"notebook_workbench\"\n",
+        "    else:\n",
+        "        # Legacy workbench, legacy colab, or other custom environments.\n",
+        "        return \"notebook_environment_unspecified\"\n",
+        "\n",
+        "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
@@ -736,6 +753,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -788,6 +806,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -826,6 +848,7 @@
         "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_3_finetuning.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
@@ -914,7 +937,6 @@
         "top_k = 1  # @param {type:\"integer\"}\n",
         "raw_response = False  # @param {type:\"boolean\"}\n",
         "\n",
-        "\n",
         "predict_vllm(\n",
         "    prompt=prompt,\n",
         "    max_tokens=max_tokens,\n",

notebooks/community/model_garden/model_garden_pytorch_llama3_deployment.ipynb

@@ -34,6 +34,11 @@
         "\n",
         "<table><tbody><tr>\n",
         "  <td style=\"text-align: center\">\n",
+        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/instances\">\n",
+        "      <img alt=\"Workbench logo\" src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" width=\"32px\"><br> Run in Workbench\n",
+        "    </a>\n",
+        "  </td>\n",
+        "  <td style=\"text-align: center\">\n",
         "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fvertex-ai-samples%2Fmain%2Fnotebooks%2Fcommunity%2Fmodel_garden%2Fmodel_garden_pytorch_llama3_deployment.ipynb\">\n",
         "      <img alt=\"Google Cloud Colab Enterprise logo\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" width=\"32px\"><br> Run in Colab Enterprise\n",
         "    </a>\n",
@@ -94,43 +99,42 @@
         "\n",
         "# @markdown 1. [Make sure that billing is enabled for your project](https://cloud.google.com/billing/docs/how-to/modify-project).\n",
         "\n",
-        "# @markdown 2. **[Optional]** [Create a Cloud Storage bucket](https://cloud.google.com/storage/docs/creating-buckets) for storing experiment outputs. Set the BUCKET_URI for the experiment environment. The specified Cloud Storage bucket (`BUCKET_URI`) should be located in the same region as where the notebook was launched. Note that a multi-region bucket (eg. \"us\") is not considered a match for a single region covered by the multi-region range (eg. \"us-central1\"). If not set, a unique GCS bucket will be created instead.\n",
-        "\n",
-        "BUCKET_URI = \"gs://\"  # @param {type:\"string\"}\n",
-        "\n",
-        "# @markdown 3. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
+        "# @markdown 2. **[Optional]** Set region. If not set, the region will be set automatically according to Colab Enterprise environment.\n",
         "\n",
         "REGION = \"\"  # @param {type:\"string\"}\n",
         "\n",
-        "# @markdown 4. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus).\n",
+        "# @markdown 3. If you want to run predictions with A100 80GB or H100 GPUs, we recommend using the regions listed below. **NOTE:** Make sure you have associated quota in selected regions. Click the links to see your current quota for each GPU type: [Nvidia A100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_a100_80gb_gpus), [Nvidia H100 80GB](https://console.cloud.google.com/iam-admin/quotas?metric=aiplatform.googleapis.com%2Fcustom_model_serving_nvidia_h100_gpus). You can request for quota following the instructions at [\"Request a higher quota\"](https://cloud.google.com/docs/quota/view-manage#requesting_higher_quota).\n",
         "\n",
         "# @markdown > | Machine Type | Accelerator Type | Recommended Regions |\n",
         "# @markdown | ----------- | ----------- | ----------- |\n",
         "# @markdown | a2-ultragpu-1g | 1 NVIDIA_A100_80GB | us-central1, us-east4, europe-west4, asia-southeast1, us-east4 |\n",
         "# @markdown | a3-highgpu-2g | 2 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
         "# @markdown | a3-highgpu-4g | 4 NVIDIA_H100_80GB | us-west1, asia-southeast1, europe-west4 |\n",
-        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, us-east5, europe-west4, us-west1, asia-southeast1 |\n",
+        "# @markdown | a3-highgpu-8g | 8 NVIDIA_H100_80GB | us-central1, europe-west4, us-west1, asia-southeast1 |\n",
         "\n",
         "# Import the necessary packages\n",
         "\n",
         "# Upgrade Vertex AI SDK.\n",
-        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.64.0'\n",
-        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "! pip3 install --upgrade --quiet 'google-cloud-aiplatform>=1.84.0'\n",
         "\n",
-        "import datetime\n",
         "import importlib\n",
         "import os\n",
         "import re\n",
-        "import uuid\n",
         "from typing import Tuple\n",
         "\n",
         "from google.cloud import aiplatform\n",
         "\n",
+        "# Upgrade Vertex AI SDK.\n",
+        "if os.environ.get(\"VERTEX_PRODUCT\") != \"COLAB_ENTERPRISE\":\n",
+        "    ! pip install --upgrade tensorflow\n",
+        "! git clone https://github.com/GoogleCloudPlatform/vertex-ai-samples.git\n",
+        "\n",
         "common_util = importlib.import_module(\n",
         "    \"vertex-ai-samples.community-content.vertex_model_garden.model_oss.notebook_util.common_util\"\n",
         ")\n",
         "\n",
         "models, endpoints = {}, {}\n",
+        "LABEL = \"vllm_gpu\"\n",
         "\n",
         "# Get the default cloud project id.\n",
         "PROJECT_ID = os.environ[\"GOOGLE_CLOUD_PROJECT\"]\n",
@@ -139,52 +143,161 @@
         "if not REGION:\n",
         "    REGION = os.environ[\"GOOGLE_CLOUD_REGION\"]\n",
         "\n",
-        "# Enable the Vertex AI API and Compute Engine API, if not already.\n",
-        "print(\"Enabling Vertex AI API and Compute Engine API.\")\n",
-        "! gcloud services enable aiplatform.googleapis.com compute.googleapis.com\n",
-        "\n",
-        "# Cloud Storage bucket for storing the experiment artifacts.\n",
-        "# A unique GCS bucket will be created for the purpose of this notebook. If you\n",
-        "# prefer using your own GCS bucket, change the value yourself below.\n",
-        "now = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
-        "BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "\n",
-        "if BUCKET_URI is None or BUCKET_URI.strip() == \"\" or BUCKET_URI == \"gs://\":\n",
-        "    BUCKET_URI = f\"gs://{PROJECT_ID}-tmp-{now}-{str(uuid.uuid4())[:4]}\"\n",
-        "    BUCKET_NAME = \"/\".join(BUCKET_URI.split(\"/\")[:3])\n",
-        "    ! gsutil mb -l {REGION} {BUCKET_URI}\n",
-        "else:\n",
-        "    assert BUCKET_URI.startswith(\"gs://\"), \"BUCKET_URI must start with `gs://`.\"\n",
-        "    shell_output = ! gsutil ls -Lb {BUCKET_NAME} | grep \"Location constraint:\" | sed \"s/Location constraint://\"\n",
-        "    bucket_region = shell_output[0].strip().lower()\n",
-        "    if bucket_region != REGION:\n",
-        "        raise ValueError(\n",
-        "            \"Bucket region %s is different from notebook region %s\"\n",
-        "            % (bucket_region, REGION)\n",
-        "        )\n",
-        "print(f\"Using this GCS Bucket: {BUCKET_URI}\")\n",
-        "\n",
-        "STAGING_BUCKET = os.path.join(BUCKET_URI, \"temporal\")\n",
-        "MODEL_BUCKET = os.path.join(BUCKET_URI, \"llama3\")\n",
-        "\n",
-        "\n",
         "# Initialize Vertex AI API.\n",
         "print(\"Initializing Vertex AI API.\")\n",
-        "aiplatform.init(project=PROJECT_ID, location=REGION, staging_bucket=STAGING_BUCKET)\n",
-        "\n",
-        "# Gets the default SERVICE_ACCOUNT.\n",
-        "shell_output = ! gcloud projects describe $PROJECT_ID\n",
-        "project_number = shell_output[-1].split(\":\")[1].strip().replace(\"'\", \"\")\n",
-        "SERVICE_ACCOUNT = f\"{project_number}-compute@developer.gserviceaccount.com\"\n",
-        "print(\"Using this default Service Account:\", SERVICE_ACCOUNT)\n",
-        "\n",
-        "\n",
-        "# Provision permissions to the SERVICE_ACCOUNT with the GCS bucket\n",
-        "! gsutil iam ch serviceAccount:{SERVICE_ACCOUNT}:roles/storage.admin $BUCKET_NAME\n",
+        "aiplatform.init(project=PROJECT_ID, location=REGION)\n",
         "\n",
         "! gcloud config set project $PROJECT_ID\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/storage.admin\"\n",
-        "! gcloud projects add-iam-policy-binding --no-user-output-enabled {PROJECT_ID} --member=serviceAccount:{SERVICE_ACCOUNT} --role=\"roles/aiplatform.user\"\n",
+        "import vertexai\n",
+        "\n",
+        "vertexai.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "vGnNkMt96bvm"
+      },
+      "outputs": [],
+      "source": [
+        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
+        "\n",
+        "# @markdown Llama 3 uses a context length of 8,192 tokens, double the context length of Llama 2. See this [Meta blog post](https://ai.meta.com/blog/meta-llama-3/) for more details.\n",
+        "\n",
+        "# @markdown Allowing for predictions with LoRA weights stored on GCS or Hugging Face is enabled. To enable serving more LoRAs in a single batch, additional GPU will be required.\n",
+        "\n",
+        "# @markdown Set the model to deploy.\n",
+        "\n",
+        "base_model_name = \"llama3-8b-chat-hf\"  # @param [\"llama3-8b-hf\", \"llama3-8b-chat-hf\", \"llama3-70b-hf\", \"llama3-70b-chat-hf\"] {isTemplate:true}\n",
+        "if base_model_name == \"llama3-8b-hf\":\n",
+        "    hf_model_id = \"meta-llama/Meta-Llama-3-8B\"\n",
+        "    version_id = \"meta-llama-3-8b\"\n",
+        "elif base_model_name == \"llama3-8b-chat-hf\":\n",
+        "    hf_model_id = \"meta-llama/Meta-Llama-3-8B-Instruct\"\n",
+        "    version_id = \"meta-llama-3-8b-instruct\"\n",
+        "elif base_model_name == \"llama3-70b-hf\":\n",
+        "    hf_model_id = \"meta-llama/Meta-Llama-3-70B\"\n",
+        "    version_id = \"meta-llama-3-70b\"\n",
+        "elif base_model_name == \"llama3-70b-chat-hf\":\n",
+        "    hf_model_id = \"meta-llama/Meta-Llama-3-70B-Instruct\"\n",
+        "    version_id = \"meta-llama-3-70b-instruct\"\n",
+        "else:\n",
+        "    raise ValueError(f\"Unsupported base model name: {base_model_name}\")\n",
+        "PUBLISHER_MODEL_NAME = f\"publishers/meta/models/llama3@{version_id}\"\n",
+        "\n",
+        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
+        "\n",
+        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_H100_80GB\"]\n",
+        "max_loras = 1\n",
+        "enforce_eager = False\n",
+        "enable_lora = True\n",
+        "\n",
+        "if \"8b\" in base_model_name:\n",
+        "    enforce_eager = False\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        # L4 serving is more cost efficient than V100 serving.\n",
+        "        machine_type = \"g2-standard-12\"\n",
+        "        accelerator_count = 1\n",
+        "        max_loras = 5\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-1g\"\n",
+        "        accelerator_count = 1\n",
+        "        max_loras = 100\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "elif \"70b\" in base_model_name:\n",
+        "    # If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA3 70B\n",
+        "    # models with 8 L4 (24G) GPUs.\n",
+        "    # Note that with the default timeout threshold of Vertex endpoints, you should\n",
+        "    # set a `max_tokens` configuration of around 1,000 tokens or fewer. If you need\n",
+        "    # longer generated sequences, file a request with Vertex to allowlist\n",
+        "    # your project for a longer timeout threshold with Vertex endpoints.\n",
+        "\n",
+        "    if accelerator_type == \"NVIDIA_L4\":\n",
+        "        machine_type = \"g2-standard-96\"\n",
+        "        accelerator_count = 8\n",
+        "        max_loras = 1\n",
+        "        enforce_eager = True\n",
+        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
+        "        machine_type = \"a2-highgpu-4g\"\n",
+        "        accelerator_count = 4\n",
+        "        max_loras = 45\n",
+        "    elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
+        "        machine_type = \"a3-highgpu-4g\"\n",
+        "        accelerator_count = 4\n",
+        "        max_loras = 45\n",
+        "    else:\n",
+        "        raise ValueError(\n",
+        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "        )\n",
+        "else:\n",
+        "    raise ValueError(\n",
+        "        f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
+        "    )\n",
+        "common_util.check_quota(\n",
+        "    project_id=PROJECT_ID,\n",
+        "    region=REGION,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    is_for_training=False,\n",
+        ")\n",
+        "\n",
+        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint). Note that [dedicated endpoint does not support VPC Service Controls](https://cloud.google.com/vertex-ai/docs/predictions/choose-endpoint-type), uncheck the box if you are using VPC-SC.\n",
+        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "z-XybZjtgF9M"
+      },
+      "source": [
+        "## Deploy prebuilt LLaMA3 models on vLLM"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ZMQ8ElQ-6bvm"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 1] Deploy with Model Garden SDK\n",
+        "\n",
+        "# @markdown Deploy with Gen AI model-centric SDK. This section uploads the prebuilt model to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model. See [use open models with Vertex AI](https://cloud.google.com/vertex-ai/generative-ai/docs/open-models/use-open-models) for documentation on other use cases.\n",
+        "from vertexai.preview import model_garden\n",
+        "\n",
+        "model = model_garden.OpenModel(PUBLISHER_MODEL_NAME)\n",
+        "endpoints[LABEL] = model.deploy(\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    use_dedicated_endpoint=use_dedicated_endpoint,\n",
+        "    accept_eula=True,  # Accept the End User License Agreement (EULA) on the model card before deploy. Otherwise, the deployment will be forbidden.\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "E8OiHHNNE_wj"
+      },
+      "outputs": [],
+      "source": [
+        "# @title [Option 2] Deploy with customized configs\n",
+        "\n",
+        "# @markdown This section uploads prebuilt LLaMA3 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
         "\n",
         "# @markdown # Access LLaMA3 models on Vertex AI for GPU based serving\n",
         "# @markdown The original models from Meta are converted into the Hugging Face format for serving in Vertex AI.\n",
@@ -193,41 +306,31 @@
         "# @markdown 2. Review and accept the agreement in the pop-up window on the model card page. If you have previously accepted the model agreement, there will not be a pop-up window on the model card page and this step is not needed.\n",
         "# @markdown 3. After accepting the agreement of LLaMA3, a `gs://` URI containing LLaMA3 pretrained and finetuned models will be shared.\n",
         "# @markdown 4. Paste the URI in the `VERTEX_AI_MODEL_GARDEN_LLAMA3` field below.\n",
-        "# @markdown 5. The LLaMA3 models will be copied into `BUCKET_URI`.\n",
-        "\n",
         "\n",
         "VERTEX_AI_MODEL_GARDEN_LLAMA3 = \"\"  # @param {type:\"string\", isTemplate:true}\n",
         "assert (\n",
         "    VERTEX_AI_MODEL_GARDEN_LLAMA3\n",
         "), \"Click the agreement of LLaMA3 in Vertex AI Model Garden, and get the GCS path of LLaMA3 model artifacts.\"\n",
-        "parsed_gcs_url = re.search(\"gs://.*?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA3)\n",
+        "parsed_gcs_url = re.search(\"gs://.*?(/)?(?=[ ]|$)\", VERTEX_AI_MODEL_GARDEN_LLAMA3)\n",
         "if parsed_gcs_url:\n",
         "    VERTEX_AI_MODEL_GARDEN_LLAMA3 = parsed_gcs_url.group()\n",
         "assert VERTEX_AI_MODEL_GARDEN_LLAMA3.startswith(\n",
         "    \"gs://\"\n",
         "), \"VERTEX_AI_MODEL_GARDEN_LLAMA3 is expected to be a GCS URI and must start with `gs://`.\"\n",
-        "print(\n",
-        "    \"Copying LLaMA3 model artifacts from\",\n",
-        "    VERTEX_AI_MODEL_GARDEN_LLAMA3,\n",
-        "    \"to \",\n",
-        "    MODEL_BUCKET,\n",
-        ")\n",
-        "\n",
-        "! gsutil -m cp -R $VERTEX_AI_MODEL_GARDEN_LLAMA3/* $MODEL_BUCKET\n",
+        "assert VERTEX_AI_MODEL_GARDEN_LLAMA3.endswith(\n",
+        "    \"/\"\n",
+        "), \"VERTEX_AI_MODEL_GARDEN_LLAMA3 must end with `/`.\"\n",
+        "model_id = os.path.join(VERTEX_AI_MODEL_GARDEN_LLAMA3, base_model_name)\n",
         "\n",
         "# The pre-built serving docker images.\n",
         "VLLM_DOCKER_URI = \"us-docker.pkg.dev/vertex-ai/vertex-vision-model-garden-dockers/pytorch-vllm-serve:20240721_0916_RC00\"\n",
         "\n",
-        "# @markdown Set use_dedicated_endpoint to False if you don't want to use [dedicated endpoint](https://cloud.google.com/vertex-ai/docs/general/deployment#create-dedicated-endpoint).\n",
-        "use_dedicated_endpoint = True  # @param {type:\"boolean\"}\n",
-        "\n",
         "\n",
         "def deploy_model_vllm(\n",
         "    model_name: str,\n",
         "    model_id: str,\n",
         "    publisher: str,\n",
         "    publisher_model_id: str,\n",
-        "    service_account: str,\n",
         "    base_model_id: str = None,\n",
         "    machine_type: str = \"g2-standard-8\",\n",
         "    accelerator_type: str = \"NVIDIA_L4\",\n",
@@ -246,6 +349,7 @@
         "    use_dedicated_endpoint: bool = False,\n",
         "    max_num_seqs: int = 256,\n",
         "    model_type: str = None,\n",
+        "    enable_llama_tool_parser: bool = False,\n",
         ") -> Tuple[aiplatform.Model, aiplatform.Endpoint]:\n",
         "    \"\"\"Deploys trained models with vLLM into Vertex AI.\"\"\"\n",
         "    endpoint = aiplatform.Endpoint.create(\n",
@@ -298,6 +402,10 @@
         "    if model_type:\n",
         "        vllm_args.append(f\"--model-type={model_type}\")\n",
         "\n",
+        "    if enable_llama_tool_parser:\n",
+        "        vllm_args.append(\"--enable-auto-tool-choice\")\n",
+        "        vllm_args.append(\"--tool-call-parser=vertex-llama-3\")\n",
+        "\n",
         "    env_vars = {\n",
         "        \"MODEL_ID\": base_model_id,\n",
         "        \"DEPLOY_SOURCE\": \"notebook\",\n",
@@ -333,117 +441,15 @@
         "        accelerator_type=accelerator_type,\n",
         "        accelerator_count=accelerator_count,\n",
         "        deploy_request_timeout=1800,\n",
-        "        service_account=service_account,\n",
         "        system_labels={\n",
         "            \"NOTEBOOK_NAME\": \"model_garden_pytorch_llama3_deployment.ipynb\",\n",
+        "            \"NOTEBOOK_ENVIRONMENT\": common_util.get_deploy_source(),\n",
         "        },\n",
         "    )\n",
         "    print(\"endpoint_name:\", endpoint.name)\n",
         "\n",
-        "    return model, endpoint"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "z-XybZjtgF9M"
-      },
-      "source": [
-        "## Deploy prebuilt LLaMA3 models on vLLM"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "cellView": "form",
-        "id": "E8OiHHNNE_wj"
-      },
-      "outputs": [],
-      "source": [
-        "# @title Deploy\n",
+        "    return model, endpoint\n",
         "\n",
-        "# @markdown This section uploads prebuilt LLaMA3 models to Model Registry and deploys it to a Vertex AI Endpoint. It takes 15 minutes to 1 hour to finish depending on the size of the model.\n",
-        "\n",
-        "# @markdown NVIDIA_L4 GPUs are used for demonstration. The serving efficiency of L4 GPUs is inferior to that of A100 GPUs, but L4 GPUs are nevertheless good serving solutions if you do not have A100 quota.\n",
-        "\n",
-        "# @markdown Llama 3 uses a context length of 8,192 tokens, double the context length of Llama 2. See this [Meta blog post](https://ai.meta.com/blog/meta-llama-3/) for more details.\n",
-        "\n",
-        "# @markdown Allowing for predictions with LoRA weights stored on GCS or Hugging Face is enabled. To enable serving more LoRAs in a single batch, additional GPU will be required.\n",
-        "\n",
-        "# @markdown Set the model to deploy.\n",
-        "\n",
-        "base_model_name = \"llama3-8b-chat-hf\"  # @param [\"llama3-8b-hf\", \"llama3-8b-chat-hf\", \"llama3-70b-hf\", \"llama3-70b-chat-hf\"] {isTemplate:true}\n",
-        "model_id = os.path.join(MODEL_BUCKET, base_model_name)\n",
-        "if base_model_name == \"llama3-8b-hf\":\n",
-        "    hf_model_id = \"meta-llama/Meta-Llama-3-8B\"\n",
-        "elif base_model_name == \"llama3-8b-chat-hf\":\n",
-        "    hf_model_id = \"meta-llama/Meta-Llama-3-8B-Instruct\"\n",
-        "elif base_model_name == \"llama3-70b-hf\":\n",
-        "    hf_model_id = \"meta-llama/Meta-Llama-3-70B\"\n",
-        "elif base_model_name == \"llama3-70b-chat-hf\":\n",
-        "    hf_model_id = \"meta-llama/Meta-Llama-3-70B-Instruct\"\n",
-        "else:\n",
-        "    raise ValueError(f\"Unsupported base model name: {base_model_name}\")\n",
-        "\n",
-        "# @markdown Find Vertex AI prediction supported accelerators and regions at https://cloud.google.com/vertex-ai/docs/predictions/configure-compute.\n",
-        "\n",
-        "accelerator_type = \"NVIDIA_L4\"  # @param [\"NVIDIA_L4\", \"NVIDIA_TESLA_A100\", \"NVIDIA_H100_80GB\"]\n",
-        "max_loras = 1\n",
-        "enforce_eager = False\n",
-        "enable_lora = True\n",
-        "\n",
-        "if \"8b\" in base_model_name:\n",
-        "    enforce_eager = False\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
-        "        # L4 serving is more cost efficient than V100 serving.\n",
-        "        machine_type = \"g2-standard-8\"\n",
-        "        accelerator_count = 1\n",
-        "        max_loras = 5\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "        machine_type = \"a2-highgpu-1g\"\n",
-        "        accelerator_count = 1\n",
-        "        max_loras = 100\n",
-        "    else:\n",
-        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
-        "        )\n",
-        "elif \"70b\" in base_model_name:\n",
-        "    # If you do not have access to 4 A100 (40G) GPUs, you may serve LLaMA3 70B\n",
-        "    # models with 8 L4 (24G) GPUs.\n",
-        "    # Note that with the default timeout threshold of Vertex endpoints, you should\n",
-        "    # set a `max_tokens` configuration of around 1,000 tokens or fewer. If you need\n",
-        "    # longer generated sequences, file a request with Vertex to allowlist\n",
-        "    # your project for a longer timeout threshold with Vertex endpoints.\n",
-        "\n",
-        "    if accelerator_type == \"NVIDIA_L4\":\n",
-        "        machine_type = \"g2-standard-96\"\n",
-        "        accelerator_count = 8\n",
-        "        max_loras = 1\n",
-        "        enforce_eager = True\n",
-        "    elif accelerator_type == \"NVIDIA_TESLA_A100\":\n",
-        "        machine_type = \"a2-highgpu-4g\"\n",
-        "        accelerator_count = 4\n",
-        "        max_loras = 45\n",
-        "    elif accelerator_type == \"NVIDIA_H100_80GB\":\n",
-        "        machine_type = \"a3-highgpu-4g\"\n",
-        "        accelerator_count = 4\n",
-        "        max_loras = 45\n",
-        "    else:\n",
-        "        raise ValueError(\n",
-        "            f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
-        "        )\n",
-        "else:\n",
-        "    raise ValueError(\n",
-        "        f\"Recommended GPU setting not found for: {accelerator_type} and {base_model_name}.\"\n",
-        "    )\n",
-        "common_util.check_quota(\n",
-        "    project_id=PROJECT_ID,\n",
-        "    region=REGION,\n",
-        "    accelerator_type=accelerator_type,\n",
-        "    accelerator_count=accelerator_count,\n",
-        "    is_for_training=False,\n",
-        ")\n",
         "\n",
         "gpu_memory_utilization = 0.9\n",
         "max_model_len = 8192  # Maximum context length.\n",
@@ -458,7 +464,6 @@
         "    publisher=\"meta\",\n",
         "    publisher_model_id=\"llama3\",\n",
         "    base_model_id=hf_model_id,\n",
-        "    service_account=SERVICE_ACCOUNT,\n",
         "    machine_type=machine_type,\n",
         "    accelerator_type=accelerator_type,\n",
         "    accelerator_count=accelerator_count,\n",
@@ -533,7 +538,7 @@
       "execution_count": null,
       "metadata": {
         "cellView": "form",
-        "id": "7d9bbf86da5e"
+        "id": "KyxwWZ1XDtMJ"
       },
       "outputs": [],
       "source": [
@@ -557,6 +562,7 @@
         "# @markdown If you encounter the issue like `ServiceUnavailable: 503 Took too long to respond when processing`, you can reduce the maximum number of output tokens, such as set `max_tokens` as 20.\n",
         "max_tokens = 50  # @param {type: \"integer\"}\n",
         "temperature = 1.0  # @param {type: \"number\"}\n",
+        "stream = False  # @param {type: \"boolean\"}\n",
         "\n",
         "# @markdown Now we can send a request.\n",
         "\n",
@@ -583,8 +589,21 @@
         "    messages=[{\"role\": \"user\", \"content\": user_message}],\n",
         "    temperature=temperature,\n",
         "    max_tokens=max_tokens,\n",
+        "    stream=stream,\n",
         ")\n",
-        "print(model_response)\n",
+        "\n",
+        "if stream:\n",
+        "    usage = None\n",
+        "    contents = []\n",
+        "    for chunk in model_response:\n",
+        "        if chunk.usage is not None:\n",
+        "            usage = chunk.usage\n",
+        "            continue\n",
+        "        print(chunk.choices[0].delta.content, end=\"\")\n",
+        "        contents.append(chunk.choices[0].delta.content)\n",
+        "    print(f\"\\n\\n{usage}\")\n",
+        "else:\n",
+        "    print(model_response)\n",
         "\n",
         "# @markdown Click \"Show Code\" to see more details."
       ]
@@ -617,11 +636,7 @@
         "\n",
         "# Delete models.\n",
         "for model in models.values():\n",
-        "    model.delete()\n",
-        "\n",
-        "delete_bucket = False  # @param {type:\"boolean\"}\n",
-        "if delete_bucket:\n",
-        "    ! gsutil -m rm -r $BUCKET_NAME"
+        "    model.delete()"
       ]
     }
   ],