Update app.py

app.py

@@ -5,7 +5,7 @@ import pandas as pd
 #         'dp', 'tp', 'pp', 'cp', 'GPU numbers', 'Batch size', 'FP8', 'Model parameters', 'Model_states', 'Activation', 'Total']
 
 col=['L', 'H', 'FFN', 'S', 'A', 'G', 
-        'dp', 'tp', 'pp', 'cp', 'GPU number', 'Batch size', 'FP8', 'Model parameters', 'Model states', 'Activation', 'Total']
+        'dp', 'tp', 'pp', 'cp', 'Number of GPUs', 'B', 'FP8', 'Model parameters (B)', 'Model states (GB)', 'Activation (GB)', 'Total (GB)']
 
 abbr = """
     <div align="center">
@@ -35,9 +35,18 @@ def Compute_Parameters_input(seq_length, hidden_size, vocab_size, act_func, tp):
 
     return num_parameters_word_embedding + num_parameters_position_embedding
 
-def Compute_Parameters_output(hidden_size, vocab_size, tp):
-    num_parameters_output_layernorm = 2 * hidden_size
-    num_parameters_output_embedding = 0 # due to sharedWordEmbedding
+def Compute_Parameters_output(hidden_size, vocab_size, is_tie_word_embedding, act_func, tp):
+    # layernorm: h/2h
+    if act_func == "LLaMA":
+        num_parameters_output_layernorm = hidden_size # RMSNorm
+    else:
+        num_parameters_output_layernorm = 2 * hidden_size # LayerNorm
+
+    if is_tie_word_embedding == "True":
+        num_parameters_output_embedding = 0 # due to sharedWordEmbedding
+    else:
+        num_parameters_output_embedding = hidden_size * vocab_size / tp
+
     return num_parameters_output_layernorm + num_parameters_output_embedding
 
 def Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, act_func, tp):
@@ -79,7 +88,7 @@ def Compute_Parameters_mlp(hidden_size, ffn_size, is_bias, act_func, tp):
     
     return num_parameters_mlp
 
-def Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, act_func, head_num, tp, pp):
+def Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func, head_num, tp, pp):
     if is_group_query == "False":
         group_query_num = head_num
     kv_hidden_size = hidden_size / head_num * group_query_num
@@ -94,7 +103,7 @@ def Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size,
     num_parameters_in_total_layers = num_parameters_in_single_layer * layer_num / pp
     
     # output part
-    parameters_output = Compute_Parameters_output(hidden_size, vocab_size, tp)    
+    parameters_output = Compute_Parameters_output(hidden_size, vocab_size, is_tie_word_embedding, act_func, tp)    
 
     if pp == 1:
         num_parameters_total = (
@@ -146,9 +155,9 @@ def Compute_Master_weight(numParametersTotal, is_dist_opt, dp, cp):
     
     return master_weight_memory
 
-def Compute_Model_states(seq_length, vocab_size, layer_num, hidden_size, ffn_size, head_num, is_group_query, group_query_num, is_bias, act_func,
+def Compute_Model_states(seq_length, vocab_size, layer_num, hidden_size, ffn_size, head_num, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func,
         dp, tp, pp, cp, is_dist_opt, is_fp8, is_fp8_init, g_ty, o_ty):
-    numParametersTotal = Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, act_func, head_num, tp, pp)
+    numParametersTotal = Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func, head_num, tp, pp)
 
     weight_memory = Compute_Weight(numParametersTotal, is_fp8, is_fp8_init)
     gradient_memory = Compute_Gradient(numParametersTotal, g_ty)
@@ -160,7 +169,7 @@ def Compute_Model_states(seq_length, vocab_size, layer_num, hidden_size, ffn_siz
 
 # activation memory:
 def compute_activation_memory_attention(activation_dtype, seq_length, b, hidden_size, kv_hidden_size, is_sp, tp):
-    # LN 2bsq
+    # LN 2bsh
     activation_mem_attn_ln = seq_length * b * hidden_size * 2
     if is_sp == "False":
         activation_mem_attn_ln *= tp
@@ -288,11 +297,13 @@ def compute_activation_memory(vocab_size, seq_length, layer_num, b, b_global, he
     return activation_memory / tp / cp
 
 # compute_btn.click.function
-def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_length, head_num, is_group_query, group_query_num, is_bias, act_func,
+def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_length, head_num, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func,
         dp, tp, pp, cp, is_sp, vp, is_dist_opt, b, b_global, is_fp8, is_fp8_init, g_ty, o_ty, record_df, count):
     # data type trans
     if is_group_query == "True":
         group_query_num = int(group_query_num)
+    else:
+        group_query_num = head_num
 
     # check input
     [result, Error_message] = check_input(dp, tp, pp, cp, hidden_size, head_num, layer_num, seq_length, vp, b, b_global)
@@ -301,13 +312,13 @@ def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_l
 
     # get model states
     numParameters, weight_memory, gradient_memory, optimizer_memory, master_weight_memory, model_states_memory = Compute_Model_states(seq_length, vocab_size, layer_num, hidden_size, 
-        ffn_size, head_num, is_group_query, group_query_num, is_bias, act_func, dp, tp, pp, cp, is_dist_opt, is_fp8, is_fp8_init, g_ty, o_ty)
+        ffn_size, head_num, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func, dp, tp, pp, cp, is_dist_opt, is_fp8, is_fp8_init, g_ty, o_ty)
 
     # get activation memory 
     activation_memory = compute_activation_memory(vocab_size, seq_length, layer_num, b, b_global, head_num, hidden_size, ffn_size, act_func, is_fp8, is_sp, is_group_query, group_query_num, tp, pp, dp, cp, vp)
 
     # get model parameters
-    numParametersTotal = Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, act_func, head_num, 1, 1)
+    numParametersTotal = Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func, head_num, 1, 1)
     # get gpu number
     gpu_num = dp * tp * pp * cp
 
@@ -349,9 +360,10 @@ def generate_csv(record_df):
 # formula string
 formula = r"""
         > **Note**🔑: In this formula, we assume LLM training with FP8 training.
