mmgp 3.3.1__tar.gz → 3.3.2__tar.gz

{mmgp-3.3.1/src/mmgp.egg-info → mmgp-3.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mmgp
-Version: 3.3.1
+Version: 3.3.2
 Summary: Memory Management for the GPU Poor
 Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -17,7 +17,7 @@ Dynamic: license-file
 
 
 <p align="center">
-  <H2>Memory Management 3.3.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.3.2 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

{mmgp-3.3.1 → mmgp-3.3.2}/README.md

@@ -1,6 +1,6 @@
 
 <p align="center">
-  <H2>Memory Management 3.3.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.3.2 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

{mmgp-3.3.1 → mmgp-3.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mmgp"
-version = "3.3.1"
+version = "3.3.2"
 authors = [
   { name = "deepbeepmeep", email = "deepbeepmeep@yahoo.com" },
 ]

{mmgp-3.3.1 → mmgp-3.3.2}/src/mmgp/offload.py

@@ -1,4 +1,4 @@
-# ------------------ Memory Management 3.3.1 for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 3.3.2 for the GPU Poor by DeepBeepMeep (mmgp)------------------
 #
 # This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  can run smoothly on a 24 GB GPU limited card. 
 # This a replacement for the accelerate library that should in theory manage offloading, but doesn't work properly with models that are loaded / unloaded several
@@ -258,11 +258,11 @@ def _move_to_pinned_tensor(source_tensor, big_tensor, offset, length):
         assert t.is_pinned()
     return t
 
-def _safetensors_load_file(file_path):
+def _safetensors_load_file(file_path, writable_tensors = True):
     from collections import OrderedDict
     sd = OrderedDict()    
 
-    with safetensors2.safe_open(file_path, framework="pt", device="cpu") as f:
+    with safetensors2.safe_open(file_path, framework="pt", device="cpu", writable_tensors =writable_tensors) as f:
         for k in f.keys():
             sd[k] = f.get_tensor(k)
         metadata = f.metadata()
@@ -504,25 +504,25 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
 
     failed_planned_allocation = False
 
-    for size in big_tensors_sizes:
-        try:
-            # if total > 7000 * ONE_MB:
-            #     raise  Exception ("test no more reserved RAM")
-            current_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True, device="cpu")
-            big_tensors.append(current_big_tensor)
-        except:
-            print(f"Unable to pin more tensors for this model as the maximum reservable memory has been reached ({total/ONE_MB:.2f})")
-            max_pinnable_bytes = total + total_pinned_bytes
-            failed_planned_allocation = True
-            break
+    # for size in big_tensors_sizes:
+    #     try:
+    #         # if total > 7000 * ONE_MB:
+    #         #     raise  Exception ("test no more reserved RAM")
+    #         current_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True, device="cpu")
+    #         big_tensors.append(current_big_tensor)
+    #     except:
+    #         print(f"Unable to pin more tensors for this model as the maximum reservable memory has been reached ({total/ONE_MB:.2f})")
+    #         max_pinnable_bytes = total + total_pinned_bytes
+    #         failed_planned_allocation = True
+    #         break
 
-        last_big_tensor += 1
-        total += size
+    #     last_big_tensor += 1
+    #     total += size
 
 
     gc.collect()
 
-        
+    last_allocated_big_tensor = -1        
     tensor_no = 0
     # prev_big_tensor = 0
     for n, (p, is_buffer) in params_dict.items():
@@ -543,37 +543,54 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
                 assert p.data.is_pinned()
             q = None
         else:
+
             big_tensor_no, offset, length = tensor_map_indexes[tensor_no]
+            if last_allocated_big_tensor <  big_tensor_no:
+                last_allocated_big_tensor += 1
+                size = big_tensors_sizes[last_allocated_big_tensor] 
+                try:
+                    # if total > 7000 * ONE_MB:
+                    #     raise  Exception ("test no more reserved RAM")
+                    current_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True, device="cpu")
+                    big_tensors.append(current_big_tensor)
+                except:
+                    print(f"Unable to pin more tensors for this model as the maximum reservable memory has been reached ({total/ONE_MB:.2f})")
+                    max_pinnable_bytes = total + total_pinned_bytes
+                    failed_planned_allocation = True
+                    break
+
+                total += size
+
             # if big_tensor_no != prev_big_tensor:
             #     gc.collect()
             #     prev_big_tensor = big_tensor_no
             # match_param, match_isbuffer = tied_weights.get(n, (None, False))
             # if match_param != None:
 
