feat: Tuning Checkpoints API

Google APIs · copybara-github · commit 3f93d1a39b09 · 2025-05-12T11:28:28.000-07:00
PiperOrigin-RevId: 757844206
diff --git a/google/cloud/aiplatform/v1/tuning_job.proto b/google/cloud/aiplatform/v1/tuning_job.proto
@@ -42,7 +42,8 @@ message TuningJob {
   };
 
   oneof source_model {
-    // The base model that is being tuned, e.g., "gemini-1.0-pro-002".
+    // The base model that is being tuned. See [Supported
+    // models](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/tuning#supported_models).
     string base_model = 4;
   }
 
@@ -118,7 +119,7 @@ message TuningJob {
     }
   ];
 
-  // Output only. The tuned model resources assiociated with this
+  // Output only. The tuned model resources associated with this
   // [TuningJob][google.cloud.aiplatform.v1.TuningJob].
   TunedModel tuned_model = 14 [(google.api.field_behavior) = OUTPUT_ONLY];
 
@@ -142,7 +143,7 @@ message TuningJob {
   string service_account = 22;
 }
 
-// The Model Registry Model and Online Prediction Endpoint assiociated with
+// The Model Registry Model and Online Prediction Endpoint associated with
 // this [TuningJob][google.cloud.aiplatform.v1.TuningJob].
 message TunedModel {
   // Output only. The resource name of the TunedModel. Format:
@@ -162,6 +163,12 @@ message TunedModel {
       type: "aiplatform.googleapis.com/Endpoint"
     }
   ];
+
+  // Output only. The checkpoints associated with this TunedModel.
+  // This field is only populated for tuning jobs that enable intermediate
+  // checkpoints.
+  repeated TunedModelCheckpoint checkpoints = 3
+      [(google.api.field_behavior) = OUTPUT_ONLY];
 }
 
 // Dataset distribution for Supervised Tuning.
@@ -245,13 +252,21 @@ message SupervisedTuningDataStats {
   repeated Content user_dataset_examples = 8
       [(google.api.field_behavior) = OUTPUT_ONLY];
 
-  // The number of examples in the dataset that have been truncated by any
-  // amount.
-  int64 total_truncated_example_count = 10;
+  // Output only. The number of examples in the dataset that have been dropped.
+  // An example can be dropped for reasons including: too many tokens, contains
+  // an invalid image, contains too many images, etc.
+  int64 total_truncated_example_count = 10
+      [(google.api.field_behavior) = OUTPUT_ONLY];
+
+  // Output only. A partial sample of the indices (starting from 1) of the
+  // dropped examples.
+  repeated int64 truncated_example_indices = 11
+      [(google.api.field_behavior) = OUTPUT_ONLY];
 
-  // A partial sample of the indices (starting from 1) of the truncated
-  // examples.
-  repeated int64 truncated_example_indices = 11;
+  // Output only. For each index in `truncated_example_indices`, the user-facing
+  // reason why the example was dropped.
+  repeated string dropped_example_reasons = 12
+      [(google.api.field_behavior) = OUTPUT_ONLY];
 }
 
 // The tuning data statistic values for
@@ -273,6 +288,9 @@ message SupervisedHyperParameters {
     // Adapter size 1.
     ADAPTER_SIZE_ONE = 1;
 
+    // Adapter size 2.
+    ADAPTER_SIZE_TWO = 6;
+
     // Adapter size 4.
     ADAPTER_SIZE_FOUR = 2;
 
@@ -281,6 +299,9 @@ message SupervisedHyperParameters {
 
     // Adapter size 16.
     ADAPTER_SIZE_SIXTEEN = 4;
+
+    // Adapter size 32.
+    ADAPTER_SIZE_THIRTY_TWO = 5;
   }
 
   // Optional. Number of complete passes the model makes over the entire
@@ -307,6 +328,11 @@ message SupervisedTuningSpec {
   // Optional. Hyperparameters for SFT.
   SupervisedHyperParameters hyper_parameters = 3
       [(google.api.field_behavior) = OPTIONAL];
+
+  // Optional. If set to true, disable intermediate checkpoints for SFT and only
+  // the last checkpoint will be exported. Otherwise, enable intermediate
+  // checkpoints for SFT. Default is false.
+  bool export_last_checkpoint_only = 6 [(google.api.field_behavior) = OPTIONAL];
 }
 
 // TunedModel Reference for legacy model migration.
@@ -331,3 +357,19 @@ message TunedModelRef {
     }];
   }
 }
+
+// TunedModelCheckpoint for the Tuned Model of a Tuning Job.
+message TunedModelCheckpoint {
+  // The ID of the checkpoint.
+  string checkpoint_id = 1;
+
+  // The epoch of the checkpoint.
+  int64 epoch = 2;
+
+  // The step of the checkpoint.
+  int64 step = 3;
+
+  // The Endpoint resource name that the checkpoint is deployed to. Format:
+  // `projects/{project}/locations/{location}/endpoints/{endpoint}`.
+  string endpoint = 4;
+}
