add mem_faction

Author: shihaobai

lightllm/common/basemodel/basemodel.py

@@ -55,6 +55,7 @@ def __init__(self, kvargs):
         self.graph_max_batch_size = kvargs.get("graph_max_batch_size", 16)
         self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192)
         self.disable_cudagraph = kvargs.get("disable_cudagraph", False)
+        self.mem_fraction = kvargs.get("mem_fraction", 0.9)
 
         self._init_datatype()
         self._init_config()
@@ -119,6 +120,7 @@ def _init_mem_manager(self):
             head_num=self.config["num_attention_heads"] // self.world_size_,
             head_dim=self.config["n_embed"] // self.config["num_attention_heads"],
             layer_num=self.config["n_layer"],
+            mem_fraction=self.mem_fraction,
         )
         self.max_total_token_num = self.mem_manager.size
         return

lightllm/common/int8kv_mem_manager.py

@@ -4,14 +4,14 @@
 
 
 class INT8KVMemoryManager(MemoryManager):
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True):
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9):
         self.kv_dtype = torch.int8
-        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=True)
+        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=mem_fraction)
 
     def get_cell_size(self):
-        return self.head_num * self.head_dim * self.layer_num * 2 * torch._utils._element_size(
+        return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(
             self.kv_dtype
-        ) + self.head_num * self.layer_num * 2 * torch._utils._element_size(self.dtype)
+        ) + 2 * self.head_num * self.layer_num * torch._utils._element_size(self.dtype)
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
         self.kv_buffer = [

lightllm/common/mem_manager.py

@@ -9,15 +9,15 @@
 
 
 class MemoryManager:
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False):
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9):
         self.size = size
         self.head_num = head_num
         self.head_dim = head_dim
         self.layer_num = layer_num
         self.always_copy = always_copy
-        self.kv_dtype = dtype
+        self.dtype = dtype
         # profile the max total token num if the size is None
-        self.profile_size()
+        self.profile_size(mem_fraction)
         # mem_state 修改为使用计数方式，方便后期实现token共享机制，实现beam search 等
         self.mem_state = torch.zeros((self.size,), dtype=torch.int32, device="cuda")
         self.indexes = torch.arange(0, self.size, dtype=torch.long, device="cuda")
@@ -33,26 +33,32 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False
         self.shared_can_use_token_num = SharedInt(f"{str(nccl_port)}_mem_manger_can_use_token_num")
 
         self.shared_can_use_token_num.set_value(self.can_use_mem_size)
-        self._init_buffers(self.size, dtype, head_num, head_dim, layer_num)
+        self._init_buffers(
+            self.size,
+            dtype,
+            head_num,
+            head_dim,
+            layer_num,
+        )
 
     def get_cell_size(self):
-        return self.head_num * self.head_dim * self.layer_num * 2 * torch._utils._element_size(self.kv_dtype)
+        return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(self.dtype)
 
-    def profile_size(self):
+    def profile_size(self, mem_fraction):
         if self.size is not None:
             return
         import torch.distributed as dist
 
         tp_rank = dist.get_rank()
         world_size = dist.get_world_size()
         total_memory = get_total_gpu_memory()
-        available_memory = get_available_gpu_memory(tp_rank, world_size) - total_memory * (1 - 0.9)
+        available_memory = get_available_gpu_memory(tp_rank, world_size) - total_memory * (1 - mem_fraction)
         cell_size = self.get_cell_size()
         self.size = int(available_memory * 1024 ** 3 / cell_size)
         logger.info(
             f"{str(available_memory)} GB space is available after load the model weight\n"
             f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n"
-            f"{self.size} is the profiled max_total_token_num with the mem_fraction 0.9\n"
+            f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n"
         )
         return
 

lightllm/common/ppl_int4kv_mem_manager.py

@@ -4,15 +4,15 @@
 
 
 class PPLINT4KVMemoryManager(MemoryManager):
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True):
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9):
         self.kv_dtype = torch.int8
         self.group_quant_size = 8
-        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=True)
+        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction)
 
     def get_cell_size(self):
-        return self.head_num * self.head_dim // 2 * self.layer_num * 2 * torch._utils._element_size(
+        return 2 * self.head_num * self.head_dim // 2 * self.layer_num * torch._utils._element_size(
             self.kv_dtype
