File size: 1,535 Bytes
d73500e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
#!/usr/bin/bash
port="21305"
GPUs="4,5,6,7"
dataset="c4_val"
prune_data_type="pt"
n_calibration_samples=64
seq_len=2048
prune_method="nbl_linearize"
num_layers_to_linearize=32
model_name=llama3-8b-base
model_name_or_path=/workspace/1016_qif/Meta-Llama-3-8B
# model_name=llama3.1-8b-base
# model_name_or_path=/workspace/1016_qif/Llama-3.1-8B
folder_name="${model_name}-${prune_method}-linearize${num_layers_to_linearize}"
nbl_metric_cache_file="../results_prune/cache/${model_name}-${prune_method}-${dataset}-${n_calibration_samples}samples.pt"
echo ${folder_name}
output_dir=../results_prune/${folder_name}
prune_model_save_path=${output_dir}/checkpoint
CUDA_VISIBLE_DEVICES=$GPUs accelerate launch --main_process_port $port \
src/compress.py \
--stage prune \
--model_name_or_path ${model_name_or_path} \
--dataset ${dataset} \
--dataset_dir ./src/llmtuner/data \
--split "train" \
--prune_data_type ${prune_data_type} \
--cutoff_len ${seq_len} \
--output_dir ${output_dir} \
--logging_steps 10 \
--bf16 \
--n_calibration_samples ${n_calibration_samples} \
--prune_method ${prune_method} \
--num_layers_to_linearize ${num_layers_to_linearize} \
--nbl_metric_cache_file ${nbl_metric_cache_file} \
--prune_model_save_path ${prune_model_save_path}
# After running this script, the pruned model with linearized layers will be saved
# in the specified `prune_model_save_path`.
# The model can be loaded as usual, and the specified attention layers will be
# replaced by the learned linear layers.
|