feat: Enable Vertex Model Garden Managed OSS Fine Tuning.

PiperOrigin-RevId: 834357700

Author: vertex-sdk-bot

tests/unit/vertexai/tuning/test_tuning.py

@@ -22,10 +22,12 @@
 import importlib
 from typing import Dict, Iterable
 from unittest import mock
+from unittest.mock import patch
 import uuid
 
 from google import auth
 from google.auth import credentials as auth_credentials
+from google.cloud import storage
 from google.cloud import aiplatform
 import vertexai
 from google.cloud.aiplatform import compat
@@ -34,26 +36,16 @@
 from google.cloud.aiplatform.metadata import experiment_resources
 from google.cloud.aiplatform_v1beta1.services import gen_ai_tuning_service
 from google.cloud.aiplatform_v1beta1.types import job_state
-from google.cloud.aiplatform_v1beta1.types import (
-    tuning_job as gca_tuning_job,
-)
+from google.cloud.aiplatform_v1beta1.types import tuning_job as gca_tuning_job
 from vertexai.preview import tuning
-from vertexai.preview.tuning import (
-    sft as preview_supervised_tuning,
-)
-from vertexai.preview.tuning._tuning import SourceModel
+from vertexai.preview.tuning import sft as preview_supervised_tuning
+from vertexai.preview.tuning._tuning import SourceModel as PreviewSourceModel
+from vertexai.preview.tuning._tuning import TuningJob as PreviewTuningJob
 from vertexai.tuning import _distillation
 from vertexai.tuning import sft as supervised_tuning
-from google.cloud import storage
-from vertexai.preview.tuning._tuning import (
-    TuningJob as PreviewTuningJob,
-)
-
-
+from vertexai.tuning._tuning import SourceModel
 import pytest
 
-from unittest.mock import patch
-
 from google.rpc import status_pb2
 
 
@@ -191,18 +183,18 @@ def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)
 
     @mock.patch.object(
-        target=PreviewTuningJob,
+        target=tuning.TuningJob,
         attribute="client_class",
         new=MockTuningJobClientWithOverride,
     )
     @pytest.mark.parametrize(
         "supervised_tuning",
-        [preview_supervised_tuning],
+        [preview_supervised_tuning, supervised_tuning],
     )
     def test_genai_tuning_service_supervised_tuning_tune_model(
         self, supervised_tuning: supervised_tuning
     ):
