Merge pull request #5034 from BerriAI/litellm_vtx_ft_param_mapping

[Feat] Vertex AI fine tuning - support translating hyperparameters

litellm/fine_tuning/main.py

@@ -255,6 +255,7 @@ def create_fine_tuning_job(
                 vertex_location=vertex_ai_location,
                 timeout=timeout,
                 api_base=api_base,
+                kwargs=kwargs,
             )
         else:
             raise litellm.exceptions.BadRequestError(

litellm/llms/fine_tuning_apis/vertex_ai.py

@@ -86,6 +86,38 @@ class VertexFineTuningAPI(VertexLLM):
             integrations=[],
         )
 
+    def convert_openai_request_to_vertex(
+        self, create_fine_tuning_job_data: FineTuningJobCreate, **kwargs
+    ) -> FineTuneJobCreate:
+        """
+        convert request from OpenAI format to Vertex format
+        https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/tuning
+        supervised_tuning_spec = FineTunesupervisedTuningSpec(
+        """
+        hyperparameters = create_fine_tuning_job_data.hyperparameters
+        supervised_tuning_spec = FineTunesupervisedTuningSpec(
+            training_dataset_uri=create_fine_tuning_job_data.training_file,
+            validation_dataset=create_fine_tuning_job_data.validation_file,
+        )
+
+        if hyperparameters:
+            if hyperparameters.n_epochs:
+                supervised_tuning_spec["epoch_count"] = int(hyperparameters.n_epochs)
+            if hyperparameters.learning_rate_multiplier:
+                supervised_tuning_spec["learning_rate_multiplier"] = float(
+                    hyperparameters.learning_rate_multiplier
+                )
+
+        supervised_tuning_spec["adapter_size"] = kwargs.get("adapter_size")
+
+        fine_tune_job = FineTuneJobCreate(
+            baseModel=create_fine_tuning_job_data.model,
+            supervisedTuningSpec=supervised_tuning_spec,
+            tunedModelDisplayName=create_fine_tuning_job_data.suffix,
+        )
+
+        return fine_tune_job
+
     async def acreate_fine_tuning_job(
         self,
         fine_tuning_url: str,
@@ -144,6 +176,7 @@ class VertexFineTuningAPI(VertexLLM):
         vertex_credentials: Optional[str],
         api_base: Optional[str],
         timeout: Union[float, httpx.Timeout],
+        **kwargs,
     ):
 
         verbose_logger.debug(
@@ -166,12 +199,8 @@ class VertexFineTuningAPI(VertexLLM):
             "Content-Type": "application/json",
         }
 
-        supervised_tuning_spec = FineTunesupervisedTuningSpec(
-            training_dataset_uri=create_fine_tuning_job_data.training_file
-        )
-        fine_tune_job = FineTuneJobCreate(
-            baseModel=create_fine_tuning_job_data.model,
-            supervisedTuningSpec=supervised_tuning_spec,
+        fine_tune_job = self.convert_openai_request_to_vertex(
+            create_fine_tuning_job_data=create_fine_tuning_job_data, **kwargs
         )
 
         fine_tuning_url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/tuningJobs"

litellm/tests/test_fine_tuning_api.py

@@ -20,6 +20,12 @@ from test_gcs_bucket import load_vertex_ai_credentials
 
 from litellm import create_fine_tuning_job
 from litellm._logging import verbose_logger
+from litellm.llms.fine_tuning_apis.vertex_ai import (
+    FineTuningJobCreate,
+    VertexFineTuningAPI,
+)
+
+vertex_finetune_api = VertexFineTuningAPI()
 
 
 def test_create_fine_tune_job():
@@ -210,3 +216,58 @@ async def test_create_vertex_fine_tune_jobs():
         assert create_fine_tuning_response.object == "fine_tuning.job"
     except:
         pass
+
+
+# Testing OpenAI -> Vertex AI param mapping
+
+
+def test_convert_openai_request_to_vertex_basic():
+    openai_data = FineTuningJobCreate(
+        training_file="gs://bucket/train.jsonl",
+        validation_file="gs://bucket/val.jsonl",
+        model="text-davinci-002",
+        hyperparameters={"n_epochs": 3, "learning_rate_multiplier": 0.1},
+        suffix="my_fine_tuned_model",
+    )
+
+    result = vertex_finetune_api.convert_openai_request_to_vertex(openai_data)
+
+    print("converted vertex ai result=", result)
+
+    assert result["baseModel"] == "text-davinci-002"
+    assert result["tunedModelDisplayName"] == "my_fine_tuned_model"
+    assert (
+        result["supervisedTuningSpec"]["training_dataset_uri"]
+        == "gs://bucket/train.jsonl"
+    )
+    assert (
+        result["supervisedTuningSpec"]["validation_dataset"] == "gs://bucket/val.jsonl"
+    )
+    assert result["supervisedTuningSpec"]["epoch_count"] == 3
+    assert result["supervisedTuningSpec"]["learning_rate_multiplier"] == 0.1
+
+
+def test_convert_openai_request_to_vertex_with_adapter_size():
+    openai_data = FineTuningJobCreate(
+        training_file="gs://bucket/train.jsonl",
+        model="text-davinci-002",
+        hyperparameters={"n_epochs": 5, "learning_rate_multiplier": 0.2},
+        suffix="custom_model",
+    )
+
+    result = vertex_finetune_api.convert_openai_request_to_vertex(
+        openai_data, adapter_size="SMALL"
+    )
+
+    print("converted vertex ai result=", result)
+
+    assert result["baseModel"] == "text-davinci-002"
+    assert result["tunedModelDisplayName"] == "custom_model"
+    assert (
+        result["supervisedTuningSpec"]["training_dataset_uri"]
+        == "gs://bucket/train.jsonl"
+    )
+    assert result["supervisedTuningSpec"]["validation_dataset"] is None
+    assert result["supervisedTuningSpec"]["epoch_count"] == 5
+    assert result["supervisedTuningSpec"]["learning_rate_multiplier"] == 0.2
+    assert result["supervisedTuningSpec"]["adapter_size"] == "SMALL"

litellm/types/router.py

@@ -179,6 +179,7 @@ class GenericLiteLLMParams(BaseModel):
         ## VERTEX AI ##
         vertex_project: Optional[str] = None,
         vertex_location: Optional[str] = None,
+        vertex_credentials: Optional[str] = None,
         ## AWS BEDROCK / SAGEMAKER ##
         aws_access_key_id: Optional[str] = None,
         aws_secret_access_key: Optional[str] = None,