#!/usr/bin/env bash
# =============================================================================
# CONTINUE the 64k b5 extension on 8x H100 -- to a FULL b5 epoch.
#
# PARENT: ext64k run3 step 695 (the 64k model: 67/65/52/29 at 8k/16k/32k/64k).
#   Recovered from strong-room-grows-fin-03 and now durable on HF at
#   checkpoints/ext64k_b5_rank0_step_695.ckpt (byte-verified).
#
# WHY cp=8 AND NOT cp=4: at 64k, cp=4 puts 16,384 tokens on each rank, which
#   MEASURED 108,493 MB on the H200 run. An H100 has 81,559 MiB. It does not
#   fit. cp=8 halves that to 8,192 tok/rank. There is no other lever: CP_B is
#   already 1 and bluscriptCP.cpp has no activation-recompute flag.
#
# WHY CP_INIT_FROM (weights only) AND NOT CP_CKPT_RESUME: ZeRO-1 shards
#   optimizer state across the GLOBAL process group (bluscriptCP.cpp:24-28), so
#   a sidecar written at world_size=4 cannot load at world_size=8 --
#   ShardedOptimizer::load_state throws ZeroCheckpointError("cross-world-size
#   resume not supported"). Weights-only is what every prior stage in this
#   project used across a topology change, including the 64k run itself.
#   TWO CONSEQUENCES, both deliberate:
#     1. the data loader restarts at token 0, so the first 42% of b5 is seen a
#        second time and the remaining 58% for the first time -- a full pass;
#     2. Adam m/v reset, which CP_REWARMUP (100 steps to 1.8e-4) absorbs.
#
# WHY THE FULL EPOCH IS ONLY NOW POSSIBLE: yesterday's box had 5 b5 train
#   shards (0.481 B tokens). All 9 are here now = 0.866 B = a real epoch.
#
# TAIL=100, NOT 2: yesterday's 2-step tail left the final LR at 1.2e-4, not the
#   ~6e-5 floor -- the checkpoint was labelled decayed but effectively was not.
#
# UPLOAD IS VERIFIED BY BYTE SIZE. Yesterday's up() helper was
#   `[ -e "$1" ] || return 0`, which silently SUCCEEDS when the source is
#   missing. That is how the step-695 checkpoint was lost: the CLI never
#   errored, the log said nothing, and 17 GB never left the box.
# =============================================================================
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
export PATH=/usr/local/cuda/bin:$PATH

PY=/root/gv/envv/bin/python3          # this box's venv (NOT /root/envv)
HF=/root/gv/envv/bin/hf
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=cont64k_logs; mkdir -p "$LOG"
M="$LOG/master.log"
REPO=unparallelled/BluScriptCP

T=65536; NP=8; GPUS=0,1,2,3,4,5,6,7
ORIG=32768; SCALE=2                   # 2 * 32768 == 65536, the parent's own cache
DATA=Data_Loader/Data/b5_65536_131072
PARENT="$CK/blumodelcp_run3_step_695.ckpt"
HAY="$PWD/Data_Loader/Data/b2_8192_16384/val_shard_00000.bin"
WARM=100; TAIL=${TAIL:-100}; FREQ=${FREQ:-100}; KEEP=2
GB=524288

say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

# ---- preflight -------------------------------------------------------------
$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing (eval would silently score nothing)"
[ -x "$EXE" ]    || die "$EXE not built"
# CP_ATTN_FUSION=split needs BOTH kernel families compiled in: CP_FUSED_ROPE
# (Stage 2 QK-norm+RoPE) and CP_ATTN_SPLIT (Stage 3 causal attention). Without
# them SplitAttnSDPA_cp.h compiles to stubs that THROW at the first forward --
# a default `make` produced exactly that and the run died 73 s in. The Makefile
# records the flag state in this stamp; refuse to start if it is not right.
STAMP=build/objects/.cp_rope_flags
[ -f "$STAMP" ] || die "no build-flag stamp: rebuild with make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp"
grep -q 'CP_FUSED_ROPE=1' "$STAMP" && grep -q 'CP_ATTN_SPLIT=1' "$STAMP" || \
  die "binary built WITHOUT the split kernels ($(cat "$STAMP")). Rebuild: make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j\$(nproc)"
[ "$EXE" -nt "$STAMP" ] || say "WARN: $EXE is older than the flag stamp -- rebuild if unsure"
[ -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"
PSZ=$(stat -c%s "$PARENT")
[ "$PSZ" = "5225752094" ] || die "parent wrong size: $PSZ (want 5225752094)"
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 ))
FREEG=$(df -BG --output=avail / | tail -1 | tr -dc 0-9)
[ "$FREEG" -ge 100 ] || die "only ${FREEG} GB free; 8-rank checkpoints need ~90 GB"

