feat: LLM - Added support for the max_context_length tuning parameter

PiperOrigin-RevId: 615362109

Author: Ark-kun

tests/unit/aiplatform/test_language_models.py

@@ -606,6 +606,11 @@ def reverse_string_2(s):""",
                     "parameterType": "NUMBER_DOUBLE",
                 },
                 "location": {"parameterType": "STRING"},
+                "max_context_length": {
+                    "defaultValue": "",
+                    "isOptional": True,
+                    "parameterType": "STRING",
+                },
                 "model_display_name": {"parameterType": "STRING"},
                 "project": {"parameterType": "STRING"},
                 "tensorboard_resource_id": {
@@ -2271,6 +2276,7 @@ def test_tune_text_generation_model_ga(
                     tensorboard=tensorboard_name,
                 ),
                 accelerator_type="TPU",
+                max_context_length="32k",
             )
             call_kwargs = mock_pipeline_service_create.call_args[1]
             pipeline_arguments = call_kwargs[
@@ -2288,6 +2294,7 @@ def test_tune_text_generation_model_ga(
             assert pipeline_arguments["tensorboard_resource_id"] == tensorboard_name
             assert pipeline_arguments["large_model_reference"] == "text-bison@001"
             assert pipeline_arguments["accelerator_type"] == "TPU"
+            assert pipeline_arguments["max_context_length"] == "32k"
             assert (
                 call_kwargs["pipeline_job"].encryption_spec.kms_key_name
                 == _TEST_ENCRYPTION_KEY_NAME

vertexai/language_models/_distillation.py

@@ -21,6 +21,7 @@ def distill_from(
         evaluation_spec: Optional[tuning.TuningEvaluationSpec] = None,
         accelerator_type: Optional[tuning._ACCELERATOR_TYPE_TYPE] = None,
         model_display_name: Optional[str] = None,
+        max_context_length: Optional[str] = None,
     ):
         """Tunes a smaller model with help from another bigger model.
 
@@ -32,6 +33,8 @@ def distill_from(
             evaluation_spec: Specification for the model evaluation during tuning.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
             model_display_name: Custom display name for the tuned model.
+            max_context_length: The max context length used for tuning.
+                Can be either '8k' or '32k'
 
         Returns:
             A tuning job for distillation.
@@ -86,6 +89,8 @@ def distill_from(
             pipeline_arguments[
                 "encryption_spec_key_name"
             ] = aiplatform_initializer.global_config.encryption_spec_key_name
+        if max_context_length is not None:
+            pipeline_arguments["max_context_length"] = max_context_length
         if model_display_name is None:
             model_display_name = (
                 f"{student_short_model_id}"
@@ -94,7 +99,6 @@ def distill_from(
         pipeline_arguments["model_display_name"] = model_display_name
         # # Not exposing these parameters:
         # temperature: Optional[float] = None,
-        # max_context_length: Optional[int] = None,
         # tpu_training_skip_cmek: Optional[bool] = None,
         # api_endpoint: Optional[str] = None,
         # version: Optional[str] = None,

vertexai/language_models/_language_models.py

@@ -228,6 +228,7 @@ def tune_model(
         tuning_evaluation_spec: Optional["TuningEvaluationSpec"] = None,
         default_context: Optional[str] = None,
         accelerator_type: Optional[_ACCELERATOR_TYPE_TYPE] = None,
+        max_context_length: Optional[str] = None,
     ) -> "_LanguageModelTuningJob":
         """Tunes a model based on training data.
 
@@ -253,6 +254,8 @@ def tune_model(
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             default_context: The context to use for all training samples by default.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
+            max_context_length: The max context length used for tuning.
+                Can be either '8k' or '32k'
 
         Returns:
             A `LanguageModelTuningJob` object that represents the tuning job.
@@ -313,6 +316,9 @@ def tune_model(
                 )
             tuning_parameters["accelerator_type"] = accelerator_type
 
+        if max_context_length:
+            tuning_parameters["max_context_length"] = max_context_length
+
         return self._tune_model(
             training_data=training_data,
             tuning_parameters=tuning_parameters,
@@ -600,6 +606,7 @@ def tune_model(
         model_display_name: Optional[str] = None,
         tuning_evaluation_spec: Optional["TuningEvaluationSpec"] = None,
         accelerator_type: Optional[_ACCELERATOR_TYPE_TYPE] = None,
+        max_context_length: Optional[str] = None,
     ) -> "_LanguageModelTuningJob":
         """Tunes a model based on training data.
 
@@ -621,6 +628,8 @@ def tune_model(
             model_display_name: Custom display name for the tuned model.
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
+            max_context_length: The max context length used for tuning.
+                Can be either '8k' or '32k'
 
         Returns:
             A `LanguageModelTuningJob` object that represents the tuning job.
@@ -641,6 +650,7 @@ def tune_model(
             model_display_name=model_display_name,
             tuning_evaluation_spec=tuning_evaluation_spec,
             accelerator_type=accelerator_type,
+            max_context_length=max_context_length,
         )
 
 
@@ -659,6 +669,7 @@ def tune_model(
         model_display_name: Optional[str] = None,
         tuning_evaluation_spec: Optional["TuningEvaluationSpec"] = None,
         accelerator_type: Optional[_ACCELERATOR_TYPE_TYPE] = None,
+        max_context_length: Optional[str] = None,
     ) -> "_LanguageModelTuningJob":
         """Tunes a model based on training data.
 
@@ -687,6 +698,8 @@ def tune_model(
             model_display_name: Custom display name for the tuned model.
             tuning_evaluation_spec: Specification for the model evaluation during tuning.
             accelerator_type: Type of accelerator to use. Can be "TPU" or "GPU".
+            max_context_length: The max context length used for tuning.
+                Can be either '8k' or '32k'
 
         Returns:
             A `LanguageModelTuningJob` object that represents the tuning job.
@@ -708,6 +721,7 @@ def tune_model(
             model_display_name=model_display_name,
             tuning_evaluation_spec=tuning_evaluation_spec,
             accelerator_type=accelerator_type,
+            max_context_length=max_context_length,
         )
         tuned_model = job.get_tuned_model()
         self._endpoint = tuned_model._endpoint