diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b128fb5338..b9643cac6a 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3057,6 +3057,14 @@ class UnslothTrainer: logger.info("Configuring DeepSeek OCR data collator...\n") FastVisionModel.for_training(self.model) + # DeepSeek OCR's (image_size, base_size, crop_mode) is a + # coupled preset; changing image_size alone desyncs the + # per-crop pixel grid from num_queries. Use Gundam. + if training_args.get("vision_image_size") is not None: + logger.info( + "Vision image resize ignored for DeepSeek OCR " + "(uses fixed Gundam preset).\n" + ) data_collator = DeepSeekOCRDataCollator( tokenizer = self.tokenizer, model = self.model, @@ -3123,7 +3131,21 @@ class UnslothTrainer: from unsloth.trainer import UnslothVisionDataCollator FastVisionModel.for_training(self.model) - data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) + vision_image_size = training_args.get("vision_image_size") + if vision_image_size is None: + data_collator = UnslothVisionDataCollator( + self.model, self.tokenizer + ) + else: + logger.info( + f"Vision image resize: {vision_image_size} (max dimension)\n" + ) + data_collator = UnslothVisionDataCollator( + self.model, + self.tokenizer, + resize = vision_image_size, + resize_dimension = "max", + ) logger.info("Vision data collator configured\n") # ========== TRAINING CONFIGURATION ========== diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d2c2316d45..0af3349c6f 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -193,6 +193,7 @@ class TrainingBackend: "hf_token": kwargs.get("hf_token", ""), "load_in_4bit": kwargs.get("load_in_4bit", True), "max_seq_length": kwargs.get("max_seq_length", 2048), + "vision_image_size": kwargs.get("vision_image_size"), "hf_dataset": kwargs.get("hf_dataset", ""), "local_datasets": kwargs.get("local_datasets"), "local_eval_datasets": kwargs.get("local_eval_datasets"), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f47a6bd599..632b38d75a 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -959,7 +959,47 @@ def _activate_transformers_version(model_name: str) -> None: activate_transformers_for_subprocess(model_name) -def _adapt_for_mlx_vlm(items): +def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]: + if width <= 0 or height <= 0 or target <= 0: + return width, height + largest_side = max(width, height) + if largest_side <= target: + return width, height + # Integer formula matches unsloth_zoo's collator (Python round() differs + # by 1px on half-pixel cases). max(1, _) avoids zero-side degenerate output. + new_w = max(1, (width * target + largest_side // 2) // largest_side) + new_h = max(1, (height * target + largest_side // 2) // largest_side) + return new_w, new_h + + +def _resize_mlx_vlm_image(image, resize): + if resize is None: + return image + try: + from PIL import Image + import numpy as np + except ImportError: + return image + if not isinstance(image, Image.Image): + return image + image = image.convert("RGB") + new_size = _mlx_vlm_max_resized_size(*image.size, int(resize)) + if new_size != image.size: + resampling = getattr(Image, "Resampling", Image).LANCZOS + image = image.resize(new_size, resampling) + # When a resize is requested, hand mlx-vlm a writable RGB ndarray so its + # PIL-path square-resize is skipped and HF processors don't warn on + # non-writable views. resize=None (Default) above keeps the original PIL. + return np.array(image, copy = True) + + +def _resize_mlx_vlm_images(value, resize): + if isinstance(value, list): + return [_resize_mlx_vlm_image(image, resize) for image in value] + return _resize_mlx_vlm_image(value, resize) + + +def _adapt_for_mlx_vlm(items, resize = None): """Adapt GPU-path VLM dataset output for mlx-vlm consumption. The GPU path embeds PIL images inside messages content as @@ -979,7 +1019,7 @@ def _adapt_for_mlx_vlm(items): if isinstance(part, dict) and part.get("type") == "image": img = part.get("image") if img is not None: - images.append(img) + images.append(_resize_mlx_vlm_image(img, resize)) new_content.append({"type": "image"}) else: new_content.append(part) @@ -990,9 +1030,9 @@ def _adapt_for_mlx_vlm(items): if images: out["image"] = images[0] if len(images) == 1 else images elif "image" in item: - out["image"] = item["image"] + out["image"] = _resize_mlx_vlm_images(item["image"], resize) elif "images" in item: - out["images"] = item["images"] + out["images"] = _resize_mlx_vlm_images(item["images"], resize) adapted.append(out) return adapted @@ -1168,6 +1208,25 @@ def _run_mlx_training(event_queue, stop_queue, config): is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False)) model._is_vlm_model = is_vlm + vision_image_size = config.get("vision_image_size") + # DeepSeek OCR uses a coupled preset tuple; skip resize like the Torch path. + _model_name_lower = str(config.get("model_name", "")).lower() + _is_deepseek_ocr = "deepseek" in _model_name_lower and "ocr" in _model_name_lower + if is_vlm and vision_image_size is not None and _is_deepseek_ocr: + _send( + "status", + status_message = ( + "MLX vision image resize ignored for DeepSeek OCR " + "(uses fixed Gundam preset)." + ), + ) + vision_image_size = None + elif is_vlm and vision_image_size is not None: + vision_image_size = int(vision_image_size) + _send( + "status", + status_message = f"MLX vision image resize: {vision_image_size} (max dimension)", + ) # ── 2. Apply LoRA / full FT ── # Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.) @@ -1302,7 +1361,10 @@ def _run_mlx_training(event_queue, stop_queue, config): progress_callback = _fmt_progress, ) if vlm_info.get("success"): - dataset = _adapt_for_mlx_vlm(vlm_info["dataset"]) + dataset = _adapt_for_mlx_vlm( + vlm_info["dataset"], + resize = vision_image_size, + ) else: errors = vlm_info.get("errors", []) raise ValueError( @@ -1317,7 +1379,10 @@ def _run_mlx_training(event_queue, stop_queue, config): dataset_name = hf_dataset or "local", ) if ev_info.get("success"): - eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"]) + eval_dataset = _adapt_for_mlx_vlm( + ev_info["dataset"], + resize = vision_image_size, + ) elif format_type: _send("status", status_message = f"Formatting dataset ({format_type})...") @@ -2248,6 +2313,7 @@ def run_training_process( eval_dataset = eval_dataset, eval_steps = eval_steps, max_seq_length = config.get("max_seq_length", 2048), + vision_image_size = config.get("vision_image_size"), optim = config.get("optim", "adamw_8bit"), lr_scheduler_type = config.get("lr_scheduler_type", "linear"), is_cpt = is_cpt, diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 7c53b0fee5..c6be1eff4e 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -5,10 +5,17 @@ Pydantic schemas for Training API """ +import re from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing import Any, Optional, List, Dict, Literal +# ASCII integer with an optional single sign. Used by _check_vision_image_size +# to reject "++512", "--256", and Unicode-digit strings ("５１２", "٥١٢") that +# would otherwise slip through str.isdigit() + int(). +_INT_RE = re.compile(r"[+-]?[0-9]+") + + _MAX_BATCH_SIZE = 4096 _MAX_GRAD_ACCUM = 4096 _MAX_STEPS = 1_000_000 @@ -18,6 +25,9 @@ _MAX_SEQ_LENGTH = 2_000_000 _MAX_LR_VALUE = 1.0 _MAX_LORA_R = 16_384 _MAX_LORA_ALPHA = 32_768 +_MIN_VISION_IMAGE_SIZE = 256 +# 2048 was the most I could get most llms to work at without getting unstable +_MAX_VISION_IMAGE_SIZE = 2048 def _parse_lr(v: Any) -> float: @@ -58,6 +68,10 @@ class TrainingStartRequest(BaseModel): hf_token: Optional[str] = Field(None, description = "HuggingFace token") load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization") max_seq_length: int = Field(2048, description = "Maximum sequence length") + vision_image_size: Optional[int] = Field( + None, + description = "Optional maximum image side length for VLM training. Null uses model default.", + ) trust_remote_code: bool = Field( False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", @@ -159,6 +173,40 @@ class TrainingStartRequest(BaseModel): ) return v + @field_validator("vision_image_size", mode = "before") + @classmethod + def _check_vision_image_size(cls, v: Any) -> Optional[int]: + # mode="before" sees True/False as bool (not 1/0) for a precise error. + if v is None: + return v + if isinstance(v, bool): + raise ValueError("vision_image_size must be an integer or null") + if isinstance(v, int): + coerced = v + elif isinstance(v, str) and _INT_RE.fullmatch(v.strip()): + coerced = int(v.strip()) + elif isinstance(v, float) and v.is_integer(): + coerced = int(v) + else: + # numpy ints / Integral subclasses, without a hard numpy import. + try: + import numbers + + if isinstance(v, numbers.Integral): + coerced = int(v) + elif isinstance(v, numbers.Real) and float(v).is_integer(): + coerced = int(v) + else: + raise TypeError + except Exception: + raise ValueError("vision_image_size must be an integer or null") + if coerced < _MIN_VISION_IMAGE_SIZE or coerced > _MAX_VISION_IMAGE_SIZE: + raise ValueError( + f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, " + f"{_MAX_VISION_IMAGE_SIZE}] (got {coerced!r})" + ) + return coerced + @field_validator("warmup_steps") @classmethod def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]: diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6e2413b3e9..41a9e15562 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -194,6 +194,7 @@ async def start_training( "hf_token": request.hf_token or "", "load_in_4bit": request.load_in_4bit, "max_seq_length": request.max_seq_length, + "vision_image_size": request.vision_image_size, "hf_dataset": request.hf_dataset or "", "local_datasets": request.local_datasets, "local_eval_datasets": request.local_eval_datasets, diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 98c7bdaa55..c36363b1ae 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -66,6 +66,7 @@ def _load_worker_module(): _worker = _load_worker_module() _normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer _normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler +_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size def test_mlx_studio_optimizer_aliases_are_explicit(): @@ -82,3 +83,14 @@ def test_mlx_studio_rejects_unknown_optimizer(): def test_mlx_studio_rejects_unknown_scheduler(): with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"): _normalize_mlx_studio_scheduler("linear_typo") + + +def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer(): + assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256) + assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512) + assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512) + assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128) + assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256) + # Half-pixel cases must match the Torch collator (not banker's round). + assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500) + assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167) diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py index 7ffa9bb384..6df491610b 100644 --- a/studio/backend/tests/test_studio_train_validation.py +++ b/studio/backend/tests/test_studio_train_validation.py @@ -18,6 +18,8 @@ from models.training import ( _MAX_LORA_ALPHA, _MAX_LORA_R, _MAX_SEQ_LENGTH, + _MAX_VISION_IMAGE_SIZE, + _MIN_VISION_IMAGE_SIZE, ) @@ -62,6 +64,52 @@ class TestBatchSizeCap: _check_field("batch_size", 0) +class TestVisionImageSizeCap: + def test_none_accepts_model_default(self): + _check_field("vision_image_size", None) + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE], + ) + def test_in_range_accepts(self, value): + _check_field("vision_image_size", value) + assert _MIN_VISION_IMAGE_SIZE == 256 + assert _MAX_VISION_IMAGE_SIZE == 2048 + + @pytest.mark.parametrize( + "value", + [_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True], + ) + def test_invalid_rejects(self, value): + with pytest.raises(ValidationError): + _check_field("vision_image_size", value) + + @pytest.mark.parametrize("value", [True, False]) + def test_bool_error_says_integer_not_range(self, value): + # Regression guard: bools must say "integer or null", not "in [256, 2048]". + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + @pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"]) + def test_multi_sign_string_says_integer_not_raw(self, value): + # Regression guard: multi-sign strings must not leak int()'s raw + # "invalid literal" message; precise contract is "integer or null". + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + assert "invalid literal" not in str(exc.value) + + @pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"]) + def test_unicode_digit_string_rejected(self, value): + # Full-width / Arabic-Indic / Devanagari digits must be rejected so the + # value reaching the backend equals the ASCII the user typed. + with pytest.raises(ValidationError) as exc: + _check_field("vision_image_size", value) + assert "integer or null" in str(exc.value) + + class TestLoraRCap: def test_at_cap_accepts(self): _check_field("lora_r", _MAX_LORA_R) diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index e9fe1d679c..c0a60d0de2 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -108,6 +108,7 @@ export const LR_DEFAULT_CPT = 5e-5; export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, + visionImageSize: null as number | null, learningRate: LR_DEFAULT_LORA, // null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT. embeddingLearningRate: null as number | null, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 9749eba9b7..f254921fa6 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -175,12 +175,20 @@ export function ParamsSection(): ReactElement { const isCpt = store.trainingMethod === "cpt"; const isRawText = isRawTextDatasetFormat(store.datasetFormat); const showVisionLora = store.isVisionModel && store.isDatasetImage === true; + // DeepSeek OCR uses a coupled preset; backend ignores user image size. + const _selectedModelLower = (store.selectedModel ?? "").toLowerCase(); + const isDeepseekOcr = + _selectedModelLower.includes("deepseek") && + _selectedModelLower.includes("ocr"); + const showVisionImageSize = showVisionLora && !isDeepseekOcr; const [loraOpen, setLoraOpen] = useState(false); const [hyperOpen, setHyperOpen] = useState(false); const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen; const [ctxInput, setCtxInput] = useState(String(store.contextLength)); const ctxAnchorRef = useRef(null); const ctxItems = CONTEXT_LENGTHS.map(String); + // Backend validator allows [256, 2048]; offer the full span. + const visionImageSizePresets = [256, 384, 512, 768, 1024, 1536, 2048]; // Keep input in sync when the store value changes externally // (e.g. model defaults being applied after model selection). @@ -947,6 +955,62 @@ export function ParamsSection(): ReactElement { + {showVisionImageSize && ( + + Resize images by maximum side length. Default uses the + model image size. Larger images use up more context. Does not upscale or change aspect ratio.{" "} + + Read more + + + } + > + + + )} { - const yamlStr = serializeConfigToYaml(store, store.isVisionModel); + // isDatasetImage is null in three windows: before a dataset check + // completes, after dataset edits, and on import. Treat all three as + // "save it" so the user's choice is never silently dropped while we + // wait to confirm the dataset type. Only a confirmed text-only dataset + // (=== false) suppresses the vision fields. + const includeVisionFields = + store.isVisionModel && store.isDatasetImage !== false; + // DeepSeek OCR ignores vision_image_size; don't emit it to YAML either, + // or a later import on a non-DeepSeek model would activate the stale value. + const selectedModelLower = (store.selectedModel ?? "").toLowerCase(); + const isDeepseekOcr = + selectedModelLower.includes("deepseek") && + selectedModelLower.includes("ocr"); + const yamlStr = serializeConfigToYaml( + store, + includeVisionFields, + includeVisionFields && !isDeepseekOcr, + ); const blob = new Blob([yamlStr], { type: "text/yaml" }); const url = URL.createObjectURL(blob); const a = document.createElement("a"); diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 5d81f0df9c..1223888b11 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -23,7 +23,12 @@ export function buildTrainingStartPayload( const isCpt = config.trainingMethod === "cpt"; const adapterMethod = config.trainingMethod !== "full"; const isQloraMethod = config.trainingMethod === "qlora"; - const isFourBitModel = (config.selectedModel ?? "").toLowerCase().includes("4bit"); + const _selectedModelLower = (config.selectedModel ?? "").toLowerCase(); + const isFourBitModel = _selectedModelLower.includes("4bit"); + // DeepSeek OCR ignores user-selected image size; do not send it. + const isDeepseekOcr = + _selectedModelLower.includes("deepseek") && + _selectedModelLower.includes("ocr"); const isEmbedding = config.isEmbeddingModel; const isRawText = isRawTextDatasetFormat(config.datasetFormat); const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null; @@ -55,6 +60,10 @@ export function buildTrainingStartPayload( hf_token: config.hfToken.trim() || null, load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel), max_seq_length: config.contextLength, + vision_image_size: + config.isVisionModel && config.isDatasetImage === true && !isDeepseekOcr + ? config.visionImageSize + : null, trust_remote_code: config.trustRemoteCode ?? false, hf_dataset: hfDataset, subset: hfDataset ? config.datasetSubset : null, diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index e512b9e28d..8fecfaf7b8 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -27,6 +27,7 @@ interface BackendTrainingDefaults { eval_steps?: number; weight_decay?: number; random_seed?: number; + vision_image_size?: number | string | null; packing?: boolean; train_on_completions?: boolean; gradient_checkpointing?: "none" | "true" | "unsloth"; diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts index 8bc9c4e064..3a0a14b150 100644 --- a/studio/frontend/src/features/training/lib/model-defaults.ts +++ b/studio/frontend/src/features/training/lib/model-defaults.ts @@ -28,6 +28,7 @@ type ModelDefaultsPatch = Partial< | "trainOnCompletions" | "gradientCheckpointing" | "randomSeed" + | "visionImageSize" | "enableWandb" | "wandbProject" | "enableTensorboard" @@ -129,6 +130,25 @@ export function mapBackendModelConfigToTrainingPatch( const randomSeed = toNumber(training?.random_seed); if (randomSeed !== undefined) patch.randomSeed = randomSeed; + // Only patch when the config carries the key; model-switch reset lives in + // setSelectedModel so same-model reloads don't wipe a user's choice. + if (Object.hasOwn(training ?? {}, "vision_image_size")) { + const raw = training?.vision_image_size; + if (raw == null) { + patch.visionImageSize = null; + } else { + // Mirror studio/backend/models/training.py:_check_vision_image_size: + // drop anything outside [_MIN_VISION_IMAGE_SIZE, _MAX_VISION_IMAGE_SIZE] + // so the store/UI never show a value the backend would reject. + const n = toNumber(raw); + if (n !== undefined && Number.isInteger(n) && n >= 256 && n <= 2048) { + patch.visionImageSize = n; + } else { + patch.visionImageSize = null; + } + } + } + const packing = toBoolean(training?.packing); if (packing !== undefined) patch.packing = packing; diff --git a/studio/frontend/src/features/training/lib/yaml-config.ts b/studio/frontend/src/features/training/lib/yaml-config.ts index d168da5347..1deabf3547 100644 --- a/studio/frontend/src/features/training/lib/yaml-config.ts +++ b/studio/frontend/src/features/training/lib/yaml-config.ts @@ -27,8 +27,27 @@ export function parseYamlConfig(text: string): BackendModelConfig { console.warn("Ignored unknown YAML keys:", unknownKeys.join(", ")); } + // File import is authoritative: forge vision_image_size = null when the + // training section is missing, malformed, or missing the key, so a stale + // store value cannot survive an import. (Same-model defaults reloads + // preserve user choice via Object.hasOwn in model-defaults.ts.) + const rawTraining = raw.training; + const isPlainTrainingObject = + rawTraining != null && + typeof rawTraining === "object" && + !Array.isArray(rawTraining); + let trainingObj: Record; + if (!isPlainTrainingObject) { + trainingObj = { vision_image_size: null }; + } else { + trainingObj = { ...(rawTraining as Record) }; + if (!Object.hasOwn(trainingObj, "vision_image_size")) { + trainingObj.vision_image_size = null; + } + } + return { - training: (raw.training ?? undefined) as BackendModelConfig["training"], + training: trainingObj as BackendModelConfig["training"], lora: (raw.lora ?? undefined) as BackendModelConfig["lora"], logging: (raw.logging ?? undefined) as BackendModelConfig["logging"], }; @@ -41,6 +60,7 @@ export function parseYamlConfig(text: string): BackendModelConfig { export function serializeConfigToYaml( state: TrainingConfigState, includeVisionFields: boolean, + includeVisionImageSize: boolean = includeVisionFields, ): string { const lora: Record = { lora_r: state.loraRank, @@ -58,25 +78,31 @@ export function serializeConfigToYaml( lora.finetune_mlp_modules = state.finetuneMLPModules; } + const training: Record = { + max_seq_length: state.contextLength, + num_epochs: state.epochs, + learning_rate: state.learningRate, + batch_size: state.batchSize, + gradient_accumulation_steps: state.gradientAccumulation, + warmup_steps: state.warmupSteps, + max_steps: state.maxSteps, + save_steps: state.saveSteps, + eval_steps: state.evalSteps, + weight_decay: state.weightDecay, + random_seed: state.randomSeed, + packing: state.packing, + train_on_completions: state.trainOnCompletions, + gradient_checkpointing: state.gradientCheckpointing, + optim: state.optimizerType, + lr_scheduler_type: state.lrSchedulerType, + }; + + if (includeVisionImageSize) { + training.vision_image_size = state.visionImageSize; + } + const config = { - training: { - max_seq_length: state.contextLength, - num_epochs: state.epochs, - learning_rate: state.learningRate, - batch_size: state.batchSize, - gradient_accumulation_steps: state.gradientAccumulation, - warmup_steps: state.warmupSteps, - max_steps: state.maxSteps, - save_steps: state.saveSteps, - eval_steps: state.evalSteps, - weight_decay: state.weightDecay, - random_seed: state.randomSeed, - packing: state.packing, - train_on_completions: state.trainOnCompletions, - gradient_checkpointing: state.gradientCheckpointing, - optim: state.optimizerType, - lr_scheduler_type: state.lrSchedulerType, - }, + training, lora, }; 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 ef16f641f5..137669f8b8 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -390,6 +390,8 @@ export const useTrainingConfigStore = create()( error instanceof Error ? error.message : "Failed to load model defaults", + // Defaults load failed; reset so no prior model's value lingers. + visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize, }); // Fallback vision check if config endpoint fails. @@ -498,7 +500,16 @@ export const useTrainingConfigStore = create()( }, setSelectedModel: (selectedModel) => { const previousModel = get().selectedModel; - set({ selectedModel, modelDefaultsError: null }); + // Reset vision_image_size on a true switch only; same-model reloads + // go through the mapper, which preserves the user's choice. + const patch: { selectedModel: string | null; modelDefaultsError: null; visionImageSize?: number | null } = { + selectedModel, + modelDefaultsError: null, + }; + if (selectedModel !== previousModel) { + patch.visionImageSize = DEFAULT_HYPERPARAMS.visionImageSize; + } + set(patch); if (!selectedModel) { _modelConfigController?.abort(); @@ -701,6 +712,7 @@ export const useTrainingConfigStore = create()( }), setEpochs: (epochs) => set({ epochs }), setContextLength: (contextLength) => set({ contextLength }), + setVisionImageSize: (visionImageSize) => set({ visionImageSize }), setLearningRate: (learningRate) => { _learningRateManuallySet = true; set({ learningRate }); @@ -755,7 +767,10 @@ export const useTrainingConfigStore = create()( resetToModelDefaults: () => { const { selectedModel } = get(); if (!selectedModel) return; - set({ modelDefaultsAppliedFor: null }); + set({ + modelDefaultsAppliedFor: null, + visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize, + }); loadAndApplyModelDefaults(selectedModel); }, applyConfigPatch: (config: BackendModelConfig) => { diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 0cb881e634..5c12e71186 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -7,6 +7,7 @@ export interface TrainingStartRequest { hf_token: string | null; load_in_4bit: boolean; max_seq_length: number; + vision_image_size?: number | null; /** Allow loading models with custom code. Only enable for repos you trust. */ trust_remote_code?: boolean; hf_dataset: string | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 5b156316ca..f40f053e8b 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -82,6 +82,7 @@ export interface TrainingConfigState { finetuneMLPModules: boolean; targetModules: string[]; maxPositionEmbeddings: number | null; + visionImageSize: number | null; } export interface TrainingConfigActions { @@ -115,6 +116,7 @@ export interface TrainingConfigActions { setUploadedEvalFile: (file: string | null) => void; setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; + setVisionImageSize: (size: number | null) => void; setLearningRate: (rate: number) => void; setEmbeddingLearningRate: (rate: number | null) => void; setOptimizerType: (value: string) => void;