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', 'Number of GPUs', 'B', 'FP8', 'Model parameters (B)', 'Model states (GB)', 'Activation (GB)', 'Total (GB)']
+        'DP', 'TP', 'PP', 'CP', 'GPUs', 'B', 'FP8', 'Model parameters (B)', 'Model states (GB)', 'Activation (GB)', 'Total (GB)']
 
 abbr = """
     <div align="center">
@@ -31,7 +31,7 @@ def Compute_Parameters_input(seq_length, hidden_size, vocab_size, act_func, tp):
     if act_func == "LLaMA":
         num_parameters_position_embedding = 0 
     else:
-        num_parameters_position_embedding = seq_length * hidden_size
+        num_parameters_position_embedding = seq_length * hidden_size / tp
 
     return num_parameters_word_embedding + num_parameters_position_embedding
 
@@ -119,13 +119,15 @@ def Compute_Parameters(seq_length, vocab_size, layer_num, hidden_size, ffn_size,
     
     return num_parameters_total
 
-def Compute_Weight(numParametersTotal, is_fp8, is_fp8_init):
-    if is_fp8 == "False":
-        weight_memory = 2 * numParametersTotal
-    elif is_fp8_init == "False":
+def Compute_Weight(numParametersTotal, precision, is_fp8, is_fp8_init):
+    weight_memory = 0
+    if precision == "FP32":
         weight_memory = 4 * numParametersTotal
     else:
-        weight_memory = 2 * numParametersTotal 
+        weight_memory = 2 * numParametersTotal
+
+    if is_fp8 == "True" and is_fp8_init == "False":
+        weight_memory += 2 * numParametersTotal
     
     return weight_memory
 
@@ -137,7 +139,7 @@ def Compute_Gradient(numParametersTotal, g_ty):
     
     return gradient_memory
 
-def Compute_Optimizer_states(numParametersTotal, o_ty, is_dist_opt, dp, cp):
+def Compute_Optimizer_states(numParametersTotal, opt_func, o_ty, is_dist_opt, dp, cp):
     if o_ty == "FP32":
         optimizer_memory = 4 * 2 * numParametersTotal 
     elif o_ty =="BF16":
@@ -146,23 +148,30 @@ def Compute_Optimizer_states(numParametersTotal, o_ty, is_dist_opt, dp, cp):
     if is_dist_opt == "True":
         optimizer_memory = optimizer_memory / (dp * cp)
 
+    # for SGD, we have no optimizer states
+    if opt_func == "SGD":
+        optimizer_memory = 0
+
     return optimizer_memory
 
-def Compute_Master_weight(numParametersTotal, is_dist_opt, dp, cp):
-    master_weight_memory = 4 * numParametersTotal
+def Compute_Master_weight(numParametersTotal, precision, is_dist_opt, dp, cp):
+    if precision == "BF16":
+        master_weight_memory = 4 * numParametersTotal
+    else:
+        master_weight_memory = 0
     if is_dist_opt == "True":
         master_weight_memory = master_weight_memory / (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, is_tie_word_embedding, act_func,
-        dp, tp, pp, cp, is_dist_opt, is_fp8, is_fp8_init, g_ty, o_ty):
+        dp, tp, pp, cp, is_dist_opt, precision, is_fp8, is_fp8_init, g_ty, opt_func, o_ty):
     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)
+    weight_memory = Compute_Weight(numParametersTotal, precision, is_fp8, is_fp8_init)
     gradient_memory = Compute_Gradient(numParametersTotal, g_ty)
-    optimizer_memory = Compute_Optimizer_states(numParametersTotal, o_ty, is_dist_opt, dp, cp)
-    master_weight_memory = Compute_Master_weight(numParametersTotal, is_dist_opt, dp, cp)
+    optimizer_memory = Compute_Optimizer_states(numParametersTotal, opt_func, o_ty, is_dist_opt, dp, cp)
+    master_weight_memory = Compute_Master_weight(numParametersTotal, precision, is_dist_opt, dp, cp)
 
     return numParametersTotal, weight_memory, gradient_memory, optimizer_memory, master_weight_memory, \
             weight_memory + gradient_memory + optimizer_memory + master_weight_memory
@@ -298,7 +307,7 @@ def compute_activation_memory(vocab_size, seq_length, layer_num, b, b_global, he
 
 # 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, 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):
+        dp, tp, pp, cp, is_sp, vp, is_dist_opt, b, b_global, precision, is_fp8, is_fp8_init, g_ty, opt_func, o_ty, record_df, count):
     # data type trans
     if is_group_query == "True":
         group_query_num = int(group_query_num)
@@ -312,7 +321,7 @@ 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, is_tie_word_embedding, 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, precision, is_fp8, is_fp8_init, g_ty, opt_func, 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)
@@ -344,7 +353,7 @@ def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_l
                 GPU numbers = {str(gpu_num)}, \n
                 Model parameters = {str(numParametersTotal)} B, \n
                 Model parameters on each device = {str(numParameters)} B, \n
-                Model_states = {str(model_states_memory)} GB, \n
+                Model_states = Weight + Gradient + Optimizer = {str(model_states_memory)} GB, \n
                 Activation = {str(activation_memory)} GB, \n
                 Total memory consumption = {str(Total)} GB \n
            """, record_df, count
@@ -389,7 +398,7 @@ formula = r"""
         $$
         {Activation} = 
         (1 + \frac{pp-1}{pp \times vp}) \times
-        \frac{(8BS + BSH) \times pp + 15BSH + 5BS \times FFN}{tp \times cp}
+        \frac{(8BS + BSH) \times pp + (15BSH + 5BS \times FFN) \times L}{tp \times cp}
         $$
 
         ***
@@ -494,7 +503,7 @@ with gr.Blocks() as demo:
             )
             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.]")
+                act_func = gr.Radio(["LLaMA", "GPT"], value="LLaMA", label="Model type", info="eg. LLaMa: SwiGLU, RoPE, RMSNorm") #, 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 (V)", value=32000)
                     layer_num = gr.Number(label="Layer number (L)", value=32)
@@ -549,13 +558,14 @@ with gr.Blocks() as demo:
                 # with gr.Row():
                 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")
-                is_fp8 = gr.Radio(["True", "False"], value="True", label="FP8 Training")
-                is_fp8_init = gr.Radio(["True", "False"], value="True", label="FP8 Initialization(will reduce memory)")
-                # with gr.Row():
+                precision = gr.Dropdown(["FP32", "BF16"], value="BF16", label="Training precision")
+                with gr.Row():
+                    is_fp8 = gr.Radio(["True", "False"], value="True", label="FP8 Training")
+                    is_fp8_init = gr.Radio(["True", "False"], value="True", label="FP8 Initialization(will reduce memory)")
                 g_ty = gr.Dropdown(["FP32", "BF16"], value="FP32", label="Gradients Dtype")
-                o_ty = gr.Dropdown(["FP32", "BF16"], value="FP32", label="Optimizer State Dtype")
+                with gr.Row():
+                    opt_func = gr.Radio(["Adam", "SGD"], value="Adam", label="Optimizer function")
+                    o_ty = gr.Dropdown(["FP32", "BF16"], value="FP32", label="Optimizer State Dtype")
 
     compute_btn = gr.Button("Compute")
     with gr.Tab("Output"):
@@ -590,7 +600,7 @@ with gr.Blocks() as demo:
     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, 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],
+                dp, tp, pp, cp, is_sp, vp, is_dist_opt, b, b_global, precision, is_fp8, is_fp8_init, g_ty, opt_func, o_ty, record_df, count],
         outputs=[output_text, record_df, count]
     )