Model dropdown

README.md

@@ -4,7 +4,7 @@ emoji: 🚀
 colorFrom: pink
 colorTo: blue
 sdk: gradio
-sdk_version: 4.36.0
+sdk_version: 3.40.1
 app_file: src/app.py
 pinned: false
 license: apache-2.0

requirements.txt

@@ -1,6 +1,6 @@
-accelerate<0.32.0
+accelerate @ git+https://github.com/huggingface/accelerate
 transformers
 timm
-huggingface_hub==0.19.4
+huggingface_hub
 tabulate
 einops

src/app.py

@@ -1,84 +1,45 @@
 import gradio as gr
 import pandas as pd
-from accelerate.utils import convert_bytes
 from hub_utils import check_for_discussion, report_results
-from huggingface_hub.utils import HfHubHTTPError
 from model_utils import calculate_memory, get_model
 
 
+# We need to store them as globals because gradio doesn't have a way for us to pass them in to the button
+MODEL = None
+
+
 def get_results(model_name: str, library: str, options: list, access_token: str):
-    model = get_model(model_name, library, access_token)
-    # try:
-    #     has_discussion = check_for_discussion(model_name)
-    # except HfHubHTTPError:
-    #     has_discussion = True
+    global MODEL
+    MODEL = get_model(model_name, library, access_token)
+    has_discussion = check_for_discussion(model_name)
     title = f"## Memory usage for '{model_name}'"
-    data = calculate_memory(model, options)
-    stages = {"model": [], "gradients": [], "optimizer": [], "step": []}
-    for i, option in enumerate(data):
-        for stage in stages:
-            stages[stage].append(option["Training using Adam (Peak vRAM)"][stage])
-        value = max(data[i]["Training using Adam (Peak vRAM)"].values())
-        if value == -1:
-            value = "N/A"
-        else:
-            value = convert_bytes(value)
-        data[i]["Training using Adam (Peak vRAM)"] = value
-
-    if any(value != -1 for value in stages["model"]):
-        out_explain = "## Training using Adam explained:\n"
-        out_explain += "When training on a batch size of 1, each stage of the training process is expected to have near the following memory results for each precision you selected:\n"
-        memory_values = pd.DataFrame(
-            columns=["dtype", "Model", "Gradient calculation", "Backward pass", "Optimizer step"]
-        )
-        for i, dtype in enumerate(options):
-            if stages["model"][i] != -1:
-                memory_values.loc[len(memory_values.index)] = [
-                    dtype,
-                    convert_bytes(stages["model"][i]),
-                    convert_bytes(stages["gradients"][i]),
-                    convert_bytes(stages["optimizer"][i]),
-                    convert_bytes(stages["step"][i]),
-                ]
-        return [
-            title,
-            gr.update(visible=True, value=pd.DataFrame(data)),
-            gr.update(visible=True, value=out_explain),
-            gr.update(visible=True, value=memory_values),
-        ]
-    else:
-        return [
-            title,
-            gr.update(visible=True, value=pd.DataFrame(data)),
-            gr.update(visible=False, value=""),
-            gr.update(visible=False, value=pd.DataFrame()),
-        ]
+    data = calculate_memory(MODEL, options)
+    return [title, gr.update(visible=True, value=pd.DataFrame(data)), gr.update(visible=not has_discussion)]
 
 
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown(
             """<img src="https://huggingface.co/spaces/hf-accelerate/model-memory-usage/resolve/main/measure_model_size.png" style="float: left;" width="250" height="250"><h1>🤗 Model Memory Calculator</h1>
+
     This tool will help you calculate how much vRAM is needed to train and perform big model inference
     on a model hosted on the 🤗 Hugging Face Hub. The minimum recommended vRAM needed for a model
     is denoted as the size of the "largest layer", and training of a model is roughly 4x its size (for Adam).
+
     These calculations are accurate within a few percent at most, such as `bert-base-cased` being 413.68 MB and the calculator estimating 413.18 MB.
+
     When performing inference, expect to add up to an additional 20% to this as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). 
     More tests will be performed in the future to get a more accurate benchmark for each model.
+
     Currently this tool supports all models hosted that use `transformers` and `timm`.
+
     To use this tool pass in the URL or model name of the model you want to calculate the memory usage for,
     select which framework it originates from ("auto" will try and detect it from the model metadata), and
     what precisions you want to use."""
         )
         out_text = gr.Markdown()
         out = gr.DataFrame(
-            headers=["dtype", "Largest Layer", "Total Size", "Training using Adam (Peak vRAM)"],
-            interactive=False,
-            visible=False,
-        )
-        out_explain = gr.Markdown()
-        memory_values = gr.DataFrame(
-            headers=["dtype", "Model", "Gradient calculation", "Backward pass", "Optimizer step"],
+            headers=["dtype", "Largest Layer", "Total Size", "Training using Adam"],
             interactive=False,
             visible=False,
         )
