use real number in model to calculate ops and para

app.py

@@ -2,7 +2,7 @@
 
 import streamlit as st
 import pandas as pd
-from model_util import fetch_dictionary_content, load_parameter
+from model_util import fetch_dictionary_content, load_parameter, get_model, classify_module, get_module_tensors
 from calc_util import *
 from render_util import create_table, header4, header5
 
@@ -15,6 +15,9 @@ if 'model_config' not in st.session_state:
 def load_model_config(model_id):
     if 'model_id' in st.session_state['model_config'] and st.session_state['model_config']['model_id'] == model_id:
         return st.session_state['model_config']
+    if 'parameter_count' in st.session_state:
+        st.session_state.pop('parameter_count')
+
     model_config = {}
     dictionary_content = fetch_dictionary_content(model_id)
     if dictionary_content:
@@ -37,6 +40,14 @@ def load_model_config(model_id):
         model_config['max_position_embeddings'] = 2048
         model_config['layernorm_operation'] = 2
 
+    try:
+        model_config['model'] = get_model(model_id, None, None)
+        module_tensors = get_module_tensors(model_config['model'])
+        model_config['module_classes'] = classify_module(module_tensors)
+    except:
+        model_config['model'] = None
+        model_config['module_classes'] = None
+
     st.session_state['model_config'] = model_config
     return model_config
 
