huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-10 20:04:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

use less memory

Browse Source
This commit is contained in:
OlivierDehaene 2023-07-19 00:42:15 +02:00
parent 05d2a77e4c
commit 0111869ad0
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
server/text_generation_server/models/flash_causal_lm.py
View File

@ -743,8 +743,9 @@ class FlashCausalLM(Model):
total_gpu_memory = torch.cuda.get_device_properties(self.device).total_memory
# 0.98 to add some wiggle room
num_blocks = (
int((total_gpu_memory - peak_memory) // total_cache_size)
int((total_gpu_memory * 0.98 - peak_memory) // total_cache_size)
# Add batch.blocks as we allocated it above, so it is included in the peak memory.
+ batch.blocks
)