optimize: Implement adaptive GPU cache clearing for BlockSwap efficiency
BlockSwap now uses minimal, pressure-based cache clearing instead of forced deep cleans: - Clear only when VRAM < 5% free (was 15%, more aggressive threshold) - Use minimal GPU cache clear (deep=False) for reduced overhead - Applied after each block/IO swap when pressure detected Also fixes peft dependency: Add explicit peft>=0.17.0 in requirements for diffusers>=0.33.1 compatibility in ComfyUI ecosystem where other plugins may install outdated peft versions.
This commit is contained in:
@@ -21,6 +21,7 @@ dependencies = [
|
||||
"einops",
|
||||
"omegaconf>=2.3.0",
|
||||
"diffusers>=0.33.1",
|
||||
"peft>=0.17.0",
|
||||
"rotary_embedding_torch>=0.5.3",
|
||||
"opencv-python",
|
||||
"gguf",
|
||||
|
||||
@@ -7,6 +7,7 @@ psutil
|
||||
einops
|
||||
omegaconf>=2.3.0
|
||||
diffusers>=0.33.1
|
||||
peft>=0.17.0
|
||||
rotary_embedding_torch>=0.5.3
|
||||
opencv-python
|
||||
gguf
|
||||
|
||||
@@ -444,7 +444,7 @@ def _wrap_block_forward(
|
||||
_log_swap_timing(debug, t_start, self._block_idx, "block")
|
||||
|
||||
# Only clear cache under memory pressure
|
||||
clear_memory(debug=debug, deep=True, force=False, timer_name="wrap_block_forward")
|
||||
clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward")
|
||||
else:
|
||||
output = original_forward(*args, **kwargs)
|
||||
|
||||
@@ -523,7 +523,7 @@ def _wrap_io_forward(
|
||||
_log_swap_timing(debug, t_start, self._module_name, "I/O")
|
||||
|
||||
# Only clear cache under memory pressure
|
||||
clear_memory(debug=debug, deep=True, force=False, timer_name="wrap_block_forward")
|
||||
clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward")
|
||||
|
||||
return output
|
||||
|
||||
|
||||
@@ -191,7 +191,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
|
||||
|
||||
Args:
|
||||
debug: Debug instance for logging (optional)
|
||||
force: If True, always clear. If False, only clear when <15% free
|
||||
force: If True, always clear. If False, only clear when <5% free
|
||||
deep: If True, perform deep cleanup including GC and OS operations.
|
||||
If False (default), only perform minimal GPU cache clearing.
|
||||
timer_name: Optional suffix for timer names to make them unique per invocation
|
||||
@@ -230,9 +230,9 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
|
||||
mem_info = get_basic_vram_info(device=None)
|
||||
|
||||
if "error" not in mem_info:
|
||||
# Check VRAM/MPS memory pressure (15% free threshold)
|
||||
# Check VRAM/MPS memory pressure (5% free threshold)
|
||||
free_ratio = mem_info["free_gb"] / mem_info["total_gb"]
|
||||
if free_ratio < 0.15:
|
||||
if free_ratio < 0.05:
|
||||
should_clear = True
|
||||
if debug:
|
||||
backend = "MPS" if torch.mps.is_available() else "VRAM"
|
||||
@@ -241,7 +241,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
|
||||
# For non-MPS systems, also check system RAM separately
|
||||
if not should_clear and not torch.mps.is_available():
|
||||
mem = psutil.virtual_memory()
|
||||
if mem.available < mem.total * 0.15:
|
||||
if mem.available < mem.total * 0.05:
|
||||
should_clear = True
|
||||
if debug:
|
||||
debug.log(f"RAM pressure: {mem.available/(1024**3):.2f}GB free of {mem.total/(1024**3):.2f}GB", category="memory")
|
||||
|
||||
Reference in New Issue
Block a user