Merge commit '3732b011219890e6d2b6935eba116dea31e10bfd'

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:
@@ -27,7 +30,7 @@ def load_model_config(model_id):
         model_config['max_position_embeddings'] = dictionary_content['max_position_embeddings']
         model_config['layernorm_operation'] = 2
     else:
-        st.warning("Model Info is not public!")
+        st.warning("Fetching information failed! Maybe model info is not public!")
         model_config['model_id'] = 'opt-1.3b'
         model_config['hidden_size'] = 2048
         model_config['num_attention_heads'] = 32
@@ -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 = [
@@ -57,7 +67,7 @@ subtotal_operations = [
 
 
 
-col1, col2, col3, col4, col5 = st.columns([1,1.5,2.5,2.5,0.1])
+col1, col2, col3, col4, col5 = st.columns([0.8, 2, 2.5, 2.5, 0.01])
 
 inference_config = {}
 parameter_count = {}
@@ -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,25 @@ 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:,}')
+
+
+    # add parameter viewer
+    if model_config['model']:
+        header4("Parameters Viewer")
+        weight_generic = st.selectbox('Select weight:', options=model_config['module_classes'])
+        modules = {}
+        for module in model_config['module_classes'][weight_generic]:
+            modules.update(module)
+        modules = {k: list(v) for k, v in modules.items()}
+        modules = pd.DataFrame(list(modules.items()), columns=["Parameter", "Shape"])
+        st.markdown(create_table(modules))
 
 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,4 +1,12 @@
 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, named_module_tensors
+from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
 
 
 def fetch_dictionary_content(model_id):
@@ -15,4 +23,119 @@ def load_parameter(model_dict, cand_keys):
     for k in cand_keys:
         if k in model_dict:
             return model_dict[k]
-    return 0
+    return 0
+
+# Reference: https://huggingface.co/spaces/hf-accelerate/model-memory-usage
+def extract_from_url(name: str):
+    "Checks if `name` is a URL, and if so converts it to a model name"
+    is_url = False
+    try:
+        result = urlparse(name)
+        is_url = all([result.scheme, result.netloc])
+    except Exception:
+        is_url = False
+    # Pass through if not a URL
+    if not is_url:
+        return name
+    else:
+        path = result.path
+        return path[1:]
+
+
+def translate_llama2(text):
+    "Translates llama-2 to its hf counterpart"
+    if not text.endswith("-hf"):
+        return text + "-hf"
+    return text
+
+
+def get_model(model_name: str, library: str, access_token: str):
+    "Finds and grabs model from the Hub, and initializes on `meta`"
+    if "meta-llama" in model_name:
+        model_name = translate_llama2(model_name)
+    if library == "auto":
+        library = None
+    model_name = extract_from_url(model_name)
+    try:
+        model = create_empty_model(model_name, library_name=library, trust_remote_code=True, access_token=access_token)
+    except GatedRepoError:
+        raise RuntimeError(
+            f"Model `{model_name}` is a gated model, please ensure to pass in your access token and try again if you have access. You can find your access token here : https://huggingface.co/settings/tokens. "
+        )
+    except RepositoryNotFoundError:
+        raise RuntimeError(f"Model `{model_name}` was not found on the Hub, please try another model name.")
+    except ValueError:
+        raise RuntimeError(
+            f"Model `{model_name}` does not have any library metadata on the Hub, please manually select a library_name to use (such as `transformers`)"
+        )
+    except (RuntimeError, OSError) as e:
+        library = check_has_model(e)
+        if library != "unknown":
+            raise RuntimeError(
+                f"Tried to load `{model_name}` with `{library}` but a possible model to load was not found inside the repo."
+            )
+        raise RuntimeError(
+            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"
+        )
+    except ImportError:
+        # hacky way to check if it works with `trust_remote_code=False`
+        model = create_empty_model(
+            model_name, library_name=library, trust_remote_code=False, access_token=access_token
+        )
+    except Exception as e:
+        raise RuntimeError(
+            f"Model `{model_name}` had an error, please open a discussion on the model's page with the error message and name: `{e}`"
+        )
+    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,
+        'mlp':0,
+        'embed':0,
+    }
+    for k, v in sizes.items():
+        for kk in size_dict:
+            if kk in k and 'weight' in k:
+                size_dict[kk] += v/1024**3
+    print(sizes)