diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 1b66306842..6b335bc139 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -63,6 +63,32 @@ async def load_checkpoint( try: # Version switching is handled automatically by the subprocess-based # export backend — no need for ensure_transformers_version() here. + + # Free GPU memory: shut down any running inference/training subprocesses + # before loading the export checkpoint (they'd compete for VRAM). + try: + from core.inference import get_inference_backend + inf = get_inference_backend() + if inf.active_model_name: + logger.info( + "Unloading inference model '%s' to free GPU memory for export", + inf.active_model_name, + ) + inf._shutdown_subprocess() + inf.active_model_name = None + inf.models.clear() + except Exception as e: + logger.warning("Could not unload inference model: %s", e) + + try: + from core.training import get_training_backend + trn = get_training_backend() + if trn.is_training_active(): + logger.info("Stopping active training to free GPU memory for export") + trn.stop_training() + except Exception as e: + logger.warning("Could not stop training: %s", e) + backend = get_export_backend() success, message = backend.load_checkpoint( checkpoint_path=request.checkpoint_path, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 521ff79d9a..156b3237a4 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -163,6 +163,19 @@ async def load_model( logger.info("Unloading GGUF model before loading Unsloth model") llama_backend.unload_model() + # Shut down any export subprocess to free VRAM + try: + from core.export import get_export_backend + exp_backend = get_export_backend() + if exp_backend.current_checkpoint: + logger.info("Shutting down export subprocess to free GPU memory for inference") + exp_backend._shutdown_subprocess() + exp_backend.current_checkpoint = None + exp_backend.is_vision = False + exp_backend.is_peft = False + except Exception as e: + logger.warning("Could not shut down export subprocess: %s", e) + # Auto-detect quantization for LoRA adapters from adapter_config.json # The training pipeline patches this file with "unsloth_training_method" # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False. diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index ae49706fc9..c0ee6deff7 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -192,7 +192,7 @@ async def start_training( "tensorboard_dir": request.tensorboard_dir or "", } - # Free GPU memory: shut down any running inference subprocess + # Free GPU memory: shut down any running inference/export subprocesses # before training starts (they'd compete for VRAM otherwise) try: from core.inference import get_inference_backend @@ -208,6 +208,18 @@ async def start_training( except Exception as e: logger.warning("Could not unload inference model: %s", e) + try: + from core.export import get_export_backend + exp_backend = get_export_backend() + if exp_backend.current_checkpoint: + logger.info("Shutting down export subprocess to free GPU memory for training") + exp_backend._shutdown_subprocess() + exp_backend.current_checkpoint = None + exp_backend.is_vision = False + exp_backend.is_peft = False + except Exception as e: + logger.warning("Could not shut down export subprocess: %s", e) + # start_training now spawns a subprocess (non-blocking) success = backend.start_training(**training_kwargs)