feat: support hyperparameters in distillation tuning

PiperOrigin-RevId: 890622141

Author: speedstorm1

api-report/genai-node.api.md

@@ -889,7 +889,9 @@ export interface DistillationDataStats {
 // @public
 export interface DistillationHyperParameters {
     adapterSize?: AdapterSize;
+    batchSize?: number;
     epochCount?: string;
+    learningRate?: number;
     learningRateMultiplier?: number;
 }
 
@@ -909,6 +911,7 @@ export interface DistillationSpec {
     studentModel?: string;
     trainingDatasetUri?: string;
     tunedTeacherModelSource?: string;
+    tuningMode?: TuningMode;
     validationDatasetUri?: string;
 }
 

api-report/genai-web.api.md

@@ -889,7 +889,9 @@ export interface DistillationDataStats {
 // @public
 export interface DistillationHyperParameters {
     adapterSize?: AdapterSize;
+    batchSize?: number;
     epochCount?: string;
+    learningRate?: number;
     learningRateMultiplier?: number;
 }
 
@@ -909,6 +911,7 @@ export interface DistillationSpec {
     studentModel?: string;
     trainingDatasetUri?: string;
     tunedTeacherModelSource?: string;
+    tuningMode?: TuningMode;
     validationDatasetUri?: string;
 }
 

api-report/genai.api.md

@@ -889,7 +889,9 @@ export interface DistillationDataStats {
 // @public
 export interface DistillationHyperParameters {
     adapterSize?: AdapterSize;
+    batchSize?: number;
     epochCount?: string;
+    learningRate?: number;
     learningRateMultiplier?: number;
 }
 
@@ -909,6 +911,7 @@ export interface DistillationSpec {
     studentModel?: string;
     trainingDatasetUri?: string;
     tunedTeacherModelSource?: string;
+    tuningMode?: TuningMode;
     validationDatasetUri?: string;
 }
 

src/converters/_tunings_converters.ts

@@ -450,6 +450,15 @@ export function createTuningJobConfigToVertex(
         fromTuningMode,
       );
     }
+  } else if (discriminatorTuningMode === 'DISTILLATION') {
+    const fromTuningMode = common.getValueByPath(fromObject, ['tuningMode']);
+    if (parentObject !== undefined && fromTuningMode != null) {
+      common.setValueByPath(
+        parentObject,
+        ['distillationSpec', 'tuningMode'],
+        fromTuningMode,
+      );
+    }
   }
 
   const fromCustomBaseModel = common.getValueByPath(fromObject, [
@@ -479,6 +488,15 @@ export function createTuningJobConfigToVertex(
         fromBatchSize,
       );
     }
+  } else if (discriminatorBatchSize === 'DISTILLATION') {
+    const fromBatchSize = common.getValueByPath(fromObject, ['batchSize']);
+    if (parentObject !== undefined && fromBatchSize != null) {
+      common.setValueByPath(
+        parentObject,
+        ['distillationSpec', 'hyperParameters', 'batchSize'],
+        fromBatchSize,
+      );
+    }
   }
 
   let discriminatorLearningRate = common.getValueByPath(rootObject, [
@@ -499,6 +517,17 @@ export function createTuningJobConfigToVertex(
         fromLearningRate,
       );
     }
+  } else if (discriminatorLearningRate === 'DISTILLATION') {
+    const fromLearningRate = common.getValueByPath(fromObject, [
+      'learningRate',
+    ]);
+    if (parentObject !== undefined && fromLearningRate != null) {
+      common.setValueByPath(
+        parentObject,
+        ['distillationSpec', 'hyperParameters', 'learningRate'],
+        fromLearningRate,
+      );
+    }
   }
 
   const fromLabels = common.getValueByPath(fromObject, ['labels']);

src/types.ts

@@ -4609,14 +4609,19 @@ export declare interface PreferenceOptimizationSpec {
   validationDatasetUri?: string;
 }
 
-/** Hyperparameters for Distillation. This data type is not supported in Gemini API. */
+/** Hyperparameters for distillation. */
 export declare interface DistillationHyperParameters {
   /** Optional. Adapter size for distillation. */
   adapterSize?: AdapterSize;
   /** Optional. Number of complete passes the model makes over the entire training dataset during training. */
   epochCount?: string;
   /** Optional. Multiplier for adjusting the default learning rate. */
   learningRateMultiplier?: number;
+  /** The batch size hyperparameter for tuning.
+      This is only supported for OSS models in Vertex. */
+  batchSize?: number;
+  /** The learning rate for tuning. OSS models only. */
+  learningRate?: number;
 }
 
 /** Distillation tuning spec for tuning. */
@@ -4637,6 +4642,8 @@ export declare interface DistillationSpec {
   tunedTeacherModelSource?: string;
   /** Optional. Cloud Storage path to file containing validation dataset for tuning. The dataset must be formatted as a JSONL file. */
   validationDatasetUri?: string;
+  /** Tuning mode for tuning. */
+  tuningMode?: TuningMode;
 }
 
 /** The `Status` type defines a logical error model that is suitable for different programming environments, including REST APIs and RPC APIs. It is used by [gRPC](https://github.com/grpc). Each `Status` message contains three pieces of data: error code, error message, and error details. You can find out more about this error model and how to work with it in the [API Design Guide](https://cloud.google.com/apis/design/errors). This data type is not supported in Gemini API. */
@@ -5203,7 +5210,7 @@ export declare interface CreateTuningJobConfig {
   preTunedModelCheckpointId?: string;
   /** Adapter size for tuning. */
   adapterSize?: AdapterSize;
-  /** Tuning mode for SFT tuning. */
+  /** Tuning mode for tuning. */
   tuningMode?: TuningMode;
   /** Custom base model for tuning. This is only supported for OSS models in Vertex. */
   customBaseModel?: string;