From 0111869ad033526124feca0029528c1543023114 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Wed, 19 Jul 2023 00:42:15 +0200
Subject: [PATCH] use less memory

---
 server/text_generation_server/models/flash_causal_lm.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
index 4a3ad6fb..517fba68 100644
--- a/server/text_generation_server/models/flash_causal_lm.py
+++ b/server/text_generation_server/models/flash_causal_lm.py
@@ -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
         )