dac2aeda1a* 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 <danielhanchen@gmail.com>
138 lines
4.6 KiB
Python
138 lines
4.6 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
||
|
||
"""Pin TrainingStartRequest hyperparameter caps at the at-cap / over-cap boundary."""
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
from pydantic import ValidationError
|
||
|
||
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
||
if str(_BACKEND_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(_BACKEND_ROOT))
|
||
|
||
from models.training import (
|
||
_MAX_BATCH_SIZE,
|
||
_MAX_LORA_ALPHA,
|
||
_MAX_LORA_R,
|
||
_MAX_SEQ_LENGTH,
|
||
_MAX_VISION_IMAGE_SIZE,
|
||
_MIN_VISION_IMAGE_SIZE,
|
||
)
|
||
|
||
|
||
def _check_field(field_name: str, value):
|
||
"""Run the field validator without constructing a full TrainingStartRequest."""
|
||
from models.training import TrainingStartRequest
|
||
|
||
schema_field = TrainingStartRequest.model_fields[field_name]
|
||
return TrainingStartRequest.__pydantic_validator__.validate_assignment(
|
||
TrainingStartRequest.model_construct(),
|
||
field_name,
|
||
value,
|
||
)
|
||
|
||
|
||
class TestSeqLengthCap:
|
||
def test_at_cap_accepts(self):
|
||
_check_field("max_seq_length", _MAX_SEQ_LENGTH)
|
||
assert _MAX_SEQ_LENGTH == 2_000_000
|
||
|
||
def test_over_cap_rejects(self):
|
||
with pytest.raises(ValidationError) as exc:
|
||
_check_field("max_seq_length", _MAX_SEQ_LENGTH + 1)
|
||
assert "max_seq_length" in str(exc.value)
|
||
|
||
def test_below_min_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("max_seq_length", 0)
|
||
|
||
|
||
class TestBatchSizeCap:
|
||
def test_at_cap_accepts(self):
|
||
_check_field("batch_size", _MAX_BATCH_SIZE)
|
||
assert _MAX_BATCH_SIZE == 4096
|
||
|
||
def test_over_cap_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("batch_size", _MAX_BATCH_SIZE + 1)
|
||
|
||
def test_below_min_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_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)
|
||
assert _MAX_LORA_R == 16_384
|
||
|
||
def test_over_cap_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("lora_r", _MAX_LORA_R + 1)
|
||
|
||
def test_below_min_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("lora_r", 0)
|
||
|
||
|
||
class TestLoraAlphaCap:
|
||
def test_at_cap_accepts(self):
|
||
_check_field("lora_alpha", _MAX_LORA_ALPHA)
|
||
assert _MAX_LORA_ALPHA == 32_768
|
||
|
||
def test_over_cap_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("lora_alpha", _MAX_LORA_ALPHA + 1)
|
||
|
||
def test_below_min_rejects(self):
|
||
with pytest.raises(ValidationError):
|
||
_check_field("lora_alpha", 0)
|