Repository: huggingface/text-generation-inference

Issue:
How to config vllm gpu_memory_utilization?
Hi team, I am trying using codegen2.5 7b model on tgi with A100 40GB and it gives me out of memory error because of vllm. I wonder if there is any way I can config gpu_memory_utilization in the code such that the vllm does not reserve too memory beforehand 


Discussion:
We're working on some Quality of life to help with that : https://github.com/huggingface/text-generation-inference/pull/630

otherwise, try to look at the error message, it should give you the name of the parameters you can tweak to fix the RAM issue 

```
--max-total-tokens
--max-batch-prefill-tokens
--max-input-length
# and maybe a couple others I usually only tweak those
```
I'm writing from memory, so just check the error message for the correct names or `text-generation-launcher --help`
Hi @Narsil, Thanks for the help. I try to tweek the number even I set --max-batch-prefill-tokens=1, and --max-batch-total-tokens=2 but it is still out of memory. What other things can I do? For the context, when I try not to use flash attention on llamamodel, it can work. When I use flash attention and not vllm, it also works. so there must be sth going on the vllm side. 
Can you share a reproducible example ? And the full stacktrace ?
I have the same issue, kv cache warmup casue OOM

Resolving patch:
diff --git a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
--- a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
@@ -154,7 +154,7 @@ def _load_gqa(config, prefix: str, weights):
     weight = weights.get_multi_weights_col(
         prefixes=[f"{prefix}.q_proj", f"{prefix}.k_proj", f"{prefix}.v_proj"],
         quantize=config.quantize,
-        dim=0
+        dim=0,
     )
 
     if config.quantize != "gptq":
@@ -168,7 +168,9 @@ def _load_gqa(config, prefix: str, weights):
             config.hidden_size,
         ], f"{list(weight.shape)} != {[(num_heads + 2 * config.num_key_value_heads) * head_size, config.hidden_size]}"
 
-    return TensorParallelColumnLinear(get_linear(weight, bias=None, quantize=config.quantize))
+    return TensorParallelColumnLinear(
+        get_linear(weight, bias=None, quantize=config.quantize)
+    )
 
 
 class FlashLlamaAttention(torch.nn.Module):
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
--- a/server/text_generation_server/models/flash_causal_lm.py
+++ b/server/text_generation_server/models/flash_causal_lm.py
@@ -39,6 +39,7 @@ def __init__(
         device: torch.device,
     ):
         self.block_size = BLOCK_SIZE
+        self.num_blocks = num_blocks
 
         element_size = torch.tensor([], dtype=dtype).element_size()
         x = self.block_size // element_size
@@ -714,7 +715,6 @@ def warmup(self, batch: FlashCausalLMBatch):
         global CACHE_MANAGER
 
         torch.cuda.empty_cache()
-        torch.cuda.reset_peak_memory_stats(self.device)
         try:
             CACHE_MANAGER = CacheManager(
                 batch.blocks,
@@ -731,23 +731,20 @@ def warmup(self, batch: FlashCausalLMBatch):
                 f"You need to decrease `--max-batch-prefill-tokens`"
             ) from e
 
-        # Inspired by the original implementation in [vllm](https://github.com/vllm-project/vllm)
-        # Calculate the number of blocks that can be allocated with the
-        # profiled peak memory.
         torch.cuda.synchronize(self.device)
-        peak_memory = torch.cuda.max_memory_reserved(self.device)
 
+        # Inspired by the original implementation in [vllm](https://github.com/vllm-project/vllm)
+        # Calculate the number of blocks that can be allocated with the free memory
         dtype_size = torch.tensor([], dtype=self.dtype).element_size()
         cache_block_size = BLOCK_SIZE * self.num_kv_heads * self.head_size
         total_cache_size = self.num_layers * cache_block_size * 2 * dtype_size
 
-        total_gpu_memory = torch.cuda.get_device_properties(self.device).total_memory
+        free_memory, _ = torch.cuda.mem_get_info(self.device)
 
-        # 0.98 to add some wiggle room
         num_blocks = (
-            int((total_gpu_memory * 0.98 - peak_memory) // total_cache_size)
+            int(free_memory // total_cache_size)
             # Add batch.blocks as we allocated it above, so it is included in the peak memory.
-            + batch.blocks
+            + CACHE_MANAGER.num_blocks
         )
 
         del CACHE_MANAGER
diff --git a/server/text_generation_server/utils/gptq/quantize.py b/server/text_generation_server/utils/gptq/quantize.py
--- a/server/text_generation_server/utils/gptq/quantize.py
+++ b/server/text_generation_server/utils/gptq/quantize.py
@@ -864,8 +864,9 @@ def quantize(
     )
 
     with init_empty_weights():
-        model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.float16,
-                                                 trust_remote_code=trust_remote_code)
+        model = AutoModelForCausalLM.from_config(
+            config, torch_dtype=torch.float16, trust_remote_code=trust_remote_code
+        )
     model = model.eval()
 
     print("LOADED model")


Test patch: