Fix @lrucache memory leaks

Author: turboderp

exllamav3/loader/safetensors.py

@@ -9,7 +9,7 @@
 import mmap
 from ..util import Timer, cuda_sync_active
 from ..ext import exllamav3_ext as ext
-from functools import lru_cache
+from functools import cached_property
 from fnmatch import fnmatch
 import time
 
@@ -192,11 +192,13 @@ def get_tensor_size(
         return bytesize
 
 
-    @lru_cache
-    def get_tensor_file_map_trie(self):
+    @cached_property
+    def _get_tensor_file_map_trie(self):
         import marisa_trie
         trie = marisa_trie.Trie(self.tensor_file_map.keys())
         return trie
+    def get_tensor_file_map_trie(self):
+        return self._get_tensor_file_map_trie
 
 
     def list_tensors(
@@ -400,12 +402,16 @@ def close(self):
                 self.handles[filename] = None
 
 
-    @lru_cache
-    def max_key_len(self):
+    @cached_property
+    def _max_key_len(self):
         l = max(len(k) for k in self.tensor_file_map.keys())
         return l
 
 
+    def max_key_len(self):
+        return self._max_key_len
+
+
     def set_new_tensors(self, new_tensors):
         self.new_tensors = new_tensors
 
@@ -513,7 +519,7 @@ def __init__(
     ):
         self.main = main
         self.stcs = []
-
+        self._get_tensor_sizes_cache = {}
 
     def compile_star_globs(self, patterns, *, flags = 0):
         # Turn list of filter globs into single, compiled regex
@@ -572,17 +578,18 @@ def has_tensor_group(
         return True
 
 
-    @lru_cache
     def get_tensor_sizes(
         self,
         prefix: str,
     ):
-        keys = [self.main.tensor_file_map.get(prefix)]
-        if keys[0] is None:
-            keys = []
-        keys += self.main.get_tensor_file_map_trie().keys(prefix + ".")
-        sizes = [self.get_tensor_size(key) for key in keys]
-        return sizes
+        if prefix not in self._get_tensor_sizes_cache:
+            keys = [self.main.tensor_file_map.get(prefix)]
+            if keys[0] is None:
+                keys = []
+            keys += self.main.get_tensor_file_map_trie().keys(prefix + ".")
+            sizes = [self.get_tensor_size(key) for key in keys]
+            self._get_tensor_sizes_cache[prefix] = sizes
+        return self._get_tensor_sizes_cache[prefix]
 
 
     def get_tensor_size(

exllamav3/model/model.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from functools import lru_cache
+from functools import cached_property
 from typing import Callable
 import torch
 from .config import Config
@@ -50,14 +50,18 @@ def find_module(self, key: str):
         return self.modules_dict[key]
 
 
-    @lru_cache
-    def get_cache_layers(self):
+    @cached_property
+    def _get_cache_layers(self):
         return [m for m in self if m.caps.get("kv_cache")]
+    def get_cache_layers(self):
+        return self._get_cache_layers
 
 
-    @lru_cache
-    def get_recurrent_layers(self):
+    @cached_property
+    def _get_recurrent_layers(self):
         return [m for m in self if m.caps.get("recurrent_cache")]
+    def get_recurrent_layers(self):
+        return self._get_recurrent_layers
 
 
     @staticmethod

exllamav3/modules/linear.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from functools import lru_cache
+from functools import cached_property
 from typing_extensions import override
 import torch
 import torch.nn.functional as F
@@ -347,13 +347,15 @@ def quant_format_id(self):
             return None
 
 
-    @lru_cache
-    def storage_size(self):
+    @cached_property
+    def _storage_size(self):
         # alt_key is only used when loading unquantized model
         if self.is_exl3_storage(self.key):
             return sum(self.config.stc.get_tensor_sizes(prefix = self.key))
         else:
             return 2 * self.in_features * self.out_features
+    def storage_size(self):
+        return self._storage_size
 
 
     def recons_size(self):

exllamav3/modules/module.py

@@ -8,7 +8,7 @@
 if TYPE_CHECKING:
     from ..model.config import Config
 from ..model.model_tp_alloc import TPAllocation
-from functools import lru_cache
+from functools import cached_property
 
 # Use host bounce when moving state from device to device in layer split
 no_p2p_copy = os.environ.get('EXLLAMA_NO_P2P_COPY', None)
@@ -133,9 +133,11 @@ def tp_import(local_context, plan, loaded):
         """
         raise NotImplementedError()
 
-    @lru_cache
-    def all_cache_modules(self) -> list[Module]:
+    @cached_property
+    def _all_cache_modules(self) -> list[Module]:
         return [m for m in self if m.caps.get("kv_cache")]
+    def all_cache_modules(self):
+        return self._all_cache_modules
 
     @abstractmethod
     def optimizer_targets(self):

exllamav3/tokenizer/tokenizer.py

@@ -5,7 +5,7 @@
 from ..util import synchronized
 from ..util.file import maybe_read_json
 from ..model.config import Config
-from functools import lru_cache
+from functools import lru_cache, cached_property
 from typing import TYPE_CHECKING
 from ..util import profile_opt
 if TYPE_CHECKING:
@@ -162,7 +162,7 @@ def get_default_token_id(config_key: str, current: int | None, default: int | No
         self.actual_vocab_size = 1 + max(
             list(self.extended_id_to_piece.keys()) + \
             list(self.unspecial_id_to_piece.keys()) + \
-            [self.raw_vocab_size() - 1]
+            [self.raw_vocab_size - 1]
         )
 
         # Useful token IDs
@@ -185,7 +185,7 @@ def get_default_token_id(config_key: str, current: int | None, default: int | No
         self.get_id_to_piece_list(False)
         self.get_piece_to_id_dict()
 
-    @lru_cache
+    @cached_property
     def raw_vocab_size(self):
         """
         Cache this function because it's suspiciously slow in HF Tokenizers
@@ -383,14 +383,14 @@ def decode_(self, seq, decode_special_tokens):
 
         if not decode_special_tokens:
 
-            max_token = self.raw_vocab_size()
+            max_token = self.raw_vocab_size
             seq = [t for t in seq if (t != self.pad_token_id and t < max_token and t != self.eos_token_id)]
             if self.eos_token_id in seq: seq = seq[:seq.index(self.eos_token_id)]
             return self.decode_unspecial(seq)
 
         else:
 
-            max_token = self.raw_vocab_size()
+            max_token = self.raw_vocab_size
             seq = [t for t in seq if (t != self.pad_token_id and t < max_token)]
             text = ""
             start = 0
@@ -486,11 +486,10 @@ def num_tokens(self, text):
 
     # Get ordinals of single-byte tokens
 
-    @synchronized
-    @lru_cache
-    def get_id_to_ord_list(self):
+    @cached_property
+    def _get_id_to_ord_list(self):
 
-        self.id_to_ord = list(range(self.raw_vocab_size()))
+        self.id_to_ord = list(range(self.raw_vocab_size))
 
         def clean_special_chars(p):
             p = p.replace(self.space_char_, " ")
@@ -508,7 +507,7 @@ def piece_to_ord(p):
                 if o <= 255: return o
             return -1
 
-        i = self.raw_vocab_size()
+        i = self.raw_vocab_size
         while True:
             if i in self.extended_id_to_piece:
                 self.id_to_ord.append(piece_to_ord(self.extended_id_to_piece[i]))
@@ -520,10 +519,15 @@ def piece_to_ord(p):
                 break
             i += 1
 
+        return self.id_to_ord
+    @synchronized
+    def get_id_to_ord_list(self):
+        return self._get_id_to_ord_list
+
     # Copy vocabulary from model
 
-    @lru_cache
-    def get_fixed_vocab(self):
+    @cached_property
+    def _get_fixed_vocab(self):
         test_enc = self.tokenizer.encode(" t", add_special_tokens = False)
         test_count = len(test_enc.ids)
         assert test_count > 0, "Tokenizer error, test string encodes to zero tokens"
@@ -532,35 +536,37 @@ def get_fixed_vocab(self):
 
         if test_count == 1 and len(test_piece) == len(" t"):
             vocab = self.tokenizer.decode_batch(
-                [[i] for i in range(self.raw_vocab_size())],
+                [[i] for i in range(self.raw_vocab_size)],
                 skip_special_tokens = False
             )
         else:
             prefix_id = self.tokenizer.encode(" ", add_special_tokens = False).ids[0]
             prefix_piece = self.tokenizer.decode([prefix_id])
             prefix_len = len(prefix_piece)
             vocab = self.tokenizer.decode_batch(
-                [[prefix_id, i] for i in range(self.raw_vocab_size())]
+                [[prefix_id, i] for i in range(self.raw_vocab_size)]
             )
             vocab = [v[prefix_len:] for v in vocab]
 
         return vocab
+    def get_fixed_vocab(self):
+        return self._get_fixed_vocab
 
-    @synchronized
-    @lru_cache
-    def get_id_to_piece_list(self, include_special_tokens = False):
+    @cached_property
+    def _get_id_to_piece_list_spc(self, include_special_tokens = False):
 
-        if include_special_tokens:
-            id_to_piece_extended = self.get_id_to_piece_list().copy()
-            for k, v in self.extended_id_to_piece.items():
-                id_to_piece_extended[k] = v
+        id_to_piece_extended = self.get_id_to_piece_list().copy()
+        for k, v in self.extended_id_to_piece.items():
+            id_to_piece_extended[k] = v
 
-            self.id_to_piece_with_special = id_to_piece_extended
-            return self.id_to_piece_with_special
+        self.id_to_piece_with_special = id_to_piece_extended
+        return self.id_to_piece_with_special
+    @cached_property
+    def _get_id_to_piece_list_nonspc(self, include_special_tokens = False):
 
         self.id_to_piece = self.get_fixed_vocab()
 
-        i = self.raw_vocab_size()
+        i = self.raw_vocab_size
         while True:
             if i in self.extended_id_to_piece:
                 self.id_to_piece.append(self.extended_id_to_piece[i])
@@ -573,13 +579,21 @@ def get_id_to_piece_list(self, include_special_tokens = False):
             i += 1
 
         return self.id_to_piece
-
     @synchronized
-    @lru_cache
-    def get_piece_to_id_dict(self):
+    def get_id_to_piece_list(self, include_special_tokens = False):
+        if include_special_tokens:
+            return self._get_id_to_piece_list_spc
+        else:
+            return self._get_id_to_piece_list_nonspc
+
+    @cached_property
+    def _get_piece_to_id_dict(self):
         all_pieces = self.get_id_to_piece_list()
         self.piece_to_id = {piece: idx for idx, piece in enumerate(all_pieces)}
         return self.piece_to_id
+    @synchronized
+    def get_piece_to_id_dict(self):
+        return self._get_piece_to_id_dict
 
     @staticmethod
     def from_config(config: Config):
@@ -603,12 +617,14 @@ def get_tokens_with_prefix_id(self, prefix_id: int):
         prefix = id_to_piece[prefix_id]
         return self.get_tokens_with_prefix_string(prefix)
 
-    @lru_cache
-    def get_vocab_dict(self):
+    @cached_property
+    def _get_vocab_dict(self):
         """
         Return tokenizer (dictionary for Formatron)
         """
         return {
             self.tokenizer.id_to_token(i): i
             for i in range(self.tokenizer.get_vocab_size())
         }
+    def get_vocab_dict(self):
+        return self._get_vocab_dict