ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #121 from unslothai/fix/vision-model-fixes

Browse source 
fix: Vision model detection and model-dataset compatibility
This commit is contained in:
Roland Tannous 2026-02-17 00:39:58 +04:00 • committed by GitHub
commit 369486246e
13 changed files with 156 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

12
studio/backend/core/training/trainer.py
View file

@ -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(

6
studio/backend/core/training/training.py
View file

@ -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:

1
studio/backend/models/training.py
View file

@ -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")

1
studio/backend/routes/training.py
View file

@ -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 "",

8
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -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
@ -693,7 +693,7 @@ export function ParamsSection(): ReactElement {
</SelectContent>
</Select>
</Row>
{store.modelType !== "vision" && (
{!showVisionLora && (
<div className="flex items-center gap-2">
<Checkbox
id="packing"

10
studio/frontend/src/features/studio/sections/training-section.tsx
View file

@ -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">

2
studio/frontend/src/features/studio/studio-page.tsx
View file

@ -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 && (

3
studio/frontend/src/features/training/api/mappers.ts
View file

@ -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.modelType === "vision") {
if (config.isVisionModel && config.isDatasetMultimodal) {
return { [input]: "image", [output]: "text" };
}

21
studio/frontend/src/features/training/api/models-api.ts Normal file
View file

@ -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;
}

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

@ -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({

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

@ -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 }),
@ -130,7 +221,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;
},
},

1
studio/frontend/src/features/training/types/api.ts
View file

@ -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;

4
studio/frontend/src/features/training/types/config.ts
View file

@ -50,6 +50,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;