Upload scripts/ext_32k.sh with huggingface_hub

scripts/ext_32k.sh

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