diff --git a/nodes/lora_merger.py b/nodes/lora_merger.py index 997b1be..7ef5109 100644 --- a/nodes/lora_merger.py +++ b/nodes/lora_merger.py @@ -48,7 +48,7 @@ def merge_multi_loras( prepare_for_large_operation(total_size_gb * 2.5, torch.device(device)) - handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths] + handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths] try: # 1. Analyze all LoRAs @@ -190,6 +190,12 @@ def merge_multi_loras( del merged_down, merged_up for d, _ in downs: del d for u, _ in ups: del u + downs.clear() + ups.clear() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() pbar.update(1) @@ -412,7 +418,7 @@ def merge_multi_loras_dare( print(f"[LoRA Multi-Merge DARE] Drop rate: {drop_rate}, Trim quantile: {trim_quantile}") prepare_for_large_operation(total_size_gb * 2.5, torch.device(device)) - handlers = [MemoryEfficientSafeOpen(p) for p in lora_paths] + handlers = [MemoryEfficientSafeOpen(p, low_memory=True) for p in lora_paths] rng = torch.Generator(device=device).manual_seed(seed) @@ -523,6 +529,16 @@ def merge_multi_loras_dare( output_sd[f"{core}.lora_up.weight"] = merged_up.to(save_dtype).cpu().contiguous() output_sd[f"{core}.alpha"] = torch.tensor(float(max_rank), dtype=save_dtype) + del merged_down, merged_up + for d, _ in downs: del d + for u, _ in ups: del u + downs.clear() + ups.clear() + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + pbar.update(1) # Final Summary diff --git a/nodes/lora_resize.py b/nodes/lora_resize.py index eef65b0..07b5876 100644 --- a/nodes/lora_resize.py +++ b/nodes/lora_resize.py @@ -461,6 +461,12 @@ def resize_lora_file( if alpha_suffix: output_sd[f"{new_block_name}{alpha_suffix}"] = torch.tensor(result["new_alpha"]).to(save_dtype) + del result + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + pbar.update(1) if verbose and fro_list: @@ -854,6 +860,10 @@ def merge_loras_to_model( del cpu_base stats["copied"] += 1 + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() pbar.update(1) diff --git a/nodes/merger.py b/nodes/merger.py index 2efa2ad..36de235 100644 --- a/nodes/merger.py +++ b/nodes/merger.py @@ -77,18 +77,17 @@ class MergerLogic: print(f"[Merger] Preparing memory for {total_size_gb:.2f}GB merge operation...") prepare_for_large_operation(total_size_gb * 1.2, torch.device(process_device)) - lazy_load = recipe_params.get('lazy_load', True) handlers = {} for name in model_names.values(): if name and name != "None": path = folder_paths.get_full_path(model_type, name) if not path: raise FileNotFoundError(f"Model '{name}' not found.") - handlers[name] = MemoryEfficientSafeOpen(path, low_memory=lazy_load) + handlers[name] = MemoryEfficientSafeOpen(path, low_memory=True) primary_handler = handlers[primary_model_name] all_keys = primary_handler.keys() - metadata = primary_handler.header.get("__metadata__", {}) + metadata = primary_handler.metadata() # Convert mismatch_mode string to enum mismatch_mode_str = recipe_params.get('mismatch_mode', 'skip') @@ -125,72 +124,78 @@ class MergerLogic: discarded_keys = 0 error_keys = [] - for key in tqdm(all_keys, desc="Merging layers", unit="layers"): - # Check discard patterns first - skip entirely - if _matches_any_pattern(key, discard_patterns): - discarded_keys += 1 + with torch.no_grad(): + for key in tqdm(all_keys, desc="Merging layers", unit="layers"): + # Check discard patterns first - skip entirely + if _matches_any_pattern(key, discard_patterns): + discarded_keys += 1 + pbar.update(1) + continue + + # Pre-load Model A's tensor with pinned memory for CUDA + cpu_tensor = primary_handler.get_tensor(key) + if process_device == 'cuda': + tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype) + else: + tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype) + del cpu_tensor + + # Check exclude patterns - use Model A only, no merge + if _matches_any_pattern(key, exclude_patterns): + merged_state_dict[key] = tensor_a.detach().to(save_torch_dtype).cpu().clone() + excluded_keys += 1 + pbar.update(1) + continue + + # Pass tensor_a metadata to recipes for zeros mode and fallback + recipe_params['_tensor_a'] = tensor_a + recipe_params['_tensor_a_shape'] = tensor_a.shape + recipe_params['_tensor_a_dtype'] = tensor_a.dtype + + try: + recipe = calc_mode_class.create_recipe(key=key, **recipe_params) + result = recipe.merge() + except MissingTensorError as e: + if mismatch_mode == MissingTensorBehavior.ERROR: + raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}") + result = None + error_keys.append(key) + + # Handle None result (mismatch occurred with skip mode) + if result is None: + result = tensor_a + skipped_keys += 1 + + if isinstance(result, dict): + for r_key, r_tensor in result.items(): + merged_state_dict[r_key] = r_tensor.detach().to(save_torch_dtype).cpu().clone() + else: + # Ensure compatibility with Model A's architecture. + # If alignment_mode is 'pad/crop', we crop results that were padded. + # If alignment_mode is 'interpolate', resizing happened during operators. + if alignment_mode == 'pad/crop': + target_shape = recipe_params['_tensor_a_shape'] + if result.shape != target_shape: + slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape)) + result = result[slices] + + merged_state_dict[key] = result.detach().to(save_torch_dtype).cpu().clone() + + # Clean up references to allow GC immediately. + # Local loop variables must be explicitly deleted to prevent PyTorch from keeping tensors in VRAM. + recipe.clean() + del recipe_params['_tensor_a'] + del tensor_a + del recipe + del result + + if recipe_params.get('force_clear_cache', True): + import gc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + pbar.update(1) - continue - - # Pre-load Model A's tensor with pinned memory for CUDA - cpu_tensor = primary_handler.get_tensor(key) - if process_device == 'cuda': - tensor_a = transfer_to_gpu_pinned(cpu_tensor, process_device, process_dtype) - else: - tensor_a = cpu_tensor.to(device=process_device, dtype=process_dtype) - del cpu_tensor - - # Check exclude patterns - use Model A only, no merge - if _matches_any_pattern(key, exclude_patterns): - merged_state_dict[key] = tensor_a.to(save_torch_dtype).cpu().clone() - excluded_keys += 1 - pbar.update(1) - continue - - # Pass tensor_a metadata to recipes for zeros mode and fallback - recipe_params['_tensor_a'] = tensor_a - recipe_params['_tensor_a_shape'] = tensor_a.shape - recipe_params['_tensor_a_dtype'] = tensor_a.dtype - - try: - recipe = calc_mode_class.create_recipe(key=key, **recipe_params) - result = recipe.merge() - except MissingTensorError as e: - if mismatch_mode == MissingTensorBehavior.ERROR: - raise ValueError(f"Layer mismatch error (mismatch_mode='error'): {e}") - result = None - error_keys.append(key) - - # Handle None result (mismatch occurred with skip mode) - if result is None: - result = tensor_a - skipped_keys += 1 - - if isinstance(result, dict): - for r_key, r_tensor in result.items(): - merged_state_dict[r_key] = r_tensor.to(save_torch_dtype).cpu().clone() - else: - # Ensure compatibility with Model A's architecture. - # If alignment_mode is 'pad/crop', we crop results that were padded. - # If alignment_mode is 'interpolate', resizing happened during operators. - if alignment_mode == 'pad/crop': - target_shape = recipe_params['_tensor_a_shape'] - if result.shape != target_shape: - slices = tuple(slice(0, min(res_s, tgt_s)) for res_s, tgt_s in zip(result.shape, target_shape)) - result = result[slices] - - merged_state_dict[key] = result.to(save_torch_dtype).cpu().clone() - - # Clean up tensor_a reference to allow GC - del recipe_params['_tensor_a'] - - if recipe_params.get('force_clear_cache', True): - import gc - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - pbar.update(1) # Log summary if excluded_keys > 0: diff --git a/nodes/merger_ops.py b/nodes/merger_ops.py index 35f2eea..3479da5 100644 --- a/nodes/merger_ops.py +++ b/nodes/merger_ops.py @@ -4,6 +4,8 @@ import torch.nn.functional as F from enum import Enum +from .merger_utils import transfer_to_gpu_pinned + class MissingTensorBehavior(Enum): """Controls behavior when a tensor key is missing from a model.""" ERROR = "error" # Raise exception (strict mode) @@ -127,12 +129,21 @@ class Operation: raise NotImplementedError def recurse(self, operation): - source_tensors = [source_oper.merge() for source_oper in operation.sources] - return operation.oper(*source_tensors) + self._source_tensors = [source_oper.merge() for source_oper in operation.sources] + return operation.oper(*self._source_tensors) def merge(self): return self.merge_func(self) + def clean(self): + if hasattr(self, '_source_tensors'): + for t in self._source_tensors: + del t + del self._source_tensors + for source in self.sources: + if hasattr(source, 'clean'): + source.clean() + class LoadTensor(Operation): def __init__(self, key, model_name, handlers, device, dtype, on_missing=MissingTensorBehavior.ERROR, fallback_shape=None, fallback_dtype=None): @@ -165,7 +176,17 @@ class LoadTensor(Operation): dtype = self.fallback_dtype if self.fallback_dtype else self.dtype return torch.zeros(self.fallback_shape, device=self.device, dtype=dtype) - return handler.get_tensor(self.key).to(device=self.device, dtype=self.dtype) + cpu_tensor = handler.get_tensor(self.key) + if self.device == 'cuda': + self._tensor = transfer_to_gpu_pinned(cpu_tensor, self.device, self.dtype) + else: + self._tensor = cpu_tensor.to(device=self.device, dtype=self.dtype) + del cpu_tensor + return self._tensor + + def clean(self): + if hasattr(self, '_tensor'): + del self._tensor class Multiply(Operation): def __init__(self, key, alpha, *sources): @@ -409,6 +430,10 @@ class PreloadedTensor(Operation): def merge(self) -> torch.Tensor: return self.tensor + def clean(self): + if hasattr(self, 'tensor'): + del self.tensor + class WeightSum(CalcMode): name = 'Weight-Sum' description = 'A * (1 - α) + B * α. Simple linear interpolation.' diff --git a/nodes/merger_utils.py b/nodes/merger_utils.py index 1656c46..d68e200 100644 --- a/nodes/merger_utils.py +++ b/nodes/merger_utils.py @@ -1,133 +1,89 @@ """ -Memory-efficient safetensors utilities with parallel I/O and pinned memory support. - -Based on benchmarking, parallel reading with ThreadPoolExecutor provides 2-4x speedup -for large models compared to sequential reads. - -Worker count is automatically optimized based on device capabilities when workers=None. -Pinned memory enables faster CPU→GPU transfers (2-3x improvement). +Memory-efficient safetensors utilities with pinned memory support. +Forces direct file reads to minimize OS page cache usage and avoids VRAM accumulation. """ import os -import mmap +import gc import json import struct import torch -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, List, Optional, Any, Tuple +from safetensors import safe_open +from typing import Dict, List, Optional, Any -# Import ComfyUI's memory pinning functions if available -try: - from comfy.model_management import pin_memory, unpin_memory, is_device_cuda - PINNING_AVAILABLE = True -except ImportError: - PINNING_AVAILABLE = False - def pin_memory(tensor): return False - def unpin_memory(tensor): return False - def is_device_cuda(device): return str(device).startswith('cuda') +# Module-level configuration for pinned transfer +_verbose = False +_pinned_transfer_stats = {"pinned": 0, "fallback": 0} +def set_verbose(enabled: bool): + """Enable/disable verbose output for pinned transfers.""" + global _verbose + _verbose = enabled -class PinnedTensor: - """Context manager for automatic pinned memory management. - - Usage: - with PinnedTensor(cpu_tensor) as pinned: - gpu_tensor = pinned.to('cuda') - # Memory automatically unpinned after transfer - """ - - def __init__(self, tensor: torch.Tensor, auto_pin: bool = True): - self.tensor = tensor - self.is_pinned = False - self.auto_pin = auto_pin and PINNING_AVAILABLE - - def __enter__(self) -> torch.Tensor: - if self.auto_pin and self.tensor.device.type == 'cpu': - self.is_pinned = pin_memory(self.tensor) - return self.tensor - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.is_pinned: - unpin_memory(self.tensor) - self.is_pinned = False +def get_pinned_transfer_stats(): + """Return pinned transfer statistics.""" + return _pinned_transfer_stats.copy() +def reset_pinned_transfer_stats(): + """Reset transfer statistics.""" + global _pinned_transfer_stats + _pinned_transfer_stats = {"pinned": 0, "fallback": 0} def transfer_to_gpu_pinned( tensor: torch.Tensor, device: str = 'cuda', dtype: Optional[torch.dtype] = None ) -> torch.Tensor: - """Transfer tensor to GPU using pinned memory for faster transfer. + """Transfer tensor to GPU using pinned memory for faster transfer.""" + global _pinned_transfer_stats - Args: - tensor: CPU tensor to transfer - device: Target GPU device - dtype: Optional dtype conversion + if tensor.device.type != 'cpu' or not torch.cuda.is_available(): + if dtype is not None: + return tensor.to(device=device, dtype=dtype) + return tensor.to(device=device) - Returns: - Tensor on GPU device - """ - if not PINNING_AVAILABLE or not is_device_cuda(torch.device(device)): - # Fall back to regular transfer - if dtype: + if not str(device).startswith('cuda'): + if dtype is not None: return tensor.to(device=device, dtype=dtype) return tensor.to(device=device) - # Pin memory for faster transfer - was_pinned = pin_memory(tensor) try: - if dtype: - result = tensor.to(device=device, dtype=dtype, non_blocking=True) + pinned = tensor.pin_memory() + if dtype is not None: + result = pinned.to(device=device, dtype=dtype, non_blocking=True) else: - result = tensor.to(device=device, non_blocking=True) - # Sync to ensure transfer is complete before unpinning - if was_pinned: - torch.cuda.current_stream().synchronize() + result = pinned.to(device=device, non_blocking=True) + + torch.cuda.current_stream().synchronize() + _pinned_transfer_stats["pinned"] += 1 return result - finally: - if was_pinned: - unpin_memory(tensor) + except Exception as e: + _pinned_transfer_stats["fallback"] += 1 + if dtype is not None: + return tensor.to(device=device, dtype=dtype) + return tensor.to(device=device) class MemoryEfficientSafeOpen: - """Memory-efficient safetensors file reader with mmap, parallel I/O, and pinned memory support. - - Features: - - mmap mode: Zero-copy tensor access via memory-mapped file - - low_memory mode: Direct file reads to minimize OS page cache usage - - Parallel loading: Multi-threaded tensor reads for 2-4x speedup - - Sorted batch reads: Keys sorted by file offset for sequential I/O - - Auto-optimized workers: Adjusts parallelism based on device capabilities - - Pinned memory: Faster CPU→GPU transfers when transferring to CUDA - - Args: - filename: Path to safetensors file - device: Target device (default 'cpu') - mmap_mode: Use memory-mapped file for zero-copy (default True) - low_memory: Use direct file reads to minimize memory (overrides mmap_mode) + """Memory-efficient safetensors file reader with streaming and pinned memory support. + + Operates strictly in streaming mode (low memory). Tensors are read byte-by-byte from + their exact offsets without using safe_open or memory mapping, ensuring tensors do not + remain persistently in OS page caches or RAM. """ - def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False): + def __init__(self, filename: str, device: str = 'cpu', **kwargs): + # Ignore mmap_mode, low_memory, or any other legacy kwargs. We always stream. self.filename = filename self.device = device - # low_memory mode forces mmap off to avoid OS page caching - self.low_memory = low_memory - self.mmap_mode = mmap_mode and not low_memory self.header, self.header_size = self._read_header() self.file = open(filename, "rb") - self.mmap_obj = None - - if self.mmap_mode: - self.mmap_obj = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): - if self.mmap_obj: - self.mmap_obj.close() self.file.close() - def keys(self) -> List[str]: """Return all tensor keys (excluding metadata).""" return [k for k in self.header.keys() if k != "__metadata__"] @@ -136,8 +92,17 @@ class MemoryEfficientSafeOpen: """Return file metadata.""" return self.header.get("__metadata__", {}) + def get_shape(self, key: str) -> tuple: + """Get tensor shape without loading tensor data.""" + if key not in self.header: + raise KeyError(f"Tensor '{key}' not found in file") + return tuple(self.header[key]["shape"]) + + def get_ndim(self, key: str) -> int: + """Get tensor ndim without loading tensor data.""" + return len(self.get_shape(key)) + def keys_sorted_by_offset(self) -> List[str]: - """Return keys sorted by file offset for optimal sequential I/O.""" keys_with_offsets = [] for key in self.keys(): offset = self.header[key]["data_offsets"][0] @@ -146,48 +111,28 @@ class MemoryEfficientSafeOpen: return [k for k, _ in keys_with_offsets] def get_tensor(self, key: str) -> torch.Tensor: - """Load a single tensor by key.""" + """Load a single tensor by key directly from file byte offset.""" if key not in self.header: raise KeyError(f"Tensor '{key}' not found in the file") metadata = self.header[key] offset_start, offset_end = metadata["data_offsets"] - if self.mmap_mode and self.mmap_obj: - if offset_start != offset_end: - start = self.header_size + 8 + offset_start - end = self.header_size + 8 + offset_end - tensor_bytes = memoryview(self.mmap_obj)[start:end] - else: - tensor_bytes = None + if offset_start != offset_end: + self.file.seek(self.header_size + 8 + offset_start) + tensor_bytes = bytearray(offset_end - offset_start) + self.file.readinto(tensor_bytes) else: - # Non-mmap mode: use pre-allocated bytearray with readinto for minimal copies tensor_bytes = None - if offset_start != offset_end: - self.file.seek(self.header_size + 8 + offset_start) - # Pre-allocate writable buffer and read directly into it - tensor_bytes = bytearray(offset_end - offset_start) - self.file.readinto(tensor_bytes) return self._deserialize_tensor(tensor_bytes, metadata) - def get_tensor_to_gpu( self, key: str, device: str = 'cuda', dtype: Optional[torch.dtype] = None ) -> torch.Tensor: - """Load a tensor and transfer to GPU using pinned memory for faster transfer. - - Args: - key: Tensor key to load - device: Target GPU device (default 'cuda') - dtype: Optional dtype conversion - - Returns: - Tensor on GPU device - """ cpu_tensor = self.get_tensor(key) return transfer_to_gpu_pinned(cpu_tensor, device, dtype) @@ -197,20 +142,9 @@ class MemoryEfficientSafeOpen: workers: Optional[int] = None, sort_by_offset: bool = True ) -> Dict[str, torch.Tensor]: - """Load multiple tensors in parallel using ThreadPoolExecutor. - - Args: - keys: List of keys to load (default: all keys) - workers: Number of worker threads (default: auto-calculated based on device) - sort_by_offset: Sort keys by file offset before loading (default: True) - - Returns: - Dict mapping key -> tensor - """ if keys is None: keys = self.keys() - # Auto-calculate workers based on device capabilities if workers is None: try: from .device_utils import get_optimal_workers, estimate_model_size @@ -220,13 +154,11 @@ class MemoryEfficientSafeOpen: workers = min(4, os.cpu_count() or 4) if sort_by_offset: - # Sort by offset for better I/O patterns keys_with_offsets = [(k, self.header[k]["data_offsets"][0]) for k in keys] keys_with_offsets.sort(key=lambda x: x[1]) keys = [k for k, _ in keys_with_offsets] result = {} - with ThreadPoolExecutor(max_workers=workers) as executor: futures = { executor.submit( @@ -251,7 +183,6 @@ class MemoryEfficientSafeOpen: return result def get_tensor_as_dict(self, key: str) -> Dict[str, Any]: - """Load a tensor and decode as JSON dict (for config tensors).""" if key not in self.header: raise KeyError(f"Tensor '{key}' not found") @@ -261,14 +192,8 @@ class MemoryEfficientSafeOpen: if offset_start == offset_end: return {} - if self.mmap_mode and self.mmap_obj: - start = self.header_size + 8 + offset_start - end = self.header_size + 8 + offset_end - tensor_bytes = bytes(self.mmap_obj[start:end]) - else: - self.file.seek(self.header_size + 8 + offset_start) - tensor_bytes = self.file.read(offset_end - offset_start) - + self.file.seek(self.header_size + 8 + offset_start) + tensor_bytes = self.file.read(offset_end - offset_start) return json.loads(tensor_bytes.decode("utf-8")) def _read_header(self): @@ -285,19 +210,13 @@ class MemoryEfficientSafeOpen: if tensor_bytes is None: byte_tensor = torch.empty(0, dtype=torch.uint8) else: - # Avoid extra copy if already a bytearray (low_memory mode) - if isinstance(tensor_bytes, bytearray): - byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8) - else: - # mmap memoryview needs copy to create writable tensor - byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8) + byte_tensor = torch.frombuffer(tensor_bytes, dtype=torch.uint8) if dtype_str in ["F8_E5M2", "F8_E4M3"]: return self._convert_float8(byte_tensor, dtype_str, shape) return byte_tensor.view(dtype).reshape(shape) - @staticmethod def _get_torch_dtype(dtype_str): dtype_map = { @@ -322,7 +241,7 @@ class MemoryEfficientSafeOpen: elif dtype_str == "F8_E4M3" and hasattr(torch, "float8_e4m3fn"): return byte_tensor.view(torch.float8_e4m3fn).reshape(shape) else: - raise ValueError(f"Unsupported float8 type: {dtype_str}. Your PyTorch version may be too old.") + raise ValueError(f"Unsupported float8 type: {dtype_str}") def _read_tensor_from_file(