say "=== CONTINUE 64k b5 to a FULL EPOCH (8x H100, cp=$NP) ==="
say "parent: ext64k run3 step 695 (64k model), weights-only via CP_INIT_FROM"
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 on disk"
say "disk: ${FREEG} GB free"

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: probe -- does 64k fit in 80 GB at cp=8? ----------------------
if [ ! -f "$LOG/probe.ok" ]; then
  say "PHASE 1: 3-step probe (8,192 tok/rank; extrapolates to ~67 GB of 81.5 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=$T CP_MAX_STEPS=3 \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$LOG/probe.log" 2>&1
  RC=$?
  # The probe must PROVE it trained, not merely fail to say "OOM". The first
  # version of this gate only grepped for OOM, so a probe that aborted on every
  # rank with `sdpa_split_rope_once_forward: build with -DCP_FUSED_ROPE` was
  # declared "PROBE OK" and the real run then died 73 s in. Three independent
  # checks now, and the peak-memory parse is one of them: an empty peak means
  # no step ever completed.
  if grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log"; then
    say "PROBE OOM -- 64k does NOT fit at cp=8 on 80 GB cards."
    say "  There is no further lever on 8 GPUs: CP_B is 1 and there is no recompute flag."
    grep -aoE 'cudaMallocAsync\(size=[0-9]+\) failed' "$LOG/probe.log" | head -2 | sed 's/^/    /' >&2
    exit 1
  fi
  if grep -qaE 'terminate called|std::runtime_error|Signal: Aborted' "$LOG/probe.log"; then
    say "PROBE CRASHED (not OOM). First error:"
    grep -aA1 'terminate called' "$LOG/probe.log" | grep -a 'what()' | head -1 | sed 's/^/    /' >&2
    exit 1
  fi
  NSTEP=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe.log")
  PEAK=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)
  [ "$RC" -eq 0 ]     || { say "PROBE exited $RC"; tail -15 "$LOG/probe.log" >&2; exit 1; }
  [ "$NSTEP" -ge 1 ]  || { say "PROBE produced NO step lines"; tail -15 "$LOG/probe.log" >&2; exit 1; }
  [ -n "$PEAK" ]      || { say "PROBE reported no peak memory"; tail -15 "$LOG/probe.log" >&2; exit 1; }
  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | sed 's/^/  /' >&2
  say "PROBE OK: $NSTEP steps, peak $PEAK of 81559 MiB/card"
  touch "$LOG/probe.ok"
fi

# ---- step count: a full epoch over what is on disk -------------------------
# The probe ran grad_accum=1. Its dt is NOT the per-step cost -- multiplying it
# by grad_accum overstated the real rate 5.8x on the 64k run (74.0 vs 12.85
# s/step) because fixed per-step costs got multiplied too. Reported for
# information only; the step count comes from the DATA, not the clock.
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)]
if not d: print("0.0"); raise SystemExit
probe_accum=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 "probe rate (UPPER BOUND, inflated by fixed costs x grad_accum): ${SPS} s/step"

STEPS=${STEPS_OVERRIDE:-$(( have / GB ))}
STABLE=$(( STEPS - WARM - TAIL ))
[ "$STABLE" -lt 1 ] && STABLE=1
# bluscriptCP.cpp:1006 dies if warmup + stable >= max_steps - 1
[ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "illegal WSD: $WARM+$STABLE >= $STEPS-1"
say "PHASE 2: $STEPS steps x $GB = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens = a full b5 epoch"
say "  WSD: warmup $WARM, stable $STABLE, decay tail $TAIL -> final checkpoint IS decayed"

if [ -f "$LOG/cont.ckpt" ] && [ -f "$(cat "$LOG/cont.ckpt")" ]; then
  say "already trained -> $(cat "$LOG/cont.ckpt")"
else
  # Reclaim leaked .zero_g* sidecars while training runs: CP_CKPT_KEEP prunes
  # .ckpt but NOT the sidecars (that leak filled a 193 GB disk on 2026-09-09).
  ( while sleep 300; do
      for s in checkpoints_bluscriptcp/rank_*/*.zero_g*; do
        [ -e "$s" ] || continue
        base=${s%.zero_g*}
        [ -f "$base.ckpt" ] || rm -f "$s"
      done
    done ) & JANITOR=$!
  trap 'kill $JANITOR 2>/dev/null' EXIT

  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/cont.log" 2>&1
  kill $JANITOR 2>/dev/null; trap - EXIT

  grep -qaiE 'out of memory|bad_alloc' "$LOG/cont.log" && die "OOM in the full run"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/cont.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/cont.ckpt"; echo "$RN" > "$LOG/runnum"
  say "cont64k -> run $RN, $(basename "$CKF"), $(grep -aE '^step ' "$LOG/cont.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
  say "  peak mem $(grep -aoE 'mem: [0-9]+MB' "$LOG/cont.log" | tail -1)"
  say "  final lr $(grep -aE '^step ' "$LOG/cont.log" | tail -1 | grep -oE 'lr [0-9.e-]+')  (want ~6e-05 = decayed)"
fi

# ---- PHASE 3: eval, 4 lengths in parallel ----------------------------------
EXT=$(basename "$(cat "$LOG/cont.ckpt")")
say "PHASE 3: eval cont64k -- 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=cont64k_len_$LEN TASKS="niah_single" \
    CKPTS="cont64k" 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="cont64k=$EXT" ./run_longeval.sh \
    >> "$LOG/eval_$LEN.log" 2>&1 ) &
done
wait
for LEN in 8192 16360 32744 65512; do
  [ -s "cont64k_len_$LEN/summary.csv" ] && say "  len $LEN -> $(sed -n 2p cont64k_len_$LEN/summary.csv)" \
                                        || say "  len $LEN -> NO RESULT"
done

# ---- PHASE 4: upload, VERIFIED ---------------------------------------------
say "PHASE 4: upload"
{
  echo "# 64k b5 continuation to a full epoch -- $(date -u +%Y-%m-%d)"
  echo
  echo "Parent: ext64k run3 step 695 (67.0/65.0/52.0/29.0 at 8k/16k/32k/64k),"
  echo "        loaded WEIGHTS-ONLY: ZeRO sidecars cannot cross world_size 4 -> 8."
  echo "Data:   b5_65536_131072, $ntr shards = $(echo "scale=3; $have/1000000000" | bc) B tokens."
  echo "Train:  $STEPS steps x $GB = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens = a full b5 epoch."
  echo "        The loader restarts at token 0, so b5's first 42% is seen a"
  echo "        second time and the rest for the first time."
  echo "        WSD warmup $WARM / stable $STABLE / decay tail $TAIL."
  echo "        8x H100, cp=$NP, 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 "cont64k_len_$LEN/summary.csv" 2>/dev/null | cut -d, -f2)
    printf "| %s | %s |\n" "$LEN" "${v:-n/a}"
  done
  echo
  echo "Parent for comparison: 67.0 / 65.0 / 52.0 / 29.0."
} > "$LOG/RESULTS.md"

# Upload, then CONFIRM the remote byte size. A missing source is a LOUD failure:
# the silent `[ -e "$1" ] || return 0` in ext64k_b5.sh is what lost the step-695
# checkpoint -- the CLI never errored and nothing was logged.
UPOK=1
up(){ # src dst
  local src=$1 dst=$2
  if [ ! -e "$src" ]; then say "  MISSING SOURCE (NOT UPLOADED): $src"; UPOK=0; return 1; fi
  if $HF upload "$REPO" "$src" "$dst" >>"$LOG/upload.log" 2>&1; then say "  sent $dst"
  else say "  UPLOAD FAILED $dst"; UPOK=0; return 1; fi
}
verify(){ # localfile remotepath
  $PY - "$REPO" "$1" "$2" <<'PYEOF'
import sys,os
from huggingface_hub import HfApi
repo,loc,rem=sys.argv[1],sys.argv[2],sys.argv[3]
want=os.path.getsize(loc)
got={f.rfilename:(f.size or 0) for f in HfApi().repo_info(repo,files_metadata=True).siblings}.get(rem)
print(f"  {'VERIFIED' if got==want else 'MISMATCH/ABSENT'} {rem} local={want:,} remote={got}")
sys.exit(0 if got==want else 1)
PYEOF
}

up "$LOG/RESULTS.md" "docs/CONT64K_B5_RESULTS_$(date -u +%Y-%m-%d).md"
for LEN in 8192 16360 32744 65512; do up "cont64k_len_$LEN" "logs/cont64k_len_$LEN/"; done
up "$LOG" logs/cont64k/
up "$(basename "$0")" "scripts/$(basename "$0")"

RN=$(cat "$LOG/runnum" 2>/dev/null)
STEP=$(basename "$(cat "$LOG/cont.ckpt")" | grep -oE 'step_[0-9]+' | cut -d_ -f2)
say "uploading the model itself (the part that was lost yesterday)"
SRC="$CK/blumodelcp_run${RN}_step_${STEP}.ckpt"
DST="checkpoints/cont64k_b5_rank0_step_${STEP}.ckpt"
if up "$SRC" "$DST"; then verify "$SRC" "$DST" | tee -a "$M" || UPOK=0; fi
for r in $(seq 0 $((NP-1))); do
  S="checkpoints_bluscriptcp/rank_$r/blumodelcp_run${RN}_step_${STEP}.zero_g1"
  D="checkpoints/cont64k_b5_rank${r}_step_${STEP}.zero_g1"
  if up "$S" "$D"; then verify "$S" "$D" | tee -a "$M" || UPOK=0; fi
done

[ "$UPOK" = 1 ] && say "=== CONT64K DONE -- all uploads verified, safe to tear down ===" \
                || say "=== CONT64K DONE but UPLOADS INCOMPLETE -- DO NOT TEAR DOWN ==="