Add run_decay_only.sh (decay-only from pretrain branch ckpts)

run_decay_only.sh

@@ -0,0 +1,132 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# ============================================================================
+# [Decay-only] 선택한 Stage1 branch 체크포인트로부터 *decay(annealing)만* 수행.
+#              (run_stage2_3_from_ckpt.sh 의 (1) decay 단계만 떼어낸 것. stage2/3 는 하지 않음.)
+#
+#  - SETTINGS  : run_stage1_sweep.sh 에서 쓴 것과 동일한 (hidden-depth-heads) 세트.
+#  - EPOCHS    : 목표 stage1 "총 budget"(decay 포함) epoch 들. SAVE_EVERY 의 배수여야
+#               stage1 에서 저장된 branch 체크포인트와 정확히 맞습니다.
+#
+#  각 (setting, budget) 조합마다:
+#     branch ckpt(= interval_epoch_BRANCH_EP.ckpt) 의 가중치에서 시작해
+#     LR 을 DECAY_LR -> DECAY_ETA_MIN 으로 DECAY_EPOCHS 만큼 decay 합니다.
+#
+#  ※ decay 길이는 *비율 기반*(stage1 의 branch 저장과 동일 공식):
+#       DECAY_EPOCHS(E) = ceil(max(DECAY_RATIO*E, MIN_DECAY_EPOCHS))
+#       BRANCH_EP(E)    = E - DECAY_EPOCHS(E)
+#    총 budget E = BRANCH_EP stable + DECAY_EPOCHS decay (길이 E 단일 WSD run 과 등가).
+#    ⚠️ DECAY_RATIO / MIN_DECAY_EPOCHS 는 run_stage1_sweep.sh 와 *반드시 동일*해야
+#       BRANCH_EP 가 stage1 이 저장한 branch epoch 과 일치합니다.
+#
+#  ※ 출력 경로/이름은 run_stage2_3_from_ckpt.sh 의 decay 단계와 동일하게 맞춰서,
+#    원하면 나중에 그 스크립트로 stage2/3 만 이어서 돌릴 수 있습니다
+#    (그 경우 decay 는 last.ckpt 자동 resume 으로 즉시 통과).
+# ============================================================================
+
+CMD="python train.py"
+CFG_DECAY="/workspace/stormer/configs/decay_finetune.yaml"
+DATA_ROOT="/workspace/stormer/wb2_h5df"
+
+S1_RESULTS_ROOT="/workspace/stormer/results/pretrain"
+DECAY_RESULTS_ROOT="/workspace/stormer/results/decay"
+
+# ---- stage1 과 동일한 (hidden-depth-heads) 세트 ----
+SETTINGS=(
+  "128-24-2"
+)
+
+# ---- ⭐ 목표 stage1 "총 budget"(decay 포함) epoch 들 (SAVE_EVERY 배수) ----
+EPOCHS=(20 40 80 100)
+
+PATCH_SIZE=4
+BATCH_SIZE=1
+
+# ---- decay(annealing) 설정 ----
+# ⚠️ DECAY_RATIO / MIN_DECAY_EPOCHS 는 run_stage1_sweep.sh 와 동일하게!
+DECAY_RATIO=0.1          # decay 길이 비율 (budget 대비)
+MIN_DECAY_EPOCHS=1       # decay 최소 길이(epoch)
+DECAY_LR=5e-4            # decay 시작 LR = stage1 의 stable LR 과 동일하게
+DECAY_TYPE=sqrt          # cosine | linear | sqrt
+DECAY_ETA_MIN=1e-8       # decay 도착 LR
+
+# DECAY_EPOCHS(E) = ceil(max(DECAY_RATIO*E, MIN_DECAY_EPOCHS))  (stage1 branch 공식과 동일)
+decay_len () {
+  awk -v e="$1" -v r="${DECAY_RATIO}" -v mn="${MIN_DECAY_EPOCHS}" \
+    'BEGIN{ d=r*e; if(d<mn) d=mn; c=int(d); if(c < d-1e-9) c++; print c }'
+}
+
+# stage1 branch 체크포인트를 epoch(=학습한 epoch 수, 1-index) 으로부터 찾는다
+resolve_s1_ckpt () {
+  local ckpt_dir="$1"; local epoch="$2"
+  local c="${ckpt_dir}/interval_epoch_$(printf '%03d' "${epoch}").ckpt"
+  if [ -f "${c}" ]; then echo "${c}"; return 0; fi
+  {
+    echo "ERROR: stage1 branch epoch '${epoch}' 체크포인트를 ${ckpt_dir} 에서 못 찾음."
+    echo "  (DECAY_RATIO/MIN_DECAY_EPOCHS 가 run_stage1_sweep.sh 와 같은지, EPOCHS 가 SAVE_EVERY 배수인지 확인)"
+    echo "사용 가능한 branch 체크포인트:"
+    ls -1 "${ckpt_dir}"/interval_epoch_*.ckpt 2>/dev/null || echo "  (없음 - stage1 을 먼저 돌렸는지 확인)"
+  } >&2
+  return 1
+}
+
+for setting in "${SETTINGS[@]}"; do
+  IFS="-" read -r h d nh <<< "${setting}"
+  RUN_ID="h${h}_d${d}_n${nh}"
+  S1_NAME="pretrain_one_step_${RUN_ID}"
+  S1_CKPT_DIR="${S1_RESULTS_ROOT}/${RUN_ID}/${S1_NAME}/checkpoints"
+
+  # 아키텍처는 stage1 체크포인트와 반드시 동일해야 하므로 모든 호출에 동일하게 전달
+  ARCH_ARGS=(
+    --model.net.init_args.hidden_size "${h}"
+    --model.net.init_args.depth "${d}"
+    --model.net.init_args.num_heads "${nh}"
+    --model.net.patch_size "${PATCH_SIZE}"
+  )
+
+  for EP in "${EPOCHS[@]}"; do
+    # 총 budget E 에서 비율 decay 를 carve-out: stable 은 (E - DECAY_EPOCHS) 에서 branch.
+    DECAY_EPOCHS="$(decay_len "${EP}")"
+    BRANCH_EP=$((EP - DECAY_EPOCHS))
+    if [ "${BRANCH_EP}" -le 0 ]; then
+      echo "SKIP: EP=${EP} <= DECAY_EPOCHS=${DECAY_EPOCHS} 라 branch 지점(${BRANCH_EP})이 유효하지 않음." >&2
+      continue
+    fi
+
+    TAG="${RUN_ID}_ep${BRANCH_EP}_d${DECAY_EPOCHS}"   # 라벨: stable branch epoch + decay 길이
+    echo "############################################################"
+    echo "# ${TAG} :  decay only"
+    echo "#   목표 총 budget=${EP}ep  =  stable@${BRANCH_EP}ep + decay ${DECAY_EPOCHS}ep"
+    echo "#   decay: ${DECAY_TYPE}, lr ${DECAY_LR} -> ${DECAY_ETA_MIN}"
+    echo "############################################################"
+
+    S1_CKPT="$(resolve_s1_ckpt "${S1_CKPT_DIR}" "${BRANCH_EP}")" || exit 1
+    echo "  [ckpt] stage1 branch @ ${BRANCH_EP}ep : ${S1_CKPT}"
+
+    ##########################################################
+    # DECAY (annealing) :  stage1 ckpt -> lr decay -> eta_min
+    ##########################################################
+    DECAY_NAME="decay_${TAG}"
+    DECAY_ROOT="${DECAY_RESULTS_ROOT}/${DECAY_NAME}"
+    echo "  --- DECAY (${DECAY_EPOCHS} ep, ${DECAY_TYPE}, lr ${DECAY_LR} -> ${DECAY_ETA_MIN}) ---"
+    ${CMD} \
+      --config "${CFG_DECAY}" \
+      --trainer.max_epochs "${DECAY_EPOCHS}" \
+      --trainer.default_root_dir "${DECAY_ROOT}" \
+      --trainer.logger.init_args.name "${DECAY_NAME}" \
+      --model.max_epochs "${DECAY_EPOCHS}" \
+      --model.lr "${DECAY_LR}" \
+      --model.eta_min "${DECAY_ETA_MIN}" \
+      --model.decay_type "${DECAY_TYPE}" \
+      --model.pretrained_path "${S1_CKPT}" \
+      "${ARCH_ARGS[@]}" \
+      --data.root_dir "${DATA_ROOT}" \
+      --data.steps 1 \
+      --data.batch_size "${BATCH_SIZE}"
+    DECAY_CKPT="${DECAY_ROOT}/${DECAY_NAME}/checkpoints/last.ckpt"
+    echo "  [done] ${TAG}: decayed ckpt -> ${DECAY_CKPT}"
+  done
+done
+
+echo "ALL DECAY-ONLY RUNS DONE"