Create app.py

app.py

@@ -0,0 +1,85 @@
+import gradio as gr
+import math
+
+# Helper function to pretty-print message sizes
+def convert_params(params):
+    if params == 0:
+        return "0"
+    size_name = ("", "K", "M", "B", "T", "P", "E", "Z", "Y")
+    i = int(math.floor(math.log(params, 1000)))
+    p = math.pow(1000, i)
+    s = round(params / p, 2)
+    return "%s %s" % (s, size_name[i])
+
+# calculates the params of a model given their hparams
+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
+
+    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
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# Transformer Model Parameter Calculator")
+
+    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)
+    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)
+    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)
+
+    result = gr.Textbox(label="Output", interactive=False)
+
+    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)
+
+    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])
+
+demo.launch()