167 lines
6.1 KiB
Bash
167 lines
6.1 KiB
Bash
#!/usr/bin/env bash
|
|
# =============================================================================
|
|
# orpo_hp_sweep.sh — ORPO Hyperparameter Sweep (200 steps each)
|
|
#
|
|
# 각 설정을 200 steps씩 돌려서 최적 조합을 찾는 스크립트.
|
|
# 결과는 sweep_results/ 디렉토리에 저장됨.
|
|
#
|
|
# Usage:
|
|
# bash scripts/orpo_hp_sweep.sh # 전체 sweep (6 runs)
|
|
# bash scripts/orpo_hp_sweep.sh --dry-run # 설정만 출력
|
|
# =============================================================================
|
|
set -uo pipefail
|
|
# NOTE: set +e — individual runs may fail; we log failures and continue the sweep
|
|
|
|
cd "$(dirname "$0")/.."
|
|
|
|
SWEEP_STEPS=200
|
|
SWEEP_DIR="checkpoints/orpo_sweep"
|
|
RESULTS_FILE="${SWEEP_DIR}/sweep_results.jsonl"
|
|
BASE_MODEL="eval/outputs/hf_3b_sft_best"
|
|
DATA_PATH="data/preference/combined_preference.jsonl"
|
|
NPROC=8
|
|
MASTER_PORT_BASE=29510
|
|
|
|
# B200 NCCL tuning (NVSwitch mesh — let NCCL auto-detect proto/channels/algo)
|
|
export NCCL_IB_DISABLE=1
|
|
export NCCL_BUFFSIZE=134217728
|
|
export OMP_NUM_THREADS=9
|
|
export MKL_NUM_THREADS=9
|
|
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
|
export NCCL_P2P_LEVEL=NVL
|
|
export PYTHONWARNINGS="ignore::UserWarning:torch.library"
|
|
|
|
mkdir -p "${SWEEP_DIR}"
|
|
declare -a FAILED_RUNS=()
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sweep configurations: (name, beta, lr, max_length, batch_size, grad_accum)
|
|
# ---------------------------------------------------------------------------
|
|
# 핵심 탐색 축:
|
|
# 1. beta: 반복 억제 강도 (0.15 vs 0.25 vs 0.35)
|
|
# 2. lr: 수렴 속도 (5e-6 vs 8e-6 vs 1.2e-5)
|
|
# 3. max_length: VRAM vs 커버리지 (1024 vs 1536)
|
|
|
|
declare -a CONFIGS=(
|
|
# name beta lr max_len bs accum
|
|
"baseline_b015_lr8e6 0.15 8e-6 1536 4 4"
|
|
"baseline_b025_lr8e6 0.25 8e-6 1536 4 4"
|
|
"strong_b035_lr8e6 0.35 8e-6 1536 4 4"
|
|
"fast_b025_lr12e6 0.25 1.2e-5 1536 4 4"
|
|
"conserv_b025_lr5e6 0.25 5e-6 1536 4 4"
|
|
"short_b025_lr8e6 0.25 8e-6 1024 4 4"
|
|
)
|
|
|
|
DRY_RUN=false
|
|
if [[ "${1:-}" == "--dry-run" ]]; then
|
|
DRY_RUN=true
|
|
fi
|
|
|
|
echo "=================================================================="
|
|
echo " ORPO Hyperparameter Sweep"
|
|
echo " Configs: ${#CONFIGS[@]}"
|
|
echo " Steps each: ${SWEEP_STEPS}"
|
|
echo " Results: ${RESULTS_FILE}"
|
|
echo "=================================================================="
|
|
|
|
for i in "${!CONFIGS[@]}"; do
|
|
read -r NAME BETA LR MAX_LEN BS ACCUM <<< "${CONFIGS[$i]}"
|
|
PORT=$((MASTER_PORT_BASE + i))
|
|
OUTPUT="${SWEEP_DIR}/${NAME}"
|
|
|
|
echo ""
|
|
echo "--- Run $((i+1))/${#CONFIGS[@]}: ${NAME} ---"
|
|
echo " beta=${BETA} lr=${LR} max_length=${MAX_LEN} bs=${BS} accum=${ACCUM}"
|
|
|
|
if [[ "${DRY_RUN}" == "true" ]]; then
|
|
echo " [DRY RUN] skipping"
|
|
continue
|
|
fi
|
|
|
|
mkdir -p "${OUTPUT}"
|
|
START_TIME=$(date +%s)
|
|
|
|
torchrun \
|
|
--nproc_per_node=${NPROC} \
|
|
--master_port=${PORT} \
|
|
train/orpo.py \
|
|
--model_path "${BASE_MODEL}" \
|
|
--custom_data_path "${DATA_PATH}" \
|
|
--output_dir "${OUTPUT}" \
|
|
--max_steps ${SWEEP_STEPS} \
|
|
--lr ${LR} \
|
|
--beta ${BETA} \
|
|
--batch_size ${BS} \
|
|
--gradient_accumulation_steps ${ACCUM} \
|
|
--max_length ${MAX_LEN} \
|
|
\
|
|
--weight_decay 0.01 \
|
|
--warmup_ratio 0.05 \
|
|
--eval_split_ratio 0.05 \
|
|
--eval_steps 100 \
|
|
--early_stopping_patience 100 \
|
|
--save_steps 200 \
|
|
--save_total_limit 1 \
|
|
--logging_steps 10 \
|
|
--report_to none \
|
|
--dataset_num_proc 64 \
|
|
--dataloader_num_workers 4 \
|
|
--no_load_best \
|
|
2>&1 | tee "${OUTPUT}/train.log"
|
|
RUN_EXIT=$?
|
|
|
|
END_TIME=$(date +%s)
|
|
ELAPSED=$((END_TIME - START_TIME))
|
|
|
|
if [[ ${RUN_EXIT} -ne 0 ]]; then
|
|
echo " [ERROR] Run ${NAME} failed with exit code ${RUN_EXIT} after ${ELAPSED}s"
|
|
echo "{\"name\":\"${NAME}\",\"beta\":${BETA},\"lr\":\"${LR}\",\"max_length\":${MAX_LEN},\"status\":\"FAILED\",\"exit_code\":${RUN_EXIT},\"elapsed_s\":${ELAPSED}}" >> "${RESULTS_FILE}"
|
|
FAILED_RUNS+=("${NAME}")
|
|
continue
|
|
fi
|
|
|
|
# Extract final metrics from log
|
|
FINAL_LOSS=$(grep -oP "'loss': '[\d.]+'" "${OUTPUT}/train.log" | tail -1 | grep -oP "[\d.]+" || echo "N/A")
|
|
EVAL_LOSS=$(grep -oP "'eval_loss': '[\d.]+'" "${OUTPUT}/train.log" | tail -1 | grep -oP "[\d.]+" || echo "N/A")
|
|
MARGIN=$(grep -oP "'rewards/margins': '[-\d.]+'" "${OUTPUT}/train.log" | tail -1 | grep -oP "[-\d.]+" || echo "N/A")
|
|
|
|
# Save result
|
|
echo "{\"name\":\"${NAME}\",\"beta\":${BETA},\"lr\":\"${LR}\",\"max_length\":${MAX_LEN},\"status\":\"OK\",\"loss\":\"${FINAL_LOSS}\",\"eval_loss\":\"${EVAL_LOSS}\",\"margin\":\"${MARGIN}\",\"elapsed_s\":${ELAPSED}}" >> "${RESULTS_FILE}"
|
|
|
|
echo " -> loss=${FINAL_LOSS} eval_loss=${EVAL_LOSS} margin=${MARGIN} time=${ELAPSED}s"
|
|
|
|
# Cleanup weights to save disk (keep logs)
|
|
rm -rf "${OUTPUT}/checkpoint-"* "${OUTPUT}/emergency_checkpoint" 2>/dev/null || true
|
|
done
|
|
|
|
echo ""
|
|
echo "=================================================================="
|
|
echo " Sweep Complete!"
|
|
echo " Results: ${RESULTS_FILE}"
|
|
if [[ -f "${RESULTS_FILE}" ]]; then
|
|
echo ""
|
|
echo " Summary:"
|
|
cat "${RESULTS_FILE}" | python3 -c "
|
|
import sys, json
|
|
results = [json.loads(l) for l in sys.stdin]
|
|
results.sort(key=lambda r: float(r.get('eval_loss', '999')))
|
|
print(f' {\"Name\":<25} {\"Beta\":>6} {\"LR\":>10} {\"Loss\":>8} {\"EvalLoss\":>10} {\"Margin\":>8} {\"Time\":>6}')
|
|
print(f' {\"-\"*25} {\"-\"*6} {\"-\"*10} {\"-\"*8} {\"-\"*10} {\"-\"*8} {\"-\"*6}')
|
|
for r in results:
|
|
print(f' {r[\"name\"]:<25} {r[\"beta\"]:>6} {r[\"lr\"]:>10} {r[\"loss\"]:>8} {r[\"eval_loss\"]:>10} {r[\"margin\"]:>8} {r[\"elapsed_s\"]:>5}s')
|
|
print()
|
|
best = results[0]
|
|
print(f' BEST: {best[\"name\"]} (eval_loss={best[\"eval_loss\"]})')
|
|
" 2>/dev/null || cat "${RESULTS_FILE}"
|
|
fi
|
|
|
|
# Report failed runs
|
|
if [[ ${#FAILED_RUNS[@]} -gt 0 ]]; then
|
|
echo ""
|
|
echo " FAILED RUNS (${#FAILED_RUNS[@]}):"
|
|
for fname in "${FAILED_RUNS[@]}"; do
|
|
echo " - ${fname}"
|
|
done
|
|
fi
|
|
echo "=================================================================="
|