Update app.py. [RoPE, RMSNorm]

app.py

@@ -25,9 +25,14 @@ def Get_BillionParameter(parameter):
     return parameter / 1000**3
 
 # model states:
-def Compute_Parameters_input(hidden_size, vocab_size, tp):
+def Compute_Parameters_input(seq_length, hidden_size, vocab_size, act_func, tp):
     num_parameters_word_embedding = hidden_size * vocab_size / tp
-    num_parameters_position_embedding = 0 #args.hidden_size * args.seq_length
+    # position embedding
+    if act_func == "LLaMA":
+        num_parameters_position_embedding = 0 
+    else:
+        num_parameters_position_embedding = seq_length * hidden_size
+
     return num_parameters_word_embedding + num_parameters_position_embedding
 
 def Compute_Parameters_output(hidden_size, vocab_size, tp):
@@ -35,10 +40,13 @@ def Compute_Parameters_output(hidden_size, vocab_size, tp):
     num_parameters_output_embedding = 0 # due to sharedWordEmbedding
     return num_parameters_output_layernorm + num_parameters_output_embedding
 
-def Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, tp):
+def Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, act_func, tp):
     # attention: 
-    # layernorm: 2h 
-    num_parameters_attention = 2 * hidden_size
+    # layernorm: h/2h
+    if act_func == "LLaMA":
+        num_parameters_mlp = hidden_size # RMSNorm
+    else:
+        num_parameters_mlp = 2 * hidden_size # LayerNorm
     # QKV weight: 3h*h/tp, bias: 3h/tp
     # output linear weight: h*h/tp, bias: h
     num_parameters_attention_Q_weight = hidden_size * hidden_size / tp
@@ -51,10 +59,13 @@ def Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, tp):
     
     return num_parameters_attention
 
-def Compute_Parameters_mlp(hidden_size, ffn_size, is_bias, act_func,  tp):
+def Compute_Parameters_mlp(hidden_size, ffn_size, is_bias, act_func, tp):
     # MLP: 
-    # layernorm: 2h
-    num_parameters_mlp = 2 * hidden_size
+    # layernorm: h/2h
+    if act_func == "LLaMA":
+        num_parameters_mlp = hidden_size # RMSNorm
+    else:
+        num_parameters_mlp = 2 * hidden_size # LayerNorm
     # mlp1 weight: h*ffn/tp, bias: ffn/tp
     # mlp2 weight: ffn*h/tp, bias: h
     if act_func == "LLaMA":
@@ -68,16 +79,16 @@ def Compute_Parameters_mlp(hidden_size, ffn_size, is_bias, act_func,  tp):
     
     return num_parameters_mlp
 
-def Compute_Parameters(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, 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
     
     # input part
-    num_parameters_input = Compute_Parameters_input(hidden_size, vocab_size, tp)
+    num_parameters_input = Compute_Parameters_input(seq_length, hidden_size, vocab_size, tp)
 
     # middle layers part
-    num_parameters_attention = Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, tp)
+    num_parameters_attention = Compute_Parameters_attention(hidden_size, kv_hidden_size, is_bias, act_func, tp)
     num_parameters_mlp = Compute_Parameters_mlp(hidden_size, ffn_size, is_bias, act_func, tp)
     num_parameters_in_single_layer = num_parameters_attention + num_parameters_mlp
     num_parameters_in_total_layers = num_parameters_in_single_layer * layer_num / pp
@@ -135,9 +146,9 @@ def Compute_Master_weight(numParametersTotal, is_dist_opt, dp, cp):
     
     return master_weight_memory
 
-def Compute_Model_states(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, act_func,
         dp, tp, pp, cp, is_dist_opt, is_fp8, is_fp8_init, g_ty, o_ty):
-    numParametersTotal = Compute_Parameters(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, act_func, head_num, tp, pp)
 
     weight_memory = Compute_Weight(numParametersTotal, is_fp8, is_fp8_init)
     gradient_memory = Compute_Gradient(numParametersTotal, g_ty)
@@ -289,14 +300,14 @@ def Compute_ALL_Model_memory(vocab_size, layer_num, hidden_size, ffn_size, seq_l
         return Error_message, 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, 
+    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)
 
     # 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(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, act_func, head_num, 1, 1)
     # get gpu number
     gpu_num = dp * tp * pp * cp