inference speed

.gitignore

@@ -0,0 +1 @@
+**__pycache__**

.vscode/launch.json

@@ -0,0 +1,20 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "streamlit",
+            "type": "python",
+            "request": "launch",
+            "module": "streamlit",
+            "console": "integratedTerminal",
+            "justMyCode": true,
+            "args": [
+                "run",
+                "app.py"
+            ]
+        }
+    ]
+}

app.py

@@ -0,0 +1,180 @@
+# Ref: Ouyang, A. (2023). Understanding the Performance of Transformer Inference (Doctoral dissertation, Massachusetts Institute of Technology).
+
+import streamlit as st
+import pandas as pd
+from model_util import fetch_dictionary_content, load_parameter
+from calc_util import *
+from render_util import create_table, header4, header5
+
+
+st.set_page_config(layout='wide')
+if 'model_config' not in st.session_state:
+    st.session_state['model_config'] = {}
+
+
+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']
+    model_config = {}
+    dictionary_content = fetch_dictionary_content(model_id)
+    if dictionary_content:
+        model_config['model_id'] = model_id
+        model_config['hidden_size'] = dictionary_content['hidden_size']
+        model_config['num_attention_heads'] = dictionary_content['num_attention_heads']
+        model_config['num_hidden_layers'] = dictionary_content['num_hidden_layers']
+        model_config['intermediate_size'] = load_parameter(dictionary_content, ['intermediate_size', 'ffn_dim'])
+        model_config['vocab_size'] = dictionary_content['vocab_size']
+        model_config['max_position_embeddings'] = dictionary_content['max_position_embeddings']
+        model_config['layernorm_operation'] = 2
+    else:
+        st.warning("Model Info is not public!")
+        model_config['model_id'] = 'opt-1.3b'
+        model_config['hidden_size'] = 2048
+        model_config['num_attention_heads'] = 32
+        model_config['num_hidden_layers'] = 24
+        model_config['intermediate_size'] = 8192
+        model_config['vocab_size'] = 50272
+        model_config['max_position_embeddings'] = 2048
+        model_config['layernorm_operation'] = 2
+
+    st.session_state['model_config'] = model_config
+    return model_config
+
+
+subtotal_parameters = [
+    'embedding_weights',
+    'attention_weights',
+    'mlp_weights',
+    'model_total_size'
+]
+
+subtotal_operations = [
+    'embeddings',
+    'attention',
+    'mlp',
+    'total',
+]
+
+
+
+col1, col2, col3, col4, col5 = st.columns(5)
+
+inference_config = {}
+parameter_count = {}
+cached_parameter_count = {}
+prefilling_operation_count = {}
+generation_operation_count = {}
+gpu_config = {}
+inference_info = {}
+
+with col1:
+    header4("Model")
+    model_id = st.text_input("huggingface model id", 'ArthurZ/opt-13b')
+    model_config = load_model_config(model_id)
+    model_config['hidden_size'] = st.number_input('hidden size', value=model_config['hidden_size'], format ="%d")
+    model_config['num_attention_heads'] = st.number_input('num attention heads', value=model_config['num_attention_heads'], format ="%d")
+    model_config['num_hidden_layers'] = st.number_input('num hidden layers', value=model_config['num_hidden_layers'], format ="%d")
+    model_config['intermediate_size'] = st.number_input('intermediate size', value=model_config['intermediate_size'], format ="%d")
+    model_config['vocab_size'] = st.number_input('vocab size', value= model_config['vocab_size'], format ="%d")
+    model_config['max_position_embeddings'] = st.number_input('max position embeddings', value=model_config['max_position_embeddings'], format ="%d")
+
+    header4("Inference Setting")
+    inference_config['batchsize'] = st.number_input('batchsize', value=1, format ="%d")
+    inference_config['input_seq_length'] = st.number_input('input seq length', value=1, format ="%d")
+    inference_config['output_seq_length'] = st.number_input('output seq length', value=1, format ="%d")
+    inference_config['byte_per_parameter'] = st.number_input('byte per parameter', value=2, format ="%d")
+    inference_config['KV_cache'] = st.checkbox("Use KV cache", value=True)
+
+    header4("GPU Setting")
+    gpu_config['Name'] = st.text_input('GPU Type', value="A6000")
+    gpu_config['TFLOP'] = st.number_input('TFLOP', value=38.7, format ="%2f")
+    gpu_config['memory_bandwidth'] = st.number_input('memory bandwidth (GB/s)', value=768, format ="%2d")
+    gpu_config['arithmetic_intensity'] = gpu_config['TFLOP']*10**12/gpu_config['memory_bandwidth']/1024**3
+    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'] = inference_config['byte_per_parameter'] * (
+        parameter_count['embedding_weights'] + 
+        parameter_count['attention_weights'] + 
+        parameter_count['mlp_weights'] + 
+        parameter_count['layernorm'])
+
+
+
+    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}
+
+    # Convert dictionaries to pandas dataframes for table display
+    df_parameters_items = pd.DataFrame(list(parameters_items.items()), columns=["Parameter", "Count"])
+    df_subtotal_parameters_items = pd.DataFrame(list(subtotal_parameters_items.items()), columns=["Parameter", "Count"])
+
+    header4("Model Parameters")
+    st.markdown(create_table(df_parameters_items))
+
+    header4("Parameters Summary")
+    st.markdown(create_table(df_subtotal_parameters_items))
+
+
+with col3: # Prefilling
+    prefilling_operation_count = prefilling_operation(model_config, inference_config)
+    inference_info['inference_prefilling_time'] = prefilling_operation_count['total'] / (gpu_config['TFLOP']*10**12)
+    inference_info['inference_prefilling_throughput'] = inference_config['input_seq_length']/inference_info['inference_prefilling_time']
+    cached_parameter_count['kv_cache'] = 2 * (inference_config['batchsize'] * (model_config['hidden_size'] * model_config['num_hidden_layers'] * inference_config['input_seq_length']))
+
+    operation_items = {key: "{:,}".format(int(prefilling_operation_count[key])) for key in prefilling_operation_count if key not in subtotal_operations}
+    subtotal_operation_items = {key: "{:,}".format(int(prefilling_operation_count[key])) for key in prefilling_operation_count if key in subtotal_operations}
+
+    ## Convert dictionaries to pandas dataframes for table display
+    df_operation_count = pd.DataFrame(list(operation_items.items()), columns=["Operation", "FLOPS"])
+    df_subtotal_operation_count = pd.DataFrame(list(subtotal_operation_items.items()), columns=["Operation", "FLOPS"])
+   
+    header4("Inference Ops: Prefilling")
+    st.markdown(create_table(df_operation_count))
+
+    header5("Summary: Prefilling")
+    st.markdown(create_table(df_subtotal_operation_count))
+    st.write(f"Prefillng throughput (tokens/s): {inference_info['inference_prefilling_throughput']:.2f}")
+
+    if inference_config['KV_cache']:
+        st.write(f"kv cache (Byte): {cached_parameter_count['kv_cache']:,}")
+        
+
+
+with col4: # Prefilling
+    generation_operation_count = generation_operation(model_config, inference_config)
+    inference_info['inference_generation_time'] = generation_operation_count['total'] / (gpu_config['TFLOP']*10**12)
+    inference_info['inference_generation_throughput'] = inference_config['output_seq_length']/inference_info['inference_generation_time']
+    inference_info['inference_client_generation_throughput'] = inference_config['output_seq_length'] / (inference_info['inference_prefilling_time'] + inference_info['inference_generation_time'])
+    cached_parameter_count['kv_cache'] = 2 * (inference_config['batchsize'] * (model_config['hidden_size'] * model_config['num_hidden_layers'] * (inference_config['input_seq_length']+inference_config['output_seq_length'])))
+
+    operation_items = {key: "{:,}".format(int(generation_operation_count[key])) for key in generation_operation_count if key not in subtotal_operations}
+    subtotal_operation_items = {key: "{:,}".format(int(generation_operation_count[key])) for key in generation_operation_count if key in subtotal_operations}
+
+    ## Convert dictionaries to pandas dataframes for table display
+    df_operation_count = pd.DataFrame(list(operation_items.items()), columns=["Operation", "FLOPS"])
+    df_subtotal_operation_count = pd.DataFrame(list(subtotal_operation_items.items()), columns=["Operation", "FLOPS"])
+   
+    header4("Inference Ops: Generation")
+    st.markdown(create_table(df_operation_count))
+
+    header5("Summary: Generation")
+    st.markdown(create_table(df_subtotal_operation_count))
+    st.write(f"Generation-only throughput (tokens/s): {inference_info['inference_generation_throughput']:.2f}")
+    st.write(f"(Client) Generation throughput (tokens/s): {inference_info['inference_client_generation_throughput']:.2f}")
+
+    if inference_config['KV_cache']:
+        st.write(f"kv cache (Byte): {cached_parameter_count['kv_cache']:,}")

