From 982a855df77d41e5f4bdc2227be4330ccaf90baa Mon Sep 17 00:00:00 2001
From: Manan17 <shahmanan170602@gmail.com>
Date: Thu, 19 Feb 2026 02:59:48 +0000
Subject: [PATCH 1/2] Setting default for the train on responses only

---
 .../training/stores/training-config-store.ts  | 41 ++++++++++++++++---
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index 3d7d352817..dbf5b3a817 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -104,8 +104,21 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             if (controller.signal.aborted) return;
             if (get().selectedModel !== modelName) return;
 
+            const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config);
+
+            // train_on_responses_only: true for LLMs, true for VLM+text dataset, false for VLM+vision dataset
+            if (!modelDetails.is_vision) {
+              patch.trainOnCompletions = true;
+            } else {
+              const datasetMultimodal = get().isDatasetMultimodal;
+              if (datasetMultimodal !== null) {
+                patch.trainOnCompletions = !datasetMultimodal;
+              }
+              // if dataset not yet checked, leave YAML default
+            }
+
             set({
-              ...mapBackendModelConfigToTrainingPatch(modelDetails.config),
+              ...patch,
               isVisionModel: modelDetails.is_vision,
               isLoadingModelDefaults: false,
               isCheckingVision: false,
@@ -129,10 +142,20 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             void checkVisionModel(modelName)
               .then((isVision) => {
                 if (get().selectedModel !== modelName) return;
-                set({
+                const updates: Record<string, unknown> = {
                   isVisionModel: isVision,
                   isCheckingVision: false,
-                });
+                };
+                // train_on_responses_only default
+                if (!isVision) {
+                  updates.trainOnCompletions = true;
+                } else {
+                  const datasetMultimodal = get().isDatasetMultimodal;
+                  if (datasetMultimodal !== null) {
+                    updates.trainOnCompletions = !datasetMultimodal;
+                  }
+                }
+                set(updates as Partial<TrainingConfigState>);
               })
               .catch(() => {
                 if (get().selectedModel !== modelName) return;
@@ -247,10 +270,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           })
             .then((res) => {
               if (controller.signal.aborted) return;
-              set({
-                isDatasetMultimodal: !!res.is_multimodal,
+              const isMultimodal = !!res.is_multimodal;
+              const updates: Record<string, unknown> = {
+                isDatasetMultimodal: isMultimodal,
                 isCheckingDataset: false,
-              });
+              };
+              // train_on_responses_only: for VLMs, true if text dataset, false if vision dataset
+              if (get().isVisionModel) {
+                updates.trainOnCompletions = !isMultimodal;
+              }
+              set(updates as Partial<TrainingConfigState>);
             })
             .catch(() => {
               if (controller.signal.aborted) return;

From b7b4cbc949e1193af5a7c91253f5e7b3609af0c4 Mon Sep 17 00:00:00 2001
From: Manan17 <shahmanan170602@gmail.com>
Date: Thu, 19 Feb 2026 05:06:41 +0000
Subject: [PATCH 2/2] Adding vision and multimodal dataset to the localstorage

---
 .../src/features/training/stores/training-config-store.ts       | 2 --
 1 file changed, 2 deletions(-)

diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts
index dbf5b3a817..260625db0d 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -46,11 +46,9 @@ let _modelConfigController: AbortController | null = null;
 const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set([
   "modelType",
   "isCheckingVision",
-  "isVisionModel",
   "isLoadingModelDefaults",
   "modelDefaultsError",
   "isCheckingDataset",
-  "isDatasetMultimodal",
 ]);
 
 function partializePersistedState(