Skip to content

Commit 1c6c012

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
feat: Support distillation tuning
PiperOrigin-RevId: 862304483
1 parent d315544 commit 1c6c012

9 files changed

Lines changed: 476 additions & 3 deletions

File tree

‎Google.GenAI.E2E.Tests/Tunings/TuneTest.cs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,4 +208,50 @@ public async Task TuneSimpleGeminiTest() {
208208

209209
StringAssert.Contains(ex.Message, "not supported in Gemini API");
210210
}
211+
212+
// TODO(ayushagra): Enable this test once the distillation e2e test is out of autopush.
213+
[TestMethod]
214+
public async Task TuneDistillationVertexTest() {
215+
Assert.Inconclusive("Vertex distillation test is currently only supported in autopush.");
216+
var baseModel = "meta/llama3_1@llama-3.1-8b-instruct";
217+
var trainingDataset = new TuningDataset {
218+
GcsUri = "gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl"
219+
};
220+
var validationDataset = new TuningValidationDataset {
221+
GcsUri = "gs://nathreya-oss-tuning-sdk-test/distillation-val-openai-opposites.jsonl"
222+
};
223+
var config = new CreateTuningJobConfig {
224+
Method = TuningMethod.DISTILLATION,
225+
BaseTeacherModel = "deepseek-ai/deepseek-v3.1-maas",
226+
EpochCount = 20,
227+
ValidationDataset = validationDataset,
228+
OutputUri = "gs://nathreya-oss-tuning-sdk-test/ayushagra-distillation-test-folder",
229+
HttpOptions = new HttpOptions {
230+
ApiVersion = "v1beta1",
231+
BaseUrl = "https://us-central1-autopush-aiplatform.sandbox.googleapis.com/"
232+
}
233+
};
234+
235+
var tuningJob = await vertexClient.Tunings.TuneAsync(baseModel, trainingDataset, config);
236+
237+
Assert.IsNotNull(tuningJob);
238+
}
239+
240+
[TestMethod]
241+
public async Task TuneDistillationGeminiTest() {
242+
var baseModel = "meta/llama3_1@llama-3.1-8b-instruct";
243+
var trainingDataset = new TuningDataset {
244+
GcsUri = "gs://nathreya-oss-tuning-sdk-test/distillation-openai-opposites.jsonl"
245+
};
246+
var config = new CreateTuningJobConfig {
247+
Method = TuningMethod.DISTILLATION,
248+
BaseTeacherModel = "deepseek-ai/deepseek-v3.1-maas",
249+
};
250+
251+
var ex = await Assert.ThrowsExceptionAsync<NotSupportedException>(async () => {
252+
await geminiClient.Tunings.TuneAsync(baseModel, trainingDataset, config);
253+
});
254+
255+
StringAssert.Contains(ex.Message, "not supported in Gemini API");
256+
}
211257
}

‎Google.GenAI/README.md‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1409,6 +1409,40 @@ public class CancelTuningJobExample {
14091409
}
14101410
```
14111411

1412+
#### Distillation
1413+
1414+
You can perform distillation by setting `Method` to `TuningMethod.DISTILLATION` in `CreateTuningJobConfig`.
1415+
1416+
```csharp
1417+
using System.Threading.Tasks;
1418+
using Google.GenAI;
1419+
using Google.GenAI.Types;
1420+
1421+
public class DistillationJob {
1422+
public static async Task main() {
1423+
// assuming credentials are set up in environment variables as instructed above.
1424+
var client = new Client(vertexAI: true);
1425+
1426+
// Prompt dataset
1427+
var trainingDataset = new TuningDataset {
1428+
GcsUri = "gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl"
1429+
};
1430+
1431+
var config = new CreateTuningJobConfig {
1432+
Method = TuningMethod.DISTILLATION,
1433+
BaseTeacherModel = "teacher-model-id",
1434+
EpochCount = 1,
1435+
};
1436+
var tuningJob = await client.Tunings.TuneAsync(
1437+
baseModel: "student-model-id",
1438+
trainingDataset: trainingDataset,
1439+
config: config,
1440+
);
1441+
Console.WriteLine(tuningJob.State);
1442+
}
1443+
}
1444+
```
1445+
14121446
### List Tuning Jobs
14131447
```csharp
14141448
using System.Threading.Tasks;

‎Google.GenAI/Tunings.cs‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,27 @@ internal JsonNode CreateTuningJobConfigToMldev(JsonNode fromObject, JsonObject p
146146
throw new NotSupportedException("beta parameter is not supported in Gemini API.");
147147
}
148148

149+
if (!Common.IsZero(Common.GetValueByPath(fromObject, new string[] { "baseTeacherModel" }))) {
150+
throw new NotSupportedException(
151+
"baseTeacherModel parameter is not supported in Gemini API.");
152+
}
153+
154+
if (!Common.IsZero(
155+
Common.GetValueByPath(fromObject, new string[] { "tunedTeacherModelSource" }))) {
156+
throw new NotSupportedException(
157+
"tunedTeacherModelSource parameter is not supported in Gemini API.");
158+
}
159+
160+
if (!Common.IsZero(
161+
Common.GetValueByPath(fromObject, new string[] { "sftLossWeightMultiplier" }))) {
162+
throw new NotSupportedException(
163+
"sftLossWeightMultiplier parameter is not supported in Gemini API.");
164+
}
165+
166+
if (!Common.IsZero(Common.GetValueByPath(fromObject, new string[] { "outputUri" }))) {
167+
throw new NotSupportedException("outputUri parameter is not supported in Gemini API.");
168+
}
169+
149170
return toObject;
150171
}
151172

@@ -175,6 +196,14 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
175196
fromObject, new string[] { "validationDataset" }))),
176197
toObject, rootObject));
177198
}
199+
} else if (discriminatorValueValidationDataset == "DISTILLATION") {
200+
if (Common.GetValueByPath(fromObject, new string[] { "validationDataset" }) != null) {
201+
Common.SetValueByPath(parentObject, new string[] { "distillationSpec" },
202+
TuningValidationDatasetToVertex(
203+
JsonNode.Parse(JsonSerializer.Serialize(Common.GetValueByPath(
204+
fromObject, new string[] { "validationDataset" }))),
205+
toObject, rootObject));
206+
}
178207
}
179208
if (Common.GetValueByPath(fromObject, new string[] { "tunedModelDisplayName" }) != null) {
180209
Common.SetValueByPath(
@@ -206,6 +235,12 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
206235
new string[] { "preferenceOptimizationSpec", "hyperParameters", "epochCount" },
207236
Common.GetValueByPath(fromObject, new string[] { "epochCount" }));
208237
}
238+
} else if (discriminatorValueEpochCount == "DISTILLATION") {
239+
if (Common.GetValueByPath(fromObject, new string[] { "epochCount" }) != null) {
240+
Common.SetValueByPath(
241+
parentObject, new string[] { "distillationSpec", "hyperParameters", "epochCount" },
242+
Common.GetValueByPath(fromObject, new string[] { "epochCount" }));
243+
}
209244
}
210245

211246
JsonNode discriminatorLearningRateMultiplier =
@@ -229,6 +264,13 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
229264
"learningRateMultiplier" },
230265
Common.GetValueByPath(fromObject, new string[] { "learningRateMultiplier" }));
231266
}
267+
} else if (discriminatorValueLearningRateMultiplier == "DISTILLATION") {
268+
if (Common.GetValueByPath(fromObject, new string[] { "learningRateMultiplier" }) != null) {
269+
Common.SetValueByPath(
270+
parentObject,
271+
new string[] { "distillationSpec", "hyperParameters", "learningRateMultiplier" },
272+
Common.GetValueByPath(fromObject, new string[] { "learningRateMultiplier" }));
273+
}
232274
}
233275

234276
JsonNode discriminatorExportLastCheckpointOnly =
@@ -252,6 +294,13 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
252294
new string[] { "preferenceOptimizationSpec", "exportLastCheckpointOnly" },
253295
Common.GetValueByPath(fromObject, new string[] { "exportLastCheckpointOnly" }));
254296
}
297+
} else if (discriminatorValueExportLastCheckpointOnly == "DISTILLATION") {
298+
if (Common.GetValueByPath(fromObject, new string[] { "exportLastCheckpointOnly" }) !=
299+
null) {
300+
Common.SetValueByPath(
301+
parentObject, new string[] { "distillationSpec", "exportLastCheckpointOnly" },
302+
Common.GetValueByPath(fromObject, new string[] { "exportLastCheckpointOnly" }));
303+
}
255304
}
256305

