From 0ac051177ab03c597af64d679b2063f05443a397 Mon Sep 17 00:00:00 2001 From: BardiaKoopah Date: Thu, 11 Jun 2026 07:29:05 -0700 Subject: [PATCH] fix(mlx): forward resume_from_checkpoint to MLXTrainer.train() (#6173) Studio's frontend exposes a Resume action and submits requests with resume_from_checkpoint set to a previous run's output_dir. The CUDA training paths in worker.py read this field from config and pass it to trainer.train() (see lines 2729-2787 and 3108-3229). The MLX path _run_mlx_training did neither: it never read config['resume_from_checkpoint'] and called trainer.train() with no args. The MLX trainer also did not accept the kwarg, so even threading it through would have been a no-op. With this PR + the unsloth-zoo companion PR adding the trainer-side support (saves optimizer_state + trainer_state, accepts and applies resume_from_checkpoint in MLXTrainer.train()), MLX Resume now works end-to-end. Verified on M2 16GB with Qwen3-0.6B + unsloth/LaTeX_OCR: loss at every post-resume step matches a fresh run bit for bit (2.168627977371216 == 2.168627977371216 at step 6, etc). Two lines: read the field near the other config.get() extractions in _run_mlx_training, pass it as a kwarg at the trainer.train() call site. Companion PR: unslothai/unsloth-zoo#751 --- studio/backend/core/training/worker.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) 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]: