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

0.98

Browse Source
This commit is contained in:
OlivierDehaene 2023-07-19 02:06:16 +02:00
parent 406b094002
commit 2934543a59
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -743,9 +743,9 @@ class FlashCausalLM(Model):
total_gpu_memory = torch.cuda.get_device_properties(self.device).total_memory
# 0.985 to add some wiggle room
# 0.98 to add some wiggle room
num_blocks = (
int((total_gpu_memory * 0.985 - 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
)