#!/usr/bin/env bash
# =============================================================================
# b4_continue.sh -- CONTINUE ext32k through the rest of b4 at a 32k window.
#
#   ext32k already trained 750 steps = 0.393 B tokens = 34% of b4's train split.
#   This runs steps 750 -> 2227, the remaining 1,477 steps = 0.775 B tokens,
#   finishing exactly one pass over b4 (1.168 B of 1.171 B available, no wrap).
#
# WHY CP_CKPT_RESUME AND NOT CP_INIT_FROM:
#   CP_INIT_FROM loads WEIGHTS ONLY and starts at step 0, so the DataLoader would
#   restart at token 0 and replay the 0.393 B ext32k already saw. CP_CKPT_RESUME
#   restores start_step=750, and bluscriptCP.cpp:2084 then calls
#   skip_batches(750 * 16 = 12,000) to fast-forward the stream to 0.393 B. That
#   is what makes this a true continuation of the epoch rather than a replay.
#
# KNOWN DEGRADATION -- the ZeRO optimizer sidecar is missing.
#   Yesterday's teardown uploaded the two .ckpt files but NOT the .zero_g1
#   sidecar, and that box is gone. Per bluscriptCP.cpp:2032, "without the sidecar
#   a bf16 run resumes with group 1's Adam m/v at zero -- it trains on, so the
#   loss curve only looks slightly wrong." Weights and group-0 optimizer state
#   restore normally; group 1's momentum rebuilds over tens of steps. Expect a
#   small loss transient right after step 750 that is an artifact, not a signal.
#
# LEARNING RATE -- CP_MAX_LR / CP_MIN_LR DO NOT EXIST.
#   bluscriptCP.cpp hardcodes max_lr=6e-4 (line 131) and min_lr=6e-5 (line 132),
#   and the env-parsing block never reads CP_MAX_LR or CP_MIN_LR -- there is not
#   even a float env helper. Every script in this project passed them as a no-op.
#   So the real schedule is 6e-4 * CP_REWARMUP_PEAK(0.3) = 1.8e-4 peak, 6e-5 floor.
#   ext32k's first 750 steps ran at exactly that (verified from its log: 1.8000e-04
#   at steps 99..649, 6.0030e-05 at 749), so this continuation MUST keep the same
#   values or it would not be a continuation. Do not "fix" the LR here.
#
# SCHEDULE: re-warm 100 steps from 750 back to 1.8e-4, flat to 2027, then a
#   200-step decay tail. decay_start = 750+100+1177 = 2027. The guard at
#   bluscriptCP.cpp:1006 needs CP_WARMUP+CP_STABLE < CP_MAX_STEPS-1:
#   100+1177 = 1277 < 2226. This is a SECOND schedule on top of ext32k's
#   completed one, so "ext32k@750 vs this" is not a steps-only contrast.
# =============================================================================
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/Adhi/BluTrain/envv/bin/python
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=b4cont_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=32768; NP=2; GPUS=0,1; ORIG=16384; SCALE=2
RUN=2; START=750; STEPS=2227; WARM=100; STABLE=1177; FREQ=50; KEEP=2
EVALS="1100 1650 2227"          # 750 is already evaluated (the ext32k result)
DATA=lendata/b4_32768_65536
HAY="$PWD/lendata/haystack/b2val_haystack.bin"
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
BASE="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 CP_REWARMUP_ANCHOR=$START CP_STABLE=$STABLE \
CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"

$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
[ -x "$EXE" ] || die "not built"
for r in 0 1; do
  [ -f "checkpoints_bluscriptcp/rank_$r/blumodelcp_run${RUN}_step_${START}.ckpt" ] || \
    die "rank_$r ext32k step-$START checkpoint missing"
done
[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant"
say "=== CONTINUE ext32k through the rest of b4 ==="
say "resume run $RUN at step $START -> $STEPS  ($((STEPS-START)) new steps, $(echo "scale=3; ($STEPS-$START)*524288/1000000000" | bc) B tokens)"
say "loader fast-forwards $((START*16)) batches = $(echo "scale=3; $START*524288/1000000000" | bc) B; reads to $(echo "scale=3; $STEPS*524288/1000000000" | bc) B of 1.171 B -- no wrap"
say "WARNING: .zero_g1 sidecar absent -- group 1 Adam moments start at zero (small transient)"
say "eval points: $EVALS"

run_seg(){
  local TGT="$1"
  [ -f "$LOG/seg_${TGT}.done" ] && { say "segment to $TGT done"; return 0; }
  say "training to step $TGT"
  ( for i in $(seq 1 6000); do
      n=$(ls checkpoints_bluscriptcp/rank_*/blumodelcp_run${RUN}_step_${TGT}.ckpt 2>/dev/null | wc -l)
      if [ "$n" -ge "$NP" ]; then
        sleep 5
        for p in $(pgrep -f 'mpirun.*bluscriptCP_exec') $(pgrep -f 'build/bluscriptCP_exec'); do kill "$p" 2>/dev/null; done
        exit 0
      fi
      sleep 20
    done ) &
  local WPID=$!
  env $ARCH $BASE CP_CKPT_RESUME=$RUN CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=$STEPS \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      >> "$LOG/train.log" 2>&1
  kill $WPID 2>/dev/null; wait $WPID 2>/dev/null
  grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM"
  local SRC="blumodelcp_run${RUN}_step_${TGT}.ckpt"
  local n; n=$(ls checkpoints_bluscriptcp/rank_*/"$SRC" 2>/dev/null | wc -l)
  [ "$n" -ge "$NP" ] || die "step $TGT checkpoint missing ($n/$NP)"
  local SAFE="b4c${TGT}_${SRC}"
  for r in 0 1; do
    [ -f "checkpoints_bluscriptcp/rank_$r/$SAFE" ] || \
      cp "checkpoints_bluscriptcp/rank_$r/$SRC" "checkpoints_bluscriptcp/rank_$r/$SAFE"
  done
  echo "$SAFE" > "$LOG/seg_${TGT}.ckpt"; touch "$LOG/seg_${TGT}.done"
  say "step $TGT reached, preserved $SAFE, loss $(grep -aE '^step ' "$LOG/train.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
}

eval_seg(){
  local TGT="$1"
  [ -f "$LOG/eval_${TGT}.done" ] && return 0
  local SAFE; SAFE=$(cat "$LOG/seg_${TGT}.ckpt")
  say "EVAL @ step $TGT"
  for LEN in 8192 16360 32744; do
    GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=b4c_${TGT}_len_$LEN TASKS="niah_single" \
    CKPTS="b4c$TGT" 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=0,1 JOBS_PER_GPU=1 BF16=1 GREEDY=1 \
    CKPT_EXTRA="b4c$TGT=$SAFE" ./run_longeval.sh >> "$LOG/eval_${TGT}_${LEN}.log" 2>&1 \
      || say "  WARN eval @$TGT len $LEN"
    [ -s "b4c_${TGT}_len_$LEN/summary.csv" ] && \
      say "  step $TGT len $LEN -> $(sed -n 2p b4c_${TGT}_len_$LEN/summary.csv)"
  done
  touch "$LOG/eval_${TGT}.done"
}

for TGT in $EVALS; do run_seg "$TGT" || die "segment $TGT failed"; eval_seg "$TGT"; done
say "=== B4 CONTINUATION DONE ==="
say "progression vs ext32k@750 (niah_single):"
for LEN in 8192 16360 32744; do
  line="  len $LEN:  750=(see ext32k)"
  for TGT in $EVALS; do
    v=$(sed -n 2p "b4c_${TGT}_len_$LEN/summary.csv" 2>/dev/null | cut -d, -f2)
    line="$line  ${TGT}=${v:-?}"
  done
  say "$line"
done