#!/usr/bin/env bash
# =============================================================================
# retrain_corrected_cache.sh -- rebuild the whole checkpoint chain under the
# CORRECTED YaRN parameterisation (YARN_ORIG_MAXPOS=4096).
#
# WHY THE FULL CHAIN, not just the arms: run7 itself was trained with the buggy
# orig_maxpos=1024 default, so arms trained from it would start from weights
# already adapted to the wrong rotation. run8 (the 4k base) is unaffected -- it
# predates any YaRN scaling -- so the chain restarts there.
#
# WHY THIS IS THE REAL TEST: switching the cache at EVAL time is not a fix.
# Measured on this box, corrected-vs-as-trained at N=100: 1k 90 vs 84, ~4k
# 14-33 vs 75, 16k 17 vs 0. The corrected ramp (low=14 high=32) is a DIFFERENT
# function from the one the weights adapted to (low=7 high=25) -- better in some
# regimes, much worse around 4k. Only retraining resolves that.
#
# Recipe recovered from run7's own log + LR trace, not from memory:
#   B=4 grad_accum=16 global_batch=1048576, Muon+AdamW ZeRO-1 symmetric, bf16
#   mask=15, max_lr 1e-3 -> min 1e-4, rewarmup 50 to peak 0.3x (=3e-4), 750 steps
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
ROOT="$(cd ../.. && pwd)"
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1

PY=${PY:-/root/Adhi/BluTrain/envv/bin/python}
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=retrain_logs; mkdir -p "$LOG"
M="$LOG/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

T=16384; YARN=4; ORIG=4096; NP=4; GPUS=0,1,2,3
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
COMMON="CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
YARN_SCALE=$YARN YARN_ORIG_MAXPOS=$ORIG \
CP_ZERO=1 CP_ZERO_V2=1 CP_ZERO_SYMMETRIC=1 CP_BF16=1 CP_OPTIMIZER=muon \
CP_MAX_LR=1e-3 CP_MIN_LR=1e-4 CP_WARMUP=1 \
CP_CKPT=1 CP_CKPT_KEEP=1 CP_CKPT_NEW_RUN=1"

[ -x "$EXE" ] || die "$EXE missing"
FREE=$(df -BG --output=avail /root | tail -1 | tr -dc '0-9')
say "free disk ${FREE}G (need ~70G: 4 checkpoint sets at ~17G)"
[ "${FREE:-0}" -ge 75 ] || die "only ${FREE}G free"

# run a training job, then VERIFY the checkpoint exists -- a zero exit code is
# not proof: bluscriptCP catches a failed save, prints [WARN] and trains on.
train(){   # $1=label $2=logfile $3...=extra env
  local label="$1" lf="$2"; shift 2
  say "training $label"
  env $ARCH $COMMON "$@" CP_DATA_ROOT="$PWD/$DATA_ROOT" \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$lf" 2>&1 || die "$label training failed -- see $lf"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$lf" | head -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || die "$label: no run number in $lf"
  CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || die "$label finished but wrote NO checkpoint (see $lf)"
  grep -qa 'checkpoint save.*failed' "$lf" && die "$label: a checkpoint save FAILED"
  say "  $label -> run $RN, $(basename "$CKF"), final loss $(grep -aE '^step ' "$lf" | tail -1 | grep -oE 'loss: [0-9.]+')"
  case "$CKF" in *$'\n'*) die "ckptpath sanity: multi-line checkpoint path";; esac
  echo "$CKF"
}

# ---- 1. arm training data (deterministic at --seed 0) -----------------------
FILL=$(ls Data_Loader/Data/flux_train_*.bin 2>/dev/null | head -12 | tr '\n' ' ')
[ -n "$FILL" ] || die "no flux_train shards for in2_gen filler"
for A in a b c; do
  if [ ! -f "in2_data/arm_$A/.done" ]; then
    say "generating arm $A data"
    $PY Tests/bluscriptcp/in2_gen.py --arm $A --T $T --examples 12816 --K 64 --shards 3 \
        --seed 0 --d-min 64 --filler-shards $FILL --out-dir in2_data/arm_$A >>"$M" 2>&1 \
        || die "arm $A train data generation failed"
    $PY Tests/bluscriptcp/in2_gen.py --arm $A --T $T --examples 64 --split val --K 64 --shards 1 \
        --seed 0 --d-min 64 --filler-shards $(ls Data_Loader/Data/flux_val*.bin | head -1) \
        --out-dir in2_data/arm_$A >>"$M" 2>&1 || die "arm $A val data generation failed"
    touch "in2_data/arm_$A/.done"
  fi
done
say "arm data ready"

# ---- 2. run7' : the YaRN 16k extension, corrected cache ---------------------
if [ -f "$LOG/run7c.ckptpath" ]; then
  R7=$(cat "$LOG/run7c.ckptpath"); say "run7' already done ($R7)"
else
  DATA_ROOT=Data_Loader/Data
  R7=$(train "run7' (YaRN 16k extension, corrected cache)" "$LOG/run7c.log" \
        CP_INIT_FROM="$PWD/$CK/blumodelcp_run8_step_1000.ckpt" \
        CP_B=4 CP_GLOBAL_BATCH=1048576 CP_MAX_STEPS=750 \
        CP_REWARMUP=50 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=250) || exit 1
  echo "$R7" > "$LOG/run7c.ckptpath"
fi

# ---- 3. arms A/B/C from run7' ----------------------------------------------
for A in a b c; do
  if [ -f "$LOG/arm_$A.ckptpath" ]; then say "arm $A already done"; continue; fi
  DATA_ROOT=in2_data/arm_$A
  P=$(train "arm $A (corrected cache)" "$LOG/arm_$A.log" \
        CP_INIT_FROM="$PWD/$R7" \
        CP_B=4 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=400 \
        CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=200) || exit 1
  echo "$P" > "$LOG/arm_$A.ckptpath"
done

say "ALL DONE"
say "  run7' : $(cat $LOG/run7c.ckptpath)"
for A in a b c; do say "  arm $A : $(cat $LOG/arm_$A.ckptpath 2>/dev/null)"; done