Upload source/files/needle_sweep.sh with huggingface_hub

source/files/needle_sweep.sh

@@ -0,0 +1,254 @@
+#!/usr/bin/env bash
+# =============================================================================
+# needle_sweep.sh -- run the retrieval probe across an extension-factor row set
+# for ONE model, then score it. Forward-only, cache-only (no fine-tuning).
+#
+# Reuses the existing CP_EVAL_PPL + CP_EVAL_TOKENS_BIN path unchanged: the .bin
+# holds synthetic trials, and per-position mean_nll in the CSV gives retrieval
+# NLL at the answer span. See needle_gen.py for the fixed-slot construction.
+#
+# GROUP-SLOT FAN-OUT: slots = NGPU / CP_SIZE, so EVERY row parallelises --
+#   T<=128k -> CP=1 -> 8 concurrent   |  256k -> CP=2 -> 4   |  512k -> CP=4 -> 2
+# Each group is pinned to a contiguous device subset so cudaSetDevice(rank) binds
+# rank->GPU correctly. Concurrency changes only WHEN work happens, never what is
+# computed (independent processes, separate CSVs), so no determinism argument is
+# needed -- unlike the search's EvaluatorPool.
+#
+# RESUMABLE: one marker per (T, cache, variant, block) in $STATE_DIR.
+#
+# Launch (114M 4k base, NTK, all five factors):
+#   MODEL_KEY=114m_run1 CKPT_RUN=1 \
+#   MODEL_ARCH="CP_N_EMBD=768 CP_N_LAYER=12 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=2048 CP_WEIGHT_TYING=1" \
+#   FACTORS="8 16 32" CACHES="ntk" \
+#     nohup bash Tests/bluscriptcp/needle_sweep.sh > needle_114m.log 2>&1 &
+#
+# The single searched cell (48M@16x) is a separate invocation:
+#   MODEL_KEY=48m_run200 CKPT_RUN=200 MODEL_ARCH="..." FACTORS="16" \
+#   CACHES="searched" SEARCHED_LAMBDA=longrope_best_48m_s16_repr.txt bash ...
+# =============================================================================
+set -euo pipefail
+cd "$(dirname "$0")/../.."   # -> repo root
+
+MODEL_KEY=${MODEL_KEY:?set MODEL_KEY (label for the results rows, e.g. 114m_run1)}
+# INIT_FROM=<path/to/base.ckpt> loads WEIGHTS ONLY instead of resuming a run.
+# Needed whenever the checkpoint's optimizer state cannot be read back in the
+# EVAL configuration: a ZeRO/Muon-trained base evaluated forward-only builds a
+# plain nn::AdamW, which would try to parse ShardedMuon's state and fail. Also
+# the only way to evaluate a base trained at a different world size, since ZeRO
+# state is sharded to the size that wrote it. A forward-only probe needs no
+# optimizer state at all, so this is the correct loader for it.
+INIT_FROM=${INIT_FROM:-}
+if [ -n "$INIT_FROM" ]; then
+  [ -f "$INIT_FROM" ] || { echo "INIT_FROM: no such file: $INIT_FROM" >&2; exit 1; }
+  CKPT_RUN=${CKPT_RUN:-0}
+  CKPT_ENV="CP_INIT_FROM=$INIT_FROM"
+else
+  CKPT_RUN=${CKPT_RUN:?set CKPT_RUN (or INIT_FROM=<path>)}
+  CKPT_ENV="CP_CKPT_RESUME=$CKPT_RUN"
+fi
+MODEL_ARCH=${MODEL_ARCH:?set MODEL_ARCH (must match how CKPT_RUN was trained)}
+
+FACTORS=${FACTORS:-"8 16 32 64 128"}     # extension factors s; T = PHYS_T * s
+CACHES=${CACHES:-"ntk"}                  # ntk | yarn | pi | searched
+# YARN_SCALE: evaluate under the trainer's YaRN cache instead of the LongRoPE lambda
+# ramp. These are DIFFERENT ways to reach the same T and a checkpoint must be scored
+# under the one it was TRAINED with, or the numbers are meaningless -- measured:
+# run7 (extended at YARN_SCALE=4, m=1.1386) scored under the lambda ramp (scale=1.0)
+# had its retrieval FLOOR collapse from 4.17 to 11.34 nats, i.e. it failed even with
+# the needle 12 tokens from the probe. bluscriptCP.cpp:914 states the intended matrix:
+#   base ckpt     + YARN_SCALE=1   (naive extrapolation @ T)
+#   extended ckpt + YARN_SCALE=4   (YaRN @ T)
+# Set YARN_SCALE to take that route; leave it empty for the LongRoPE route.
+YARN_SCALE=${YARN_SCALE:-}
+SEARCHED_LAMBDA=${SEARCHED_LAMBDA:-}     # required when CACHES contains "searched"
+PHYS_T=${PHYS_T:-4096}
+HEAD_DIM=${HEAD_DIM:-64}
+ROPE_THETA=${ROPE_THETA:-500000}
+BLOCKS=${BLOCKS:-5}
+NEEDLE_LEN=${NEEDLE_LEN:-8}
+SEED=${SEED:-0}
+GPUS=${GPUS:-0,1,2,3,4,5,6,7}
+FILLER=${FILLER:-Data_Loader/Data/edufineweb_val_00000.bin}
+DATA_ROOT=${DATA_ROOT:-$PWD/Data_Loader/Data}
+NEEDLE_DIR=${NEEDLE_DIR:-needles}
+# Absolute offset for the probe+answer span (needle_gen --probe-pos). Unset = the
+# end of the trial, the historical layout. Set it to decouple needle-to-probe
+# DISTANCE from the needle's ABSOLUTE POSITION, which the default layout welds
+# into a single variable. Filenames do NOT encode it, so give each probe position
+# its own NEEDLE_DIR/CSV_DIR/STATE_DIR.
+PROBE_POS=${PROBE_POS:-}
+CSV_DIR=${CSV_DIR:-needle_csv}
+STATE_DIR=${STATE_DIR:-needle_state}
+RESULTS=${RESULTS:-needle_results.csv}
+# Per-rank forward-only logits are (T/CP)*vocab*4 bytes. Keep <= this so a 49 GB
+# card holds logits + activations. Pre-flight 2b must confirm it (and that logits
+# are fp32 -- a bf16 model triggers an as_type copy that doubles the peak).
+MAX_TOK_PER_RANK=${MAX_TOK_PER_RANK:-131072}
+
+EXE=./build/bluscriptCP_exec
+export LD_LIBRARY_PATH=BluTrain/Tensor-Implementations/lib:BluTrain/Profiler/lib:${LD_LIBRARY_PATH:-}
+IFS=',' read -r -a GPUARR <<< "$GPUS"
+NGPU=${#GPUARR[@]}
+KV=$(sed -E 's/.*CP_N_KVHEAD=([0-9]+).*/\1/' <<< "$MODEL_ARCH")
+mkdir -p "$NEEDLE_DIR" "$CSV_DIR" "$STATE_DIR"
+log(){ echo "[$(date '+%F %H:%M:%S')] $*"; }
+
+# ---- pre-flight -------------------------------------------------------------
+[ -x "$EXE" ] || { echo "ERROR: $EXE missing -- build first"; exit 1; }
+[ -f "$FILLER" ] || { echo "ERROR: filler stream $FILLER not found"; exit 1; }
+# Detect a live evaluator by RESOLVED EXECUTABLE, not by command-line text. `pgrep -f
+# bluscriptCP_exec` also matches any shell whose argv merely CONTAINS that string --
+# including the wrapper that launches this script -- which produced a false "already
+# running" abort. `pgrep -x` is no good either: comm truncates to 15 chars and
+# "bluscriptCP_exec" is 16. Comparing /proc/PID/exe is exact.
+exec_running(){
+  local want; want=$(readlink -f "$EXE") || return 1
+  local p
+  for p in /proc/[0-9]*; do
+    [ "$(readlink -f "$p/exe" 2>/dev/null)" = "$want" ] && { echo "${p#/proc/}"; return 0; }
+  done
+  return 1
+}
+if pids=$(exec_running); then
+  # ALLOW_CONCURRENT_EXE=1 permits a DELIBERATE overlap -- e.g. probing a mid-training
+  # checkpoint while the pretrain keeps running. Only use it when the memory headroom is
+  # known: a forward-only probe is ~10-15 GB against a training job's ~47 GB on an 80 GB
+  # card. The context_rot_2026-08-19 copies of this script already carried this flag; the
+  # tracked one never got it.
+  if [ "${ALLOW_CONCURRENT_EXE:-0}" = "1" ]; then
+    log "WARNING: $EXE already running (pid $pids); ALLOW_CONCURRENT_EXE=1 -- proceeding"
+  else
+    echo "ERROR: $EXE already running (pid $pids) -- this sweep would oversubscribe GPUs $GPUS"
+    echo "  set ALLOW_CONCURRENT_EXE=1 to override when the overlap is intended"
+    exit 1
+  fi
+fi
+case " $CACHES " in *" searched "*)
+  [ -n "$SEARCHED_LAMBDA" ] || { echo "ERROR: CACHES has 'searched' but SEARCHED_LAMBDA is unset"; exit 1; }
+  [ -f "$SEARCHED_LAMBDA" ] || { echo "ERROR: SEARCHED_LAMBDA '$SEARCHED_LAMBDA' not found"; exit 1; };;
+esac
+
+# CP degree per row, from the plan's CP table (memory-driven, not a free choice).
+cp_for_T(){
+  local T=$1 cp=1
+  while [ $(( T / cp )) -gt "$MAX_TOK_PER_RANK" ]; do cp=$(( cp * 2 )); done
+  echo "$cp"
+}
+# Verified topology only: pure ring>=4 is known-broken in this repo.
+topo_for(){
+  local cp=$1
+  if   [ "$cp" = "1" ]; then echo ""
+  elif [ $(( KV % cp )) -eq 0 ]; then echo "CP_ATTN_MODE=ulysses"
+  elif [ $(( cp % KV )) -eq 0 ]; then echo "CP_ATTN_MODE=hybrid CP_ULYSSES_SIZE=$KV"
+  else echo "UNSUPPORTED"; fi
+}
+
+# lambda file per (cache, head_dim, theta, s). NTK/PI are theta-independent closed
+# forms; YaRN depends on base via find_correction_range, so it MUST be regenerated
+# per theta. All are head_dim-dependent (vector length = head_dim/2).
+lambda_file(){
+  local cache=$1 s=$2
+  if [ "$cache" = "searched" ]; then echo "$SEARCHED_LAMBDA"; return; fi
+  local out="$NEEDLE_DIR/lam_${cache}_hd${HEAD_DIM}_th${ROPE_THETA}_s${s}.txt"
+  if [ ! -f "$out" ]; then
+    python3 - "$cache" "$s" "$out" "$HEAD_DIM" "$ROPE_THETA" "$PHYS_T" <<'PY'
+import sys, importlib.util
+cache, s, out, hd, theta, orig = (sys.argv[1], int(sys.argv[2]), sys.argv[3],
+                                  int(sys.argv[4]), float(sys.argv[5]), float(sys.argv[6]))
+spec = importlib.util.spec_from_file_location('lrs', 'Tests/bluscriptcp/longrope_search.py')
+lrs = importlib.util.module_from_spec(spec); spec.loader.exec_module(lrs)
+half, lam_max = hd // 2, s * 1.25
+if   cache == 'ntk':  lam = lrs.seed_NTK(hd, s, lam_max)
+elif cache == 'pi':   lam = lrs.seed_PI(half, s, lam_max)
+elif cache == 'yarn': lam = lrs.seed_YaRN(hd, s, orig, 32.0, 1.0, theta, lam_max)
+else: sys.exit(f"unknown cache '{cache}'")
+lrs.write_factors(out, lam, 0, s, int(orig * s))
+PY
+  fi
+  echo "$out"
+}
+
+# Overridable so a REFINED distance grid can be added around a crossing. Set
+# VARIANTS to the UNION of old and new points: already-completed ones are skipped
+# by their state markers, so only the new points actually run, and the scorer then
+# reports the whole grid in one table.
+VARIANTS=${VARIANTS:-"p0 p0.25 p0.5 p0.75 p1 absent null"}
+log "needle sweep: model=$MODEL_KEY run=$CKPT_RUN factors=[$FACTORS] caches=[$CACHES] blocks=$BLOCKS"
+log "  GPUS=$GPUS (n=$NGPU) kv=$KV head_dim=$HEAD_DIM theta=$ROPE_THETA"
+
+for S in $FACTORS; do
+  T=$(( PHYS_T * S ))
+  CP=$(cp_for_T "$T")
+  TOPO=$(topo_for "$CP")
+  [ "$TOPO" = "UNSUPPORTED" ] && { echo "ERROR: no verified topology for CP=$CP with kv=$KV at T=$T"; exit 1; }
+  [ "$CP" -le "$NGPU" ] || { echo "ERROR: T=$T needs CP=$CP but only $NGPU GPU(s) given"; exit 1; }
+  N=$([ "$T" -le 131072 ] && echo 8 || echo 4)      # pinned trials-per-block (plan A1)
+  SLOTS=$(( NGPU / CP ))
+  log "== s=$S T=$T CP=$CP ${TOPO:-(single-GPU)} N=$N slots=$SLOTS =="
+
+  # .bin files depend only on (T, variant, block) -- shared across models and caches.
+  for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do
+    [ -f "$NEEDLE_DIR/needle_T${T}_${V}_b${B}.bin" ] && continue
+    python3 Tests/bluscriptcp/needle_gen.py --T "$T" --variant "$V" --block "$B" \
+        --trials "$N" --needle-len "$NEEDLE_LEN" --filler "$FILLER" \
+        --seed "$SEED" --out-dir "$NEEDLE_DIR" \
+        ${PROBE_POS:+--probe-pos "$PROBE_POS"} >/dev/null
+  done; done
+
+  for CACHE in $CACHES; do
+    LAM=$(lambda_file "$CACHE" "$S")
+    # The lambda vector is length-specific: reject a mismatch loudly rather than
+    # scoring a cache tuned for another length (the T == S_search invariant).
+    SS=$(awk '/^S_search/{print $2}' "$LAM")
+    [ "$SS" = "$T" ] || { echo "ERROR: $LAM has S_search=$SS but this row is T=$T"; exit 1; }
+
+    # ---- job list, then a group-slot loop ----
+    jobs=()
+    for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do jobs+=("$V:$B"); done; done
+    running=0; g=0
+    for job in "${jobs[@]}"; do
+      V=${job%%:*}; B=${job##*:}
+      MK="$STATE_DIR/${MODEL_KEY}.T${T}.${CACHE}.${V}.b${B}.done"
+      OUTCSV="$CSV_DIR/ppl_needle_T${T}_${V}_b${B}.csv"
+      [ -f "$MK" ] && continue
+      # pin this group to its own contiguous device subset
+      devs=$(IFS=,; echo "${GPUARR[*]:$((g*CP)):$CP}")
+      if [ -n "$YARN_SCALE" ]; then
+        ROPE_ENV="YARN_SCALE=$YARN_SCALE"          # trainer's YaRN cache
+      else
+        ROPE_ENV="CP_LONGROPE_FACTORS=$LAM"        # LongRoPE lambda ramp
+      fi
+      (
+        env $MODEL_ARCH CP_SIZE=$CP $TOPO \
+            CP_EVAL_PPL=1 CP_T=$T CP_B=1 \
+            CP_EVAL_WINDOWS=$N CP_EVAL_SKIP_WINDOWS=0 \
+            CP_EVAL_TOKENS_BIN="$NEEDLE_DIR/needle_T${T}_${V}_b${B}.bin" \
+            $ROPE_ENV $CKPT_ENV \
+            CP_EVAL_OUT="$OUTCSV" CP_DATA_ROOT="$DATA_ROOT" \
+            CUDA_VISIBLE_DEVICES=$devs mpirun -np $CP "$EXE" \
+          > "$CSV_DIR/${MODEL_KEY}_T${T}_${CACHE}_${V}_b${B}.log" 2>&1 \
+          && touch "$MK"
+      ) &
+      g=$(( (g + 1) % SLOTS )); running=$(( running + 1 ))
+      if [ "$running" -ge "$SLOTS" ]; then wait -n; running=$(( running - 1 )); fi
+    done
+    wait
+
+    # every job must have produced its marker, or scoring would read a stale CSV
+    miss=0
+    for V in $VARIANTS; do for B in $(seq 0 $((BLOCKS-1))); do
+      [ -f "$STATE_DIR/${MODEL_KEY}.T${T}.${CACHE}.${V}.b${B}.done" ] || {
+        echo "  MISSING: $V b$B -- see $CSV_DIR/${MODEL_KEY}_T${T}_${CACHE}_${V}_b${B}.log"; miss=1; }
+    done; done
+    [ "$miss" = "0" ] || { echo "ERROR: s=$S cache=$CACHE incomplete; not scoring"; exit 1; }
+
+    log "  scoring s=$S cache=$CACHE -> $RESULTS"
+    python3 Tests/bluscriptcp/needle_score.py \
+        --needle-dir "$NEEDLE_DIR" --csv-dir "$CSV_DIR" --T "$T" --blocks "$BLOCKS" \
+        --model "$MODEL_KEY" --cache "$CACHE" --factor "$S" \
+        --head-dim "$HEAD_DIM" --rope-theta "$ROPE_THETA" --out "$RESULTS" \
+        --p-fracs "$(tr ' ' '\n' <<< "$VARIANTS" | grep '^p' | paste -sd,)"
+  done
+done
+
+log "SWEEP DONE -> $RESULTS  (CSVs in $CSV_DIR, markers in $STATE_DIR)"