@@ -45,7 +56,6 @@ subtotal_parameters = [
     'embedding_weights',
     'attention_weights',
     'mlp_weights',
-    'model_total_size (Byte)'
 ]
 
 subtotal_operations = [
@@ -98,27 +108,16 @@ with col1:
     st.write(f"arithmetic_intensity: {gpu_config['arithmetic_intensity']:.3f}")
 
 with col2:
-    parameter_count['word_embedding'] = model_config['vocab_size']*model_config['hidden_size']
-    parameter_count['positional_embedding'] = model_config['max_position_embeddings']*model_config['hidden_size']
-    
-    parameter_count['attention_Q']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
-    parameter_count['attention_K']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
-    parameter_count['attention_V']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
-    parameter_count['attention_out'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
-    
-    parameter_count['layernorm'] = 2*model_config['layernorm_operation']*model_config['num_hidden_layers']*model_config['hidden_size']
-    parameter_count['mlp1'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['intermediate_size']
-    parameter_count['mlp2'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['intermediate_size']
-    parameter_count['embedding_weights'] = parameter_count['word_embedding'] + parameter_count['positional_embedding']
-    parameter_count['attention_weights'] = parameter_count['attention_out'] + parameter_count['attention_Q'] + parameter_count['attention_K'] + parameter_count['attention_V']
-    parameter_count['mlp_weights'] = parameter_count['mlp1'] + parameter_count['mlp2']
-    parameter_count['model_total_size (Byte)'] = inference_config['byte_per_parameter'] * (
-        parameter_count['embedding_weights'] + 
-        parameter_count['attention_weights'] + 
-        parameter_count['mlp_weights'] + 
-        parameter_count['layernorm'])
-
-
+    if 'parameter_count' not in st.session_state:
+        if model_config['model']:
+            st.info("Model info fetcted!")
+            parameter_count = calc_model_size_from_model(model_config, inference_config)
+        else:
+            st.info("Fail to fetch model info. Using estimation!")
+            parameter_count = model_size_estimate(model_config, inference_config)
+        st.session_state.parameter_count = parameter_count 
+    else:
+        parameter_count = st.session_state.parameter_count
 
     parameters_items = {key: "{:,}".format(int(parameter_count[key])) for key in parameter_count if key not in subtotal_parameters}
     subtotal_parameters_items = {key: "{:,}".format(int(parameter_count[key])) for key in parameter_count if key in subtotal_parameters}
@@ -133,6 +132,14 @@ with col2:
     header4("Parameters Summary")
     st.markdown(create_table(df_subtotal_parameters_items))
 
+    model_total_size_in_byte = inference_config['byte_per_parameter'] * (
+                                                                            parameter_count['embedding_weights'] + 
+                                                                            parameter_count['attention_weights'] + 
+                                                                            parameter_count['mlp_weights'] + 
+                                                                            parameter_count['layernorm']
+                                                                        )
+    st.write(f'model_total_size (Byte): {model_total_size_in_byte:,}')
+
 
 with col3: # Prefilling
     prefilling_operation_count = prefilling_operation(model_config, inference_config)

calc_util.py

@@ -1,5 +1,47 @@
 import numpy as np
-
+from collections import defaultdict
+from functools import partial
+from typing import List
+from model_util import get_module_tensors_matched
+
+def calc_model_size_from_model(model_config, inference_config):
+    get_module_tensors_matched_partial = partial(get_module_tensors_matched, module_classes_dict = model_config['module_classes'])
+
+    parameter_count = defaultdict(float)
+    parameter_count['word_embedding'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'embed' in x and 'pos' not in x)])
+    parameter_count['positional_embedding'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'embed' in x and 'pos' in x)])
+
+    parameter_count['attention_Q'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'att' in x and 'q' in x)])
+    parameter_count['attention_K'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'att' in x and 'k' in x)])
+    parameter_count['attention_V'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'att' in x and 'v' in x)])
+    parameter_count['attention_out'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'att' in x and ('out_' in x or 'o_' in x))])
+    
+    parameter_count['layernorm'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'norm' in x)])
+    parameter_count['mlp_weights'] = sum([v.numel() for v in get_module_tensors_matched_partial(lambda x: 'fc' in x or 'mlp' in x)])
+       
+    parameter_count['embedding_weights'] = parameter_count['word_embedding'] + parameter_count['positional_embedding']
+    parameter_count['attention_weights'] = parameter_count['attention_out'] + parameter_count['attention_Q'] + parameter_count['attention_K'] + parameter_count['attention_V']
+        
+    return parameter_count
+
+def model_size_estimate(model_config, inference_config):
+    parameter_count = {}
+    parameter_count['word_embedding'] = model_config['vocab_size']*model_config['hidden_size']
+    parameter_count['positional_embedding'] = model_config['max_position_embeddings']*model_config['hidden_size']
+    
+    parameter_count['attention_Q']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
+    parameter_count['attention_K']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
+    parameter_count['attention_V']   = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
+    parameter_count['attention_out'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['hidden_size']/model_config['num_attention_heads']*model_config['num_attention_heads']
+    
+    parameter_count['layernorm'] = 2*model_config['layernorm_operation']*model_config['num_hidden_layers']*model_config['hidden_size']
+    parameter_count['mlp1'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['intermediate_size']
+    parameter_count['mlp2'] = model_config['num_hidden_layers']*model_config['hidden_size']*model_config['intermediate_size']
+    parameter_count['embedding_weights'] = parameter_count['word_embedding'] + parameter_count['positional_embedding']
+    parameter_count['attention_weights'] = parameter_count['attention_out'] + parameter_count['attention_Q'] + parameter_count['attention_K'] + parameter_count['attention_V']
+    parameter_count['mlp_weights'] = parameter_count['mlp1'] + parameter_count['mlp2']
+    
+    return parameter_count
 
 def multiplication_in_int64(array):
     return np.cumprod(np.array(array, dtype=np.int64))[-1]
@@ -19,28 +61,76 @@ def word_embedding_operation(model_config, inference_config):
     #The resultant matrix after the multiplication will be of size \( B \times s \times d_{model} \).
     #For each element in this resultant matrix, the number of FLOPs required is \( 2 \times n_{vocab} \). This is because for a single element in the output matrix, we have \( 2N \) FLOPs (with \( N \) being the common dimension), leading to the matrix multiplication FLOP count as:
     #\begin{equation}
-    #2 \times B \times s \times n_{vocab} \times d_{model}
+    #2 \times B \times s \times n_{v ocab} \times d_{model}
     #\end{equation}
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' not in x, model_config['module_classes'])
+        A = [inference_config['batchsize'], inference_config['input_seq_length'], modules[0][0]]
+        B = modules[0]
+        op_count = matrix_operation(A, B)
+        return op_count
+
     A = [inference_config['batchsize'], inference_config['input_seq_length'], model_config['vocab_size']]
     B = [model_config['vocab_size'], model_config['hidden_size']]
-    return matrix_operation(A, B)
+    op_count = matrix_operation(A, B)
+    return op_count
 
 
 def positional_embedding_operation(model_config, inference_config):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'embed' in x and 'pos' in x, model_config['module_classes'])
+        return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], modules[0][-1]])
+
     return multiplication_in_int64([inference_config['batchsize'], inference_config['input_seq_length'], model_config['hidden_size']])
 
 ### Below three are the same
 def attention_K_operation(model_config, inference_config, seq_length):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
