#!/usr/bin/env bash
# =============================================================================
# ext_32k.sh -- YaRN-extend the best 16k model to 32k, then evaluate.
#
# WHY THIS RUN EXISTS
#   The length ladder ended in a tie: l3 (b3, 16-32k docs TRUNCATED to the 16k
#   window) matched l2 (b2, 8-16k docs that FIT) to the decimal at 16,360 --
#   35.0 vs 35.0, and -6.0/+2.0/+0.0 pp across 4k/8k/16k, none resolvable.
#   Conclusion: document length helps only while documents fit inside the window.
#   So widen the window first, THEN feed b3, whose documents now fit.
#
# PARENT: l2, the best 16k arm, not run7'. l2 is where the 16k ladder topped out.
#
# THE YaRN INVARIANT that cost 27pp when it was wrong:
#   YARN_SCALE * YARN_ORIG_MAXPOS must equal CP_T EXACTLY, and ORIG_MAXPOS must
#   be the length the weights were ACTUALLY trained at (16384 here, not the
#   1024 default at YARNOps.cpp:57).
#     2 * 16384 = 32768 == CP_T   OK
#
# MEMORY -- 2 GPUS IS TIGHT. cp cannot exceed the GPU count, so cp=2 puts
# T/2 = 16,384 tokens on each rank:
#     projected 130.8 GB/GPU = 91.4% of an H200's 143 GB
#   That is close enough that fragmentation may still OOM, so STEP 1 is a 3-step
#   probe and the script aborts rather than wasting hours. This does NOT fit an
#   80 GB H100 (measured: it OOMs).
#
# TOKEN BUDGET: 750 steps x 1,048,576 = 786 M tokens, matching run7's original
# 4k->16k extension exactly, so the two extension stages are comparable.
#
# GATE ON SHORT CONTEXT, NOT PERPLEXITY. The 2026-09-02 mismatch condition
# IMPROVED 16k (0 -> 14-17%) while COLLAPSING 4k from 74% to 16%. Perplexity at
# the new length would have looked fine. The eval therefore spans 1k..32k and a
# regression at 1k/4k is a failure even if 32k improves.
# =============================================================================
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=ext32k_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=${NP:-2}; GPUS=${GPUS:-0,1}
ORIG=16384; SCALE=2
STEPS=${STEPS:-750}
PARENT="$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt"
DATA=lendata/b3_16384_32768
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"

[ $(( SCALE * ORIG )) -eq $T ] || die "YaRN invariant broken: $SCALE * $ORIG != $T"
$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing: $PY -m pip install tiktoken"
[ -x "$EXE" ]    || die "$EXE not built (make libtensor SM_ARCH=90 && make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp)"
[ -f "$PARENT" ] || die "parent l2 missing: $PARENT"
[ -d "$DATA" ]   || die "b3 data missing: $DATA"

say "=== 16k -> 32k YaRN extension ==="
say "CP_T=$T YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG (product $(( SCALE*ORIG )) == $T)"
say "parent $(basename $PARENT), data $DATA, cp=$NP, $STEPS steps x 1,048,576 = 786 M tokens"

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_MAX_LR=1e-3 CP_MIN_LR=1e-4"

# ---- STEP 1: 3-step memory probe -------------------------------------------
if [ ! -f "$LOG/probe.ok" ]; then
  say "probe: 3 steps at T=$T to confirm it fits before committing hours"
  env $ARCH $COMMON CP_REWARMUP=1 CP_REWARMUP_PEAK=0.3 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 at T=$T with cp=$NP. Projected 130.8 GB/GPU."
    say "  This needs more ranks: cp=4 -> 67.5 GB/GPU. Not runnable on $NP GPUs."
    exit 1
  fi
  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | tail -1 | sed 's/^/  probe: /' | tee -a "$M"
  say "probe peak mem: $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
  touch "$LOG/probe.ok"
else
  say "probe already passed"
fi

# ---- STEP 2: the extension --------------------------------------------------
if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
  say "extension already done -> $(cat "$LOG/ext.ckpt")"
else
  say "training $STEPS steps"
  env $ARCH $COMMON CP_REWARMUP=50 CP_REWARMUP_PEAK=0.3 \
      CP_CKPT=1 CP_CKPT_KEEP=2 CP_CKPT_FREQ=50 CP_CKPT_NEW_RUN=1 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=1048576 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 during training"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/ext.log" | tail -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 "no checkpoint written"
  bn=$(basename "$CKF"); n=$(ls checkpoints_bluscriptcp/rank_*/"$bn" 2>/dev/null | wc -l)
  [ "$n" -ge "$NP" ] || die "only $n/$NP rank shards for $bn"
  echo "$CKF" > "$LOG/ext.ckpt"
  say "ext32k -> run $RN, $bn, final loss $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
fi

# ---- STEP 3: eval 1k..32k, matched cache -----------------------------------
EXT=$(basename "$(cat "$LOG/ext.ckpt")")
EXTRA="l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt ext32k=$EXT"
say "eval: l2 (16k parent) vs ext32k, matched YARN_ORIG=$ORIG SCALE=$SCALE"
# PROMPT_LEN + MAX_NEW must pad up to <= CP_T (PAD=256): 32744+24 = 32768 exactly
for LEN in 1024 4096 16360 32744; do
  say "PROMPT_LEN=$LEN"
  GEN_DATA_ROOT=$DATA HAY="$PWD/lendata/haystack/b2val_haystack.bin" \
  PY=$PY OUT=ext32k_len_$LEN TASKS="niah_single" CKPTS="l2 ext32k" N=100 \
  PROMPT_LEN=$LEN MAX_NEW=24 T=$T DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN=$SCALE YARN_ORIG=$ORIG GPUS=$GPUS JOBS_PER_GPU=1 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN nonzero at $LEN"
  say "  -> $(sed -n 1p ext32k_len_$LEN/summary.csv 2>/dev/null) / $(sed -n 2p ext32k_len_$LEN/summary.csv 2>/dev/null)"
done
say "=== 32k EXTENSION DONE ==="
say "REGRESSION CHECK: compare 1024 and 4096 against l2. A 32k gain bought with a"
say "  short-context collapse is the 2026-09-02 mismatch failure repeating."