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:
Adrien Toupet
2025-11-03 17:29:58 -05:00
parent baf0058924
commit a900af96fd
4 changed files with 8 additions and 6 deletions
+1
View File
@@ -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",
+1
View File
@@ -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
+2 -2
View File
@@ -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
+4 -4
View File
@@ -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")