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:
@@ -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
@@ -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 = {
|
||||
|
||||
Reference in New Issue
Block a user