#!/usr/bin/env bash
# =============================================================================
# lenbucket_experiment.sh -- does TRAINING DOCUMENT LENGTH change 16k retrieval?
#
# THE QUESTION
#   The 2026-09-03 result left retrieval at 27% at the full 16,384 window. One
#   standing hypothesis for the ceiling is the training corpus: flux documents
#   have median 826 tokens, p90 2,721, and only 0.27% reach 16k -- so almost no
#   training example ever required attending across the window. This tests it by
#   holding everything else fixed and changing only document length.
#
# THREE ARMS, one variable
#   armA_c  flux plain text     median 826, p90 2,721 tok   ALREADY TRAINED (run 28)
#   L1      b1_4096_8192        docs in [4096, 8192)        trained here
#   L2      b2_8192_16384       docs in [8192, 16384)       trained here
#   All three: 400 steps, global batch 524,288, from run7' (run 6 step 750),
#   YARN_ORIG_MAXPOS=4096 matched, Muon+ZeRO-1, bf16. Only the data differs.
#
# WHY armA_c IS RE-EVALUATED HERE rather than reusing its published numbers:
#   those were measured with the flux val shard as NIAH filler, and flux is not
#   on this box (nor in the HF model repo). A different haystack is a different
#   test, so all three arms are re-scored against ONE haystack instead.
#
# HAYSTACK CHOICE: b3_16384_32768 val shard -- documents of 16-32k tokens, held
#   out from every arm above. b1's or b2's own val shard would be in-distribution
#   for exactly one arm and would flatter it.
#
# MEMORY: the H200 runs used CP_B=4 and 108 GB/GPU. These are 80 GB H100s, so
#   CP_B=2 with grad_accum 16 holds the global batch at 524,288 while roughly
#   halving activation memory. If step 0 OOMs, the run retries at CP_B=1/accum 32
#   -- same global batch, same optimizer trajectory, just slower.
# =============================================================================
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=lenbucket_logs; mkdir -p "$LOG"
M="$LOG/master.log"

# stderr, never stdout: train_one()'s stdout is its return value.
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

T=16384; NP=4; GPUS=0,1,2,3; ORIG=4096
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=4 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"

PARENT="$CK/blumodelcp_run6_step_750.ckpt"      # run7'
HAY_ABS="$PWD/lendata/haystack/b3val_haystack.bin"
LENS="1024 2048 4096 8192 16360"
N=${N:-100}
TASKS="niah_single val_backward"

[ -x "$EXE" ]     || die "$EXE not built"
[ -f "$PARENT" ]  || die "parent run7' missing: $PARENT"
[ -f "$HAY_ABS" ] || die "haystack missing: $HAY_ABS"

say "=== length-bucket experiment ==="
say "free disk $(df -BG --output=avail /root | tail -1 | tr -dc '0-9')G"

# ---------------------------------------------------------------------------
train_one(){            # $1 = arm label (l1/l2)   $2 = data dir
  local A="$1"
  local DD="$2"
  local lf="$LOG/${A}.log"
  if [ -f "$LOG/${A}.ckpt" ] && [ -f "$(cat "$LOG/${A}.ckpt")" ]; then
    say "$A already trained -> $(cat "$LOG/${A}.ckpt")"; return 0
  fi
  [ -d "$DD" ] || { say "$A: data dir $DD missing -- SKIP"; return 1; }
  local ntr; ntr=$(ls "$DD"/*train*.bin 2>/dev/null | wc -l)
  local nva; nva=$(ls "$DD"/*val*.bin   2>/dev/null | wc -l)
  say "$A: $ntr train shards, $nva val shards in $DD"
  [ "$ntr" -ge 1 ] && [ "$nva" -ge 1 ] || { say "$A: need >=1 train and >=1 val shard"; return 1; }

  # CP_B x grad_accum must hold global batch at 524,288 to match armA_c exactly.
  local B ACC
  for cfg in "2 16" "1 32"; do
    B=${cfg% *}; ACC=${cfg#* }
    say "$A: attempting CP_B=$B (grad_accum=$ACC, global 524288)"
    env $ARCH $COMMON CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DD" \
        CP_B=$B CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=400 \
        CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=200 \
        CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
        > "$lf" 2>&1
    if grep -qaiE 'out of memory|CUDA error: out of memory|bad_alloc' "$lf"; then
      say "$A: OOM at CP_B=$B -- retrying smaller"; continue
    fi
    break
  done

  local RN; RN=$(grep -aoE 'CP_Training_log[0-9]+' "$lf" | head -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || { say "$A: no run number in $lf -- FAILED"; return 1; }
  local CKF; CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || { say "$A: run $RN wrote NO checkpoint -- FAILED"; return 1; }
  # all four ZeRO-1 shards or the load is silently partial
  local bn n; bn=$(basename "$CKF")
  n=$(ls checkpoints_bluscriptcp/rank_*/"$bn" 2>/dev/null | wc -l)
  [ "$n" -ge 4 ] || { say "$A: only $n/4 rank shards for $bn -- FAILED"; return 1; }
  echo "$CKF" > "$LOG/${A}.ckpt"
  say "$A -> run $RN, $bn, B=$B, final loss $(grep -aE '^step ' "$lf" | tail -1 | grep -oE 'loss: [0-9.]+')"
  say "$A  peak mem $(grep -aoE 'mem: [0-9]+MB' "$lf" | tail -1)"
  return 0
}

train_one l1 lendata/b1_4096_8192 || say "L1 incomplete -- continuing"
train_one l2 lendata/b2_8192_16384 || say "L2 incomplete -- continuing"

# ---------------------------------------------------------------------------
# evaluate all three arms against ONE haystack
# ---------------------------------------------------------------------------
EXTRA=""
[ -f "$CK/blumodelcp_run28_step_400.ckpt" ] && EXTRA="armA_c=blumodelcp_run28_step_400.ckpt"
for A in l1 l2; do
  f="$LOG/${A}.ckpt"; [ -f "$f" ] || continue
  p=$(cat "$f"); [ -f "$p" ] || continue
  EXTRA="$EXTRA ${A}=$(basename "$p")"
done
[ -n "$EXTRA" ] || die "no checkpoints to evaluate"
MODELS=$(echo $EXTRA | tr ' ' '\n' | cut -d= -f1 | tr '\n' ' ')
say "eval models: $MODELS"
say "haystack: $HAY_ABS ($(stat -Lc%s "$HAY_ABS") B)"

for LEN in $LENS; do
  say "eval PROMPT_LEN=$LEN N=$N"
  # GEN_DATA_ROOT: generation never reads training data, but bluscriptCP builds
  # the DataLoader before reaching the CP_GENERATE branch, and this box has no
  # Data_Loader/Data (no flux). Any dir with a *train*.bin and a *val*.bin works.
  GEN_DATA_ROOT=lendata/b1_4096_8192 \
  HAY="$HAY_ABS" PY=$PY OUT=lenb_len_$LEN TASKS="$TASKS" CKPTS="$MODELS" N=$N \
  PROMPT_LEN=$LEN MAX_NEW=24 T=$T DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN_ORIG=$ORIG GPUS=0,1,2,3 JOBS_PER_GPU=2 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN: eval at $LEN returned nonzero"
  [ -f "lenb_len_$LEN/summary.csv" ] && say "  -> $(sed -n 2p lenb_len_$LEN/summary.csv)"
done
say "eval done"

# ---------------------------------------------------------------------------
{
  echo "# Training document length vs 16k retrieval -- $(date -u +%Y-%m-%d)"
  echo
  echo "All arms: 400 steps, global batch 524,288, from run7' (run 6 step 750),"
  echo "YARN_ORIG_MAXPOS=4096 matched, Muon+ZeRO-1 bf16, 4xH100."
  echo "Only the training data differs. NIAH scored generatively (substring match),"
  echo "greedy, N=$N per cell, one shared haystack held out from every arm."
  echo
  echo "| arm | training docs |"
  echo "|---|---|"
  echo "| armA_c | flux plain text, median 826 / p90 2,721 tok |"
  echo "| l1 | b1_4096_8192, docs in [4096, 8192) |"
  echo "| l2 | b2_8192_16384, docs in [8192, 16384) |"
  echo
  for LEN in $LENS; do
    echo "## PROMPT_LEN=$LEN"; echo '```'
    cat "lenb_len_$LEN/summary.csv" 2>/dev/null || echo "(missing)"
    echo '```'; echo
  done
  echo "## Training losses"; echo '```'
  for A in l1 l2; do
    [ -f "$LOG/${A}.log" ] && \
      echo "$A: $(grep -aE '^step ' "$LOG/${A}.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
  done
  echo '```'
  echo
  echo "NOTE: losses are NOT comparable across arms -- each is measured on its own"
  echo "data distribution. Longer documents are a different prediction problem, so"
  echo "a higher loss on b2 than on flux says nothing about model quality. Only the"
  echo "retrieval numbers compare."
} > "$LOG/RESULTS.md"
say "wrote $LOG/RESULTS.md"
say "=== LENBUCKET DONE ==="