From 3b16e317b893d197c1bb48ae2e27f328d6c88c8a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 21 Jul 2026 03:52:21 -0700 Subject: [PATCH] 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. --- studio/backend/core/training/worker.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 2df4fa58c6..5ae838a2b2 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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