#!/usr/bin/env bash
# =============================================================================
# ext64k_b5.sh -- YaRN extension 32k -> 64k on b5, deadline-aware.
#
# PHASE 1  3-step probe: does 64k fit at cp=4 on 143 GB, and how fast is a step?
# PHASE 2  train as many steps as fit before HARD_STOP, WSD tail sized so the
#          FINAL checkpoint is DECAYED and comparable to the 32k results
# PHASE 3  eval 4 lengths IN PARALLEL on the 4 GPUs (wall clock = slowest job)
# PHASE 4  upload ckpt + sidecar + logs
#
# WHY b5: documents are 65,536-131,072 tokens, so the MINIMUM equals the 64k
#   window and every training sample is one contiguous document. b4 (32-64k)
#   would straddle in roughly half its samples. Same reasoning that made b4 the
#   right choice at 32k.
#
# PARENT: b4 full-epoch step 2227 -- the best 32k model (70.0/58.0/43.0 at
#   8k/16k/32k). CP_INIT_FROM loads WEIGHTS ONLY, so its cp=2 shard layout is
#   irrelevant to this cp=4 run, and no sidecar is needed (none is read).
#
# YaRN INVARIANT: SCALE * ORIG_MAXPOS must equal seq_len EXACTLY and ORIG_MAXPOS
#   must be the length the parent was actually trained at. 2 * 32768 = 65536.
#
# LR: CP_MAX_LR / CP_MIN_LR DO NOT EXIST -- bluscriptCP.cpp:131-132 hardcodes
#   6e-4 / 6e-5 and the env block never reads them. Real schedule is
#   6e-4 * CP_REWARMUP_PEAK(0.3) = 1.8e-4 peak, 6e-5 floor -- the same values
#   every previous run used, so results stay comparable.
#
# DEADLINE: HARD_STOP is a UTC epoch. Step count is derived from the MEASURED
#   probe rate, not an estimate, so the run lands inside the window.
# =============================================================================
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=/root/envv/bin/python3
HF=/root/envv/bin/hf
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=ext64k_logs; mkdir -p "$LOG"
M="$LOG/master.log"
REPO=unparallelled/BluScriptCP
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

T=65536; NP=4; GPUS=0,1,2,3
ORIG=32768; SCALE=2                  # 2 * 32768 == 65536
DATA=lendata/b5_65536_131072
PARENT="$CK/b4epoch32k_step2227.ckpt"
HAY="$PWD/lendata/haystack/b2val_haystack.bin"
WARM=100; TAIL=${TAIL:-2}; FREQ=50; KEEP=2
GB=524288
# 23:30 IST = 18:00 UTC today
HARD_STOP=${HARD_STOP:-$(date -u -d "$(date -u +%Y-%m-%d) 18:00:00" +%s)}
RESERVE_MIN=${RESERVE_MIN:-50}       # eval (28) + upload (20) + slack

