feat: Support accelerator_type in RLHF tuning

PiperOrigin-RevId: 617572810

Author: vertex-sdk-bot

tests/unit/aiplatform/test_language_models.py

@@ -723,6 +723,11 @@ def reverse_string_2(s):""",
                     "isOptional": True,
                     "parameterType": "NUMBER_INTEGER",
                 },
+                "accelerator_type": {
+                    "defaultValue": "",
+                    "isOptional": True,
+                    "parameterType": "STRING",
+                },
                 "tensorboard_resource_id": {
                     "isOptional": True,
                     "parameterType": "STRING",
@@ -2696,6 +2701,7 @@ def test_tune_text_generation_model_rlhf(
         kl_coeff = 0.3
         tensorboard_resource_id = _get_test_tensorboard_resource_id()
         eval_dataset = "gs://bucket/eval.jsonl"
+        accelerator_type = "TPU"
 
         with mock.patch.object(
             target=model_garden_service_client.ModelGardenServiceClient,
@@ -2718,6 +2724,7 @@ def test_tune_text_generation_model_rlhf(
                 reward_model_train_steps=reward_model_train_steps,
                 reinforcement_learning_train_steps=reinforcement_learning_train_steps,
                 kl_coeff=kl_coeff,
+                accelerator_type=accelerator_type,
                 tuning_evaluation_spec=preview_language_models.TuningEvaluationSpec(
                     tensorboard=tensorboard_resource_id,
                     evaluation_data=eval_dataset,
@@ -2756,6 +2763,7 @@ def test_tune_text_generation_model_rlhf(
                 pipeline_arguments["tensorboard_resource_id"] == tensorboard_resource_id
             )
             assert pipeline_arguments["eval_dataset"] == eval_dataset
+            assert pipeline_arguments["accelerator_type"] == "TPU"
 
     @pytest.mark.parametrize(
         "job_spec",

vertexai/language_models/_language_models.py

@@ -413,6 +413,7 @@ class _RlhfTuningParameters:
     deploy_model: Optional[bool] = None
     eval_dataset: Optional[str] = None
     project: Optional[str] = None
+    accelerator_type: Optional[_ACCELERATOR_TYPE_TYPE] = None
     tensorboard_resource_id: Optional[str] = None
 
     def asdict(self) -> Dict[str, Any]:
@@ -439,6 +440,7 @@ def tune_model_rlhf(
         kl_coeff: Optional[float] = None,
         default_context: Optional[str] = None,
         tuning_job_location: Optional[str] = None,
+        accelerator_type: Optional[_ACCELERATOR_TYPE_TYPE] = None,
         tuning_evaluation_spec: Optional["TuningEvaluationSpec"] = None,
     ) -> "_LanguageModelTuningJob":
         """Tunes a model using reinforcement learning from human feedback.
@@ -491,6 +493,7 @@ def tune_model_rlhf(
                 negative" or "Translate this sentence to Danish". Do not specify this
                 if your dataset already prepends the instruction to the inputs field.
             tuning_job_location: GCP location where the tuning job should be run.
+            accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
             tuning_evaluation_spec: Evaluation settings to use during tuning.
 
         Returns:
@@ -527,6 +530,13 @@ def tune_model_rlhf(
             model_id=self._model_id,
         )
 
+        if accelerator_type:
+            if accelerator_type not in _ACCELERATOR_TYPES:
+                raise ValueError(
+                    f"Unsupported accelerator type: {accelerator_type}."
+                    f" Supported types: {_ACCELERATOR_TYPES}"
+                )
+
         tuning_parameters = _RlhfTuningParameters(
             prompt_dataset=prompt_dataset_uri,
             preference_dataset=preference_dataset_uri,
@@ -542,6 +552,7 @@ def tune_model_rlhf(
             kl_coeff=kl_coeff,
             instruction=default_context,
             eval_dataset=eval_dataset,
+            accelerator_type=accelerator_type,
             tensorboard_resource_id=tensorboard_resource_id,
         )
 
@@ -574,7 +585,7 @@ def _tune_model_rlhf(
             raise ValueError(
                 _get_invalid_tuning_location_msg(
                     requested_location=tuning_parameters.location,
-                    valid_locations=_SUPPORTED_RLHF_LOCATIONS,
+                    valid_locations=_TUNING_LOCATIONS,
                 )
             )
         if self._model_id not in _SUPPORTED_RLHF_MODELS:
@@ -3433,13 +3444,6 @@ class _PreviewCodeGenerationModel(CodeGenerationModel, _CountTokensCodeGeneratio
 # Currently, deployment can only work in these locations
 _TUNED_MODEL_LOCATIONS = _SUPPORTED_LOCATIONS
 
-# TODO(b/318874365): Use _SUPPORTED_LOCATIONS defined above once DRZ for RLHF is
-# implemented.
-_SUPPORTED_RLHF_LOCATIONS = {
-    "us-central1",
-    "europe-west4",
-}
-
 # All models supported by RLHF that can also be used for online and batch prediction:
 _SUPPORTED_RLHF_MODELS = {
     "text-bison@001",