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

fix: unload inference model before training to free GPU memory

Browse source 
When starting training, shut down the inference subprocess first
so the training subprocess has full GPU memory available.
This commit is contained in:
Roland Tannous 2026-03-05 22:28:11 +00:00
commit cbe2896705
1 changed files with 16 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

16
studio/backend/routes/training.py
View file

@ -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)