mismatch default value and choice list of gradio dropdown

app.py

@@ -16,6 +16,7 @@ PRECISION_TO_BYTES = {"float32": 4,
                       "int8": 1}
 
 ZERO_STAGES = [0, 1, 2, 3]
+BATCH_SIZES = [1,2,4,8,16,32,64]
 OPTIMIZERS = ["adam", "adamw", "sgd"]
 HUGGINGFACE_URL_CONFIG = "https://huggingface.co/{}/resolve/main/config.json"
 
@@ -151,33 +152,32 @@ def activations_memory(num_layers, sequence_length, micro_batch_size, hidden_siz
 
 def vram_required(model_size, hidden_size, sequence_length, num_layers, num_heads, micro_batch_size, num_gpus, optimizer, zero_stage, gradient_checkpointing, mixed_precision):
     # Reference: https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/
+
     model_vram = model_memory(model_size, mixed_precision=mixed_precision)
     gradients_vram = gradients_memory(model_size)
     optimizer_vram = optimizer_memory(model_size, optimizer=optimizer)
 
     # Baseline
     if zero_stage == 0:
-        aggregated_vram = model_vram + gradients_vram + optimizer_vram
+        pass
     # Optimizer state partitioning
-    if zero_stage == 1:
-        aggregated_vram = model_vram + gradients_vram + (optimizer_vram / num_gpus)
+    if zero_stage >= 1:
+        optimizer_vram = optimizer_vram / num_gpus
     # Gradient + Optimzer state partitioning
-    if zero_stage == 2:
-        aggregated_vram = model_vram + ((gradients_vram + optimizer_vram) / num_gpus)
+    if zero_stage >= 2:
+        gradients_vram = gradients_vram / num_gpus
     # Parameter partitioning + Gradient + Optimizer partitioning
     if zero_stage == 3:
-        aggregated_vram = (model_vram / num_gpus) + (gradients_vram / num_gpus) + (optimizer_vram / num_gpus)
+        aggregated_vram = model_vram / num_gpus
 
-    print(f"ZeRO stage {zero_stage} takes {aggregated_vram} GB")
+    aggregated_vram = round(model_vram, 2) + gradients_vram + optimizer_vram
 
     activations_vram = activations_memory(num_layers, sequence_length, micro_batch_size, hidden_size, num_heads)
     if gradient_checkpointing:
-        activations_vram = activations_vram ** 0.5
+        activations_vram = round(activations_vram ** 0.5, 2)
     
-    print(f"Activations require {activations_vram} GB with gradient checkpointing: {gradient_checkpointing}")
     total_vram = aggregated_vram + activations_vram
-    print(f"Estimated 'minimal' VRAM requirement on {num_gpus} GPUs per GPU is {total_vram} GB")
-    return total_vram
+    return {"total": total_vram, "model": model_vram, "gradients": gradients_vram, "optimizer": optimizer_vram, "activations": activations_vram}
 
 def build_interface(estimate_vram_fn):
     training_params = []
@@ -190,11 +190,11 @@ def build_interface(estimate_vram_fn):
 
 
         with gr.Row(visible=False) as model_params_row:
-            model_params = [gr.Slider(label="Model Size", minimum=0.1, maximum=1000, step=0.1, value=7, info="Model size (in billion parameters)"),
+            model_params = [gr.Slider(label="Model Size", minimum=0.1, maximum=400, step=0.1, value=7, info="Model size (in billion parameters)"),
                             gr.Slider(label="Hidden size", minimum=256, maximum=8192, step=128, value=4096, info="Hidden size"),
-                            gr.Slider(label="Sequence length", minimum=256, maximum=1_000_000, step=256, value=8192, info="Sequence length"),
-                            gr.Slider(label="Num layers", minimum=1, maximum=64, step=1, value=32, info="Number of layers"),
-                            gr.Slider(label="Num heads", minimum=1, maximum=64, step=1, value=32, info="Number of attention heads")
+                            gr.Slider(label="Sequence length", minimum=256, maximum=128_000, step=256, value=8192, info="Sequence length"),
+                            gr.Slider(label="Num layers", minimum=8, maximum=64, step=1, value=32, info="Number of layers"),
+                            gr.Slider(label="Num heads", minimum=8, maximum=64, step=1, value=32, info="Number of attention heads")
             ]
 
 
@@ -212,16 +212,17 @@ def build_interface(estimate_vram_fn):
         
         
         with gr.Row(equal_height=True):
-            training_params = [gr.Dropdown(label="Micro batch size", choices=[1,2,4,8,16,32,64], value=4, info="Micro batch size (batch size per device/GPU)"),
+            training_params = [gr.Dropdown(label="Micro batch size", choices=BATCH_SIZES, value=4, info="Micro batch size (batch size per device/GPU)"),
                                gr.Dropdown(label="ZeRO stage", choices=ZERO_STAGES, value=0, info="ZeRO optimization stage"),
-                               gr.Dropdown(label="Gradient checkpointing", choices=[True, False], value=False, info="Enable gradient checkpointing"),
-                               gr.Dropdown(label="Mixed preision", choices=[True, False], value=False, info="Enable mixed precision for model training"),
+                               gr.Dropdown(label="Gradient checkpointing", choices=[True, False], value=True, info="Enable gradient checkpointing"),
+                               gr.Dropdown(label="Mixed precision", choices=[False, True], value=False, info="Enable mixed precision for model training"),
                                gr.Dropdown(label="Optimizer", choices=OPTIMIZERS, value="adamw", info="Type of optimizer"),
                                gr.Slider(label="Num GPUs", minimum=1, maximum=64, step=1, value=4, info="Number of GPUs. Necessary for estimating ZeRO stages"),
                                gr.Textbox(label="Cache dir", value=None, placeholder=".huggingface_configs", info="HuggingFace cache directory to download config from")
             ]
 
         submit_btn = gr.Button("Estimate!")
+        
         output = gr.Textbox(label="Total estimated VRAM per device/GPU (in GB)")
 
         submit_btn.click(
@@ -235,22 +236,24 @@ def build_interface(estimate_vram_fn):
 
 def estimate_vram(arg_keys, *args):
     params = dict(zip(arg_keys, args))
-    print(params)
+    print("Parameters: ", params)
 
     model_config = ModelConfig(params["model_size"], params["hidden_size"], params["sequence_length"], params["num_layers"], params["num_heads"])
     training_config = TrainingConfig(params["micro_batch_size"], params["num_gpus"], params["optimizer"], params["zero_stage"], params["gradient_checkpointing"], params["mixed_precision"])
-    if params["repo_id"]:
-        # If cache directory set, then download config
-        if params["cache_dir"]:
-            config = scrape_config_from_hub(params["repo_id"])
-            model_config.overwrite_with_hf_config(config)
-        # By default, scrape config.json from hub
-        else:
-            config = download_config_from_hub(params["repo_id"], params["cache_dir"])
-            model_config.overwrite_with_hf_config(config.to_dict())
+    if not params["repo_id"]:
+        return "No model selected!"
+    # If cache directory set, then download config
+    if params["cache_dir"]:
+        config = scrape_config_from_hub(params["repo_id"])
+        model_config.overwrite_with_hf_config(config)
+    # By default, scrape config.json from hub
+    else:
+        config = download_config_from_hub(params["repo_id"], params["cache_dir"])
+        model_config.overwrite_with_hf_config(config.to_dict())
 
-    total_vram = vram_required(**vars(model_config), **vars(training_config))
-    return total_vram
+    total_vram_dict = vram_required(**vars(model_config), **vars(training_config))
+    output_str = f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['gradients']}GB (gradients) + {total_vram_dict['optimizer']}GB (optimizer) + {total_vram_dict['activations']}GB activations"
+    return output_str
 
 if __name__ == "__main__":
     parser = parse_args()
@@ -276,4 +279,5 @@ if __name__ == "__main__":
                 config = scrape_config_from_hub(args.repo_id)
             model_config.overwrite_with_hf_config(config)
 
-        total_vram = vram_required(**vars(model_config), **vars(training_config))
+        total_vram_dict = vram_required(**vars(model_config), **vars(training_config))
+        print(f"Total {total_vram_dict['total']}GB = {total_vram_dict['model']}GB (model) + {total_vram_dict['gradients']}GB (gradients) + {total_vram_dict['optimizer']}GB (optimizer) + {total_vram_dict['activations']}GB activations")