calc_util.py

@@ -0,0 +1,156 @@
+import numpy as np
+
+
+def multiplication_in_int64(array):
+    return np.cumprod(np.array(array, dtype=np.int64))[-1]
+
+def matrix_operation(shapeA, shapeB):
+    assert(shapeA[-1] == shapeB[0])
+    op = np.cumprod(np.array(shapeA[:-1], np.float64))    
+    return multiplication_in_int64([2, op[-1], shapeA[-1], shapeB[-1]])
+
+def word_embedding_operation(model_config, inference_config):
+    #Given:
+    #\begin{itemize}
+    #    \item Matrix \( X \) of size \( B \times s \) (representing the batch size and sequence length respectively).
+    #    \item Embedding matrix \( W_e \) of size \( n_{vocab} \times d_{model} \).
+    #\end{itemize}
+    
+    #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}
+    #\end{equation}
+    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)
+
+
+def positional_embedding_operation(model_config, inference_config):
+    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):
+    A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    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):
+    A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    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):
+    A = [inference_config['batchsize'], seq_length, model_config['hidden_size']]
+    B = [model_config['hidden_size'], model_config['hidden_size']/model_config['num_attention_heads']]
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * matrix_operation(A, B)
+
+## 
+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']/model_config['num_attention_heads']]
+    B = [model_config['hidden_size']/model_config['num_attention_heads'], seq_length_K]
+    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).
+    # 3 is a modeled value
+    softmax_operation = (3*inference_config['batchsize']*seq_length*seq_length)
+    return model_config['num_hidden_layers'] * model_config['num_attention_heads'] * softmax_operation
+
+def attention_multV_operation(model_config, inference_config, seq_length_Q, seq_length_V):
+    A = [inference_config['batchsize'], seq_length_Q, seq_length_V]
+    B = [seq_length_V, model_config['hidden_size']/model_config['num_attention_heads']]
+    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):
+    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)
+
+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
+    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):
+    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)
+
+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 = {}
+    prefilling_operation_count['word_embedding'] = word_embedding_operation(model_config, inference_config)
+    prefilling_operation_count['positional_embedding'] = positional_embedding_operation(model_config, inference_config)
+    
+    prefilling_operation_count['attention_Q'] = attention_Q_operation(model_config, inference_config, inference_config['input_seq_length'])
+    prefilling_operation_count['attention_K'] = attention_K_operation(model_config, inference_config, inference_config['input_seq_length'])
+    prefilling_operation_count['attention_V'] = attention_V_operation(model_config, inference_config, inference_config['input_seq_length'])
+    prefilling_operation_count['attention_QK'] = attention_QK_operation(model_config, inference_config, inference_config['input_seq_length'], inference_config['input_seq_length'])
+    prefilling_operation_count['attention_softmax'] = attention_softmax_operation(model_config, inference_config, inference_config['input_seq_length'])
+    prefilling_operation_count['attention_multV'] = attention_multV_operation(model_config, inference_config, inference_config['input_seq_length'], inference_config['input_seq_length'])
+    prefilling_operation_count['attention_out'] = attention_out_operation(model_config, inference_config, inference_config['input_seq_length'])
+
+    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['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
+
+def generation_operation(model_config, inference_config):
+    generation_operation_count = {}
+    generation_operation_count['word_embedding'] = 0
+    generation_operation_count['positional_embedding'] = 0
+    generation_operation_count['attention_K'] = 0
+    generation_operation_count['attention_V'] = 0
+    generation_operation_count['attention_Q'] = 0
+    generation_operation_count['attention_QK'] = 0
+    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['layernorm'] = 0
+
+    for t in range(inference_config['output_seq_length']):
+        if inference_config['KV_cache']:
+            generation_operation_count['attention_K'] += attention_K_operation(model_config, inference_config, 1)
+            generation_operation_count['attention_V'] += attention_V_operation(model_config, inference_config, 1)
+            generation_operation_count['attention_Q'] += attention_Q_operation(model_config, inference_config, 1)
+            generation_operation_count['attention_QK'] += attention_QK_operation(model_config, inference_config, seq_length_Q=1, seq_length_K=(t+1)+inference_config['input_seq_length'])
+            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)
+        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'])
+            generation_operation_count['attention_Q'] += attention_Q_operation(model_config, inference_config, (t+1)+inference_config['input_seq_length'])
+            generation_operation_count['attention_QK'] += attention_QK_operation(model_config, inference_config, seq_length_Q=(t+1)+inference_config['input_seq_length'], seq_length_K=(t+1)+inference_config['input_seq_length'])
+            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['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

model_util.py

@@ -0,0 +1,18 @@
+import requests
+
+
+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))
+    
+    # Check if the request was successful
+    if response.status_code == 200:
+        return response.json()  # Parse the JSON content into a Python dictionary
+    else:
+        return None
+    
+def load_parameter(model_dict, cand_keys):
+    for k in cand_keys:
+        if k in model_dict:
+            return model_dict[k]
+    return 0

render_util.py

@@ -0,0 +1,22 @@
+import streamlit as st
+
+def create_table(df):
+    # Table header based on df columns
+    header = "| " + " | ".join(df.columns) + " |"
+    # Number of columns in df to set table divider accordingly
+    divider = "|:---" * len(df.columns[:-1]) + "|-----:|"
+    rows = [header, divider]
+
+    for _, row in df.iterrows():
+        rows.append("| " + " | ".join(row.astype(str)) + " |")
+
+    return "\n".join(rows)
+
+def header3(text):
+    st.markdown(f"### {text}")
+
+def header4(text):
+    st.markdown(f"#### {text}")
+
+def header5(text):
+    st.markdown(f"##### {text}")