Upload scripts/lenbucket_experiment.sh with huggingface_hub

scripts/lenbucket_experiment.sh

@@ -0,0 +1,181 @@
+#!/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 ==="