+        total = 0
+        for module in modules:
+            if len(module) > 1:
+                A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+                B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
+                total += model_config['num_attention_heads']*matrix_operation(A, B)
+            else:
+                total += model_config['hidden_size']
+        return total
+
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
     B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 def attention_Q_operation(model_config, inference_config, seq_length):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'att' in x and 'q' in x , model_config['module_classes'])
+        total = 0
+        for module in modules:
+            if len(module) > 1:
+                A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+                B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
+                total += model_config['num_attention_heads']*matrix_operation(A, B)
+            else:
+                total += model_config['hidden_size']
+        return total
+     
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
     B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 def attention_V_operation(model_config, inference_config, seq_length):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'att' in x and 'v' in x , model_config['module_classes'])
+        total = 0
+        for module in modules:
+            if len(module) > 1:
+                A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+                B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
+                total += model_config['num_attention_heads']*matrix_operation(A, B)
+            else:
+                total += model_config['hidden_size']
+        return total
+     
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
     B = [model_config['hidden_size'], model_config['hidden_size_per_head']]
     return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
@@ -49,7 +139,7 @@ def attention_V_operation(model_config, inference_config, seq_length):
 def attention_QK_operation(model_config, inference_config, seq_length_Q, seq_length_K):
     A = [inference_config['batchsize'], seq_length_Q, model_config['hidden_size_per_head']]
     B = [model_config['hidden_size_per_head'], seq_length_K]
-    return model_config['num_hidden_layers'] * model_config['num_attention_heads']* matrix_operation(A, B)
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
 
 def attention_softmax_operation(model_config, inference_config,seq_length):
     # Ref: Ouyang, A. (2023). Understanding the Performance of Transformer Inference (Doctoral dissertation, Massachusetts Institute of Technology).
@@ -63,6 +153,18 @@ def attention_multV_operation(model_config, inference_config, seq_length_Q, seq_
     return model_config['num_hidden_layers'] * model_config['num_attention_heads']* matrix_operation(A, B)
 
 def attention_out_operation(model_config, inference_config, seq_length):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'att' in x and 'k' in x , model_config['module_classes'])
+        total = 0
+        for module in modules:
+            if len(module) > 1:
+                A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+                B = [model_config['hidden_size'], model_config['hidden_size']]
+                total += matrix_operation(A, B)
+            else:
+                total += model_config['hidden_size']
+        return total
+    
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
     B = [model_config['hidden_size'], model_config['hidden_size']]
     return model_config['num_hidden_layers'] * matrix_operation(A, B)
@@ -70,19 +172,34 @@ def attention_out_operation(model_config, inference_config, seq_length):
 def layernorm_operation(model_config, inference_config, seq_length):
     # Ref: Ouyang, A. (2023). Understanding the Performance of Transformer Inference (Doctoral dissertation, Massachusetts Institute of Technology).
     # 5 is a modeled value
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'norm' in x, model_config['module_classes'])
+        total = 0
+        for module in modules:
+            total += model_config['hidden_size']
+        return 5*total
+    
     layernorm_operation = (5*inference_config['batchsize']*seq_length*model_config['hidden_size'])
     return model_config['num_hidden_layers'] * model_config['layernorm_operation'] * layernorm_operation
 
 