-            if big_tensor_no>=0 and big_tensor_no < last_big_tensor:
-                current_big_tensor = big_tensors[big_tensor_no]
-                if is_buffer :
-                    _force_load_buffer(p) # otherwise potential memory leak
-                if isinstance(p, QTensor):
-                    if p._qtype == qint4:
-                        length1 = torch.numel(p._data._data) * p._data._data.element_size()
-                        p._data._data =  _move_to_pinned_tensor(p._data._data, current_big_tensor, offset, length1)
-                        if hasattr(p,"_scale_shift"):
-                            length2 = torch.numel(p._scale_shift) * p._scale_shift.element_size() 
-                            p._scale_shift = _move_to_pinned_tensor(p._scale_shift, current_big_tensor, offset + length1, length2)
-                        else:
-                            length2 = torch.numel(p._scale) * p._scale.element_size() 
-                            p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
-                            length3 = torch.numel(p._shift) * p._shift.element_size() 
-                            p._shift = _move_to_pinned_tensor(p._shift, current_big_tensor, offset + length1 + length2, length3)
+            # if big_tensor_no>=0 and big_tensor_no < last_big_tensor:
+            current_big_tensor = big_tensors[big_tensor_no]
+            if is_buffer :
+                _force_load_buffer(p) # otherwise potential memory leak
+            if isinstance(p, QTensor):
+                if p._qtype == qint4:
+                    length1 = torch.numel(p._data._data) * p._data._data.element_size()
+                    p._data._data =  _move_to_pinned_tensor(p._data._data, current_big_tensor, offset, length1)
+                    if hasattr(p,"_scale_shift"):
+                        length2 = torch.numel(p._scale_shift) * p._scale_shift.element_size() 
+                        p._scale_shift = _move_to_pinned_tensor(p._scale_shift, current_big_tensor, offset + length1, length2)
                     else:
-                        length1 = torch.numel(p._data) * p._data.element_size() 
-                        p._data = _move_to_pinned_tensor(p._data, current_big_tensor, offset, length1)
                         length2 = torch.numel(p._scale) * p._scale.element_size() 
                         p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
+                        length3 = torch.numel(p._shift) * p._shift.element_size() 
+                        p._shift = _move_to_pinned_tensor(p._shift, current_big_tensor, offset + length1 + length2, length3)
                 else:
-                    length = torch.numel(p.data) * p.data.element_size() 
-                    p.data = _move_to_pinned_tensor(p.data, current_big_tensor, offset, length)
+                    length1 = torch.numel(p._data) * p._data.element_size() 
+                    p._data = _move_to_pinned_tensor(p._data, current_big_tensor, offset, length1)
+                    length2 = torch.numel(p._scale) * p._scale.element_size() 
+                    p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
+            else:
+                length = torch.numel(p.data) * p.data.element_size() 
+                p.data = _move_to_pinned_tensor(p.data, current_big_tensor, offset, length)
             tensor_no += 1
         del p
     model._pinned_bytes = total
@@ -598,7 +615,7 @@ def _welcome():
     if welcome_displayed:
          return 
     welcome_displayed = True
-    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.3.1) by DeepBeepMeep ************{ENDC}{UNBOLD}")
+    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.3.3) by DeepBeepMeep ************{ENDC}{UNBOLD}")
 
 def _extract_num_from_str(num_in_str):
     size = len(num_in_str)
@@ -1151,7 +1168,7 @@ def move_loras_to_device(model, device="cpu" ):
         if ".lora_" in k:
             m.to(device)
 
-def fast_load_transformers_model(model_path: str, do_quantize = False, quantizationType =  qint8, pinToMemory = False, partialPinning = False, forcedConfigPath = None, modelClass=None, modelPrefix = None, verboseLevel = -1):
+def fast_load_transformers_model(model_path: str, do_quantize = False, quantizationType =  qint8, pinToMemory = False, partialPinning = False, forcedConfigPath = None, modelClass=None, modelPrefix = None, writable_tensors = True, verboseLevel = -1):
     """
     quick version of .LoadfromPretrained of  the transformers library
     used to build a model and load the corresponding weights (quantized or not)
@@ -1167,7 +1184,7 @@ def fast_load_transformers_model(model_path: str, do_quantize = False, quantizat
     
     verboseLevel = _compute_verbose_level(verboseLevel)
 
-    with safetensors2.safe_open(model_path) as f:
+    with safetensors2.safe_open(model_path, writable_tensors =writable_tensors) as f:
         metadata = f.metadata() 
 
     if metadata is None:
@@ -1231,13 +1248,13 @@ def fast_load_transformers_model(model_path: str, do_quantize = False, quantizat
 
     model._config = transformer_config
             
-    load_model_data(model,model_path, do_quantize = do_quantize, quantizationType = quantizationType, pinToMemory= pinToMemory, partialPinning= partialPinning, modelPrefix = modelPrefix, verboseLevel=verboseLevel )
+    load_model_data(model,model_path, do_quantize = do_quantize, quantizationType = quantizationType, pinToMemory= pinToMemory, partialPinning= partialPinning, modelPrefix = modelPrefix, writable_tensors =writable_tensors ,verboseLevel=verboseLevel )
 
     return model
 
 
 
-def load_model_data(model, file_path: str, do_quantize = False, quantizationType = qint8, pinToMemory = False, partialPinning = False, modelPrefix = None, verboseLevel = -1):
+def load_model_data(model, file_path: str, do_quantize = False, quantizationType = qint8, pinToMemory = False, partialPinning = False, modelPrefix = None, writable_tensors = True,  verboseLevel = -1):
     """
     Load a model, detect if it has been previously quantized using quanto and do the extra setup if necessary
     """
@@ -1275,7 +1292,7 @@ def load_model_data(model, file_path: str, do_quantize = False, quantizationType
         if "module" in state_dict:
             state_dict = state_dict["module"]
     else:
-        state_dict, metadata = _safetensors_load_file(file_path)
+        state_dict, metadata = _safetensors_load_file(file_path, writable_tensors =writable_tensors)
             
         if metadata is None:
             quantization_map = None
@@ -1470,7 +1487,6 @@ class offload:
     def __init__(self):
         self.active_models = []
         self.active_models_ids = []
-        self.active_subcaches = {}        
         self.models = {}
         self.cotenants_map = { 
                             "text_encoder": ["vae", "text_encoder_2"],
@@ -1732,7 +1748,6 @@ class offload:
  
         self.active_models = []
         self.active_models_ids = []
-        self.active_subcaches = []
         torch.cuda.empty_cache()
         gc.collect()
         self.last_reserved_mem_check = time.time()
@@ -2051,20 +2066,23 @@ class offload:
             last_offload_obj = None
 
         self.unload_all()
-        self.default_stream = None
+        self.active_models = None
+        self.default_stream = None 
+        self.transfer_stream = None
+        self.parameters_ref = None
         keys= [k for k in self.blocks_of_modules.keys()]
         for k in keys:
             del self.blocks_of_modules[k]
 
         self.blocks_of_modules = None
 
-
         for model_id, model in self.models.items():
             move_loras_to_device(model, "cpu")
             if hasattr(model, "_pinned_bytes"):
                 total_pinned_bytes -= model._pinned_bytes
             if hasattr(model, "_loras_model_data"):
                 unload_loras_from_model(model)
+            model = None
 
         self.models = None            
 
@@ -2074,7 +2092,7 @@ class offload:
 
 
 
-def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, partialPinning = False, loras = None, quantizeTransformer = True,  extraModelsToQuantize = None, quantizationType = qint8, budgets= 0, workingVRAM = None, asyncTransfers = True, compile = False, perc_reserved_mem_max = 0, coTenantsMap = None, verboseLevel = -1):
+def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, partialPinning = False, loras = None, quantizeTransformer = True,  extraModelsToQuantize = None, quantizationType = qint8, budgets= 0, workingVRAM = None, asyncTransfers = True, compile = False, convertFloatToBfloat16 = True, perc_reserved_mem_max = 0, coTenantsMap = None, verboseLevel = -1):
     """Hook to a pipeline or a group of modules in order to reduce their VRAM requirements:
     pipe_or_dict_of_modules : the pipeline object or a dictionary of modules of the model
     quantizeTransformer: set True by default will quantize on the fly the video / image model
@@ -2181,7 +2199,6 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
         modelPinned = (pinAllModels or model_id in modelsToPin) and not hasattr(current_model,"_already_pinned")
 
         current_model_size = 0
-
         for n, p in current_model.named_parameters():
             p.requires_grad = False
             if isinstance(p, QTensor):
@@ -2201,7 +2218,7 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
                     current_model_size +=  torch.numel(p._data) * p._data.element_size()
 
             else:
-                if p.data.dtype == torch.float32:
+                if convertFloatToBfloat16 and p.data.dtype == torch.float32:
                     # convert any left overs float32 weight to bloat16 to divide by 2 the model memory footprint
                     p.data = p.data.to(torch.bfloat16)
                 current_model_size +=  torch.numel(p.data) * p.data.element_size()

{mmgp-3.3.1 → mmgp-3.3.2}/src/mmgp/safetensors2.py

@@ -1,4 +1,4 @@
-# ------------------ Safetensors2 1.0 by DeepBeepMeep (mmgp)------------------
+# ------------------ Safetensors2 1.1 by DeepBeepMeep (mmgp)------------------
 #
 # This module entirely written in Python is a replacement for the safetensor library which requires much less RAM to load models.
 # It can be conveniently used to keep a low RAM consumption when handling  transit data (for instance when quantizing or transferring tensors to reserver RAM)
@@ -16,12 +16,14 @@ import safetensors
 import accelerate
 import os
 from collections import OrderedDict
+import warnings
 
+warnings.filterwarnings("ignore", ".*The given buffer is not writable, and PyTorch does not support non-writable tensors*")
 
 _old_torch_load_file = None
 _old_safe_open = None
 
-
+all_tensors_are_read_only = False 
 
 mmm = {}
 verboseLevel = 1
