Merge branch 'nightly' into feature/eval-split-auto-detection

studio/backend/core/training/trainer.py

@@ -111,17 +111,21 @@ class UnslothTrainer:
                    model_name: str,
                    max_seq_length: int = 2048,
                    load_in_4bit: bool = True,
-                   hf_token: Optional[str] = None) -> bool:
+                   hf_token: Optional[str] = None,
+                   is_dataset_multimodal: bool = False) -> bool:
         """Load model for training (supports both text and vision models)"""
         try:
             print("\nClearing GPU memory before training...")
             clear_gpu_cache()
 
-            # Detect if this is a vision model first
-            self.is_vlm = is_vision_model(model_name)
+            # Detect if this is a vision model AND dataset is multimodal
+            # A vision-capable model with a text-only dataset should use FastLanguageModel
+            self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal
             self.model_name = model_name
 
-            logger.info(f"Model type detected: {'Vision' if self.is_vlm else 'Text'}")
+            logger.info(f"Model architecture is vision: {is_vision_model(model_name)}")
+            logger.info(f"Dataset is multimodal: {is_dataset_multimodal}")
+            logger.info(f"Using VLM path: {self.is_vlm}")
 
             # Reset training state for new run
             self._update_progress(

studio/backend/core/training/training.py

@@ -104,7 +104,8 @@ class TrainingBackend:
                        subset: str = None,
                        train_split: str = "train",
                        eval_split: str = None,
-                       eval_steps: float = 0.01) -> bool:
+                       eval_steps: float = 0.01,
+                       is_dataset_multimodal: bool = False) -> bool:
         """
         Start training.
 
@@ -161,7 +162,8 @@ class TrainingBackend:
                 model_name=model_name,
                 max_seq_length=max_seq_length,
                 load_in_4bit=load_in_4bit if use_lora_actual else False,  # Only 4bit for LoRA
-                hf_token=hf_token if hf_token.strip() else None
+                hf_token=hf_token if hf_token.strip() else None,
+                is_dataset_multimodal=is_dataset_multimodal,
             )
 
             if not success or self.trainer.should_stop:

studio/backend/models/training.py

@@ -65,6 +65,7 @@ class TrainingStartRequest(BaseModel):
     finetune_language_layers: bool = Field(False, description="Finetune language layers")
     finetune_attention_modules: bool = Field(False, description="Finetune attention modules")
     finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules")
+    is_dataset_multimodal: bool = Field(False, description="Whether the dataset contains multimodal (image) data")
 
     # Logging parameters
     enable_wandb: bool = Field(False, description="Enable Weights & Biases logging")

studio/backend/routes/training.py

@@ -178,6 +178,7 @@ async def start_training(
             "finetune_language_layers": request.finetune_language_layers,
             "finetune_attention_modules": request.finetune_attention_modules,
             "finetune_mlp_modules": request.finetune_mlp_modules,
+            "is_dataset_multimodal": request.is_dataset_multimodal,
             "enable_wandb": request.enable_wandb,
             "wandb_token": request.wandb_token or "",
             "wandb_project": request.wandb_project or "",

studio/backend/utils/datasets/chat_templates.py

@@ -60,7 +60,24 @@ def get_tokenizer_chat_template(tokenizer, model_name):
             print(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
             print(f"   Falling back to tokenizer's default chat template")
     else:
-        print(f"📝 Using tokenizer's default chat template (no Unsloth template match)")
+        # Check if tokenizer actually has a chat_template set
+        has_chat_template = (
+            hasattr(tokenizer, 'chat_template')
+            and tokenizer.chat_template is not None
+        )
+        if has_chat_template:
+            print(f"📝 Using tokenizer's own chat template (no Unsloth template match)")
+        else:
+            # Base model with no chat template — apply default ChatML
+            print(f"📝 No chat template found — applying default ChatML template (base model)")
+            try:
+                tokenizer = get_chat_template(
+                    tokenizer,
+                    chat_template="chatml",
+                )
+            except Exception as e:
+                print(f"⚠️ Failed to apply default ChatML template: {e}")
+                print(f"   Falling back to tokenizer as-is")
 
     return tokenizer
 
@@ -227,6 +244,16 @@ def apply_chat_template_to_dataset(
     # ALPACA FORMAT
     if final_format == "alpaca":
 
+        # Set alpaca chat template on tokenizer for saving (if not already set)
+        # This ensures the template is saved with the model for inference
+        if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
+            try:
+                from unsloth.chat_templates import get_chat_template
+                tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
+                print(f"📝 Set alpaca chat template on tokenizer for model saving")
+            except Exception as e:
+                print(f"⚠️ Could not set alpaca template on tokenizer: {e}")
+
         # Use custom template if provided
         def _format_alpaca_custom(examples):
             texts = []

studio/frontend/src/features/studio/sections/params-section.tsx

@@ -109,7 +109,7 @@ function SliderRow({
 export function ParamsSection(): ReactElement {
   const store = useTrainingConfigStore();
   const isLora = store.trainingMethod !== "full";
-  const isVision = store.modelType === "vision";
+  const showVisionLora = store.isVisionModel && store.isDatasetMultimodal === true;
   const [loraOpen, setLoraOpen] = useState(false);
   const [hyperOpen, setHyperOpen] = useState(false);
 
@@ -350,7 +350,7 @@ export function ParamsSection(): ReactElement {
               />
 
               {/* Vision checkboxes */}
-              {isVision && (
+              {showVisionLora && (
                 <div className="flex flex-col gap-2 pt-1">
                   {(
                     [
@@ -400,7 +400,7 @@ export function ParamsSection(): ReactElement {
               )}
 
               {/* Text target modules */}
-              {!isVision && (
+              {!showVisionLora && (
                 <div className="flex flex-col gap-2 pt-1">
                   <span className="text-xs font-medium text-muted-foreground">
                     Target Modules
@@ -707,7 +707,7 @@ export function ParamsSection(): ReactElement {
                     </SelectContent>
                   </Select>
                 </Row>
-                {store.modelType !== "vision" && (
+                {!showVisionLora && (
                   <div className="flex items-center gap-2">
                     <Checkbox
                       id="packing"

studio/frontend/src/features/studio/sections/training-section.tsx

@@ -38,6 +38,9 @@ export function TrainingSection() {
   const store = useTrainingConfigStore();
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
   const [logOpen, setLogOpen] = useState(false);
+  const isIncompatible =
+    !store.isVisionModel && store.isDatasetMultimodal === true;
+
 
   return (
     <div data-tour="studio-training" className="lg:col-span-4">
@@ -98,7 +101,7 @@ export function TrainingSection() {
           data-tour="studio-start"
           className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
           onClick={() => void startTrainingRun()}
-          disabled={isStarting}
+          disabled={isStarting || isIncompatible}
         >
           <HugeiconsIcon icon={Rocket01Icon} className="size-4" />
           {isStarting ? "Starting..." : "Start Training"}
@@ -106,6 +109,11 @@ export function TrainingSection() {
         {startError && (
           <p className="text-xs text-red-500 leading-relaxed">{startError}</p>
         )}
+        {isIncompatible && (
+          <p className="text-xs text-red-500 leading-relaxed">
+            Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset.
+          </p>
+        )}
 
         {/* Save / Clear */}
         <div className="grid grid-cols-2 gap-2">

studio/frontend/src/features/studio/studio-page.tsx

@@ -36,7 +36,7 @@ export function StudioPage(): ReactElement {
   const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
   const closeDialog = useDatasetPreviewDialogStore((s) => s.close);
 
-  const canGoBack = runtimePhase === "stopped" || runtimePhase === "error";
+  const canGoBack = runtimePhase === "stopped" || runtimePhase === "error" || runtimePhase === "completed";
   const tourEnabled = hasHydratedRuntime && !isHydratingRuntime;
   const isConfigTour = !showTrainingView;
   const tourSteps = showTrainingView ? studioTrainingTourSteps : studioTourSteps;
@@ -70,7 +70,7 @@ export function StudioPage(): ReactElement {
           datasetSplit={config.datasetSplit}
           mode={dialogMode}
           initialData={dialogInitial}
-          isVlm={config.modelType === "vision"}
+          isVlm={config.isVisionModel && config.isDatasetMultimodal === true}
         />
 
         {canGoBack && (

studio/frontend/src/features/training/api/mappers.ts

@@ -55,6 +55,7 @@ export function buildTrainingStartPayload(
     finetune_language_layers: config.finetuneLanguageLayers,
     finetune_attention_modules: config.finetuneAttentionModules,
     finetune_mlp_modules: config.finetuneMLPModules,
+    is_dataset_multimodal: !!config.isDatasetMultimodal,
     enable_wandb: config.enableWandb,
     wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null,
     wandb_project: config.enableWandb
@@ -73,7 +74,7 @@ function buildCustomFormatMapping(
   const { input, output } = config.datasetManualMapping;
   if (!input || !output) return undefined;
 
-  if (config.modelType === "vision") {
+  if (config.isVisionModel && config.isDatasetMultimodal) {
     return { [input]: "image", [output]: "text" };
   }
 

studio/frontend/src/features/training/api/models-api.ts

@@ -0,0 +1,21 @@
+import { authFetch } from "@/features/auth";
+
+interface VisionCheckResponse {
+    model_name: string;
+    is_vision: boolean;
+}
+
+/**
+ * Check whether a model is a vision model by asking the backend.
+ * Calls GET /api/models/check-vision/{model_name}.
+ */
+export async function checkVisionModel(modelName: string): Promise<boolean> {
+    const encoded = encodeURIComponent(modelName);
+    const response = await authFetch(`/api/models/check-vision/${encoded}`);
+    if (!response.ok) {
+        // If the check fails (e.g. network error), default to non-vision
+        return false;
+    }
+    const data = (await response.json()) as VisionCheckResponse;
+    return data.is_vision;
+}

studio/frontend/src/features/training/hooks/use-training-actions.ts

@@ -29,7 +29,7 @@ export function useTrainingActions() {
 
     try {
       const datasetName = getDatasetName(config);
-      const isVlm = config.modelType === "vision";
+      const isVlm = config.isVisionModel && config.isDatasetMultimodal === true;
 
       if (datasetName) {
         const check = await checkDatasetFormat({

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -3,6 +3,8 @@ import type { StepNumber } from "@/types/training";
 import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import type { TrainingConfigState, TrainingConfigStore } from "../types/config";
+import { checkVisionModel } from "../api/models-api";
+import { checkDatasetFormat } from "../api/datasets-api";
 
 const MIN_STEP: StepNumber = 1;
 const MAX_STEP: StepNumber = STEPS.length as StepNumber;
@@ -24,9 +26,20 @@ const initialState: TrainingConfigState = {
   datasetSplit: null,
   datasetManualMapping: emptyManualMapping(),
   uploadedFile: null,
+  isCheckingVision: false,
+  isVisionModel: false,
+  isCheckingDataset: false,
+  isDatasetMultimodal: null,
   ...DEFAULT_HYPERPARAMS,
 };
 
+// AbortController for in-flight vision checks so rapid model changes
+// cancel stale requests.
+let _visionCheckController: AbortController | null = null;
+
+// AbortController for in-flight dataset multimodal checks.
+let _datasetCheckController: AbortController | null = null;
+
 function clampStep(step: number): StepNumber {
   return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
 }
@@ -57,26 +70,104 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
       nextStep: () => set({ currentStep: clampStep(get().currentStep + 1) }),
       prevStep: () => set({ currentStep: clampStep(get().currentStep - 1) }),
       setModelType: (modelType) => set({ modelType, selectedModel: null }),
-      setSelectedModel: (selectedModel) => set({ selectedModel }),
+      setSelectedModel: (selectedModel) => {
+        set({ selectedModel });
+
+        // Cancel any in-flight vision check
+        _visionCheckController?.abort();
+        _visionCheckController = null;
+
+        if (!selectedModel) {
+          set({ isCheckingVision: false });
+          return;
+        }
+
+        // Fire async backend check to determine if model is vision
+        const controller = new AbortController();
+        _visionCheckController = controller;
+        set({ isCheckingVision: true });
+
+        checkVisionModel(selectedModel)
+          .then((isVision) => {
+            // Only apply if this is still the active check
+            if (controller.signal.aborted) return;
+            set({
+              isVisionModel: isVision,
+              isCheckingVision: false,
+            });
+          })
+          .catch(() => {
+            if (controller.signal.aborted) return;
+            // On error, default to text and stop loading
+            set({ isCheckingVision: false });
+          });
+      },
       setTrainingMethod: (trainingMethod) => set({ trainingMethod }),
       setHfToken: (hfToken) => set({ hfToken }),
       setDatasetSource: (datasetSource) => set({ datasetSource }),
       setDatasetFormat: (datasetFormat) => set({ datasetFormat }),
-      setDataset: (dataset) =>
+      setDataset: (dataset) => {
+        // Cancel any in-flight dataset check
+        _datasetCheckController?.abort();
+        _datasetCheckController = null;
         set({
           dataset,
           datasetSubset: null,
           datasetSplit: null,
           datasetManualMapping: emptyManualMapping(),
-        }),
-      setDatasetSubset: (datasetSubset) =>
+          isDatasetMultimodal: null,
+          isCheckingDataset: false,
+        });
+      },
+      setDatasetSubset: (datasetSubset) => {
+        _datasetCheckController?.abort();
+        _datasetCheckController = null;
         set({
           datasetSubset,
           datasetSplit: null,
           datasetManualMapping: emptyManualMapping(),
-        }),
-      setDatasetSplit: (datasetSplit) =>
-        set({ datasetSplit, datasetManualMapping: emptyManualMapping() }),
+          isDatasetMultimodal: null,
+          isCheckingDataset: false,
+        });
+      },
+      setDatasetSplit: (datasetSplit) => {
+        _datasetCheckController?.abort();
+        _datasetCheckController = null;
+        set({
+          datasetSplit,
+          datasetManualMapping: emptyManualMapping(),
+          isDatasetMultimodal: null,
+          isCheckingDataset: false,
+        });
+        // Trigger async dataset multimodal check
+        const state = get();
+        const datasetName = state.datasetSource === "huggingface"
+          ? state.dataset
+          : state.uploadedFile;
+        if (!datasetName) return;
+
+        const controller = new AbortController();
+        _datasetCheckController = controller;
+        set({ isCheckingDataset: true });
+
+        checkDatasetFormat({
+          datasetName,
+          hfToken: state.hfToken.trim() || null,
+          subset: state.datasetSubset,
+          split: datasetSplit || "train",
+        })
+          .then((res) => {
+            if (controller.signal.aborted) return;
+            set({
+              isDatasetMultimodal: !!res.is_multimodal,
+              isCheckingDataset: false,
+            });
+          })
+          .catch(() => {
+            if (controller.signal.aborted) return;
+            set({ isDatasetMultimodal: null, isCheckingDataset: false });
+          });
+      },
       setDatasetManualMapping: (datasetManualMapping) =>
         set({ datasetManualMapping }),
       setUploadedFile: (uploadedFile) => set({ uploadedFile }),
@@ -131,7 +222,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         return s as unknown as TrainingConfigStore;
       },
       partialize: (state) => {
-        const { modelType, ...rest } = state;
+        const { modelType, isCheckingVision, isVisionModel, isCheckingDataset, isDatasetMultimodal, ...rest } = state;
         return rest;
       },
     },

studio/frontend/src/features/training/types/api.ts

@@ -37,6 +37,7 @@ export interface TrainingStartRequest {
   finetune_language_layers: boolean;
   finetune_attention_modules: boolean;
   finetune_mlp_modules: boolean;
+  is_dataset_multimodal: boolean;
   enable_wandb: boolean;
   wandb_token: string | null;
   wandb_project: string | null;

studio/frontend/src/features/training/types/config.ts

@@ -51,6 +51,10 @@ export interface TrainingConfigState {
   enableTensorboard: boolean;
   tensorboardDir: string;
   logFrequency: number;
+  isCheckingVision: boolean;
+  isVisionModel: boolean;
+  isCheckingDataset: boolean;
+  isDatasetMultimodal: boolean | null;
   finetuneVisionLayers: boolean;
   finetuneLanguageLayers: boolean;
   finetuneAttentionModules: boolean;