-def mlp1_operation(model_config, inference_config, seq_length):
+def mlp_operation(model_config, inference_config, seq_length):
+    if model_config['module_classes']:
+        modules = get_module_tensors_matched(lambda x: 'fc' in x or 'mlp' in x, model_config['module_classes'])
+        total = 0
+        for module in modules:
+            if len(module) > 1:
+                A = [inference_config['batchsize'], seq_length, module[1]]
+                B = [module[1], module[0]]
+                total += matrix_operation(A, B)
+            else:
+                total += modules[-1][0]
+        return total
+    
     A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
     B = [model_config['hidden_size'], model_config['intermediate_size']]
-    return model_config['num_hidden_layers'] * matrix_operation(A, B)
+    return model_config['num_hidden_layers'] * (2*matrix_operation(A, B))
 
-def mlp2_operation(model_config, inference_config, seq_length):
-    A = [inference_config['batchsize'], seq_length, model_config['intermediate_size']]
-    B = [model_config['intermediate_size'], model_config['hidden_size']]
-    return model_config['num_hidden_layers'] * matrix_operation(A, B)
 
 def prefilling_operation(model_config, inference_config):
     prefilling_operation_count = {}
@@ -99,12 +216,10 @@ def prefilling_operation(model_config, inference_config):
 
     prefilling_operation_count['layernorm'] =layernorm_operation(model_config, inference_config, inference_config['input_seq_length'])
 
-    prefilling_operation_count['mlp1'] = mlp1_operation(model_config, inference_config, inference_config['input_seq_length'])
-    prefilling_operation_count['mlp2'] = mlp2_operation(model_config, inference_config, inference_config['input_seq_length'])
+    prefilling_operation_count['mlp'] = mlp_operation(model_config, inference_config, inference_config['input_seq_length'])
     
     prefilling_operation_count['embeddings'] = prefilling_operation_count['word_embedding'] + prefilling_operation_count['positional_embedding']
     prefilling_operation_count['attention'] = sum([v for k,v in prefilling_operation_count.items() if 'attention' in k])
-    prefilling_operation_count['mlp'] = prefilling_operation_count['mlp1'] + prefilling_operation_count['mlp2']
     prefilling_operation_count['total'] = (prefilling_operation_count['embeddings'] + prefilling_operation_count['attention'] + prefilling_operation_count['mlp'] + prefilling_operation_count['layernorm'])
     
     return prefilling_operation_count
@@ -120,8 +235,7 @@ def generation_operation(model_config, inference_config):
     generation_operation_count['attention_softmax'] = 0
     generation_operation_count['attention_multV'] = 0
     generation_operation_count['attention_out'] = 0
-    generation_operation_count['mlp1'] = 0
-    generation_operation_count['mlp2'] = 0
+    generation_operation_count['mlp'] = 0
     generation_operation_count['layernorm'] = 0
 
     for t in range(inference_config['output_seq_length']):
@@ -133,8 +247,7 @@ def generation_operation(model_config, inference_config):
             generation_operation_count['attention_softmax'] += attention_softmax_operation(model_config, inference_config, 1)
             generation_operation_count['attention_multV'] += attention_multV_operation(model_config, inference_config, seq_length_Q=1, seq_length_V=(t+1)+inference_config['input_seq_length'])
             generation_operation_count['attention_out'] += attention_out_operation(model_config, inference_config, 1)
-            generation_operation_count['mlp1'] += mlp1_operation(model_config, inference_config, 1)
-            generation_operation_count['mlp2'] += mlp2_operation(model_config, inference_config, 1)
+            generation_operation_count['mlp'] += mlp_operation(model_config, inference_config, 1)
         else:
             generation_operation_count['attention_K'] += attention_K_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
             generation_operation_count['attention_V'] += attention_V_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
