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

studio: fix 'Standard' gradient checkpointing leaving the smart offloader active

Browse source 
The training UI's 'Standard' option sends the string "true". The worker
converted "false"/"none"/"" to False but passed "true" through as a string.
unsloth's _configure_gradient_checkpointing only unpatches the smart
offloader on the (True, False) boolean branch; the string "true" matches
neither that nor "unsloth", so it returns without unpatching and the
offloader stays active. On memory-constrained / unified-memory GPUs that
offload then OOM-crashes even though the user asked for standard
checkpointing (reported on gfx1201 R9700, applies to Strix Halo too).

Normalize "true"/"1"/"yes" -> True so unsloth gets a real bool and
unpatches, matching the existing normalization in trainer.py. "unsloth"
and "mlx" stay strings; booleans pass through unchanged.
This commit is contained in:
Daniel Han 2026-07-21 03:52:21 -07:00
commit 3b16e317b8
1 changed files with 13 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -1607,9 +1607,19 @@ def _run_mlx_training(event_queue, stop_queue, config):
# get_peft_model and MLXTrainer both accept and handle strings.
gc_setting = config.get("gradient_checkpointing", "mlx")
if isinstance(gc_setting, str):
use_grad_checkpoint = (
gc_setting if gc_setting.lower() not in ("false", "none", "") else False
)
_gc = gc_setting.lower()
if _gc in ("false", "none", ""):
use_grad_checkpoint = False
elif _gc in ("true", "1", "yes"):
# HF-standard GC (UI "Standard"): pass a real bool so unsloth
# unpatches the smart offloader. Leaving it the string "true" hits
# neither the "unsloth" nor the (True, False) branch in
# _configure_gradient_checkpointing, so the offloader stays patched
# and active against the user's choice (memory crash on
# constrained/unified GPUs).
use_grad_checkpoint = True
else:
use_grad_checkpoint = gc_setting # "unsloth" / "mlx" stay strings
else:
use_grad_checkpoint = gc_setting