Upload scripts/ext64k_b5.sh with huggingface_hub

scripts/ext64k_b5.sh

@@ -0,0 +1,216 @@
+#!/usr/bin/env bash
+# =============================================================================
+# ext64k_b5.sh -- YaRN extension 32k -> 64k on b5, deadline-aware.
+#
+# PHASE 1  3-step probe: does 64k fit at cp=4 on 143 GB, and how fast is a step?
+# PHASE 2  train as many steps as fit before HARD_STOP, WSD tail sized so the
+#          FINAL checkpoint is DECAYED and comparable to the 32k results
+# PHASE 3  eval 4 lengths IN PARALLEL on the 4 GPUs (wall clock = slowest job)
+# PHASE 4  upload ckpt + sidecar + logs
+#
+# WHY b5: documents are 65,536-131,072 tokens, so the MINIMUM equals the 64k
+#   window and every training sample is one contiguous document. b4 (32-64k)
+#   would straddle in roughly half its samples. Same reasoning that made b4 the
+#   right choice at 32k.
+#
+# PARENT: b4 full-epoch step 2227 -- the best 32k model (70.0/58.0/43.0 at
+#   8k/16k/32k). CP_INIT_FROM loads WEIGHTS ONLY, so its cp=2 shard layout is
+#   irrelevant to this cp=4 run, and no sidecar is needed (none is read).
+#
+# YaRN INVARIANT: SCALE * ORIG_MAXPOS must equal seq_len EXACTLY and ORIG_MAXPOS
+#   must be the length the parent was actually trained at. 2 * 32768 = 65536.
+#
+# LR: CP_MAX_LR / CP_MIN_LR DO NOT EXIST -- bluscriptCP.cpp:131-132 hardcodes
+#   6e-4 / 6e-5 and the env block never reads them. Real schedule is
+#   6e-4 * CP_REWARMUP_PEAK(0.3) = 1.8e-4 peak, 6e-5 floor -- the same values
+#   every previous run used, so results stay comparable.
+#
+# DEADLINE: HARD_STOP is a UTC epoch. Step count is derived from the MEASURED
+#   probe rate, not an estimate, so the run lands inside the window.
+# =============================================================================
+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
+HF=/root/envv/bin/hf
+EXE=./build/bluscriptCP_exec
+CK=checkpoints_bluscriptcp/rank_0
+LOG=ext64k_logs; mkdir -p "$LOG"
+M="$LOG/master.log"
+REPO=unparallelled/BluScriptCP
+say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
+die(){ say "FATAL: $*"; exit 1; }
+
+T=65536; NP=4; GPUS=0,1,2,3
+ORIG=32768; SCALE=2                  # 2 * 32768 == 65536
+DATA=lendata/b5_65536_131072
+PARENT="$CK/b4epoch32k_step2227.ckpt"
+HAY="$PWD/lendata/haystack/b2val_haystack.bin"
+WARM=100; TAIL=${TAIL:-2}; FREQ=50; KEEP=2
+GB=524288
+# 23:30 IST = 18:00 UTC today
+HARD_STOP=${HARD_STOP:-$(date -u -d "$(date -u +%Y-%m-%d) 18:00:00" +%s)}
+RESERVE_MIN=${RESERVE_MIN:-50}       # eval (28) + upload (20) + slack
+
+$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
+[ -x "$EXE" ]    || die "$EXE not built"
+[ -f "$PARENT" ] || die "parent missing: $PARENT"
+[ -f "$HAY" ]    || die "haystack missing: $HAY"
+[ -d "$DATA" ]   || die "$DATA missing"
+[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant: $SCALE * $ORIG != $T"
+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 train+val shards in $DATA (have $ntr/$nva)"
+have=$(( $(du -cb "$DATA"/*train*.bin | tail -1 | cut -f1) / 2 ))
+
+say "=== 32k -> 64k YaRN extension on b5 (4x H200, cp=$NP) ==="
+say "parent: b4 full epoch step 2227 (best 32k model)"
+say "YaRN SCALE=$SCALE ORIG=$ORIG -> $((SCALE*ORIG)) == CP_T=$T  OK"
+say "b5: $ntr train shards = $(echo "scale=3; $have/1000000000" | bc) B tokens available"
+say "hard stop $(date -u -d @$HARD_STOP +%H:%M) UTC = $(date -u -d @$((HARD_STOP+19800)) +%H:%M) IST, reserving ${RESERVE_MIN}m"
+
+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_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=0.3"
+ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
+
+# ---- PHASE 1: fit + speed --------------------------------------------------
+if [ ! -f "$LOG/probe.ok" ]; then
+  say "PHASE 1: 3-step probe (64k at cp=4 projects to 108.8 GB of 143 GB)"
+  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 -- 64k does not fit at cp=$NP even on 143 GB cards."
+    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" | sed 's/^/  /' >&2
+  say "PROBE OK: peak $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
+  touch "$LOG/probe.ok"
+fi
+
+# ---- derive the step count from the MEASURED rate --------------------------
+# The probe used a tiny global batch, so its dt is NOT the real per-step cost.
+# Scale it: real grad_accum is GB/T vs the probe's 65536/T.
+SPS=$($PY - "$LOG/probe.log" $GB $T <<'PYEOF'
+import re,sys,statistics
+t=open(sys.argv[1],errors='ignore').read()
+d=[float(x) for x in re.findall(r'dt: ([0-9.]+)ms', t)]
+probe_accum=65536//int(sys.argv[3]) or 1
+real_accum=int(sys.argv[2])//int(sys.argv[3])
+per=statistics.median(d[1:]) if len(d)>2 else d[-1]
+print(f"{per/1000.0*real_accum/probe_accum:.3f}")
+PYEOF
+)
+say "measured step cost at real batch: ${SPS} s/step"
+if [ -n "${STEPS_OVERRIDE:-}" ]; then
+  say "STEPS_OVERRIDE=$STEPS_OVERRIDE -- the probe ran grad_accum=1 so its dt x8"
+  say "  overstated the real step cost 5.8x (74.0 vs a measured 12.85 s/step)."
+fi
+STEPS=$($PY - "$SPS" "$HARD_STOP" "$RESERVE_MIN" "$have" "$GB" "$WARM" "$TAIL" <<'PYEOF'
+import sys,time,math
+sps=float(sys.argv[1]); stop=int(sys.argv[2]); res=int(sys.argv[3])
+have=int(sys.argv[4]); gb=int(sys.argv[5]); warm=int(sys.argv[6]); tail=int(sys.argv[7])
+budget=stop-time.time()-res*60
+n=int(budget/sps)
+n=min(n, have//gb)                       # never wrap the data
+n=max(n, warm+tail+10)                   # schedule must be legal
+print(n)
+PYEOF
+)
+STEPS=${STEPS_OVERRIDE:-$STEPS}
+STABLE=$(( STEPS - WARM - TAIL ))
+[ "$STABLE" -lt 1 ] && STABLE=1
+say "PHASE 2: $STEPS steps = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens = $(echo "scale=0; $STEPS*$GB*100/866000000" | bc)% of a b5 epoch"
+say "  WSD: warmup $WARM, stable $STABLE, decay tail $TAIL -> FINAL checkpoint IS decayed"
+
+if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
+  say "already trained -> $(cat "$LOG/ext.ckpt")"
+else
+  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=$GB 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 in 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 for $(basename "$CKF")"
+  echo "$CKF" > "$LOG/ext.ckpt"; echo "$RN" > "$LOG/runnum"
+  say "ext64k -> 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)"
+  say "  final lr $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'lr [0-9.e-]+')  (should be ~6e-05 = decayed)"
+fi
+
+# ---- PHASE 3: eval, 4 lengths in PARALLEL on 4 GPUs ------------------------
+EXT=$(basename "$(cat "$LOG/ext.ckpt")")
+say "PHASE 3: eval ext64k -- 4 lengths, one GPU each"
+# 65512 + MAX_NEW 24 = 65536 exactly; 65528 would pad to 65792 and abort.
+for LEN in 8192 16360 32744 65512; do
+  ( GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=ext64k_len_$LEN TASKS="niah_single" \
+    CKPTS="ext64k" 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=$(( (LEN==8192)?0:(LEN==16360)?1:(LEN==32744)?2:3 )) JOBS_PER_GPU=1 \
+    BF16=1 GREEDY=1 CKPT_EXTRA="ext64k=$EXT" ./run_longeval.sh \
+    >> "$LOG/eval_$LEN.log" 2>&1 ) &
+done
+wait
+for LEN in 8192 16360 32744 65512; do
+  [ -s "ext64k_len_$LEN/summary.csv" ] && say "  len $LEN -> $(sed -n 2p ext64k_len_$LEN/summary.csv)" \
+                                       || say "  len $LEN -> NO RESULT"
+done
+
+# ---- PHASE 4: upload -------------------------------------------------------
+say "PHASE 4: upload"
+{
+  echo "# 32k -> 64k YaRN extension on b5 -- $(date -u +%Y-%m-%d)"
+  echo
+  echo "Parent: b4 full-epoch step 2227 (best 32k model: 70.0/58.0/43.0 at 8k/16k/32k)."
+  echo "Data:   b5_65536_131072 -- min doc 65,536 == the window, so every sample is"
+  echo "        ONE contiguous document."
+  echo "YaRN:   SCALE=2 ORIG_MAXPOS=32768 -> 65536 exactly."
+  echo "Train:  $STEPS steps x $GB = $(echo "scale=3; $STEPS*$GB/1000000000" | bc) B tokens"
+  echo "        = $(echo "scale=0; $STEPS*$GB*100/866000000" | bc)% of b5's 0.866 B train split (no wrap)."
+  echo "        WSD warmup $WARM / stable $STABLE / decay tail $TAIL -- final ckpt DECAYED."
+  echo "        LR peak 1.8e-4, floor 6e-5 (CP_MAX_LR is a no-op; see B4 results doc)."
+  echo "        4x H200, cp=4, B=1 x grad_accum $((GB/T))."
+  echo "Eval:   niah_single, N=100, greedy, shared haystack b2_8192_16384 val."
+  echo
+  echo "## Results"; echo
+  echo "| context | niah_single |"; echo "|---|---|"
+  for LEN in 8192 16360 32744 65512; do
+    v=$(sed -n 2p "ext64k_len_$LEN/summary.csv" 2>/dev/null | cut -d, -f2)
+    printf "| %s | %s |\n" "$LEN" "${v:-n/a}"
+  done
+  echo
+  echo "For comparison, the 32k parent at the same lengths: 70.0 / 58.0 / 43.0 / n/a"
+  echo "(it cannot run at 64k -- its RoPE cache is sized 32,768)."
+  echo
+  echo "NOTE: this is a PARTIAL b5 epoch, sized to a wall-clock deadline. It is a"
+  echo "first 64k data point, not a converged extension."
+} > "$LOG/RESULTS.md"
+up(){ [ -e "$1" ] || return 0
+  $HF upload "$REPO" "$1" "$2" >>"$LOG/upload.log" 2>&1 && say "  ok $2" || say "  FAIL $2"; }
+up "$LOG/RESULTS.md" docs/EXT64K_B5_RESULTS_$(date -u +%Y-%m-%d).md
+for LEN in 8192 16360 32744 65512; do up "ext64k_len_$LEN" "logs/ext64k_len_$LEN/"; done
+up "$LOG" logs/ext64k/
+up ext64k_b5.sh scripts/ext64k_b5.sh
+RN=$(cat "$LOG/runnum" 2>/dev/null)
+for r in 0 1 2 3; do
+  for e in ckpt zero_g1; do
+    up "checkpoints_bluscriptcp/rank_$r/blumodelcp_run${RN}_step_${STEPS}.$e" \
+       "checkpoints/ext64k_b5_rank${r}_step_${STEPS}.$e"
+  done
+done
+say "=== EXT64K DONE ==="