#!/bin/bash
# Submit the 16 compression jobs for the LLaDA-8B-INSTRUCT table.
#
# Calib is REUSED from the Base table -- not rebuilt. Verified equivalent:
# the Instruct tokenizer is a strict superset of Base (adds only
# <|start_header_id|>/<|end_header_id|>/<|eot_id|> at ids 126346-126348; no
# id remapping, same MASK_ID 126336), and build_windows/apply_noise depend on
# (tokenizer, corpus, seed) alone -- no model forward. Checked on 4000 real C4
# docs plus an end-to-end rebuild: the shipped calib's window hashes reproduce
# exactly under the Instruct tokenizer.
#
# Node policy differs from the Base run: compress.sbatch hardcodes
# --exclude=node1..node4 to force L40S (5.7x faster). node6 is currently full
# and 48 foreign jobs are queued, so we override to exclude only the
# drained/draining nodes and take whichever GPU frees first. A5000 is proven
# for this path (base/all_m08_c4 finished there in 8.57h) and fits --time=24h.
set -euo pipefail
cd /home/tl356/LLaDA/trajmc_main

MODEL=/home/tl356/LLaDA-8B-Instruct
CLEAN=results/calib/base_c4_n1400_s42_full_2c9a327_calib.pt
TRAJ=results/calib/ours_c4_n1400_s42_full_2c9a327_calib.pt
W=weights_instruct
AVOID="node3,node5"     # drain / draining

for f in "$CLEAN" "$TRAJ"; do [ -s "$f" ] || { echo "missing calib: $f" >&2; exit 1; }; done
[ -d "$MODEL" ] || { echo "missing model: $MODEL" >&2; exit 1; }

# NOTE: label is an explicit arg, not $(basename save_path) -- gptq and awq
# share the w4g128_c4_clean basename and would collide in squeue.
# Skip a target that is already finished OR already in flight, so this script can be
# re-run safely after a partial failure. compress*.py are themselves idempotent on the
# summary file, but that only helps AFTER a job has burned a GPU slot; the squeue check
# is what stops a re-run from double-submitting the jobs still running from last time.
_skip () {  # <label> <save_path>  -> 0 = skip
  if [ -s "$2/compression_summary.json" ]; then
    echo "$1  SKIP (done: $2)"; return 0
  fi
  if squeue -u "$USER" -h -o "%j" | grep -qx "tmi_$1"; then
    echo "$1  SKIP (in flight)"; return 0
  fi
  return 1
}

sub_lowrank () {  # <label> <calib> <model_ratio> <decomp> <save_path>
  _skip "$1" "$5" && return 0
  sbatch --parsable --requeue --exclude="$AVOID" --job-name="tmi_$1" \
    jobs/compress.sbatch \
    --calib "$2" --model_ratio "$3" --layer_type all --decomp "$4" \
    --model_path "$MODEL" --save_path "$5"
}
sub_baseline () { # <label> <calib> <method> <bits> <sparsity> <save_path>
  _skip "$1" "$6" && return 0
  sbatch --parsable --requeue --time=24:00:00 --exclude="$AVOID" --job-name="tmi_$1" \
    jobs/compress_baselines.sbatch \
    --calib "$2" --method "$3" --bits "$4" --groupsize 128 --sparsity "$5" \
    --layer_type all --model_path "$MODEL" --save_path "$6"
}

echo "== low-rank (compress.py) =="
# SVD-LLM 20% / 40% reduction  == model_ratio 0.8 / 0.6 (cholesky)
echo "svdllm_m08_clean  $(sub_lowrank svdllm_m08_clean $CLEAN 0.8 cholesky $W/base/all_m08_c4)"
echo "svdllm_m08_traj   $(sub_lowrank svdllm_m08_traj $TRAJ  0.8 cholesky $W/ours/all_m08_c4)"
echo "svdllm_m06_clean  $(sub_lowrank svdllm_m06_clean $CLEAN 0.6 cholesky $W/base/all_m06_c4)"
echo "svdllm_m06_traj   $(sub_lowrank svdllm_m06_traj $TRAJ  0.6 cholesky $W/ours/all_m06_c4)"
# ASVD 20% reduction
echo "asvd_m08_clean    $(sub_lowrank asvd_m08_clean $CLEAN 0.8 asvd     $W/asvd/all_m08_c4)"
echo "asvd_m08_traj     $(sub_lowrank asvd_m08_traj $TRAJ  0.8 asvd     $W/asvd/all_m08_c4_traj)"

echo "== quant / pruning (compress_baselines.py) =="
echo "gptq_w4_clean     $(sub_baseline gptq_w4_clean $CLEAN gptq 4 0.23 $W/gptq/w4g128_c4_clean)"
echo "gptq_w4_traj      $(sub_baseline gptq_w4_traj $TRAJ  gptq 4 0.23 $W/gptq/w4g128_c4_traj)"
echo "gptq_w3_clean     $(sub_baseline gptq_w3_clean $CLEAN gptq 3 0.23 $W/gptq/w3g128_c4_clean)"
echo "gptq_w3_traj      $(sub_baseline gptq_w3_traj $TRAJ  gptq 3 0.23 $W/gptq/w3g128_c4_traj)"
echo "awq_w4_clean      $(sub_baseline awq_w4_clean $CLEAN awq  4 0.23 $W/awq/w4g128_c4_clean)"
echo "awq_w4_traj       $(sub_baseline awq_w4_traj $TRAJ  awq  4 0.23 $W/awq/w4g128_c4_traj)"
echo "awq_w3_clean      $(sub_baseline awq_w3_clean $CLEAN awq  3 0.23 $W/awq/w3g128_c4_clean)"
echo "awq_w3_traj       $(sub_baseline awq_w3_traj $TRAJ  awq  3 0.23 $W/awq/w3g128_c4_traj)"
echo "sparsegpt50_clean $(sub_baseline sparsegpt50_clean $CLEAN sparsegpt 4 0.5 $W/sparsegpt/unstr50_c4_clean)"
echo "sparsegpt50_traj  $(sub_baseline sparsegpt50_traj $TRAJ  sparsegpt 4 0.5 $W/sparsegpt/unstr50_c4_traj)"