Update record func, add count = gr.Number.

app.py

@@ -4,8 +4,8 @@ import pandas as pd
 col=['Layer number', 'Hidden size', 'FFN Hidden size', 'Sequence length', 'Head number', 'Group number', 
         'dp', 'tp', 'pp', 'cp', 'GPU numbers', 'Batch size', 'FP8', 'Model parameters', 'Model_states', 'Activation', 'Total']
 
-# global data
-table_data = pd.DataFrame(columns=col)
+# # global data
+# table_data = pd.DataFrame(columns=col)
 
 def Get_GigaByte(memory):
     return memory / 1024**3
@@ -267,7 +267,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, act_func,
-        dp, tp, pp, cp, is_sp, is_ip, vp, is_dist_opt, b, b_global, is_fp8, is_fp8_init, g_ty, o_ty):
+        dp, tp, pp, cp, is_sp, is_ip, vp, is_dist_opt, b, b_global, is_fp8, is_fp8_init, g_ty, o_ty, record_df, count):
     # get model states
     numParameters, weight_memory, gradient_memory, optimizer_memory, master_weight_memory, model_states_memory = Compute_Model_states(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)
@@ -288,11 +288,14 @@ def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_l
     Total = round(model_states_memory + activation_memory, 3)
 
     # record
-    global table_data
     new_row = pd.DataFrame([[layer_num, hidden_size, ffn_size, seq_length, head_num, group_query_num, dp, tp, pp, cp, gpu_num, b, is_fp8, 
                             numParametersTotal, model_states_memory, activation_memory, Total]], 
                             columns=col)
-    table_data = table_data._append(new_row, ignore_index=True)
+    if count == 1:
+        record_df = new_row
+    else:    
+        record_df = record_df._append(new_row, ignore_index=True)
+    count = count + 1
 
     # return str(gpu_num), str(model_states) + " GB", str(activation) + " GB", str(total) + " GB", table_data
     return f"""
@@ -302,25 +305,16 @@ def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_l
                 Model_states = {str(model_states_memory)} GB, \n
                 Activation = {str(activation_memory)} GB, \n
                 Total memory consumption = {str(Total)} GB \n
-           """, table_data
+           """, record_df, count
 
-def generate_csv():
-    # 创建示例 DataFrame
-    df = table_data
-    
+def generate_csv(record_df):
     # 将 DataFrame 保存为 CSV 文件
     csv_filename = "data.csv"
-    df.to_csv(csv_filename, index=False)
+    record_df.to_csv(csv_filename, index=False)
     
     # 返回 CSV 文件路径
     return csv_filename
 
-        # P = P_{input} + P_{middle} + 
-        # \begin{cases} 
-        # P_{output},  & \text{if }pp = 1 \\\\
-        # 0, & \text{if }pp > 1
-        # \end{cases} \\\\
-
 # formula string
 formula = r"""
         > **Note**🔑: In this formula, we assume LLM training with FP32 Gradient and Optimizer state, and bias = False, Zero1 = False, SP = True. 
@@ -491,18 +485,19 @@ with gr.Blocks() as demo:
         label="Record Table",
         headers=col
     )
-
+    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,
-                dp, tp, pp, cp, is_sp, is_ip, vp, is_dist_opt, b, b_global, is_fp8, is_fp8_init, g_ty, o_ty],
-        outputs=[output_text, record_df]
+                dp, tp, pp, cp, is_sp, is_ip, 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]
     )
 
     output_file=gr.File(label="When you click the download button, the downloaded form will be displayed here.")
     # download func
     download_btn.click(
         fn=generate_csv,
+        inputs=record_df,
         outputs=output_file
     )