From 3d5a7484a2694a9fef949e8947d1d75d8e97317c Mon Sep 17 00:00:00 2001 From: Yishuo Wang Date: Fri, 11 Aug 2023 11:18:19 +0800 Subject: [PATCH] [LLM] fix bloom and starcoder memory release (#8728) --- python/llm/src/bigdl/llm/ggml/model/bloom/bloom.py | 8 +++++--- .../llm/src/bigdl/llm/ggml/model/starcoder/starcoder.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/python/llm/src/bigdl/llm/ggml/model/bloom/bloom.py b/python/llm/src/bigdl/llm/ggml/model/bloom/bloom.py index 0eea6528..ec33ddaa 100644 --- a/python/llm/src/bigdl/llm/ggml/model/bloom/bloom.py +++ b/python/llm/src/bigdl/llm/ggml/model/bloom/bloom.py @@ -313,9 +313,6 @@ class Bloom(GenerationMixin): } } - def free(self): - bloom_free(self.ctx) - def _tokenize(self, text: bytes, add_bos: bool = False) -> List[int]: """Tokenize a string. @@ -427,3 +424,8 @@ class Bloom(GenerationMixin): seed=self.seed, n_threads=self.n_threads, n_batch=self.n_batch) + + def __del__(self): + if self.ctx is not None: + bloom_free(self.ctx) + self.ctx = None diff --git a/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder.py b/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder.py index e40fc991..e54392d7 100644 --- a/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder.py +++ b/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder.py @@ -315,9 +315,6 @@ class Starcoder(GenerationMixin): } } - def free(self): - starcoder_free(self.ctx) - def _tokenize(self, text: bytes, add_bos: bool = False) -> List[int]: """Tokenize a string. @@ -433,3 +430,8 @@ class Starcoder(GenerationMixin): seed=self.seed, n_threads=self.n_threads, n_batch=self.n_batch) + + def __del__(self): + if self.ctx is not None: + starcoder_free(self.ctx) + self.ctx = None