@@ -143,14 +256,12 @@ def generation_operation(model_config, inference_config):
             generation_operation_count['attention_softmax'] += attention_softmax_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
             generation_operation_count['attention_multV'] += attention_multV_operation(model_config, inference_config, seq_length_Q=(t+1)+inference_config['input_seq_length'], seq_length_V=(t+1)+inference_config['input_seq_length'])
             generation_operation_count['attention_out'] += attention_out_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
-            generation_operation_count['mlp1'] += mlp1_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
-            generation_operation_count['mlp2'] += mlp2_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            generation_operation_count['mlp'] += mlp_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
 
         generation_operation_count['layernorm'] += layernorm_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
 
     generation_operation_count['embeddings'] = generation_operation_count['word_embedding'] + generation_operation_count['positional_embedding'] 
     generation_operation_count['attention'] = sum([v for k,v in generation_operation_count.items() if 'attention' in k])
-    generation_operation_count['mlp'] = generation_operation_count['mlp1'] + generation_operation_count['mlp2']
     generation_operation_count['total'] = (generation_operation_count['attention'] + generation_operation_count['mlp'] + generation_operation_count['layernorm'])
 
     return generation_operation_count
@@ -196,12 +307,9 @@ def layernorm_activation_memory(model_config, inference_config, seq_length):
     per_layernorm_per_layer = 2 * inference_config['batchsize'] * seq_length * model_config['hidden_size']
     return model_config['num_hidden_layers'] * model_config['layernorm_operation'] * per_layernorm_per_layer
 
-def mlp1_activation_memory(model_config, inference_config, seq_length):
-    per_layer = inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['intermediate_size'])
-    return model_config['num_hidden_layers'] * per_layer
-
-def mlp2_activation_memory(model_config, inference_config, seq_length):
-    per_layer = inference_config['batchsize'] * seq_length * (model_config['intermediate_size'] + model_config['hidden_size'])
+def mlp_activation_memory(model_config, inference_config, seq_length):
+    # two mlp layer
+    per_layer = 2 * inference_config['batchsize'] * seq_length * (model_config['hidden_size'] + model_config['intermediate_size'])
     return model_config['num_hidden_layers'] * per_layer
 
 def prefilling_activation_memory(model_config, inference_config):
@@ -220,8 +328,7 @@ def prefilling_activation_memory(model_config, inference_config):
 
     activation_memory['layernorm'] = layernorm_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
 
-    activation_memory['mlp1'] = mlp1_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
-    activation_memory['mlp2'] = mlp2_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
+    activation_memory['mlp'] = mlp_activation_memory(model_config, inference_config, inference_config['input_seq_length'])
     
     activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding']
     activation_memory['attention'] = (
@@ -230,7 +337,6 @@ def prefilling_activation_memory(model_config, inference_config):
         activation_memory['attention_softmax'] + activation_memory['attention_multV'] +
         activation_memory['attention_out']
     )
