feat: support hyperparameters in distillation tuning

PiperOrigin-RevId: 890622141

Author: speedstorm1

google/genai/tests/tunings/test_tune.py

@@ -245,6 +245,25 @@
         ),
         exception_if_mldev="not supported in Gemini API",
     ),
+    pytest_helper.TestTableItem(
+        name="test_tune_oss_distillation_hyperparams",
+        parameters=genai_types.CreateTuningJobParameters(
+            base_model="qwen/qwen3@qwen3-4b",
+            training_dataset=genai_types.TuningDataset(
+                gcs_uri="gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl",
+            ),
+            config=genai_types.CreateTuningJobConfig(
+                method="DISTILLATION",
+                base_teacher_model="deepseek-ai/deepseek-r1-0528-maas",
+                learning_rate=1e-4,
+                batch_size=4,
+                output_uri="gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test",
+                tuning_mode="TUNING_MODE_FULL",
+                http_options=VERTEX_HTTP_OPTIONS,
+            ),
+        ),
+        exception_if_mldev="not supported in Gemini API",
+    ),
     pytest_helper.TestTableItem(
         name="test_tune_encryption_spec",
         parameters=genai_types.CreateTuningJobParameters(

google/genai/tunings.py

@@ -409,6 +409,13 @@ def _CreateTuningJobConfig_to_vertex(
           ['supervisedTuningSpec', 'tuningMode'],
           getv(from_object, ['tuning_mode']),
       )
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['tuning_mode']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'tuningMode'],
+          getv(from_object, ['tuning_mode']),
+      )
 
   if getv(from_object, ['custom_base_model']) is not None:
     setv(
@@ -427,6 +434,13 @@ def _CreateTuningJobConfig_to_vertex(
           ['supervisedTuningSpec', 'hyperParameters', 'batchSize'],
           getv(from_object, ['batch_size']),
       )
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['batch_size']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'hyperParameters', 'batchSize'],
+          getv(from_object, ['batch_size']),
+      )
 
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:
@@ -438,6 +452,13 @@ def _CreateTuningJobConfig_to_vertex(
           ['supervisedTuningSpec', 'hyperParameters', 'learningRate'],
           getv(from_object, ['learning_rate']),
       )
+  elif discriminator == 'DISTILLATION':
+    if getv(from_object, ['learning_rate']) is not None:
+      setv(
+          parent_object,
+          ['distillationSpec', 'hyperParameters', 'learningRate'],
+          getv(from_object, ['learning_rate']),
+      )
 
   discriminator = getv(root_object, ['config', 'method'])
   if discriminator is None:

google/genai/types.py

@@ -11649,10 +11649,7 @@ class PreferenceOptimizationSpecDict(TypedDict, total=False):
 
 
 class DistillationHyperParameters(_common.BaseModel):
-  """Hyperparameters for Distillation.
-
-  This data type is not supported in Gemini API.
-  """
+  """Hyperparameters for distillation."""
 
   adapter_size: Optional[AdapterSize] = Field(
       default=None, description="""Optional. Adapter size for distillation."""
@@ -11665,13 +11662,19 @@ class DistillationHyperParameters(_common.BaseModel):
       default=None,
       description="""Optional. Multiplier for adjusting the default learning rate.""",
   )
+  batch_size: Optional[int] = Field(
+      default=None,
+      description="""The batch size hyperparameter for tuning.
+      This is only supported for OSS models in Vertex.""",
+  )
+  learning_rate: Optional[float] = Field(
+      default=None,
+      description="""The learning rate for tuning. OSS models only.""",
+  )
 
 
 class DistillationHyperParametersDict(TypedDict, total=False):
-  """Hyperparameters for Distillation.
-
-  This data type is not supported in Gemini API.
-  """
+  """Hyperparameters for distillation."""
 
   adapter_size: Optional[AdapterSize]
   """Optional. Adapter size for distillation."""
@@ -11682,6 +11685,13 @@ class DistillationHyperParametersDict(TypedDict, total=False):
   learning_rate_multiplier: Optional[float]
   """Optional. Multiplier for adjusting the default learning rate."""
 
+  batch_size: Optional[int]
+  """The batch size hyperparameter for tuning.
+      This is only supported for OSS models in Vertex."""
+
+  learning_rate: Optional[float]
+  """The learning rate for tuning. OSS models only."""
+
 
 DistillationHyperParametersOrDict = Union[
     DistillationHyperParameters, DistillationHyperParametersDict
@@ -11723,6 +11733,9 @@ class DistillationSpec(_common.BaseModel):
       default=None,
       description="""Optional. Cloud Storage path to file containing validation dataset for tuning. The dataset must be formatted as a JSONL file.""",
   )
+  tuning_mode: Optional[TuningMode] = Field(
+      default=None, description="""Tuning mode for tuning."""
+  )
 
 
 class DistillationSpecDict(TypedDict, total=False):
@@ -11752,6 +11765,9 @@ class DistillationSpecDict(TypedDict, total=False):
   validation_dataset_uri: Optional[str]
   """Optional. Cloud Storage path to file containing validation dataset for tuning. The dataset must be formatted as a JSONL file."""
 
+  tuning_mode: Optional[TuningMode]
+  """Tuning mode for tuning."""
+
 
 DistillationSpecOrDict = Union[DistillationSpec, DistillationSpecDict]
 
@@ -13933,7 +13949,7 @@ class CreateTuningJobConfig(_common.BaseModel):
       default=None, description="""Adapter size for tuning."""
   )
   tuning_mode: Optional[TuningMode] = Field(
-      default=None, description="""Tuning mode for SFT tuning."""
+      default=None, description="""Tuning mode for tuning."""
   )
   custom_base_model: Optional[str] = Field(
       default=None,
@@ -14014,7 +14030,7 @@ class CreateTuningJobConfigDict(TypedDict, total=False):
   """Adapter size for tuning."""
 
   tuning_mode: Optional[TuningMode]
-  """Tuning mode for SFT tuning."""
+  """Tuning mode for tuning."""
 
   custom_base_model: Optional[str]
   """Custom base model for tuning. This is only supported for OSS models in Vertex."""