@@ -232,7 +234,7 @@ def torch_write_file(sd, file_path, quantization_map = None, config = None, extr
 class SafeTensorFile:
     """Main class for accessing safetensors files that provides memory-efficient access"""
     
-    def __init__(self, file_path, metadata, catalog, skip_bytes, lazy_loading = True):
+    def __init__(self, file_path, metadata, catalog, skip_bytes, lazy_loading = True, writable_tensors = True):
         self._file_path = file_path
         self._metadata = metadata
         self._catalog = catalog
@@ -241,19 +243,20 @@ class SafeTensorFile:
         self.sd = None
         self.mtracker = None
         self.lazy_loading = lazy_loading
+        self.writable_tensors = writable_tensors
 
     @classmethod
-    def load_metadata(cls, file_path, lazy_loading = True):
+    def load_metadata(cls, file_path, lazy_loading = True, writable_tensors = True):
         with open(file_path, 'rb') as f:
             catalog, metadata, skip_bytes = _read_safetensors_header(file_path, f)
 
-        return cls(file_path, metadata, catalog, skip_bytes, lazy_loading)
+        return cls(file_path, metadata, catalog, skip_bytes, lazy_loading, writable_tensors )
 
-    def init_tensors(self, lazyTensors = True):
+    def init_tensors(self, lazyTensors = True, writable_tensors = True):
         if self.sd is None:
             self.lazy_loading = lazyTensors
             if lazyTensors:
-                self.sd = self.create_tensors_with_mmap()
+                self.sd = self.create_tensors_with_mmap(writable_tensors)
             else:
                 self.sd = self.create_tensors_without_mmap()
         # else:
@@ -263,7 +266,7 @@ class SafeTensorFile:
         return self.sd
     
             
-    def create_tensors_with_mmap(self):
+    def create_tensors_with_mmap(self, writable_tensors = True):
  
         self.mtracker = MmapTracker(self._file_path)
         import mmap
@@ -302,7 +305,7 @@ class SafeTensorFile:
         with open(self._file_path, 'rb') as f:
             i = 0
             for map_start, map_size in maps_info:
-                mm = mmap.mmap(f.fileno(), map_size, offset=map_start, access=mmap.ACCESS_COPY) #.ACCESS_READ
+                mm = mmap.mmap(f.fileno(), map_size, offset=map_start, access=  mmap.ACCESS_COPY  if writable_tensors else mmap.ACCESS_READ) 
                 maps.append((mm, map_start, map_size))
                 self.mtracker.register(mm, i, map_start, map_size)
                 i = i+ 1
@@ -359,7 +362,7 @@ class SafeTensorFile:
     def get_tensor(self, name: str) -> torch.tensor:
         """Get a tensor by name"""
         # To do : switch to a JIT tensor creation per tensor
-        self.init_tensors()
+        self.init_tensors(self.lazy_loading, writable_tensors= self.writable_tensors)
         return self.sd[name]
  
     def keys(self) -> List[str]:
@@ -374,7 +377,7 @@ class SafeTensorFile:
         
     def tensors(self) -> Dict[str, torch.tensor]:
         """Get dictionary of all tensors"""
-        self.init_tensors(self.lazy_loading)
+        self.init_tensors(self.lazy_loading, writable_tensors= self.writable_tensors)
         return self.sd
         
     def metadata(self) -> Optional[Dict[str, str]]:
@@ -383,7 +386,7 @@ class SafeTensorFile:
         
     def __len__(self) -> int:
         """Get number of tensors"""
-        self.init_tensors(self.lazy_loading)
+        self.init_tensors(self.lazy_loading, writable_tensors= self.writable_tensors)
         return len(self.keys())
         
     def __contains__(self, key: str) -> bool:
@@ -401,17 +404,22 @@ class SafeTensorFile:
 class _SafeTensorLoader:
     """Context manager for loading SafeTensorFile"""
     
-    def __init__(self, filename: str ):
+    def __init__(self, filename: str, writable_tensors = True ):
         self.filename = Path(filename)
+        self.writable_tensors = writable_tensors
         self.sft = None
         if not self.filename.exists():
             raise FileNotFoundError(f"File not found: {filename}")
             
     def __enter__(self) -> SafeTensorFile:
         """Open file and return SafeTensorFile instance"""
-        
+        writable_tensors = self.writable_tensors
+
+        if all_tensors_are_read_only:
+            writable_tensors = False
+
         try:
-            self.sft = SafeTensorFile.load_metadata(self.filename)
+            self.sft = SafeTensorFile.load_metadata(self.filename, writable_tensors= writable_tensors)
             return self.sft 
             
         except Exception as e:
@@ -428,14 +436,14 @@ class _SafeTensorLoader:
         pass
 
 
-def safe_open(filename: str, framework: str = "pt",device = "cpu") -> _SafeTensorLoader:
+def safe_open(filename: str, framework: str = "pt",device = "cpu", writable_tensors = True) -> _SafeTensorLoader:
     if device != "cpu" or framework !="pt":
         return _old_safe_open(filename =filename, framework=framework, device=device)
-    return _SafeTensorLoader(filename)
+    return _SafeTensorLoader(filename, writable_tensors = writable_tensors)
 
-def torch_load_file( filename, device = 'cpu' ) -> Dict[str, torch.Tensor]:
+def torch_load_file( filename, device = 'cpu', writable_tensors = True) -> Dict[str, torch.Tensor]:
     sd = {}
-    with safe_open(filename, framework="pt", device = device ) as f:
+    with safe_open(filename, framework="pt", device = device, writable_tensors =writable_tensors  ) as f:
         for k in f.keys():
             sd[k] = f.get_tensor(k)
         return sd

{mmgp-3.3.1 → mmgp-3.3.2/src/mmgp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mmgp
-Version: 3.3.1
+Version: 3.3.2
 Summary: Memory Management for the GPU Poor
 Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -17,7 +17,7 @@ Dynamic: license-file
 
 
 <p align="center">
-  <H2>Memory Management 3.3.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.3.2 for the GPU Poor by DeepBeepMeep</H2>	
 </p>