script/run_abprune_smoke.sh

#!/usr/bin/env bash
set -euo pipefail

export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-1}"

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"

if [[ $# -lt 1 ]]; then
  cat <<'USAGE'
Usage:
  script/run_abprune_smoke.sh <model> [output_dir] [extra fuse_layers args...]

Examples:
  script/run_abprune_smoke.sh Qwen/Qwen3-1.7B
  script/run_abprune_smoke.sh /path/to/model /path/to/output --num_progressive 3
USAGE
  exit 1
fi

model="$1"
shift

dataset="${DATASET:-slimpajama}"
dataset_config="${DATASET_CONFIG:-none}"
num_progressive="${NUM_PROGRESSIVE:-4}"
seq_len="${SEQ_LEN:-128}"
target_tokens="${TARGET_TOKENS:-8192}"
calib_sequences="${CALIB_SEQUENCES:-8}"
distill_batch_size="${DISTILL_BATCH_SIZE:-1}"
eval_batch_size="${EVAL_BATCH_SIZE:-1}"
eval_num_samples="${EVAL_NUM_SAMPLES:-8}"
distill_seq_len="${DISTILL_SEQ_LEN:-128}"
lora_epochs="${LORA_EPOCHS:-0}"
distill_epochs="${DISTILL_EPOCHS:-0.1}"
distill_kl_weight="${DISTILL_KL_WEIGHT:-0.01}"
distill_kl_temp="${DISTILL_KL_TEMP:-4.0}"
dtype="${DTYPE:-bfloat16}"
batch_size="${BATCH_SIZE:-1}"
use_pertensor_fisher="${USE_PERTENSOR_FISHER:-0}"
save_full_model_cycles="${SAVE_FULL_MODEL_CYCLES:-3}"

fisher_args=(--fisher_mode param)
if [[ "$use_pertensor_fisher" == "1" ]]; then
  fisher_args=(--fisher_mode tensor)
fi

output_dir_suffix="progressive_common_smoke"
if [[ "$use_pertensor_fisher" == "1" ]]; then
  output_dir_suffix="${output_dir_suffix}_pertensor"
fi

model_slug="$(echo "$model" | tr '/:@' '___' | tr -cs '[:alnum:]_.-' '_' | sed 's/^_\\+//; s/_\\+$//')"
output_dir_default="$repo_root/results/${model_slug}_${output_dir_suffix}"
output_dir=""
if [[ $# -gt 0 && "${1:0:2}" != "--" ]]; then
  output_dir="$1"
  shift
elif [[ -n "${OUTDIR:-}" ]]; then
  output_dir="${OUTDIR}"
else
  output_dir="${output_dir_default}"
fi
if [[ -n "${RUN_NAME:-}" ]]; then
  output_dir="${output_dir}_${RUN_NAME}"
fi

python_args=(
  --model "$model" \
  --dataset "$dataset" \
  --dataset_config "$dataset_config" \
  --target_tokens "$target_tokens" \
  --num_samples "$calib_sequences" \
  --seq_len "$seq_len" \
  --batch_size "$batch_size" \
  --distill_batch_size "$distill_batch_size" \
  --distill_seq_len "$distill_seq_len" \
  --distill_epochs "$distill_epochs" \
  --eval_batch_size "$eval_batch_size" \
  --eval_seq_len "$seq_len" \
  --eval_num_samples "$eval_num_samples" \
  --distill_kl_weight "$distill_kl_weight" \
  --distill_kl_temp "$distill_kl_temp" \
  --distill_weight_decay 0.0 \
  --distill_max_grad_norm 1.0 \
  --distill_grad_accum_steps 1 \
  --distill_eval_every 0 \
  --lora_eval_every 0 \
  --lora_epochs "$lora_epochs" \
)

python_args+=("${fisher_args[@]}")
if [[ -n "$save_full_model_cycles" ]]; then
  python_args+=(--save_full_model_cycles "$save_full_model_cycles")
fi

python_args+=(
  --distill_method reparam \
  --redistrib_teacher_source previous_cycle \
  --comm_enabled \
  --comm_mu_auto \
  --layer auto \
  --exclude_pairs -1 \
  --num_progressive "$num_progressive" \
  --output_dir "$output_dir" \
  --dtype "$dtype" \
)
python_args+=("$@")

mkdir -p "$output_dir"
run_args_file="$output_dir/run_args.txt"
start_epoch=$(date +%s)
start_time=$(date --iso-8601=seconds)
{
  echo "start_time=$start_time"
  echo "command:"
  printf '%q ' python "$repo_root/src/fuse_layers.py" "${python_args[@]}"
  echo
} > "$run_args_file"

write_run_summary() {
  local exit_code=$?
  local end_epoch end_time elapsed_seconds
  end_epoch=$(date +%s)
  end_time=$(date --iso-8601=seconds)
  elapsed_seconds=$((end_epoch - start_epoch))
  {
    echo "end_time=$end_time"
    echo "elapsed_seconds=$elapsed_seconds"
    echo "exit_code=$exit_code"
  } >> "$run_args_file"
}
trap write_run_summary EXIT

python "$repo_root/src/fuse_layers.py" "${python_args[@]}"