refactor: simplify model retention logic in multi-GPU unload

- Renamed `keep_loaded` variable to `should_retain` for improved clarity
- Simplified assignment by directly retrieving `_mgpu_keep_loaded` attribute with default False
- Updated logging accordingly; may alter behavior for non-DisTorch models to no longer retain automatically
This commit is contained in:
John Pollock
2025-09-28 13:22:19 -05:00
parent 0d056141c0
commit 18493f5277
+2 -3
View File
@@ -395,14 +395,13 @@ if hasattr(mm, 'unload_all_models') and not hasattr(mm.unload_all_models, '_mgpu
mp = lm.model # weakref call to ModelPatcher
if mp is not None and hasattr(mp, 'model'):
# Check if this is a DisTorch model with keep_loaded flag
keep_loaded = getattr(mp.model, '_mgpu_keep_loaded', False)
should_retain = getattr(mp.model, '_mgpu_keep_loaded', False)
model_name = type(getattr(mp, 'model', mp)).__name__
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={keep_loaded}")
logger.mgpu_mm_log(f"[UNLOAD_DEBUG] Model {i}: {model_name}, keep_loaded={should_retain}")
# 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)