-        > 1. Interleaved pipeline.
-        > 2. bias = False.
-        > 3. SP = True. 
+        > 1. LlaMA-family Model.
+        > 2. Interleaved pipeline.
+        > 3. bias = False.
+        > 4. SP = True. 
 
         <div align="center">
         <img src=file/T1.jpg width=50%/>
@@ -359,7 +371,7 @@ formula = r"""
 
         $$
         {Total\ Model\ parameters} = 
-        HV + HS + (4H^2 + 3H \times FFN + 2H) \times L + 2H + HV
+        HV + (4H^2 + 3H \times FFN + 2H) \times L + H
         $$
 
         ***
@@ -371,7 +383,7 @@ formula = r"""
         $$
         {Model\ states} = 
         (6 + \frac{12}{dp \times cp}) \times
-        (\frac{(\frac{4H^2 + 3H \times FFN}{tp} + 2H) \times L}{pp} + \frac{HV}{tp} + HS)
+        (\frac{(\frac{4H^2 + 3H \times FFN}{tp} + 2H) \times L}{pp} + \frac{HV}{tp})
         $$
 
         $$
@@ -477,26 +489,27 @@ with gr.Blocks() as demo:
             # Input 1.[Model Parameters]
             gr.Markdown(
                 """
-                <h1>Model Parameters:</h1>
+                <h2>Model Parameters:</h2>
                 """
             )
             with gr.Accordion("Model Parameters"):
                 # with gr.Row():
                 act_func = gr.Radio(["LLaMA", "GPT"], value="LLaMA", label="Model type") #, info="Action Function in MLP, whether to use GLU (Gated Linear Unit). [e.g \"True\" for LlaMA, \"False\" for GPT.]")
                 with gr.Row():
-                    vocab_size = gr.Number(label="Vocab size", value=32000)
-                    layer_num = gr.Number(label="Layer number", value=32)
+                    vocab_size = gr.Number(label="Vocab size (V)", value=32000)
+                    layer_num = gr.Number(label="Layer number (L)", value=32)
                 with gr.Row():
-                    hidden_size = gr.Number(label="Hidden size", value=4096)
-                    ffn_size = gr.Number(label="FFN Hidden size", value=11008)
+                    hidden_size = gr.Number(label="Hidden size (H)", value=4096)
+                    ffn_size = gr.Number(label="FFN Hidden size (FFN)", value=11008)
                 with gr.Row():