-        sft_tuning_job = supervised_tuning.preview_train(
+        sft_tuning_job = supervised_tuning.train(
             source_model="gemini-1.0-pro-001",
             train_dataset="gs://some-bucket/some_dataset.jsonl",
             # Optional:
@@ -237,42 +229,42 @@ def test_genai_tuning_service_supervised_tuning_tune_model(
         assert sft_tuning_job.tuned_model_endpoint_name
 
     @mock.patch.object(
-        target=PreviewTuningJob,
+        target=tuning.TuningJob,
         attribute="client_class",
         new=MockTuningJobClientWithOverride,
     )
     @pytest.mark.parametrize(
         "supervised_tuning",
-        [preview_supervised_tuning],
+        [supervised_tuning],
     )
     def test_genai_tuning_service_encryption_spec(
         self, supervised_tuning: supervised_tuning
     ):
         """Test that the global encryption spec propagates to the tuning job."""
         vertexai.init(encryption_spec_key_name="test-key")
 
-        sft_tuning_job = supervised_tuning.preview_train(
+        sft_tuning_job = supervised_tuning.train(
             source_model="gemini-1.0-pro-001",
             train_dataset="gs://some-bucket/some_dataset.jsonl",
         )
         assert sft_tuning_job.encryption_spec.kms_key_name == "test-key"
 
     @mock.patch.object(
-        target=PreviewTuningJob,
+        target=tuning.TuningJob,
         attribute="client_class",
         new=MockTuningJobClientWithOverride,
     )
     @pytest.mark.parametrize(
         "supervised_tuning",
-        [preview_supervised_tuning],
+        [supervised_tuning],
     )
     def test_genai_tuning_service_service_account(
         self, supervised_tuning: supervised_tuning
     ):
         """Test that the service account propagates to the tuning job."""
         vertexai.init(service_account="[email protected]")
 
-        sft_tuning_job = supervised_tuning.preview_train(
+        sft_tuning_job = supervised_tuning.train(
             source_model="gemini-1.0-pro-002",
             train_dataset="gs://some-bucket/some_dataset.jsonl",
         )
@@ -331,19 +323,35 @@ def test_genai_tuning_service_distillation_distill_model(self):
         attribute="client_class",
         new=MockTuningJobClientWithOverride,
     )
+    @mock.patch.object(
+        target=tuning.TuningJob,
+        attribute="client_class",
+        new=MockTuningJobClientWithOverride,
+    )
     @pytest.mark.parametrize(
-        "supervised_tuning",
-        [preview_supervised_tuning],
+        "sft_train_method, source_model",
+        [
+            (
+                preview_supervised_tuning.preview_train,
+                PreviewSourceModel(
+                    base_model="meta/[email protected]",
+                    custom_base_model="gs://test-bucket/custom-weights",
+                ),
+            ),
+            (
+                supervised_tuning.train,
+                SourceModel(
+                    base_model="meta/[email protected]",
+                    custom_base_model="gs://test-bucket/custom-weights",
+                ),
+            ),
+        ],
     )
     def test_create_tuning_job_success(
-        self, supervised_tuning: preview_supervised_tuning
+        self, sft_train_method: supervised_tuning.train, source_model: SourceModel
     ):
-        model = SourceModel(
-            base_model="meta/[email protected]",
-            custom_base_model="gs://test-bucket/custom-weights",
-        )
-        sft_tuning_job = supervised_tuning.preview_train(
-            source_model=model,
+        sft_tuning_job = sft_train_method(
+            source_model=source_model,
             epochs=1,
             train_dataset="gs://test-bucket/test_train_dataset/",
             validation_dataset="gs://test-bucket/test_validation_dataset/",

vertexai/tuning/__init__.py

@@ -16,8 +16,10 @@
 
 # We just want to re-export certain classes
 # pylint: disable=g-multiple-import,g-importing-member
+from vertexai.tuning._tuning import SourceModel
 from vertexai.tuning._tuning import TuningJob
 
 __all__ = [
+    "SourceModel",
     "TuningJob",
 ]

vertexai/tuning/_supervised_tuning.py

@@ -21,19 +21,25 @@
     tuning_job as gca_tuning_job_types,
 )
 from vertexai import generative_models
-from vertexai.tuning import _tuning
+from vertexai.tuning import (
+    SourceModel,
+    TuningJob,
+)
 
 
 def train(
     *,
-    source_model: Union[str, generative_models.GenerativeModel],
+    source_model: Union[str, generative_models.GenerativeModel, SourceModel],
     train_dataset: Union[str, datasets.MultimodalDataset],
     validation_dataset: Optional[Union[str, datasets.MultimodalDataset]] = None,
     tuned_model_display_name: Optional[str] = None,
+    tuning_mode: Optional[Literal["FULL", "PEFT_ADAPTER"]] = None,
     epochs: Optional[int] = None,
+    learning_rate: Optional[float] = None,
     learning_rate_multiplier: Optional[float] = None,
     adapter_size: Optional[Literal[1, 4, 8, 16, 32]] = None,
     labels: Optional[Dict[str, str]] = None,
+    output_uri: Optional[str] = None,
 ) -> "SupervisedTuningJob":
     """Tunes a model using supervised training.
 
@@ -44,14 +50,41 @@ def train(
         tuned_model_display_name: The display name of the
           [TunedModel][google.cloud.aiplatform.v1.Model]. The name can be up to
           128 characters long and can consist of any UTF-8 characters.
+        tuning_mode: Tuning mode for this tuning job. Can only be used with OSS
+          models.
         epochs: Number of training epoches for this tuning job.
-        learning_rate_multiplier: Learning rate multiplier for tuning.
+        learning_rate: Learning rate for tuning. Can only be used with OSS
+          models. Mutually exclusive with `learning_rate_multiplier`.
+        learning_rate_multiplier: Learning rate multiplier for tuning. Can only
+          be used with 1P models. Mutually exclusive with `learning_rate`.
         adapter_size: Adapter size for tuning.
         labels: User-defined metadata to be associated with trained models
+        output_uri: The Google Cloud Storage URI to write the tuned model to.
+          Can only be used with OSS models.
 
     Returns:
         A `TuningJob` object.
     """
+    if tuning_mode is None:
+        tuning_mode_value = None
+    elif tuning_mode == "FULL":
+        tuning_mode_value = (
+            gca_tuning_job_types.SupervisedTuningSpec.TuningMode.TUNING_MODE_FULL
+        )
+    elif tuning_mode == "PEFT_ADAPTER":
+        tuning_mode_value = (
+            gca_tuning_job_types.SupervisedTuningSpec.TuningMode.TUNING_MODE_PEFT_ADAPTER
+        )
+    else:
+        raise ValueError(
+            f"Unsupported tuning mode: {tuning_mode}. The supported tuning modes are [FULL, PEFT_ADAPTER]"
+        )
+
+    if learning_rate and learning_rate_multiplier:
+        raise ValueError(
+            "Only one of `learning_rate` and `learning_rate_multiplier` can be set."
+        )
+
     if adapter_size is None:
         adapter_size_value = None
     elif adapter_size == 1:
@@ -83,10 +116,12 @@ def train(
     if isinstance(validation_dataset, datasets.MultimodalDataset):
         validation_dataset = validation_dataset.resource_name
     supervised_tuning_spec = gca_tuning_job_types.SupervisedTuningSpec(
+        tuning_mode=tuning_mode_value,
         training_dataset_uri=train_dataset,
         validation_dataset_uri=validation_dataset,
         hyper_parameters=gca_tuning_job_types.SupervisedHyperParameters(
             epoch_count=epochs,
+            learning_rate=learning_rate,
             learning_rate_multiplier=learning_rate_multiplier,
             adapter_size=adapter_size_value,
         ),
@@ -95,20 +130,26 @@ def train(
     if isinstance(source_model, generative_models.GenerativeModel):
         source_model = source_model._prediction_resource_name.rpartition("/")[-1]
 
+    if labels is None:
+        labels = {}
+    if "mg-source" not in labels and output_uri:
+        labels["mg-source"] = "sdk"
+
     supervised_tuning_job = (
         SupervisedTuningJob._create(  # pylint: disable=protected-access
             base_model=source_model,
             tuning_spec=supervised_tuning_spec,
             tuned_model_display_name=tuned_model_display_name,
             labels=labels,
+            output_uri=output_uri,
         )
     )
     _ipython_utils.display_model_tuning_button(supervised_tuning_job)
 
     return supervised_tuning_job
 
 
-class SupervisedTuningJob(_tuning.TuningJob):
+class SupervisedTuningJob(TuningJob):
     def __init__(self, tuning_job_name: str):
         super().__init__(tuning_job_name=tuning_job_name)
         _ipython_utils.display_model_tuning_button(self)

vertexai/tuning/_tuning.py

@@ -43,6 +43,42 @@
 _LOGGER = aiplatform_base.Logger(__name__)
 
 
+class SourceModel:
+    r"""A model that is used in managed OSS supervised tuning.
+
+    Usage:
+        ```
+        model = SourceModel(
+            base_model="meta/[email protected]",
+            custom_base_model="gs://user-bucket/custom-weights",
+        )
+        sft_tuning_job = sft.train(
+            source_model=model,
+            train_dataset="gs://my-bucket/train.jsonl",
+            validation_dataset="gs://my-bucket/validation.jsonl",
+            epochs=4,
+            tuned_model_display_name="my-tuned-model",
+            output_uri="gs://user-bucket/tuned-model"
+        )
+
+        while not sft_tuning_job.has_ended:
+            time.sleep(60)
+            sft_tuning_job.refresh()
+
+        tuned_model = aiplatform.Model(sft_tuning_job.tuned_model_name)
+        ```
+    """
+
+    def __init__(
+        self,
+        base_model: str,
+        custom_base_model: str = "",
+    ):
+        r"""Initializes SourceModel."""
+        self.base_model = base_model
+        self.custom_base_model = custom_base_model
+
+
 class TuningJobClientWithOverride(aiplatform_utils.ClientWithOverride):
     _is_temporary = True
     _default_version = compat.V1BETA1
@@ -133,7 +169,7 @@ def tuning_data_statistics(self) -> gca_tuning_job_types.TuningDataStats:
     def _create(
         cls,
         *,
-        base_model: str,
+        base_model: Union[str, SourceModel],
         tuning_spec: Union[
             gca_tuning_job_types.SupervisedTuningSpec,
             gca_tuning_job_types.DistillationSpec,
@@ -144,15 +180,13 @@ def _create(
         project: Optional[str] = None,
         location: Optional[str] = None,
         credentials: Optional[auth_credentials.Credentials] = None,
+        output_uri: Optional[str] = None,
     ) -> "TuningJob":
         r"""Submits TuningJob.
 
         Args:
-            base_model (str):
-                Model name for tuning, e.g., "gemini-1.0-pro"
-                or "gemini-1.0-pro-001".
-
-                This field is a member of `oneof`_ ``source_model``.
+            base_model: Model for tuning.
+                Supported types: str, SourceModel.
             tuning_spec: Tuning Spec for Fine Tuning.
                 Supported types: SupervisedTuningSpec, DistillationSpec.
             tuned_model_display_name: The display name of the
@@ -179,6 +213,7 @@ def _create(
                 Overrides location set in aiplatform.init.
             credentials: Custom credentials to use to call tuning job service.
                 Overrides credentials set in aiplatform.init.
+            output_uri: The Google Cloud Storage location to write the artifacts. This is only used for OSS models.
 
         Returns:
             Submitted TuningJob.
@@ -192,17 +227,25 @@ def _create(
             tuned_model_display_name = cls._generate_display_name()
 
         gca_tuning_job = gca_tuning_job_types.TuningJob(
-            base_model=base_model,
             tuned_model_display_name=tuned_model_display_name,
             description=description,
             labels=labels,
-            # The tuning_spec one_of is set later
+            # The tuning_spec one_of is set later.
+            output_uri=output_uri,
         )
 
         if isinstance(tuning_spec, gca_tuning_job_types.SupervisedTuningSpec):
             gca_tuning_job.supervised_tuning_spec = tuning_spec
+            if isinstance(base_model, SourceModel):
+                gca_tuning_job.base_model = base_model.base_model
+                gca_tuning_job.custom_base_model = base_model.custom_base_model
+            else:
+                gca_tuning_job.base_model = base_model
         elif isinstance(tuning_spec, gca_tuning_job_types.DistillationSpec):
             gca_tuning_job.distillation_spec = tuning_spec
+            if isinstance(base_model, SourceModel):
+                raise RuntimeError("Distillation is not supported for custom models.")
+            gca_tuning_job.base_model = base_model
         else:
             raise RuntimeError(f"Unsupported tuning_spec kind: {tuning_spec}")