File size: 2,084 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 | #!/usr/bin/env bash
set -euo pipefail
# Prune-only runner for LLM-Pruner on Qwen-family checkpoints.
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
workdir="$repo_root/compare_model/LLM-Pruner"
base_model="${BASE_MODEL:-Qwen/Qwen3-1.7B}"
prune_ckpt_path="${PRUNE_CKPT_PATH:-qwen3_1_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:-24}"
block_attention_layer_start="${BLOCK_ATTENTION_LAYER_START:-4}"
block_attention_layer_end="${BLOCK_ATTENTION_LAYER_END:-24}"
block_attention_roots="${BLOCK_ATTENTION_ROOTS:-q_proj,k_proj}"
block_mlp_roots="${BLOCK_MLP_ROOTS:-gate_proj,up_proj}"
pruner_type="${PRUNER_TYPE:-taylor}"
taylor_mode="${TAYLOR_MODE:-param_first}"
device="${DEVICE:-cuda}"
eval_device="${EVAL_DEVICE:-cuda}"
script_name="${PRUNE_SCRIPT:-llama3.py}"
output_dir="${OUTDIR:-$workdir/prune_log/$prune_ckpt_path}"
python_args=(
--base_model "$base_model"
--pruning_ratio "$pruning_ratio"
--block_wise
--block_attention_roots "$block_attention_roots"
--block_mlp_roots "$block_mlp_roots"
--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"
--test_after_train
--device "$device"
--eval_device "$eval_device"
--save_ckpt_log_name "$prune_ckpt_path"
--save_model
)
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[@]}"
|