Rework to calculate context size rather than relying on linear regression

app.py

@@ -1,78 +1,113 @@
+from typing import Any
 import gradio as gr
 import pandas as pd
-
-quants = (
-    pd.read_csv("quants.csv")
-    .applymap(str)
-    .groupby("quant")["bpw"]
-    .apply(float)
-    .to_dict()
-)
-models = (
-    pd.read_csv("models.csv")
-    .applymap(str)
-    .groupby("model")["params"]
-    .apply(float)
-    .to_dict()
-)
-
-
-def context_sizes(model):
-    return pd.read_csv(
-        "context_sizes/" + model.replace("/", "_") + ".csv",
-        header=None,
-        names=["context", "size"],
+import requests
+
+quants = {
+    "Q2_K": 3.35,
+    "Q3_K_S": 3.5,
+    "Q3_K_M": 3.91,
+    "Q3_K_L": 4.27,
+    "Q4_0": 4.55,
+    "Q4_K_S": 4.58,
+    "Q4_K_M": 4.85,
+    "Q5_0": 5.54,
+    "Q5_K_S": 5.54,
+    "Q5_K_M": 5.69,
+    "Q6_K": 6.59,
+    "Q8_0": 8.5,
+}
+
+
+def calc_model_size(parameters: int, quant: float) -> int:
+    return parameters * quant // 8
+
+
+def get_model_config(hf_model: str) -> dict[str, Any]:
+    config = requests.get(
+        f"https://huggingface.co/{hf_model}/raw/main/config.json"
+    ).json()
+    model_index = {}
+    try:
+        model_index = requests.get(
+            f"https://huggingface.co/{hf_model}/raw/main/model.safetensors.index.json"
+        ).json()
+    except:
+        model_index = requests.get(
+            f"https://huggingface.co/{hf_model}/raw/main/pytorch_model.bin.index.json"
+        ).json()
+
+    # assume fp16 weights
+    config["parameters"] = model_index["metadata"]["total_size"] / 2
+    return config
+
+
+def calc_input_buffer_size(model_config, context: int) -> float:
+    return 4096 + 2048 * model_config["hidden_size"] + context * 4 + context * 2048
+
+
+def calc_compute_buffer_size(model_config, context: int) -> float:
+    return (
+        (context / 1024 * 2 + 0.75) * model_config["num_attention_heads"] * 1024 * 1024
     )
 
 
-def linear_regression(xs, ys) -> tuple[float, float]:
-    sum_y = ys.sum()
-    sum_x = sum(xs)
-    sum_xy = sum([x * y for x, y in zip(xs, ys)])
-    sum_x2 = sum([x**2 for x in xs])
-    n = len(xs)
-
-    a = (sum_y * sum_x2 - sum_x * sum_xy) / (n * sum_x2 - sum_x**2)
-    b = (n * sum_xy - sum_x * sum_y) / (n * sum_x2 - sum_x**2)
-    return a, b
-
-
-def calc_model_size(parameters: float, quant: float) -> float:
-    return round(parameters * quant / 8, 2)
-
-
-def calc_context_size(context, model) -> float:
-    sizes = context_sizes(model)
-    a, b = linear_regression(sizes["context"], sizes["size"])
-    return round((a + b * context) / 1024, 2)
+def calc_context_size(model_config, context: int) -> float:
+    n_gqa = model_config["num_attention_heads"] / model_config["num_key_value_heads"]
+    n_embd_gqa = model_config["hidden_size"] / n_gqa
+    n_elements = n_embd_gqa * (model_config["num_hidden_layers"] * context)
+    return 2 * n_elements * 2
 
 
 def calc(model_base, context, quant_size):
-    model_params = models[model_base]
-    quant_bpw = quants[quant_size]
-
-    model_size = calc_model_size(model_params, quant_bpw)
-    context_size = calc_context_size(context, model_base)
+    model_config = get_model_config(model_base)
+    quant_bpw = 0
+    try:
+        quant_bpw = float(quant_size)
+    except:
+        quant_bpw = quants[quant_size]
+
+    model_size = round(
+        calc_model_size(model_config["parameters"], quant_bpw) / 1024 / 1024 / 1024, 2
+    )
+    context_size = round(
+        (
+            calc_input_buffer_size(model_config, context)
+            + calc_context_size(model_config, context)
+            + calc_compute_buffer_size(model_config, context)
+        )
+        / 1024
+        / 1024
+        / 1024,
+        2,
+    )
 
-    return model_size, context_size, model_size + context_size
+    return model_size, context_size, round(model_size + context_size, 2)
 
 
 title = "GGUF VRAM Calculator"
 
 with gr.Blocks(title=title, theme=gr.themes.Monochrome()) as app:
-    default_model = "Mistral 7B"
+    default_model = "mistralai/Mistral-7B-v0.1"
     default_quant = "Q4_K_S"
     default_context = 8192
-    default_model_size = calc_model_size(models[default_model], quants[default_quant])
-    default_context_size = calc_context_size(default_context, default_model)
+    default_size = calc(default_model, default_context, default_quant)
+    default_model_size = default_size[0]
+    default_context_size = default_size[1]
 
-    gr.Markdown(f"# {app.title}")
-    model = gr.Dropdown(
-        list(models.keys()), value=default_model, label="Select Model Base"
+    gr.Markdown(
+        f"# {app.title}\nThis is meant only as a guide and is will not be 100% accurate, this also does not account for anything that might be running in the background on your system or CUDA system memory fallback on Windows"
+    )
+    model = gr.Textbox(
+        value=default_model,
+        label="Enter Unquantized HF Model Name (e.g. mistralai/Mistral-7B-v0.1)",
+    )
+    context = gr.Number(
+        minimum=1, value=default_context, label="Desired Context Size (Tokens)"
     )
-    context = gr.Number(minimum=1, value=default_context, label="Context Size (Tokens)")
-    quant = gr.Dropdown(
-        list(quants.keys()), value=default_quant, label="Select Quant Size"
+    quant = gr.Textbox(
+        value=default_quant,
+        label="Enter GGUF Quant (Q4_K_S) or BPW for other quantization schemes such as exl2",
     )
     btn = gr.Button(value="Submit", variant="primary")
     btn.click(
@@ -85,15 +120,15 @@ with gr.Blocks(title=title, theme=gr.themes.Monochrome()) as app:
         outputs=[
             gr.Number(
                 label="Model Size (GB)",
-                value=default_model_size,
+                value=default_size[0],
             ),
             gr.Number(
                 label="Context Size (GB)",
-                value=default_context_size,
+                value=default_size[1],
             ),
             gr.Number(
                 label="Total Size (GB)",
-                value=default_model_size + default_context_size,
+                value=default_size[2],
             ),
         ],
     )

context_sizes/Llama2 13B.csv

@@ -1,12 +0,0 @@
-512,475
-1024,912
-2048,1794
-3072,2676
-4096,3558
-6144,5322
-8192,7086
-12288,10614
-16384,14142
-24576,21198
-32768,28254
-65536,56508

context_sizes/Llama2 20B.csv

@@ -1,12 +0,0 @@
-512,695
-1024,1352
-2048,2674
-3072,3996
-4096,5318
-6144,7962
-8192,10606
-12288,15894
-16384,21182
-24576,31782.52
-32768,42335.26
-65536,84670.52

context_sizes/Llama2 70B.csv

@@ -1,12 +0,0 @@
-512,305
-1024,498
-2048,948
-3072,1398
-4096,1848
-6144,2748
-8192,3648
-12288,5448
-16384,7248
-24576,10848
-32768,14448
-65536,28896

context_sizes/Llama2 7B.csv

@@ -1,12 +0,0 @@
-512,326.5
-1024,602
-2048,1180
-3072,1758
-4096,2336
-6144,3492
-8192,4648
-12288,6960
-16384,9272
-24576,13896
-32768,18520
-65536,37016

context_sizes/Mistral 7B.csv

@@ -1,12 +0,0 @@
-512,137
-1024,218
-2048,412
-3072,606
-4096,800
-6144,1188
-8192,1576
-12288,2352
-16384,3128
-24576,4680
-32768,6232
-65536,12440

context_sizes/Mixtral 8x7B.csv

@@ -1,12 +0,0 @@
-512,181.72
-1024,249.22
-2048,443.22
-3072,637.22
-4096,831.22
-6144,1219.22
-8192,1607.22
-12288,2383.22
-16384,3159.22
-24576,4711.22
-32768,6263.22
-65536,12471.22

context_sizes/Solar 10.7B_11B.csv

@@ -1,12 +0,0 @@
-512,172.19
-1024,285.19
-2048,543.19
-3072,801.19
-4096,1059.19
-6144,1575.19
-8192,2091.19
-12288,3123.19
-16384,4155.19
-24576,6219.19
-32768,8283.19
-65536,16539.19

context_sizes/Yi 34B.csv

@@ -1,12 +0,0 @@
-512,262.19
-1024,399.19
-2048,753.19
-3072,1107.19
-4096,1461.19
-6144,2169.19
-8192,2877.19
-12288,4293.19
-16384,5709.19
-24576,8541.19
-32768,11373.19
-65536,22701.19

models.csv

@@ -1,9 +0,0 @@
-model,params
-Llama2 7B,7
-Llama2 13B,13
-Llama2 70B,70
-Mistral 7B,7
-Llama2 20B,20
-Mixtral 8x7B,46.7
-Yi 34B,34
-Solar 10.7B/11B,10.7

quants.csv

@@ -1,13 +0,0 @@
-quant,bpw
-Q2_K,3.35
-Q3_K_S,3.5
-Q3_K_M,3.91
-Q3_K_L,4.27
-Q4_0,4.55
-Q4_K_S,4.58
-Q4_K_M,4.85
-Q5_0,5.54
-Q5_K_S,5.54
-Q5_K_M,5.69
-Q6_K,6.59
-Q8_0,8.5