Compare commits

...
Sign in to create a new pull request.

1 commit

Author SHA1 Message Date
Daniel Han
3b16e317b8 studio: fix 'Standard' gradient checkpointing leaving the smart offloader active
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.
2026-07-21 03:52:21 -07:00

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