-    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
     activation_memory['total'] = (
         activation_memory['embeddings'] + activation_memory['attention'] +
         activation_memory['mlp'] + activation_memory['layernorm']
@@ -238,7 +344,6 @@ def prefilling_activation_memory(model_config, inference_config):
     
     activation_memory['embeddings'] = activation_memory['word_embedding'] + activation_memory['positional_embedding'] 
     activation_memory['attention'] = sum([v for k,v in activation_memory.items() if 'attention' in k])
-    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
     activation_memory['total'] = (activation_memory['attention'] + activation_memory['mlp'] + activation_memory['layernorm'])
 
     return activation_memory
@@ -255,8 +360,7 @@ def generation_activation_memory(model_config, inference_config):
     activation_memory['attention_softmax'] = 0
     activation_memory['attention_multV'] = 0
     activation_memory['attention_out'] = 0
-    activation_memory['mlp1'] = 0
-    activation_memory['mlp2'] = 0
+    activation_memory['mlp'] = 0
     activation_memory['layernorm'] = 0
 
     for t in range(inference_config['output_seq_length']):
@@ -268,8 +372,7 @@ def generation_activation_memory(model_config, inference_config):
             activation_memory['attention_softmax'] += attention_softmax_activation_memory(model_config, inference_config, 1)
             activation_memory['attention_multV'] += attention_multV_activation_memory(model_config, inference_config, seq_length_Q=1, seq_length_V=(t+1)+inference_config['input_seq_length'])
             activation_memory['attention_out'] += attention_out_activation_memory(model_config, inference_config, 1)
-            activation_memory['mlp1'] += mlp1_activation_memory(model_config, inference_config, 1)
-            activation_memory['mlp2'] += mlp2_activation_memory(model_config, inference_config, 1)
+            activation_memory['mlp'] += mlp_activation_memory(model_config, inference_config, 1)
         else:
             activation_memory['attention_K'] += attention_K_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
             activation_memory['attention_V'] += attention_V_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
@@ -278,8 +381,7 @@ def generation_activation_memory(model_config, inference_config):
             activation_memory['attention_softmax'] += attention_softmax_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
             activation_memory['attention_multV'] += attention_multV_activation_memory(model_config, inference_config, seq_length_Q=(t+1)+inference_config['input_seq_length'], seq_length_V=(t+1)+inference_config['input_seq_length'])
             activation_memory['attention_out'] += attention_out_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
-            activation_memory['mlp1'] += mlp1_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
-            activation_memory['mlp2'] += mlp2_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            activation_memory['mlp'] += mlp_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
 
         activation_memory['layernorm'] += layernorm_activation_memory(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
 
@@ -290,7 +392,6 @@ def generation_activation_memory(model_config, inference_config):
         activation_memory['attention_softmax'] + activation_memory['attention_multV'] +
         activation_memory['attention_out']
     )
-    activation_memory['mlp'] = activation_memory['mlp1'] + activation_memory['mlp2']
     activation_memory['total'] = (
         activation_memory['embeddings'] + activation_memory['attention'] +
         activation_memory['mlp'] + activation_memory['layernorm']

model_util.py

@@ -1,11 +1,14 @@
 import requests
+import re
+from collections import defaultdict
 # Utilities related to loading in and working with models/specific models
 from urllib.parse import urlparse
 import torch
 from accelerate.commands.estimate import check_has_model, create_empty_model
-from accelerate.utils import compute_module_sizes
+from accelerate.utils import compute_module_sizes, named_module_tensors
 from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
 
+
 def fetch_dictionary_content(model_id):
     MODEL_URL = "https://huggingface.co/{model_id}/raw/main/config.json"
     response = requests.get(MODEL_URL.format(model_id=model_id))
@@ -85,10 +88,46 @@ def get_model(model_name: str, library: str, access_token: str):
         )
     return model
 
+def get_module_tensors(model):
+    module_tensors = {}
+    for name, tensor in named_module_tensors(model, recurse=True):
+        module_tensors[name] = tensor.shape
+
+    return module_tensors
+
+
+def classify_module(module_tensors):
+    # A dictionary to store counts for each generic layer type
+    module_classes = defaultdict(list)
+
+    # This function removes all numbers from a given string
+    def remove_numbers(s):
+        return re.sub(r'\d+', '', s)
+    
+    # Loop through all named parameters of the model
+    for name in module_tensors:
+        # Remove numbers from the name
+        generic_name = remove_numbers(name)
+        generic_name = generic_name.replace('..', '.')
+        
+        # If the name already exists in the dictionary, increase the count, else set it to 1
+        module_classes[generic_name].append({name: module_tensors[name]})
+
+    return module_classes
+
+def get_module_tensors_matched(filter_fn, module_classes_dict):
+    matched = []
+    for generic, module_list in module_classes_dict.items():
+        if filter_fn(generic.lower()):
+            matched.extend([v for module in module_list for v in module.values()])
+       
+    return matched
 
 
 if __name__ == '__main__':
     model = get_model('NousResearch/Nous-Hermes-Llama2-13b', None, None)
+    module_tensors = get_module_tensors(model)
+    module_classes = classify_module(module_tensors)
     sizes = compute_module_sizes(model, dtype=torch.int8)
     size_dict = {
         'attn':0,