From cbe289670532148d11c085852dd3d99fa688a91b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 5 Mar 2026 22:28:11 +0000 Subject: [PATCH] fix: unload inference model before training to free GPU memory When starting training, shut down the inference subprocess first so the training subprocess has full GPU memory available. --- studio/backend/routes/training.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 58ee3a075c..ae49706fc9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -192,6 +192,22 @@ async def start_training( "tensorboard_dir": request.tensorboard_dir or "", } + # Free GPU memory: shut down any running inference subprocess + # before training starts (they'd compete for VRAM otherwise) + try: + from core.inference import get_inference_backend + inf_backend = get_inference_backend() + if inf_backend.active_model_name: + logger.info( + "Unloading inference model '%s' to free GPU memory for training", + inf_backend.active_model_name, + ) + inf_backend._shutdown_subprocess() + inf_backend.active_model_name = None + inf_backend.models.clear() + except Exception as e: + logger.warning("Could not unload inference model: %s", e) + # start_training now spawns a subprocess (non-blocking) success = backend.start_training(**training_kwargs)