diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index fb5e9b5137..427687451c 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1389,6 +1389,12 @@ def _run_mlx_training(event_queue, stop_queue, config): # Force text-only for non-image datasets even on vision-capable models # (e.g. Qwen3.5-VL trained on plain alpaca text). _send("status", status_message = f"Loading {model_name}...") + # Pull through resume_from_checkpoint so MLXTrainer.train() can restore + # optimizer + step state and continue cleanly. Was previously dropped on + # the floor for the MLX path, so the Resume UI button silently restarted + # from step 0 (the CUDA path at lines 2729 / 3108 has been forwarding + # this all along). + resume_from_checkpoint = config.get("resume_from_checkpoint") or None is_dataset_image = bool(config.get("is_dataset_image", False)) training_type = config.get("training_type", "LoRA/QLoRA") use_lora = training_type == "LoRA/QLoRA" @@ -1852,7 +1858,7 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train() + trainer.train(resume_from_checkpoint=resume_from_checkpoint) # ── 12. Save and finalize ── if trainer.stop_requested and not _stop_save[0]: