Update app.py

app.py

@@ -1,7 +1,8 @@
 import gradio as gr
 import math
+from transformers import AutoConfig  # Required for Hugging Face integration
 
-# Helper function to pretty-print message sizes
+# ---- Helper Functions ---- #
 def convert_params(params):
     if params == 0:
         return "0"
@@ -11,57 +12,78 @@ def convert_params(params):
     s = round(params / p, 2)
     return "%s %s" % (s, size_name[i])
 
-# ---- Transformer Parameter Calculation ---- #
-def calc_params(vocab_size, tied_embeddings, hidden_size, sequence_length, num_layers, moe, num_experts, expert_interval, topk, ffn_expansion_factor, num_mlp_linears, kv_size_ratio):
-    if tied_embeddings:
-        embedding_params = hidden_size * vocab_size
-    else:
-        embedding_params = 2 * hidden_size * vocab_size
-    position_embedding_params = hidden_size * sequence_length
-    attention_params = int(2 * (1 + kv_size_ratio) * num_layers * hidden_size * hidden_size)
-    layernorm_params = 13 * num_layers * hidden_size
-
-    if moe:
-        num_expert_layers = num_layers / expert_interval
-        ffn_expert_params = num_mlp_linears * ffn_expansion_factor * num_expert_layers * num_experts * hidden_size * hidden_size
-        ffn_dense_params = num_mlp_linears * ffn_expansion_factor * (num_layers - num_expert_layers) * hidden_size * hidden_size
-        ffn_params = ffn_expert_params + ffn_dense_params
-        gating_params = num_expert_layers * hidden_size * num_experts
-    else:
-        ffn_params = num_mlp_linears * ffn_expansion_factor * num_layers * hidden_size * hidden_size
-
-    total_params = embedding_params + attention_params + ffn_params + position_embedding_params + layernorm_params
-
-    if moe:
-        total_params += gating_params
-
-    result = f"""
-    Embedding parameters: {convert_params(embedding_params)}
-    Attention parameters: {convert_params(attention_params)}
-    FFN parameters: {convert_params(ffn_params)}
-    {'Gating parameters: ' + convert_params(gating_params) if moe else ''}
-    Total Params in the Model: {convert_params(total_params)}
-    """
-    return result
-
-# ---- Memory Calculation Code (from the second script) ---- #
-def calc_mem(args):
+# Set defaults for missing arguments
+def set_defaults(args, defaults):
+    for key, value in defaults.items():
+        if getattr(args, key) is None:
+            setattr(args, key, value)
+    return args
+
+# Set value if it's None, else use the config value
+def set_if_none(args, key, config, config_key, defaults):
+    if getattr(args, key) is None:
+        setattr(args, key, config.get(config_key, defaults[key]))
+    return args
+
+# Get Hugging Face model arguments
+def get_hf_model_args(args, defaults):
+    if args.hf_model_name_or_path:
+        try:
+            config = AutoConfig.from_pretrained(args.hf_model_name_or_path, trust_remote_code=True).to_dict()
+        except Exception as e:
+            raise gr.Error(f"Error fetching Hugging Face model: {str(e)}")
+        
+        # Update arguments with Hugging Face model config values
+        args.num_layers = config.get("num_hidden_layers", defaults["num_layers"])
+        args.hidden_size = config.get("hidden_size", defaults["hidden_size"])
+        args.num_attention_heads = config.get("num_attention_heads", defaults["num_attention_heads"])
+        args.vocab_size = config.get("vocab_size", defaults["vocab_size"])
+        args.sequence_length = config.get("max_position_embeddings", defaults["sequence_length"])
+
+    return set_defaults(args, defaults)
+
+# ---- Memory Calculation ---- #
+def calc_mem(hf_model_name_or_path, num_gpus, tensor_parallel_size, pipeline_parallel_size, batch_size_per_gpu, sequence_length, vocab_size, hidden_size, num_attention_heads, num_layers, ffn_expansion_factor, is_mixed_precision, misc_mem_gib):
+    
+    # Define defaults
+    defaults = {
+        "num_layers": 44,
+        "hidden_size": 6144,
+        "num_attention_heads": 64,
+        "vocab_size": 51200,
+        "sequence_length": 2048,
+        "ffn_expansion_factor": 4,
+    }
+    
+    # Create a simple args object to simulate parsed arguments
+    class Args:
+        def __init__(self, **kwargs):
+            for key, value in kwargs.items():
+                setattr(self, key, value)
+
+    args = Args(hf_model_name_or_path=hf_model_name_or_path, num_gpus=num_gpus, tensor_parallel_size=tensor_parallel_size,
+                pipeline_parallel_size=pipeline_parallel_size, batch_size_per_gpu=batch_size_per_gpu, sequence_length=sequence_length,
+                vocab_size=vocab_size, hidden_size=hidden_size, num_attention_heads=num_attention_heads, num_layers=num_layers,
+                ffn_expansion_factor=ffn_expansion_factor, is_mixed_precision=is_mixed_precision, misc_mem_gib=misc_mem_gib)
+
+    # Fetch Hugging Face model args if a model is provided
+    args = get_hf_model_args(args, defaults)
+
     dp_degree = args.num_gpus / (args.tensor_parallel_size * args.pipeline_parallel_size)
     embed_params = 2 * args.vocab_size * args.hidden_size
     positional_params = args.hidden_size * args.sequence_length
     ln_params = 8 * args.hidden_size * args.num_layers + (2 * args.hidden_size)