$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
[ -x "$EXE" ]    || die "$EXE not built"
[ -f "$PARENT" ] || die "parent missing: $PARENT"
[ -f "$HAY" ]    || die "haystack missing: $HAY"
[ -d "$DATA" ]   || die "$DATA missing"
[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant: $SCALE * $ORIG != $T"
ntr=$(ls "$DATA"/*train*.bin 2>/dev/null | wc -l)
nva=$(ls "$DATA"/*val*.bin 2>/dev/null | wc -l)
[ "$ntr" -ge 1 ] && [ "$nva" -ge 1 ] || die "need train+val shards in $DATA (have $ntr/$nva)"
have=$(( $(du -cb "$DATA"/*train*.bin | tail -1 | cut -f1) / 2 ))

say "=== 32k -> 64k YaRN extension on b5 (4x H200, cp=$NP) ==="
say "parent: b4 full epoch step 2227 (best 32k model)"
say "YaRN SCALE=$SCALE ORIG=$ORIG -> $((SCALE*ORIG)) == CP_T=$T  OK"
say "b5: $ntr train shards = $(echo "scale=3; $have/1000000000" | bc) B tokens available"
say "hard stop $(date -u -d @$HARD_STOP +%H:%M) UTC = $(date -u -d @$((HARD_STOP+19800)) +%H:%M) IST, reserving ${RESERVE_MIN}m"

COMMON="CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG \
CP_ZERO=1 CP_ZERO_V2=1 CP_ZERO_SYMMETRIC=1 CP_BF16=1 CP_OPTIMIZER=muon \
CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=0.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"

# ---- PHASE 1: fit + speed --------------------------------------------------
if [ ! -f "$LOG/probe.ok" ]; then
  say "PHASE 1: 3-step probe (64k at cp=4 projects to 108.8 GB of 143 GB)"
  env $ARCH $COMMON CP_STABLE=0 CP_CKPT=0 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=65536 CP_MAX_STEPS=3 \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$LOG/probe.log" 2>&1
  if grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log"; then
    say "PROBE OOM -- 64k does not fit at cp=$NP even on 143 GB cards."
    grep -aoE 'cudaMallocAsync\(size=[0-9]+\) failed' "$LOG/probe.log" | head -2 | sed 's/^/    /' >&2
    exit 1
  fi
  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | sed 's/^/  /' >&2
  say "PROBE OK: peak $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
  touch "$LOG/probe.ok"
fi

# ---- derive the step count from the MEASURED rate --------------------------
# The probe used a tiny global batch, so its dt is NOT the real per-step cost.
# Scale it: real grad_accum is GB/T vs the probe's 65536/T.
SPS=$($PY - "$LOG/probe.log" $GB $T <<'PYEOF'
import re,sys,statistics
t=open(sys.argv[1],errors='ignore').read()
d=[float(x) for x in re.findall(r'dt: ([0-9.]+)ms', t)]
probe_accum=65536//int(sys.argv[3]) or 1
real_accum=int(sys.argv[2])//int(sys.argv[3])
per=statistics.median(d[1:]) if len(d)>2 else d[-1]
print(f"{per/1000.0*real_accum/probe_accum:.3f}")
PYEOF
)
say "measured step cost at real batch: ${SPS} s/step"
if [ -n "${STEPS_OVERRIDE:-}" ]; then
  say "STEPS_OVERRIDE=$STEPS_OVERRIDE -- the probe ran grad_accum=1 so its dt x8"
  say "  overstated the real step cost 5.8x (74.0 vs a measured 12.85 s/step)."
fi
STEPS=$($PY - "$SPS" "$HARD_STOP" "$RESERVE_MIN" "$have" "$GB" "$WARM" "$TAIL" <<'PYEOF'
import sys,time,math
sps=float(sys.argv[1]); stop=int(sys.argv[2]); res=int(sys.argv[3])
have=int(sys.argv[4]); gb=int(sys.argv[5]); warm=int(sys.argv[6]); tail=int(sys.argv[7])
budget=stop-time.time()-res*60
n=int(budget/sps)
n=min(n, have//gb)                       # never wrap the data
n=max(n, warm+tail+10)                   # schedule must be legal
print(n)
PYEOF
)
STEPS=${STEPS_OVERRIDE:-$STEPS}
STABLE=$(( STEPS - WARM - TAIL ))
[ "$STABLE" -lt 1 ] && STABLE=1
say "PHASE 2: $STEPS steps = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens = $(echo "scale=0; $STEPS*$GB*100/866000000" | bc)% of a b5 epoch"
say "  WSD: warmup $WARM, stable $STABLE, decay tail $TAIL -> FINAL checkpoint IS decayed"

if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
  say "already trained -> $(cat "$LOG/ext.ckpt")"
else
  env $ARCH $COMMON CP_STABLE=$STABLE \
      CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP CP_CKPT_NEW_RUN=1 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=$STEPS \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$LOG/ext.log" 2>&1
  grep -qaiE 'out of memory|bad_alloc' "$LOG/ext.log" && die "OOM in the full run"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/ext.log" | head -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || die "no run number"
  CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || die "run $RN wrote NO checkpoint"
  n=$(ls checkpoints_bluscriptcp/rank_*/"$(basename "$CKF")" 2>/dev/null | wc -l)
  [ "$n" -ge "$NP" ] || die "only $n/$NP rank shards for $(basename "$CKF")"
  echo "$CKF" > "$LOG/ext.ckpt"; echo "$RN" > "$LOG/runnum"
  say "ext64k -> run $RN, $(basename "$CKF"), loss $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
  say "  peak mem $(grep -aoE 'mem: [0-9]+MB' "$LOG/ext.log" | tail -1)"
  say "  final lr $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'lr [0-9.e-]+')  (should be ~6e-05 = decayed)"
fi

# ---- PHASE 3: eval, 4 lengths in PARALLEL on 4 GPUs ------------------------
EXT=$(basename "$(cat "$LOG/ext.ckpt")")
say "PHASE 3: eval ext64k -- 4 lengths, one GPU each"
# 65512 + MAX_NEW 24 = 65536 exactly; 65528 would pad to 65792 and abort.
for LEN in 8192 16360 32744 65512; do
  ( GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=ext64k_len_$LEN TASKS="niah_single" \
    CKPTS="ext64k" N=100 PROMPT_LEN=$LEN MAX_NEW=24 T=$T \
    DEPTHS="0.0,0.25,0.5,0.75,1.0" YARN_ORIG=$ORIG YARN=$SCALE \
    GPUS=$(( (LEN==8192)?0:(LEN==16360)?1:(LEN==32744)?2:3 )) JOBS_PER_GPU=1 \
    BF16=1 GREEDY=1 CKPT_EXTRA="ext64k=$EXT" ./run_longeval.sh \
    >> "$LOG/eval_$LEN.log" 2>&1 ) &
done
wait
for LEN in 8192 16360 32744 65512; do
  [ -s "ext64k_len_$LEN/summary.csv" ] && say "  len $LEN -> $(sed -n 2p ext64k_len_$LEN/summary.csv)" \
                                       || say "  len $LEN -> NO RESULT"
done

# ---- PHASE 4: upload -------------------------------------------------------
say "PHASE 4: upload"
{
  echo "# 32k -> 64k YaRN extension on b5 -- $(date -u +%Y-%m-%d)"
  echo
  echo "Parent: b4 full-epoch step 2227 (best 32k model: 70.0/58.0/43.0 at 8k/16k/32k)."
  echo "Data:   b5_65536_131072 -- min doc 65,536 == the window, so every sample is"
  echo "        ONE contiguous document."
  echo "YaRN:   SCALE=2 ORIG_MAXPOS=32768 -> 65536 exactly."
  echo "Train:  $STEPS steps x $GB = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens"
  echo "        = $(echo "scale=0; $STEPS*$GB*100/866000000" | bc)% of b5's 0.866 B train split (no wrap)."
  echo "        WSD warmup $WARM / stable $STABLE / decay tail $TAIL -- final ckpt DECAYED."
  echo "        LR peak 1.8e-4, floor 6e-5 (CP_MAX_LR is a no-op; see B4 results doc)."
  echo "        4x H200, cp=4, B=1 x grad_accum $((GB/T))."
  echo "Eval:   niah_single, N=100, greedy, shared haystack b2_8192_16384 val."
  echo
  echo "## Results"; echo
  echo "| context | niah_single |"; echo "|---|---|"
  for LEN in 8192 16360 32744 65512; do
    v=$(sed -n 2p "ext64k_len_$LEN/summary.csv" 2>/dev/null | cut -d, -f2)
    printf "| %s | %s |\n" "$LEN" "${v:-n/a}"
  done
  echo
  echo "For comparison, the 32k parent at the same lengths: 70.0 / 58.0 / 43.0 / n/a"
  echo "(it cannot run at 64k -- its RoPE cache is sized 32,768)."
  echo
  echo "NOTE: this is a PARTIAL b5 epoch, sized to a wall-clock deadline. It is a"
  echo "first 64k data point, not a converged extension."
} > "$LOG/RESULTS.md"
up(){ [ -e "$1" ] || return 0
  $HF upload "$REPO" "$1" "$2" >>"$LOG/upload.log" 2>&1 && say "  ok $2" || say "  FAIL $2"; }
up "$LOG/RESULTS.md" docs/EXT64K_B5_RESULTS_$(date -u +%Y-%m-%d).md
for LEN in 8192 16360 32744 65512; do up "ext64k_len_$LEN" "logs/ext64k_len_$LEN/"; done
up "$LOG" logs/ext64k/
up ext64k_b5.sh scripts/ext64k_b5.sh
RN=$(cat "$LOG/runnum" 2>/dev/null)
for r in 0 1 2 3; do
  for e in ckpt zero_g1; do
    up "checkpoints_bluscriptcp/rank_$r/blumodelcp_run${RN}_step_${STEPS}.$e" \
       "checkpoints/ext64k_b5_rank${r}_step_${STEPS}.$e"
  done
done
say "=== EXT64K DONE ==="