perf: Add low_memory mode to reduce RAM usage in LoRA merge

- MemoryEfficientSafeOpen now supports low_memory=True which disables mmap
- Uses direct file reads with readinto() to avoid OS page caching
- _deserialize_tensor avoids bytearray copy when input is already bytearray
- merge_loras_to_model now uses low_memory=True for base model

This should reduce peak RAM from 2x model size to 1x + overhead
This commit is contained in:
silveroxides
2026-01-15 03:13:51 +01:00
parent e9f7cda706
commit 8deb4fd739
2 changed files with 22 additions and 6 deletions
+3 -2
View File
@@ -1625,10 +1625,11 @@ def merge_loras_to_model(
print(f"[LoRA Merge To Model] Merging {len(lora_paths)} LoRAs with weights: {lora_weights}")
prepare_for_large_operation(total_size_gb * 1.5, torch.device(device))
# Open all files
base_handler = MemoryEfficientSafeOpen(base_model_path)
# Open all files - use low_memory for base model to avoid OS page caching
base_handler = MemoryEfficientSafeOpen(base_model_path, low_memory=True)
lora_handlers = [MemoryEfficientSafeOpen(lp) for lp in lora_paths]
try:
# Detect format and extract pairs for each LoRA
lora_infos = []
+19 -4
View File
@@ -93,6 +93,7 @@ class MemoryEfficientSafeOpen:
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
@@ -102,12 +103,15 @@ class MemoryEfficientSafeOpen:
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)
"""
def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True):
def __init__(self, filename: str, device: str = 'cpu', mmap_mode: bool = True, low_memory: bool = False):
self.filename = filename
self.device = device
self.mmap_mode = mmap_mode
# 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
@@ -123,6 +127,7 @@ class MemoryEfficientSafeOpen:
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__"]
@@ -156,13 +161,17 @@ class MemoryEfficientSafeOpen:
else:
tensor_bytes = None
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)
tensor_bytes = self.file.read(offset_end - 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,
@@ -276,13 +285,19 @@ class MemoryEfficientSafeOpen:
if tensor_bytes is None:
byte_tensor = torch.empty(0, dtype=torch.uint8)
else:
byte_tensor = torch.frombuffer(bytearray(tensor_bytes), dtype=torch.uint8)
# 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)
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 = {