-    attention_params = int(2 * (1 + args.kv_size_ratio) * args.num_layers * args.hidden_size * args.hidden_size)
-    mlp_params = args.num_mlp_linears * args.num_layers * args.hidden_size * args.ffn_expansion_factor * args.hidden_size
+    attention_params = int(2 * (1 + args.ffn_expansion_factor) * args.num_layers * args.hidden_size * args.hidden_size)
+    mlp_params = args.ffn_expansion_factor * args.num_layers * args.hidden_size * args.hidden_size
     total_params = embed_params + positional_params + ln_params + attention_params + mlp_params
 
-    bytes_per_param = args.low_prec_bytes_per_val if args.is_mixed_precision else args.high_prec_bytes_per_val
+    bytes_per_param = 2 if args.is_mixed_precision else 4
     model_mem = total_params * bytes_per_param
-    per_gpu_model_mem = model_mem / (args.tensor_parallel_size * args.pipeline_parallel_size)
-    per_gpu_mem_gib = per_gpu_model_mem / 1024**3 + args.misc_mem_gib
+    per_gpu_mem_gib = (model_mem / (args.tensor_parallel_size * args.pipeline_parallel_size)) / 1024**3 + args.misc_mem_gib
 
     return f"Per-GPU Memory Required for Training: {per_gpu_mem_gib:.2f} GiB"
 
-# Gradio Interface
+# ---- Gradio Interface ---- #
 with gr.Blocks() as demo:
     with gr.Tabs():
         with gr.TabItem("Parameter Calculation"):
@@ -101,6 +123,6 @@ with gr.Blocks() as demo:
 
             memory_result = gr.Textbox(label="Memory Calculation Result", interactive=False)
             calc_memory_button = gr.Button("Calculate Memory")
-            calc_memory_button.click(calc_mem, inputs=[num_gpus, tensor_parallel_size, pipeline_parallel_size, batch_size_per_gpu, sequence_length, vocab_size, hidden_size, num_attention_heads, num_layers, ffn_expansion_factor, is_mixed_precision, misc_mem_gib], outputs=memory_result)
+            calc_memory_button.click(calc_mem, inputs=[hf_model_name_or_path, num_gpus, tensor_parallel_size, pipeline_parallel_size, batch_size_per_gpu, sequence_length, vocab_size, hidden_size, num_attention_heads, num_layers, ffn_expansion_factor, is_mixed_precision, misc_mem_gib], outputs=memory_result)
 
 demo.launch()