Prepare

pip install -r requirements.txt

Run

Step 1: Pick data

You need to change the model path and the input file in pickdata.py and run it.

python run.py

Then you get all the .pt data in ./savedata.

Step 2: Change data format from .pt to .dat and blockwise it

Activation

You need to change the from_dir and to_dir in reshape.py and run it.

python reshape.py

Then you get all the .dat data in ./data_llama-3.1-8B-bf16-prefill-520-len-5.

Weight

You need to change the filedir and to_dir in reshape_weight.py and run it.

python reshape_weight.py

Then you get all the .dat data in ./data_llama-3.1-8B-bf16-prefill-520-len-5.

Config

Model and Dtype and Generate Length

This is the running config.

run.py:

# config
config.max_new_tokens = 5 # decoding len
max_length = 36 # prefill len
dtype = torch.bfloat16 # model dtype
model_path = "/data/coding/llama-3.1-8b" # model path
input_file = "./data_for_gem5_prompt.txt" # input file, inside is a list of strings

Pick Tensor

This is the data you want to pick out and the layers you want to pick.

pickdata.py:

save_dir = "./savedata" # dir you want to save .pt data

data_need = [
    "activation_attention_norm_input",
    "activation_attention_input",
    "activation_attention_output_after_outproj",
    "activation_ffn_norm_input",
    "activation_ffn_input",
    "activation_ffn_output",
    "constant_ROPE_cos_pos",
    "constant_ROPE_sin_pos",
    "activation_logit_output_after_outproj",
    "activation_logit_output_before_outproj",
    "activation_transformerblock_output",
    "activation_token_embedding_tokens",
    "kv_cache_k_after_update",
    "kv_cache_v_after_update",
    # "activation_attention_output_before_reshape",
    # "activation_attention_xq_after_ROPE",
    # "activation_attention_xk_after_ROPE",
    # "activation_attention_xq_before_ROPE",
    # "activation_attention_xk_before_ROPE",
    # "activation_attention_xv",
    # "activation_attention_scores_before_div_sqrt_dim",
    # "kv_cache_k_quant_after_update",
    # "kv_cache_v_quant_after_update",
    # "kv_cache_k_rest_after_update",
    # "kv_cache_v_rest_after_update",
]

layer_need = range(32) # pick out all layers

Change format and blockwise

Activation

reshape.py:

from_dir = r'./savedata' # dir of .pt data
to_dir = r'./data_llama-3.1-8B-bf16-prefill-520-len-5' # dir you want to save .dat

Weight

reshape_weight.py:

filedir = 'D:\data-tanyifan\modeldata\llama-3.1-8b' # dir of the model
to_dir = r'./data_llama-3.1-8B-bf16-prefill-520-len-5' # dir you want to save .dat