-                    sequence_len = gr.Number(label="Sequence length", value=2048)
-                    head_num = gr.Number(label="Number of Attention Heads", value=32)
+                    sequence_len = gr.Number(label="Sequence length (S)", value=2048)
+                    head_num = gr.Number(label="Number of Attention Heads (A)", value=32)
                 with gr.Row():
                     is_group_query = gr.Radio(["True", "False"], value="False", label="Use Group Query Attention")
-                    group_query_num = gr.Textbox(label="Number of Query Groups", max_lines=1, value=None, interactive=False)
-                is_bias = gr.Radio(["True", "False"], value="False", label="Use Bias")
-
+                    group_query_num = gr.Textbox(label="Number of Query Groups (G)", max_lines=1, value=None, interactive=False)
+                with gr.Row():
+                    is_bias = gr.Radio(["True", "False"], value="False", label="Use Bias")
+                    is_tie_word_embedding = gr.Radio(["True", "False"], value="False", label="Tie word embeddings")
                 # change editable function
                 def toggle_textbox_editable(radio_value):
                     # 根据 radio_value 的值来决定 textbox 是否可编辑
@@ -511,30 +524,30 @@ with gr.Blocks() as demo:
             # Input 2.[Parallelism]
             gr.Markdown(
                 """
-                <h1>Parallelism config:</h1>
+                <h2>Parallelism config:</h2>
                 """
             )
             with gr.Accordion("Parallelism config"):
                 # with gr.Row():
-                dp = gr.Number(label="Data parallelism", value=2)
-                tp = gr.Number(label="Tensor parallelism", value=2)
-                pp = gr.Number(label="Pipeline parallelism", value=2)
-                cp = gr.Number(label="Context parallelism", value=2)
+                dp = gr.Number(label="Data parallelism (dp)", value=2)
+                tp = gr.Number(label="Tensor parallelism (tp)", value=2)
+                pp = gr.Number(label="Pipeline parallelism (pp)", value=2)
+                cp = gr.Number(label="Context parallelism (cp)", value=1)
                 # with gr.Row():
                 is_sp = gr.Radio(["True", "False"], value="True", label="Sequence parallelism")
-                vp = gr.Number(label="Virtual Pipeline Size")
+                vp = gr.Number(label="Virtual Pipeline Size (vp)")
                 is_dist_opt = gr.Radio(["True", "False"], value="True", label="Use Distributed Optimizer(Zero1)")
 
         with gr.Column():
             # Input 3.[Training Settings]
             gr.Markdown(
                 """
-                <h1>Training Config:</h1>
+                <h2>Training Config:</h2>
                 """
             )
             with gr.Accordion("Training Config"):
                 # with gr.Row():
-                b = gr.Number(label="Micro Batch size", value=4)
+                b = gr.Number(label="Micro Batch size (B)", value=4)
                 b_global = gr.Number(label="Global Batch size", value=64)
                 # with gr.Row():
                 gr.Checkbox(label="True", value=True, info="BF16 Training")
@@ -547,11 +560,11 @@ with gr.Blocks() as demo:
     compute_btn = gr.Button("Compute")
     with gr.Tab("Output"):
         with gr.Column():
-            gr.Markdown(
-                """
-                <h1>Output Data:</h1>
-                """
-            )
+            # gr.Markdown(
+            #     """
+            #     <h1>Output Data:</h1>
+            #     """
+            # )
             output_text = gr.Textbox(
                 label="Compute result", 
                 interactive=False, 
@@ -565,7 +578,7 @@ with gr.Blocks() as demo:
             , latex_delimiters=[{ "left": "$$", "right": "$$", "display": True }]
         )
 
-    gr.Markdown(abbr)
+    # gr.Markdown(abbr)
 
     record_df = gr.Dataframe(
         label="Record Table",
@@ -576,7 +589,7 @@ with gr.Blocks() as demo:
     count = gr.Number(label="Row count", value=1, visible=False)
     compute_btn.click(
         fn=Compute_ALL_Model_memory, 
-        inputs=[vocab_size, layer_num, hidden_size, ffn_size, sequence_len, head_num, is_group_query, group_query_num, is_bias, act_func,
+        inputs=[vocab_size, layer_num, hidden_size, ffn_size, sequence_len, head_num, is_group_query, group_query_num, is_bias, is_tie_word_embedding, act_func,
                 dp, tp, pp, cp, is_sp, vp, is_dist_opt, b, b_global, is_fp8, is_fp8_init, g_ty, o_ty, record_df, count],
         outputs=[output_text, record_df, count]
     )