diff --git a/pyproject.toml b/pyproject.toml index e884ad8..9058012 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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", diff --git a/requirements.txt b/requirements.txt index fc82b69..73d2db0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/src/optimization/blockswap.py b/src/optimization/blockswap.py index 6cb1d57..e3b967c 100644 --- a/src/optimization/blockswap.py +++ b/src/optimization/blockswap.py @@ -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 diff --git a/src/optimization/memory_manager.py b/src/optimization/memory_manager.py index 10cd5d1..71fa501 100644 --- a/src/optimization/memory_manager.py +++ b/src/optimization/memory_manager.py @@ -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")