Add memory analyzer & utomatically configure KV cache size (#6)

cacheflow/models/model_utils.py

@@ -1,21 +1,20 @@
-import random
 from typing import Union
 
-import numpy as np
 import torch
 import torch.nn as nn
 
+from cacheflow.models.memory_analyzer import CacheFlowMemoryAnalyzer
+from cacheflow.models.memory_analyzer import OPTMemoryAnalyzer
 from cacheflow.models.opt import OPTForCausalLM
+from cacheflow.models.utils import get_torch_dtype
 
-MODEL_CLASSES = {
+
+_MODELS = {
     'opt': OPTForCausalLM,
 }
 
-STR_DTYPE_TO_TORCH_DTYPE = {
-    'half': torch.half,
-    'float': torch.float,
-    'float16': torch.float16,
-    'float32': torch.float32,
+_MEMORY_ANALYZERS = {
+    'opt': OPTMemoryAnalyzer,
 }
 
 
@@ -23,20 +22,23 @@ def get_model(
     model_name: str,
     dtype: Union[torch.dtype, str],
 ) -> nn.Module:
-    if isinstance(dtype, str):
-        torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype.lower()]
-    else:
-        torch_dtype = dtype
-    for model_class, hf_model in MODEL_CLASSES.items():
+    torch_dtype = get_torch_dtype(dtype)
+    for model_class, hf_model in _MODELS.items():
         if model_class in model_name:
-            model = hf_model.from_pretrained(model_name, torch_dtype=torch_dtype)
+            model = hf_model.from_pretrained(
+                model_name, torch_dtype=torch_dtype)
             return model.eval()
-    raise ValueError(f'Invalid model name: {model_name}')
+    raise ValueError(f'Unsupported model name: {model_name}')
 
 
-def set_seed(seed: int) -> None:
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
+def get_memory_analyzer(
+    model_name: str,
+    block_size: int,
+    dtype: Union[torch.dtype, str],
+) -> CacheFlowMemoryAnalyzer:
+    torch_dtype = get_torch_dtype(dtype)
+    for model_class, memory_analyzer in _MEMORY_ANALYZERS.items():
+        if model_class in model_name:
+            return memory_analyzer(
+                model_name, block_size, torch_dtype)
+    raise ValueError(f'Unsupported model name: {model_name}')