feat: support hyperparameters in distillation tuning

PiperOrigin-RevId: 882708166

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

@@ -13421,7 +13421,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,
@@ -13502,7 +13502,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."""