-        ) + self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size(
+        ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size(
             self.dtype
         )
 

lightllm/common/ppl_int8kv_mem_manager.py

@@ -4,14 +4,27 @@
 
 
 class PPLINT8KVMemoryManager(MemoryManager):
-    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True):
-        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=True)
+    def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=True, mem_fraction=0.9):
+        self.kv_dtype = torch.int8
+        self.group_quant_size = 8
+        super().__init__(size, dtype, head_num, head_dim, layer_num, always_copy=always_copy, mem_fraction=mem_fraction)
+
+    def get_cell_size(self):
+        return 2 * self.head_num * self.head_dim * self.layer_num * torch._utils._element_size(
+            self.kv_dtype
+        ) + 2 * self.head_num * self.head_dim // self.group_quant_size * self.layer_num * torch._utils._element_size(
+            self.dtype
+        )
 
     def _init_buffers(self, size, dtype, head_num, head_dim, layer_num):
-        group_quant_size = 8
-        self.kv_buffer = [torch.empty((size, 2 * head_num, head_dim), dtype=torch.int8, device="cuda") for _ in range(layer_num)]
-        self.scale_buffer = [torch.empty((size, 2 * head_num, head_dim // group_quant_size), dtype=dtype, device="cuda") for _ in range(layer_num)]
-    
+        self.kv_buffer = [
+            torch.empty((size, 2 * head_num, head_dim), dtype=torch.int8, device="cuda") for _ in range(layer_num)
+        ]
+        self.scale_buffer = [
+            torch.empty((size, 2 * head_num, head_dim // self.group_quant_size), dtype=dtype, device="cuda")
+            for _ in range(layer_num)
+        ]
+
     def _free_buffers(self):
         self.kv_buffer = None
-        self.scale_buffer = None
+        self.scale_buffer = None

lightllm/models/deepseek2/model.py

@@ -48,6 +48,7 @@ def _init_mem_manager(self):
             head_num=1,
             head_dim=self.config["kv_lora_rank"] + self.config["qk_rope_head_dim"],
             layer_num=self.config["num_hidden_layers"],
+            mem_fraction=self.mem_fraction,
         )
         return
 

lightllm/models/gemma_2b/model.py

@@ -12,6 +12,7 @@
 
 from lightllm.common.mem_utils import MemoryManager
 
+
 class Gemma_2bTpPartModel(TpPartBaseModel):
     # weight class
     pre_and_post_weight_class = Gemma_2bPreAndPostLayerWeight
@@ -38,20 +39,22 @@ def _verify_params(self):
         # assert self.config["num_key_value_heads"] % self.world_size_ == 0
         assert self.config["num_attention_heads"] % self.world_size_ == 0
         return
-    
+
     def _init_custom(self):
         self._init_to_get_rotary()
         return
-    
+
     def _init_mem_manager(self):
-        self.mem_manager = MemoryManager(self.max_total_token_num,
-                                        dtype=self.data_type,
-                                        head_num=self.config["num_key_value_heads"], # [SYM] always == 1
-                                        head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
-                                        layer_num=self.config["num_hidden_layers"])       
+        self.mem_manager = MemoryManager(
+            self.max_total_token_num,
+            dtype=self.data_type,
+            head_num=self.config["num_key_value_heads"],  # [SYM] always == 1
+            head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
+            layer_num=self.config["num_hidden_layers"],
+            mem_fraction=self.mem_fraction,
+        )
         return
 
-
     def _init_to_get_rotary(self, default_base=10000):
         if self.config.get("rope_scaling", {}) is None:
             rope_scaling_factor = 1.0
@@ -64,16 +67,16 @@ def _init_to_get_rotary(self, default_base=10000):
             max_seq_len = self.config["max_sequence_length"]
         else:
             max_position_embeddings = self.config.get(
-                "max_position_embeddings",
-                2048 if base <= 10000.0 + 1e-5 else 16384
+                "max_position_embeddings", 2048 if base <= 10000.0 + 1e-5 else 16384
             )
             max_seq_len = max_position_embeddings * rope_scaling_factor
 
-        inv_freq = 1.0 / (base ** (torch.arange(0, self.head_dim_, 2, device="cpu", dtype=torch.float32) / self.head_dim_))
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, self.head_dim_, 2, device="cpu", dtype=torch.float32) / self.head_dim_)
+        )
         t = torch.arange(max_seq_len + 1024 * 64, device="cpu", dtype=torch.float32) / rope_scaling_factor
         freqs = torch.outer(t, inv_freq)
 
         self._cos_cached = torch.cos(freqs).to(self.data_type).cuda()
         self._sin_cached = torch.sin(freqs).to(self.data_type).cuda()
         return
-    

lightllm/models/internlm2_wquant/model.py

@@ -3,31 +3,36 @@
 import torch
 
 from lightllm.models.internlm2.layer_weights.pre_and_post_layer_weight import Internlm2PreAndPostLayerWeight
-from lightllm.models.internlm2_wquant.layer_weights.transformer_layer_weight import Internlm2TransformerLayerWeightQuantized
+from lightllm.models.internlm2_wquant.layer_weights.transformer_layer_weight import (
+    Internlm2TransformerLayerWeightQuantized,
+)
 from lightllm.models.internlm_wquant.model import InternlmTpPartModelWQuant
 from lightllm.common.mem_utils import select_mem_manager_class
 
 
 class Internlm2TpPartModelWQuant(InternlmTpPartModelWQuant):
     # weight class
-    pre_and_post_weight_class = Internlm2PreAndPostLayerWeight 
+    pre_and_post_weight_class = Internlm2PreAndPostLayerWeight
     transformer_weight_class = Internlm2TransformerLayerWeightQuantized
 
     def __init__(self, kvargs):
         super().__init__(kvargs)
-    
+
     def _verify_params(self):
         assert self.load_way in ["HF", "DS"], "llama only supports HF and DS format to load Now!"
         assert any("w4a16" in mode_ or "w8a16" in mode_ for mode_ in self.mode), "only for weight quant model"
         assert self.config["num_key_value_heads"] % self.world_size_ == 0
         assert self.config["num_attention_heads"] % self.world_size_ == 0
         return
-    
+
     def _init_mem_manager(self):
-        self.mem_manager = select_mem_manager_class(self.mode)(self.max_total_token_num, 
-                                                     dtype=torch.float16,
-                                                     head_num=self.config["num_key_value_heads"] // self.world_size_,
-                                                     head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
-                                                     layer_num=self.config["num_hidden_layers"],
-                                                     always_copy=True)
-        return
+        self.mem_manager = select_mem_manager_class(self.mode)(
+            self.max_total_token_num,
+            dtype=torch.float16,
+            head_num=self.config["num_key_value_heads"] // self.world_size_,
+            head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
+            layer_num=self.config["num_hidden_layers"],
+            always_copy=True,
+            mem_fraction=self.mem_fraction,
+        )
+        return

lightllm/models/internlm_wquant/model.py

@@ -3,7 +3,9 @@
 import torch
 
 from lightllm.models.internlm_wquant.layer_infer.transformer_layer_infer import InternlmTransformerLayerInferWquant
-from lightllm.models.internlm_wquant.layer_weights.transformer_layer_weight import InternlmTransformerLayerWeightQuantized
+from lightllm.models.internlm_wquant.layer_weights.transformer_layer_weight import (
+    InternlmTransformerLayerWeightQuantized,
+)
 from lightllm.models.llama.model import LlamaTpPartModel
 from lightllm.common.mem_utils import select_mem_manager_class
 
@@ -17,19 +19,24 @@ class InternlmTpPartModelWQuant(LlamaTpPartModel):
 
     def __init__(self, kvargs):
         super().__init__(kvargs)
-    
+
     def _verify_params(self):
         assert self.load_way in ["HF", "DS"], "llama only supports HF and DS format to load Now!"
-        assert any("w6a16" in mode_ or "w4a16" in mode_ or "w8a16" in mode_ for mode_ in self.mode), "only for weight quant model"
+        assert any(
+            "w6a16" in mode_ or "w4a16" in mode_ or "w8a16" in mode_ for mode_ in self.mode
+        ), "only for weight quant model"
         assert self.config["num_key_value_heads"] % self.world_size_ == 0
         assert self.config["num_attention_heads"] % self.world_size_ == 0
         return
-    
+
     def _init_mem_manager(self):
-        self.mem_manager = select_mem_manager_class(self.mode)(self.max_total_token_num, 
-                                                     dtype=torch.float16,
-                                                     head_num=self.config["num_key_value_heads"] // self.world_size_,
-                                                     head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
-                                                     layer_num=self.config["num_hidden_layers"],
-                                                     always_copy=True)
-        return
+        self.mem_manager = select_mem_manager_class(self.mode)(
+            self.max_total_token_num,
+            dtype=torch.float16,
+            head_num=self.config["num_key_value_heads"] // self.world_size_,
+            head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
+            layer_num=self.config["num_hidden_layers"],
+            always_copy=True,
+            mem_fraction=self.mem_fraction,
+        )
+        return

lightllm/models/llama/model.py

@@ -62,6 +62,7 @@ def _init_mem_manager(self):
             head_num=self.config["num_key_value_heads"] // self.world_size_,
             head_dim=self.config["hidden_size"] // self.config["num_attention_heads"],
             layer_num=self.config["num_hidden_layers"],
+            mem_fraction=self.mem_fraction,
         )
         return