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

Studio: add VLM image-size control for training

Browse source 
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.
This commit is contained in:
Dariton4000 2026-05-23 20:48:55 +02:00
commit 7cc8beb2e6
16 changed files with 216 additions and 26 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -3123,7 +3123,17 @@ 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 ==========

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

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

62
studio/backend/core/training/worker.py
View file

@ -959,7 +959,43 @@ 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
scale = float(target) / float(largest_side)
return max(1, int(round(width * scale))), max(1, int(round(height * scale)))
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)
# mlx-vlm's internal collator square-resizes PIL images. Return an ndarray
# so Studio's max-dimension resize is the final resize, like trainer.py.
return np.asarray(image)
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 +1015,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 +1026,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 +1204,13 @@ 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")
if 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 +1345,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 +1363,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 +2297,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,

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

@ -18,6 +18,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 +61,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 +166,20 @@ class TrainingStartRequest(BaseModel):
)
return v
@field_validator("vision_image_size")
@classmethod
def _check_vision_image_size(cls, v: Optional[int]) -> Optional[int]:
if v is None:
return v
if isinstance(v, bool) or not isinstance(v, int):
raise ValueError("vision_image_size must be an integer or null")
if v < _MIN_VISION_IMAGE_SIZE or v > _MAX_VISION_IMAGE_SIZE:
raise ValueError(
f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, "
f"{_MAX_VISION_IMAGE_SIZE}] (got {v!r})"
)
return v
@field_validator("warmup_steps")
@classmethod
def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]:

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

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

9
studio/backend/tests/test_mlx_training_worker_config.py
View file

@ -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,11 @@ 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)

24
studio/backend/tests/test_studio_train_validation.py
View file

@ -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,28 @@ 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)
class TestLoraRCap:
def test_at_cap_accepts(self):
_check_field("lora_r", _MAX_LORA_R)

1
studio/frontend/src/config/training.ts
View file

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

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

@ -139,6 +139,7 @@ export function ParamsSection(): ReactElement {
const [ctxInput, setCtxInput] = useState(String(store.contextLength));
const ctxAnchorRef = useRef<HTMLDivElement>(null);
const ctxItems = CONTEXT_LENGTHS.map(String);
const visionImageSizePresets = [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).
@ -915,6 +916,52 @@ export function ParamsSection(): ReactElement {
</TabsContent>
<TabsContent value="memory" className="mt-3 flex flex-col gap-3">
{showVisionLora && (
<Row
label="Image Size"
tooltip={
<>
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.{" "}
<a
href="https://unsloth.ai/docs/basics/vision-fine-tuning"
target="_blank"
rel="noopener noreferrer"
className="text-primary underline"
>
Read more
</a>
</>
}
>
<Select
value={
store.visionImageSize == null
? "default"
: String(store.visionImageSize)
}
onValueChange={(value) => {
if (value === "default") {
store.setVisionImageSize(null);
return;
}
store.setVisionImageSize(Number(value));
}}
>
<SelectTrigger className="w-32">
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="default">Default</SelectItem>
{visionImageSizePresets.map((size) => (
<SelectItem key={size} value={String(size)}>
{size}
</SelectItem>
))}
</SelectContent>
</Select>
</Row>
)}
<Row
label="Grad Checkpoint"
tooltip={

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

@ -55,6 +55,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
? config.visionImageSize
: null,
trust_remote_code: config.trustRemoteCode ?? false,
hf_dataset: hfDataset,
subset: hfDataset ? config.datasetSubset : null,

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

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

8
studio/frontend/src/features/training/lib/model-defaults.ts
View file

@ -28,6 +28,7 @@ type ModelDefaultsPatch = Partial<
| "trainOnCompletions"
| "gradientCheckpointing"
| "randomSeed"
| "visionImageSize"
| "enableWandb"
| "wandbProject"
| "enableTensorboard"
@ -129,6 +130,13 @@ export function mapBackendModelConfigToTrainingPatch(
const randomSeed = toNumber(training?.random_seed);
if (randomSeed !== undefined) patch.randomSeed = randomSeed;
if (Object.hasOwn(training ?? {}, "vision_image_size")) {
const visionImageSize = training?.vision_image_size == null
? null
: toNumber(training.vision_image_size);
if (visionImageSize !== undefined) patch.visionImageSize = visionImageSize;
}
const packing = toBoolean(training?.packing);
if (packing !== undefined) patch.packing = packing;

42
studio/frontend/src/features/training/lib/yaml-config.ts
View file

@ -58,25 +58,31 @@ export function serializeConfigToYaml(
lora.finetune_mlp_modules = state.finetuneMLPModules;
}
const training: Record<string, unknown> = {
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 (includeVisionFields) {
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,
};

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

@ -701,6 +701,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
}),
setEpochs: (epochs) => set({ epochs }),
setContextLength: (contextLength) => set({ contextLength }),
setVisionImageSize: (visionImageSize) => set({ visionImageSize }),
setLearningRate: (learningRate) => {
_learningRateManuallySet = true;
set({ learningRate });
@ -755,7 +756,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
resetToModelDefaults: () => {
const { selectedModel } = get();
if (!selectedModel) return;
set({ modelDefaultsAppliedFor: null });
set({
modelDefaultsAppliedFor: null,
visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize,
});
loadAndApplyModelDefaults(selectedModel);
},
applyConfigPatch: (config: BackendModelConfig) => {

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

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

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

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