@@ -94,12 +55,18 @@ with gr.Blocks() as demo:
             access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
         with gr.Row():
             btn = gr.Button("Calculate Memory Usage")
+            post_to_hub = gr.Button(
+                value="Report results in this model repo's discussions!\n(Will open in a new tab)", visible=False
+            )
 
     btn.click(
         get_results,
         inputs=[inp, library, options, access_token],
-        outputs=[out_text, out, out_explain, memory_values],
-        api_name=False,
+        outputs=[out_text, out, post_to_hub],
+    )
+
+    post_to_hub.click(report_results, inputs=[inp, library, access_token]).then(
+        lambda: gr.Button.update(visible=False), outputs=post_to_hub
     )
 
 

src/hub_utils.py

@@ -13,7 +13,7 @@ def check_for_discussion(model_name: str):
     model_name = extract_from_url(model_name)
     discussions = list(api.get_repo_discussions(model_name))
     return any(
-        discussion.author == "model-sizer-bot"
+        discussion.title == "[AUTOMATED] Model Memory Requirements" and discussion.author == "model-sizer-bot"
         for discussion in discussions
     )
 
@@ -27,9 +27,9 @@ def report_results(model_name, library, access_token):
     post = f"""# Model Memory Requirements\n
 
 You will need about {data[1]} VRAM to load this model for inference, and {data[3]} VRAM to train it using Adam.
-
-These calculations were measured from the [Model Memory Utility Space](https://huggingface.co/spaces/hf-accelerate/model-memory-usage) on the Hub.
-
+    
+These calculations were measured from the [Model Memory Utility Space](https://hf.co/spaces/hf-accelerate/model-memory-utility) on the Hub.
+    
 The minimum recommended vRAM needed for this model assumes using [Accelerate or `device_map="auto"`](https://huggingface.co/docs/accelerate/usage_guides/big_modeling) and is denoted by the size of the "largest layer". 
 When performing inference, expect to add up to an additional 20% to this, as found by [EleutherAI](https://blog.eleuther.ai/transformer-math/). More tests will be performed in the future to get a more accurate benchmark for each model.
 

src/model_utils.py

@@ -3,7 +3,7 @@ from urllib.parse import urlparse
 
 import gradio as gr
 import torch
-from accelerate.commands.estimate import check_has_model, create_empty_model, estimate_training_usage
+from accelerate.commands.estimate import check_has_model, create_empty_model
 from accelerate.utils import calculate_maximum_sizes, convert_bytes
 from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
 
@@ -27,8 +27,8 @@ def extract_from_url(name: str):
         return path[1:]
 
 
-def translate_llama(text):
-    "Translates Llama-2 and CodeLlama to its hf counterpart"
+def translate_llama2(text):
+    "Translates llama-2 to its hf counterpart"
     if not text.endswith("-hf"):
         return text + "-hf"
     return text
@@ -36,8 +36,8 @@ def translate_llama(text):
 
 def get_model(model_name: str, library: str, access_token: str):
     "Finds and grabs model from the Hub, and initializes on `meta`"
-    if "meta-llama/Llama-2-" in model_name or "meta-llama/CodeLlama-" in model_name:
-        model_name = translate_llama(model_name)
+    if "meta-llama" in model_name:
+        model_name = translate_llama2(model_name)
     if library == "auto":
         library = None
     model_name = extract_from_url(model_name)
@@ -84,12 +84,10 @@ def calculate_memory(model: torch.nn.Module, options: list):
         dtype_largest_layer = largest_layer[0]
 
         modifier = DTYPE_MODIFIER[dtype]
-        dtype_training_size = estimate_training_usage(
-            dtype_total_size, dtype if dtype != "float16/bfloat16" else "float16"
-        )
         dtype_total_size /= modifier
         dtype_largest_layer /= modifier
 
+        dtype_training_size = convert_bytes(dtype_total_size * 4)
         dtype_total_size = convert_bytes(dtype_total_size)
         dtype_largest_layer = convert_bytes(dtype_largest_layer)
         data.append(
@@ -97,7 +95,7 @@ def calculate_memory(model: torch.nn.Module, options: list):
                 "dtype": dtype,
                 "Largest Layer or Residual Group": dtype_largest_layer,
                 "Total Size": dtype_total_size,
-                "Training using Adam (Peak vRAM)": dtype_training_size,
+                "Training using Adam": dtype_training_size,
             }
         )
     return data