File size: 2,200 Bytes
3738140 | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | #!/usr/bin/env bash
set -euo pipefail
# meta-llama/Llama-2-7b-hf, meta-llama/Llama-3.1-8B
# Prune-only runner for LLM-Pruner on Llama-family checkpoints.
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-2}"
repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
workdir="$repo_root/compare_model/LLM-Pruner"
base_model="${BASE_MODEL:-meta-llama/Llama-2-7b-hf}"
prune_ckpt_path="${PRUNE_CKPT_PATH:-llama2_7b_prune}"
pruning_ratio="${PRUNING_RATIO:-0.25}"
block_mlp_layer_start="${BLOCK_MLP_LAYER_START:-4}"
block_mlp_layer_end="${BLOCK_MLP_LAYER_END:-30}"
block_attention_layer_start="${BLOCK_ATTENTION_LAYER_START:-4}"
block_attention_layer_end="${BLOCK_ATTENTION_LAYER_END:-30}"
pruner_type="${PRUNER_TYPE:-taylor}"
taylor_mode="${TAYLOR_MODE:-param_first}"
device="${DEVICE:-cpu}"
eval_device="${EVAL_DEVICE:-cuda}"
default_script="hf_prune.py"
skip_eval_flag="--skip_post_eval"
if [[ "$base_model" == *"Llama-3"* ]] || [[ "$base_model" == *"Llama-3."* ]] || [[ "$base_model" == *"llama-3"* ]]; then
default_script="llama3.py"
skip_eval_flag="--skip_eval_after_prune"
fi
script_name="${PRUNE_SCRIPT:-$default_script}"
output_dir="${OUTDIR:-$workdir/prune_log/$prune_ckpt_path}"
python_args=(
--base_model "$base_model"
--pruning_ratio "$pruning_ratio"
--block_wise
--block_mlp_layer_start "$block_mlp_layer_start"
--block_mlp_layer_end "$block_mlp_layer_end"
--block_attention_layer_start "$block_attention_layer_start"
--block_attention_layer_end "$block_attention_layer_end"
--pruner_type "$pruner_type"
--taylor "$taylor_mode"
--device "$device"
--eval_device "$eval_device"
--save_ckpt_log_name "$prune_ckpt_path"
--save_model
"$skip_eval_flag"
)
python_args+=("$@")
mkdir -p "$output_dir"
git_commit="unknown"
if git -C "$repo_root" rev-parse --is-inside-work-tree >/dev/null 2>&1; then
git_commit=$(git -C "$repo_root" rev-parse HEAD)
fi
{
echo "git_commit=$git_commit"
echo "command:"
printf '%q ' python "$repo_root/compare_model/LLM-Pruner/$script_name" "${python_args[@]}"
echo
} > "$output_dir/run_args.txt"
cd "$workdir"
PYTHONPATH="$workdir:$repo_root${PYTHONPATH:+:$PYTHONPATH}" \
python "$script_name" "${python_args[@]}"
|