257306
JsonNode discriminatorAdapterSize =
@@ -273,6 +322,12 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
273322
new string[] { "preferenceOptimizationSpec", "hyperParameters", "adapterSize" },
274323
Common.GetValueByPath(fromObject, new string[] { "adapterSize" }));
275324
}
325+
} else if (discriminatorValueAdapterSize == "DISTILLATION") {
326+
if (Common.GetValueByPath(fromObject, new string[] { "adapterSize" }) != null) {
327+
Common.SetValueByPath(
328+
parentObject, new string[] { "distillationSpec", "hyperParameters", "adapterSize" },
329+
Common.GetValueByPath(fromObject, new string[] { "adapterSize" }));
330+
}
276331
}
277332
if (!Common.IsZero(Common.GetValueByPath(fromObject, new string[] { "batchSize" }))) {
278333
throw new NotSupportedException("batchSize parameter is not supported in Vertex AI.");
@@ -293,6 +348,30 @@ internal JsonNode CreateTuningJobConfigToVertex(JsonNode fromObject, JsonObject
293348
Common.GetValueByPath(fromObject, new string[] { "beta" }));
294349
}
295350

351+
if (Common.GetValueByPath(fromObject, new string[] { "baseTeacherModel" }) != null) {
352+
Common.SetValueByPath(
353+
parentObject, new string[] { "distillationSpec", "baseTeacherModel" },
354+
Common.GetValueByPath(fromObject, new string[] { "baseTeacherModel" }));
355+
}
356+
357+
if (Common.GetValueByPath(fromObject, new string[] { "tunedTeacherModelSource" }) != null) {
358+
Common.SetValueByPath(
359+
parentObject, new string[] { "distillationSpec", "tunedTeacherModelSource" },
360+
Common.GetValueByPath(fromObject, new string[] { "tunedTeacherModelSource" }));
361+
}
362+
363+
if (Common.GetValueByPath(fromObject, new string[] { "sftLossWeightMultiplier" }) != null) {
364+
Common.SetValueByPath(
365+
parentObject,
366+
new string[] { "distillationSpec", "hyperParameters", "sftLossWeightMultiplier" },
367+
Common.GetValueByPath(fromObject, new string[] { "sftLossWeightMultiplier" }));
368+
}
369+
370+
if (Common.GetValueByPath(fromObject, new string[] { "outputUri" }) != null) {
371+
Common.SetValueByPath(parentObject, new string[] { "outputUri" },
372+
Common.GetValueByPath(fromObject, new string[] { "outputUri" }));
373+
}
374+
296375
return toObject;
297376
}
298377

@@ -586,6 +665,12 @@ internal JsonNode TuningDatasetToVertex(JsonNode fromObject, JsonObject parentOb
586665
new string[] { "preferenceOptimizationSpec", "trainingDatasetUri" },
587666
Common.GetValueByPath(fromObject, new string[] { "gcsUri" }));
588667
}
668+
} else if (discriminatorValueGcsUri == "DISTILLATION") {
669+
if (Common.GetValueByPath(fromObject, new string[] { "gcsUri" }) != null) {
670+
Common.SetValueByPath(parentObject,
671+
new string[] { "distillationSpec", "promptDatasetUri" },
672+
Common.GetValueByPath(fromObject, new string[] { "gcsUri" }));
673+
}
589674
}
590675

591676
JsonNode discriminatorVertexDatasetResource =
@@ -606,6 +691,12 @@ internal JsonNode TuningDatasetToVertex(JsonNode fromObject, JsonObject parentOb
606691
parentObject, new string[] { "preferenceOptimizationSpec", "trainingDatasetUri" },
607692
Common.GetValueByPath(fromObject, new string[] { "vertexDatasetResource" }));
608693
}
694+
} else if (discriminatorValueVertexDatasetResource == "DISTILLATION") {
695+
if (Common.GetValueByPath(fromObject, new string[] { "vertexDatasetResource" }) != null) {
696+
Common.SetValueByPath(
697+
parentObject, new string[] { "distillationSpec", "promptDatasetUri" },
698+
Common.GetValueByPath(fromObject, new string[] { "vertexDatasetResource" }));
699+
}
609700
}
610701
if (!Common.IsZero(Common.GetValueByPath(fromObject, new string[] { "examples" }))) {
611702
throw new NotSupportedException("examples parameter is not supported in Vertex AI.");
@@ -756,6 +847,12 @@ internal JsonNode TuningJobFromVertex(JsonNode fromObject, JsonObject parentObje
756847
Common.GetValueByPath(fromObject, new string[] { "preferenceOptimizationSpec" }));
757848
}
758849

