Update app.py

app.py

@@ -11,34 +11,23 @@ def convert_params(params):
     s = round(params / p, 2)
     return "%s %s" % (s, size_name[i])
 
-# calculates the params of a model given their hparams
+# ---- 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):
-    # Calculate embedding and unembedding params. If tied, re-use the same params
     if tied_embeddings:
         embedding_params = hidden_size * vocab_size
     else:
         embedding_params = 2 * hidden_size * vocab_size
     position_embedding_params = hidden_size * sequence_length
-    # Each QKVO matrix is (hxh)
-    # Unless using GQA/MQA which makes K/V smaller
     attention_params = int(2 * (1 + kv_size_ratio) * num_layers * hidden_size * hidden_size)
-    # (4*2)lh from the layernorm weights and biases for each of the QKV and mlp_in layernorms, 1h for the final layernorm.
-    # the extra 4lh is a mystery but we include it here
     layernorm_params = 13 * num_layers * hidden_size
-    #ffn_params = 12 * num_layers * hidden_size * hidden_size
 
     if moe:
-        # the number of layers that are MoE. (e.g. interval is 2 for GShard)
         num_expert_layers = num_layers / expert_interval
-        # the number of FFN params for each MoE layer
         ffn_expert_params = num_mlp_linears * ffn_expansion_factor * num_expert_layers * num_experts * hidden_size * hidden_size
-        # the number of FFN params for every dense layer
         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
-        # the number of gating layer params assuming it's implemented as a simple linear layer
         gating_params = num_expert_layers * hidden_size * num_experts
     else:
-        # num_mlp_layers * (h x [ffn_expansion_factor * h]) FFN matrices
         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
@@ -55,34 +44,63 @@ def calc_params(vocab_size, tied_embeddings, hidden_size, sequence_length, num_l
     """
     return result
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Transformer Model Parameter Calculator")
+# ---- Memory Calculation Code (from the second script) ---- #
+def calc_mem(args):
+    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
+    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
+    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
 
-    vocab_size = gr.Number(label="Vocab Size", value=51200)
-    tied_embeddings = gr.Checkbox(label="Tied Embeddings", value=False)
-    hidden_size = gr.Number(label="Hidden Size", value=6144)
-    sequence_length = gr.Number(label="Sequence Length", value=2048)
-    num_layers = gr.Number(label="Number of Layers", value=44)
-    ffn_expansion_factor = gr.Number(label="FFN Expansion Factor", value=4)
-    num_mlp_linears = gr.Number(label="Number of Linear Layers per MLP Block", value=2)
-    kv_size_ratio = gr.Number(label="KV Size Ratio", value=1.0)
+    return f"Per-GPU Memory Required for Training: {per_gpu_mem_gib:.2f} GiB"
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    with gr.Tabs():
+        with gr.TabItem("Parameter Calculation"):
+            vocab_size = gr.Number(label="Vocab Size", value=51200)
+            tied_embeddings = gr.Checkbox(label="Tied Embeddings", value=False)
+            hidden_size = gr.Number(label="Hidden Size", value=6144)
+            sequence_length = gr.Number(label="Sequence Length", value=2048)
+            num_layers = gr.Number(label="Number of Layers", value=44)
+            ffn_expansion_factor = gr.Number(label="FFN Expansion Factor", value=4)
+            num_mlp_linears = gr.Number(label="Number of Linear Layers per MLP Block", value=2)
+            kv_size_ratio = gr.Number(label="KV Size Ratio", value=1.0)
 
-    # MoE Parameters inside an accordion
-    with gr.Accordion("MoE Parameters", open=False):
-        moe = gr.Checkbox(label="MoE", value=False)
-        num_experts = gr.Number(label="Number of Experts", value=8)
-        expert_interval = gr.Number(label="Expert Interval", value=1)
-        topk = gr.Number(label="Top k Routing", value=1)
+            with gr.Accordion("MoE Parameters", open=False):
+                moe = gr.Checkbox(label="MoE", value=False)
+                num_experts = gr.Number(label="Number of Experts", value=8)
+                expert_interval = gr.Number(label="Expert Interval", value=1)
+                topk = gr.Number(label="Top k Routing", value=1)
 
-    result = gr.Textbox(label="Output", interactive=False)
+            result = gr.Textbox(label="Output", interactive=False)
+            calculate_button = gr.Button("Calculate")
+            calculate_button.click(calc_params, inputs=[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], outputs=result)
 
-    def run_calculation(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):
-        return 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)
+        with gr.TabItem("Memory Calculation"):
+            hf_model_name_or_path = gr.Textbox(label="HuggingFace Model Name or Path", value="")
+            num_gpus = gr.Number(label="Number of GPUs", value=1)
+            tensor_parallel_size = gr.Number(label="Tensor Parallel Size", value=1)
+            pipeline_parallel_size = gr.Number(label="Pipeline Parallel Size", value=1)
+            batch_size_per_gpu = gr.Number(label="Batch Size per GPU", value=8)
+            sequence_length = gr.Number(label="Sequence Length", value=2048)
+            vocab_size = gr.Number(label="Vocab Size", value=51200)
+            hidden_size = gr.Number(label="Hidden Size", value=6144)
+            num_attention_heads = gr.Number(label="Number of Attention Heads", value=64)
+            num_layers = gr.Number(label="Number of Layers", value=44)
+            ffn_expansion_factor = gr.Number(label="FFN Expansion Factor", value=4)
+            is_mixed_precision = gr.Checkbox(label="Mixed Precision", value=True)
+            misc_mem_gib = gr.Number(label="Misc Memory Overhead (GiB)", value=5)
 
-    calculate_button = gr.Button("Calculate")
-    calculate_button.click(run_calculation, 
-                           inputs=[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],
-                           outputs=[result])
+            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)
 
 demo.launch()