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

fix(mlx): forward resume_from_checkpoint to MLXTrainer.train() (#6173)

Browse source 
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
This commit is contained in:
BardiaKoopah 2026-06-11 07:29:05 -07:00 • committed by GitHub
commit 0ac051177a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -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]: