feat: Enable Vertex Model Garden Managed OSS Fine Tuning.

PiperOrigin-RevId: 815806517

Author: vertex-sdk-bot

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 OSS 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/llama3_1@llama-3.1-8b",
+            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}")