Upload scripts/ext32k_b4.sh with huggingface_hub

scripts/ext32k_b4.sh

@@ -0,0 +1,134 @@
+#!/usr/bin/env bash
+# =============================================================================
+# ext32k_b4.sh -- YaRN extension 16k -> 32k on b4, then evaluate 1k..32k.
+#
+# WHY b4 AND NOT b3, at a 32,768 window:
+#   b3 documents are 16,384-32,768 tokens -- mean 21,224 = 65% of the window, and
+#   the MINIMUM is 16,384. So every b3 sample carries an EOT boundary plus part of
+#   a second document, and no dependency can span 32k. That is the same defect as
+#   the original flux corpus (median 826 tok), one scale up.
+#   b4 documents are 32,768-65,536 -- the minimum EQUALS the window, so every
+#   sample is one contiguous document and a 32k-spanning dependency is possible.
+#   Truncating a 44k document to 32k is fine; a straddled window is not.
+#
+# YaRN INVARIANT (the bug that cost 27pp at 16k):
+#   SCALE * ORIG_MAXPOS must equal seq_len EXACTLY, and ORIG_MAXPOS must be the
+#   length the parent was actually TRAINED at.  2 * 16384 = 32768. The default
+#   ORIG_MAXPOS is 1024 (YARNOps.cpp:57) and would silently be wrong.
+#
+# PARENT: l2 = the best 16k arm (b2 docs, which fit the 16k window). Extending
+#   from the strongest 16k model, not from run7'.
+#
+# GLOBAL BATCH held at 524,288 (B=1, T=32768, grad_accum=16) so the optimizer
+#   trajectory matches every 16k arm. Only the window changed.
+#
+# MEMORY: 2x H200 = 143 GB. cp=2 puts 16,384 tokens on each rank -> ~130.8 GB,
+#   91% utilisation. That is close enough that fragmentation could still OOM, so
+#   PHASE 1 is a 3-step probe and the full run only starts if it survives.
+# =============================================================================
+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/envv/bin/python3
+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=2; GPUS=0,1
+ORIG=16384; SCALE=2                       # 2 * 16384 == 32768 exactly
+STEPS=${STEPS:-750}
+STABLE=${STABLE:-$(( STEPS - 200 ))}      # WSD: flat LR, minimal legal decay tail
+FREQ=${FREQ:-50}
+KEEP=${KEEP:-1}                           # disk is tight on this box (~45 GB free)
+DATA=lendata/b4_32768_65536
+PARENT="$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt"
+HAY="$PWD/lendata/haystack/b2val_haystack.bin"
+ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
+
+$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
+[ -x "$EXE" ]    || die "$EXE not built"
+[ -f "$PARENT" ] || die "parent l2 missing: $PARENT"
+[ -f "$HAY" ]    || die "haystack missing"
+[ -d "$DATA" ]   || die "$DATA missing"
+ntr=$(ls "$DATA"/*train*.bin 2>/dev/null | wc -l); nva=$(ls "$DATA"/*val*.bin 2>/dev/null | wc -l)
+[ "$ntr" -ge 1 ] && [ "$nva" -ge 1 ] || die "need >=1 train and >=1 val shard in $DATA (have $ntr/$nva)"
+have=$(( $(du -cb "$DATA"/*train*.bin | tail -1 | cut -f1) / 2 ))
+need=$(( STEPS*524288 ))
+say "=== 16k -> 32k YaRN extension on b4 ==="
+say "$ntr train shards = $(echo "scale=3; $have/1000000000" | bc) B tokens; need $(echo "scale=3; $need/1000000000" | bc) B for $STEPS steps"
+[ "$have" -lt "$need" ] && say "WARNING: data repeats $(echo "scale=2; $need/$have" | bc)x"
+say "YaRN: SCALE=$SCALE ORIG_MAXPOS=$ORIG -> $((SCALE*ORIG)) (must equal CP_T=$T)"
+[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant violated"
+
+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 CP_WARMUP=100 CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3"
+
+# ---- PHASE 1: does 32k fit at cp=2 on 143 GB? -----------------------------
+if [ ! -f "$LOG/probe.ok" ]; then
+  say "PHASE 1: 3-step memory probe (projected 130.8 GB of 143 GB = 91%)"
+  env $ARCH $COMMON CP_STABLE=0 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 -- 32k does not fit at cp=2 on this box."
+    say "  32k needs cp=4 (67.5 GB/GPU). Options: a 4-GPU box, or stop here."
+    grep -aoE 'cudaMallocAsync\(size=[0-9]+\) failed' "$LOG/probe.log" | head -2 | sed 's/^/    /' >&2
+    exit 1
+  fi
+  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | tail -1 | sed 's/^/  /' >&2
+  say "PROBE OK: peak $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
+  touch "$LOG/probe.ok"
+fi
+
+# ---- PHASE 2: the extension ------------------------------------------------
+if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
+  say "extension already trained -> $(cat "$LOG/ext.ckpt")"
+else
+  say "PHASE 2: $STEPS steps, stable=$STABLE, ckpt every $FREQ keep $KEEP"
+  env $ARCH $COMMON CP_STABLE=$STABLE \
+      CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP CP_CKPT_NEW_RUN=1 \
+      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
+      CP_B=1 CP_GLOBAL_BATCH=524288 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 the full run"
+  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/ext.log" | head -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 "run $RN wrote NO checkpoint"
+  n=$(ls checkpoints_bluscriptcp/rank_*/"$(basename "$CKF")" 2>/dev/null | wc -l)
+  [ "$n" -ge "$NP" ] || die "only $n/$NP rank shards"
+  echo "$CKF" > "$LOG/ext.ckpt"
+  say "ext32k -> run $RN, $(basename "$CKF"), loss $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
+  say "  peak mem $(grep -aoE 'mem: [0-9]+MB' "$LOG/ext.log" | tail -1)"
+fi
+
+# ---- PHASE 3: evaluate, including the new 32k point ------------------------
+EXT=$(basename "$(cat "$LOG/ext.ckpt")")
+EXTRA="ext32k=$EXT"
+[ -f "$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt" ] && \
+  EXTRA="$EXTRA l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt"
+MODELS=$(echo $EXTRA | tr ' ' '\n' | cut -d= -f1 | tr '\n' ' ')
+say "PHASE 3: eval models $MODELS"
+# 32744 + MAX_NEW 24 = 32768 exactly; 32760 would pad to 33024 and abort.
+for LEN in 4096 8192 16360 32744; do
+  say "eval PROMPT_LEN=$LEN"
+  GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=ext32k_len_$LEN TASKS="niah_single" \
+  CKPTS="$MODELS" N=100 PROMPT_LEN=$LEN MAX_NEW=24 T=$T \
+  DEPTHS="0.0,0.25,0.5,0.75,1.0" YARN_ORIG=$ORIG YARN=$SCALE \
+  GPUS=0,1 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"
+  [ -s "ext32k_len_$LEN/summary.csv" ] && say "  -> $(sed -n 1p ext32k_len_$LEN/summary.csv) / $(sed -n 2p ext32k_len_$LEN/summary.csv)"
+done
+say "=== EXT32K DONE ==="