From ae8bb7cf2c345a13c2cf975fab20b9249bf359d9 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 28 Sep 2025 11:14:34 -0500 Subject: [PATCH] feat: Refine model retention logic in multi-GPU unloading - Modified condition to retain models lacking `_mgpu_keep_loaded` attribute or with `keep_loaded=True` - Improves reliability of unloading by distinguishing DisTorch and non-DisTorch models - Addresses potential premature unloading of intended persistent models in multi-GPU setups --- model_management_mgpu.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/model_management_mgpu.py b/model_management_mgpu.py index 64bf315..c2f3a93 100644 --- a/model_management_mgpu.py +++ b/model_management_mgpu.py @@ -399,7 +399,12 @@ if hasattr(mm, 'unload_all_models') and not hasattr(mm.unload_all_models, '_mgpu model_name = type(getattr(mp, 'model', mp)).__name__ logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={keep_loaded}") - if keep_loaded: + # Retain models that either: + # 1. Are non-DisTorch models (missing _mgpu_keep_loaded attribute) + # 2. Are DisTorch models with keep_loaded=True + should_retain = not hasattr(mp.model, '_mgpu_keep_loaded') or keep_loaded + + if should_retain: kept_models.append(lm) logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Adding to kept_models: {model_name}") # GC ANCHOR TEST: Prevent premature GC of clone patchers