feat: auto-detect model+dataset compatibility to select VLM vs LLM training path

2026-02-16 19:18:49 +00:00 · 2026-02-16 19:18:49 +00:00 · fa0ca59215
commit fa0ca59215
parent 09f3b6bce5
12 changed files with 94 additions and 21 deletions
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@ -110,17 +110,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(
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@ -96,7 +96,8 @@ class TrainingBackend:
                       # Optional parameters
                       custom_format_mapping: dict = None,
                       subset: str = None,
-                       split: str = "train") -> bool:
+                       split: str = "train",
+                       is_dataset_multimodal: bool = False) -> bool:
        """
        Start training.

@ -150,7 +151,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:
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@ -55,6 +55,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")
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@ -176,6 +176,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 "",
--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/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.isVisionModel;
+  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
@ -693,7 +693,7 @@ export function ParamsSection(): ReactElement {
                    </SelectContent>
                  </Select>
                </Row>
-                {!store.isVisionModel && (
+                {!showVisionLora && (
                  <div className="flex items-center gap-2">
                    <Checkbox
                      id="packing"
--- a/studio/frontend/src/features/studio/sections/training-section.tsx
+++ b/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">
--- a/studio/frontend/src/features/studio/studio-page.tsx
+++ b/studio/frontend/src/features/studio/studio-page.tsx
@ -70,7 +70,7 @@ export function StudioPage(): ReactElement {
          datasetSplit={config.datasetSplit}
          mode={dialogMode}
          initialData={dialogInitial}
-          isVlm={config.isVisionModel}
+          isVlm={config.isVisionModel && config.isDatasetMultimodal === true}
        />

        {canGoBack && (
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@ -54,6 +54,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
@ -72,7 +73,7 @@ function buildCustomFormatMapping(
  const { input, output } = config.datasetManualMapping;
  if (!input || !output) return undefined;

-  if (config.isVisionModel) {
+  if (config.isVisionModel && config.isDatasetMultimodal) {
    return { [input]: "image", [output]: "text" };
  }

--- a/studio/frontend/src/features/training/hooks/use-training-actions.ts
+++ b/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.isVisionModel;
+      const isVlm = config.isVisionModel && config.isDatasetMultimodal === true;

      if (datasetName) {
        const check = await checkDatasetFormat({
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@ -4,6 +4,7 @@ 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;
@ -27,6 +28,8 @@ const initialState: TrainingConfigState = {
  uploadedFile: null,
  isCheckingVision: false,
  isVisionModel: false,
+  isCheckingDataset: false,
+  isDatasetMultimodal: null,
  ...DEFAULT_HYPERPARAMS,
 };

@ -34,6 +37,9 @@ const initialState: TrainingConfigState = {
 // 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;
 }
@ -100,21 +106,68 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
      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 }),
@ -168,7 +221,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
        return s as unknown as TrainingConfigStore;
      },
      partialize: (state) => {
-        const { modelType, isCheckingVision, isVisionModel, ...rest } = state;
+        const { modelType, isCheckingVision, isVisionModel, isCheckingDataset, isDatasetMultimodal, ...rest } = state;
        return rest;
      },
    },
--- a/studio/frontend/src/features/training/types/api.ts
+++ b/studio/frontend/src/features/training/types/api.ts
@ -36,6 +36,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;
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@ -52,6 +52,8 @@ export interface TrainingConfigState {
  logFrequency: number;
  isCheckingVision: boolean;
  isVisionModel: boolean;
+  isCheckingDataset: boolean;
+  isDatasetMultimodal: boolean | null;
  finetuneVisionLayers: boolean;
  finetuneLanguageLayers: boolean;
  finetuneAttentionModules: boolean;