From dac2aeda1aa45271351d55fe90db77e8ea2195b2 Mon Sep 17 00:00:00 2001 From: Dariton4000 <184837391+Dariton4000@users.noreply.github.com> Date: Wed, 27 May 2026 14:01:24 +0200 Subject: [PATCH] Studio: expose image size setting in training UI (#5743) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Studio: add VLM image-size control for training Studio vision fine-tuning had no explicit way to cap image resolution, so users could not trade visual detail against context and memory use from the training UI, YAML config, or API payload. :) Add a nullable `vision_image_size` setting that keeps the current model default when unset and applies a max-side resize when provided. - Add `vision_image_size` to the training request model, route payload, backend training config, and frontend API/types plumbing. - Validate the value server-side as either null or an integer in the supported 256-2048 range. - Surface an Image Size selector for vision LoRA training with Default plus common preset sizes. - Include the value in training start payloads only for image-dataset vision models, and serialize it into vision-aware YAML configs. - Map backend model defaults back into the training store and reset the value when reapplying model defaults. - Pass the resize through the Torch trainer via `UnslothVisionDataCollator` using max-dimension semantics. - Apply the same max-dimension resize in the MLX VLM path before mlx-vlm's internal collation, preserving aspect ratio and avoiding upscaling. - Add backend validation coverage and MLX resize-size tests for the new behavior. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: thread vision_image_size into DeepSeek OCR + writable MLX ndarray - trainer.py: DeepSeek OCR collator now honors the new vision_image_size setting as image_size. Falls back to 640 when null. base_size stays at 1024 and crop_mode stays True so the Gundam preset's dynamic cropping of large documents keeps working. - worker.py: _resize_mlx_vlm_image returns np.array(image, copy=True) instead of np.asarray(image). The PIL view from np.asarray is not writable, which makes HF VLM processors emit "The given NumPy array is not writable, and PyTorch does not support non-writable tensors..." when they call torch.from_numpy. copy=True keeps the same shape and dtype but produces a writable buffer. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: align YAML export gate with API mapper + extend Image Size dropdown - training-section.tsx: handleSaveConfig now passes isVisionModel && isDatasetImage === true to serializeConfigToYaml, matching buildTrainingStartPayload. Stops vision_image_size from leaking into exported YAML for text-only datasets where the API would have sent null. - params-section.tsx: add 256 to visionImageSizePresets so the dropdown spans the validator's full [256, 2048] range. Also render a synthetic SelectItem for the current value when it was loaded from YAML or model defaults and is not in the preset list, so the controlled Select always shows the active size. * Studio: validate vision_image_size in YAML/model-default loader mapBackendModelConfigToTrainingPatch now mirrors the backend validator at studio/backend/models/training.py:169 by dropping any value that is not an integer in [256, 2048]. Pre-fix, an imported YAML like vision_image_size: 4096 or 640.5 would land in the store and the UI would happily display it, only to fail when Start Training posted to the backend. With this guard the store never holds a value the backend would reject. * Studio: precise error messages for invalid vision_image_size inputs Switch the field_validator to mode="before" so True/False surface as bool (not Pydantic's coerced 1/0) and give a precise "must be an integer or null" message instead of the misleading "must be in [256, 2048] (got 1)". Also explicitly accepts numpy Integral and integral Real scalars so YAML or programmatic callers using numpy ints keep working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: test that bool inputs yield the precise 'integer or null' error Regression guard for the validator switch to mode="before". Pre-fix, vision_image_size: True was rejected with "must be in [256, 2048] (got 1)" because Pydantic coerced before our check ran. New test asserts the message now reads "integer or null". * Studio: tighten vision_image_size loader + YAML save + MLX rounding Round 2 of follow-up review surfaced three usability issues: - model-defaults.ts: switching to a model whose backend YAML omits vision_image_size now explicitly resets the store value to null. Pre-fix, a stale 2048 from a previous model would silently apply to the new run because every checked-in model-default file omits the key. - training-section.tsx: handleSaveConfig now includes vision fields unless isDatasetImage is definitively false. isDatasetImage is null during dataset checks, after dataset edits, and on import; treating unknown as "drop" would silently lose the user's selection in those windows. Confirmed-text-only datasets still drop the value. - worker.py: _mlx_vlm_max_resized_size now mirrors the Torch collator's integer formula (w * size + size_func // 2) // size_func instead of Python round(), which uses banker's rounding and disagreed by 1px on half-pixel inputs like 333x1000 with target 500 (was 166, now 167). Test_mlx_training_worker_config gains parity assertions. * Studio: reset vision_image_size in the model-config error fallback path mapBackendModelConfigToTrainingPatch resets stale image size on the success path, but if the /api/models/config endpoint throws, training-config-store.ts falls through to checkVisionModel and only updates capability flags. Pre-fix that left a stale 2048 (or any prior selection) in the store, so once dataset detection marked the new dataset as image, the next training start would silently apply the previous model's size. The error branch now also resets to the DEFAULT_HYPERPARAMS.visionImageSize sentinel. * Studio: revert DeepSeek OCR Image Size knob + move missing-key reset Round 3 of the parallel-reviewer pass surfaced two issues that I had introduced earlier in this PR's follow-ups. - trainer.py: my prior change threaded vision_image_size into the DeepSeek OCR collator's image_size argument. The collator's (image_size, base_size, crop_mode) is a single preset (Tiny / Small / Base / Large / Gundam); changing image_size in isolation desynchronizes the per-crop pixel grid from num_queries downstream and produces wrong token grids on documents larger than the per-crop tile. The fix pins the collator back at the Gundam preset and logs a clear "ignored for DeepSeek OCR" notice when the user has selected a non-default Image Size. - model-defaults.ts + training-config-store.ts: the round 4 fix that reset visionImageSize when a model YAML omitted the key also fired on same-model reloads (ensureModelDefaultsLoaded re-fires on page refresh), wiping a value the user had just selected. The reset is now in setSelectedModel, gated on selectedModel != previousModel, so true model switches still clear stale values while reloads keep the user's selection. * Studio: extend DeepSeek OCR Image Size exclusion to MLX + frontend Round 4 of the parallel-reviewer pass flagged that the Torch trainer exclusion I added did not have a matching MLX guard, and that the UI still offered the dropdown for DeepSeek OCR even though the backend ignores it. - worker.py: _run_mlx_training now mirrors the Torch exclusion. When the model name matches DeepSeek OCR, vision_image_size is forced back to None before _adapt_for_mlx_vlm sees it, so dataset images pass through unchanged just like the Torch path. Emits a clear status line when this happens. - params-section.tsx: the Image Size Row is now gated on showVisionImageSize (showVisionLora && !isDeepseekOcr) instead of showVisionLora alone, so DeepSeek OCR users no longer see a control that silently has no effect. - mappers.ts: buildTrainingStartPayload sends null for vision_image_size whenever the selected model is DeepSeek OCR, so the backend log line about ignoring the value never fires from a UI-driven start. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: tighten YAML import/save for vision_image_size Two YAML-path asymmetries that could leak a stale image size into training: - parseYamlConfig now treats a missing training.vision_image_size as null. Without this, importing a YAML saved before this feature (or any config that omits the key) preserved whatever value the user had previously set on a different model. The model-defaults reload path still uses Object.hasOwn so same-model defaults reloads do not wipe a manual selection; only file import normalises the missing key. - handleSaveConfig now passes a DeepSeek-OCR-specific guard to serializeConfigToYaml so saved YAML matches what the API mapper actually sends. Previously a state with visionImageSize set could emit the key even though Studio ignored it at training time for DeepSeek OCR, and a later import for a non-DeepSeek vision model would activate the stale value. serializeConfigToYaml gains an optional third parameter includeVisionImageSize defaulting to includeVisionFields, preserving the existing 2-arg call signature for backwards compatibility. * Studio: also reset vision_image_size when YAML lacks a training section Round 9's parseYamlConfig normalization only fired when the YAML had a training mapping that omitted vision_image_size. A lora-only or logging-only YAML (or one with `training: null`) still left trainingObj unset, the mapper saw no vision_image_size key, and the previously selected store value persisted into the next training run. Now an absent or null training section is synthesised as { vision_image_size: null } so model-defaults.ts always patches visionImageSize back to Default on file import. Same-model defaults reloads still preserve manual choices via the existing Object.hasOwn gate in mapBackendModelConfigToTrainingPatch. * Studio: unify parseYamlConfig non-object training handling A fresh static review (Opus subagent) flagged P3-1: parseYamlConfig only synthesised vision_image_size: null when raw.training was either absent or a plain object missing the key. If raw.training is a scalar or an array (malformed but still parseable), the value was passed through unchanged, the mapper's Object.hasOwn returned false, and any previously selected visionImageSize persisted - the same stale-state leak the lora-only fallback was added to close. Treat any non-plain-object raw.training (null, array, scalar) as a malformed/missing section and reset to { vision_image_size: null }. * Studio: tighten code comments for vision_image_size path * Studio: tighten vision_image_size validator + restore lost comment context Two issues surfaced by a fresh adversarial review of the validator: 1. v.strip().lstrip("+-").isdigit() let "++512" / "--256" / "+-+512" slip past the gate, then int("++512") raised an uncaught ValueError and Pydantic surfaced "invalid literal for int() with base 10: '++512'" instead of the contracted "vision_image_size must be an integer or null". 2. str.isdigit() returns True for Unicode digit families (full-width '５１２', Arabic-Indic '٥١٢', Devanagari '१०२४'), and int() coerces them, so the value reaching the backend wasn't the ASCII the user typed. Replaced the lstrip+isdigit pair with re.fullmatch(r'[+-]?[0-9]+', stripped), which rejects both shapes with the precise error and accepts the documented ones ('256', '+512', ' 1024 '). Added 8 regression test cases covering multi-sign strings, lone sign, and the three Unicode digit families. Also restored comment context lost in f9c39331: - model-defaults.ts: name studio/backend/models/training.py:_check_vision_image_size as the spec the [256, 2048] range mirrors, so a maintainer changing the cap in one file can find the other. - training-section.tsx: enumerate the three windows in which isDatasetImage is null (before a check, after dataset edits, on import) so a future maintainer doesn't simplify the gate to `isCheckingDataset`. - worker.py: qualify the writable-ndarray comment with "when a resize is requested" so it doesn't misadvertise the resize=None early-return. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/training/trainer.py | 24 +++++- studio/backend/core/training/training.py | 1 + studio/backend/core/training/worker.py | 78 +++++++++++++++++-- studio/backend/models/training.py | 48 ++++++++++++ studio/backend/routes/training.py | 1 + .../tests/test_mlx_training_worker_config.py | 12 +++ .../tests/test_studio_train_validation.py | 48 ++++++++++++ studio/frontend/src/config/training.ts | 1 + .../studio/sections/params-section.tsx | 64 +++++++++++++++ .../studio/sections/training-section.tsx | 19 ++++- .../src/features/training/api/mappers.ts | 11 ++- .../src/features/training/api/models-api.ts | 1 + .../features/training/lib/model-defaults.ts | 20 +++++ .../src/features/training/lib/yaml-config.ts | 64 ++++++++++----- .../training/stores/training-config-store.ts | 19 ++++- .../src/features/training/types/api.ts | 1 + .../src/features/training/types/config.ts | 2 + 17 files changed, 384 insertions(+), 30 deletions(-) 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;