850+
if (Common.GetValueByPath(fromObject, new string[] { "distillationSpec" }) != null) {
851+
Common.SetValueByPath(
852+
toObject, new string[] { "distillationSpec" },
853+
Common.GetValueByPath(fromObject, new string[] { "distillationSpec" }));
854+
}
855+
759856
if (Common.GetValueByPath(fromObject, new string[] { "tuningDataStats" }) != null) {
760857
Common.SetValueByPath(
761858
toObject, new string[] { "tuningDataStats" },

‎Google.GenAI/types/CreateTuningJobConfig.cs‎

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ public record CreateTuningJobConfig {
3535
public HttpOptions ? HttpOptions { get; set; }
3636

3737
/// <summary>
38-
/// The method to use for tuning (SUPERVISED_FINE_TUNING or PREFERENCE_TUNING). If not set, the
39-
/// default method (SFT) will be used.
38+
/// The method to use for tuning (SUPERVISED_FINE_TUNING or PREFERENCE_TUNING or DISTILLATION).
39+
/// If not set, the default method (SFT) will be used.
4040
/// </summary>
4141
[JsonPropertyName("method")]
4242
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
@@ -173,6 +173,46 @@ public double
173173
get; set;
174174
}
175175

176+
/// <summary>
177+
/// The base teacher model that is being distilled. Distillation only.
178+
/// </summary>
179+
[JsonPropertyName("baseTeacherModel")]
180+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
181+
public string
182+
? BaseTeacherModel {
183+
get; set;
184+
}
185+
186+
/// <summary>
187+
/// The resource name of the Tuned teacher model. Distillation only.
188+
/// </summary>
189+
[JsonPropertyName("tunedTeacherModelSource")]
190+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
191+
public string
192+
? TunedTeacherModelSource {
193+
get; set;
194+
}
195+
196+
/// <summary>
197+
/// Multiplier for adjusting the weight of the SFT loss. Distillation only.
198+
/// </summary>
199+
[JsonPropertyName("sftLossWeightMultiplier")]
200+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
201+
public double
202+
? SftLossWeightMultiplier {
203+
get; set;
204+
}
205+
206+
/// <summary>
207+
/// The Google Cloud Storage location where the tuning job outputs are written.
208+
/// </summary>
209+
[JsonPropertyName("outputUri")]
210+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
211+
public string
212+
? OutputUri {
213+
get; set;
214+
}
215+
176216
/// <summary>
177217
/// Deserializes a JSON string to a CreateTuningJobConfig object.
178218
/// </summary>

‎Google.GenAI/types/DistillationHyperParameters.cs‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
// Auto-generated code. Do not edit.
18+
19+
using System;
20+
using System.Text.Json;
21+
using System.Text.Json.Serialization;
22+
using Google.GenAI.Serialization;
23+
24+
namespace Google.GenAI.Types {
25+
/// <summary>
26+
/// Hyperparameters for Distillation. This data type is not supported in Gemini API.
27+
/// </summary>
28+
29+
public record DistillationHyperParameters {
30+
/// <summary>
31+
/// Optional. Adapter size for distillation.
32+
/// </summary>
33+
[JsonPropertyName("adapterSize")]
34+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
35+
public AdapterSize ? AdapterSize { get; set; }
36+
37+
/// <summary>
38+
/// Optional. Number of complete passes the model makes over the entire training dataset during
39+
/// training.
40+
/// </summary>
41+
[JsonPropertyName("epochCount")]
42+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
43+
[JsonConverter(typeof(StringToNullableLongConverter))]
44+
public long
45+
? EpochCount {
46+
get; set;
47+
}
48+
49+
/// <summary>
50+
/// Optional. Multiplier for adjusting the default learning rate.
51+
/// </summary>
52+
[JsonPropertyName("learningRateMultiplier")]
53+
[JsonIgnore(Condition = JsonIgnoreCondition.WhenWritingNull)]
54+
public double
55+
? LearningRateMultiplier {
56+
get; set;
57+
}
58+
59+
/// <summary>
60+
/// Deserializes a JSON string to a DistillationHyperParameters object.
61+
/// </summary>
62+
/// <param name="jsonString">The JSON string to deserialize.</param>
63+
/// <param name="options">Optional JsonSerializerOptions.</param>
64+
/// <returns>The deserialized DistillationHyperParameters object, or null if deserialization
65+
/// fails.</returns>
66+
public static DistillationHyperParameters
67+
? FromJson(string jsonString, JsonSerializerOptions? options = null) {
68+
try {
69+
return JsonSerializer.Deserialize<DistillationHyperParameters>(jsonString, options);
70+
} catch (JsonException e) {
71+
Console.Error.WriteLine($"Error deserializing JSON: {e.ToString()}");
72+
return null;
73+
}
74+
}
75+
}
76+
}

0 commit comments

Comments
 (0)