Upload scripts/retrain_corrected_cache.sh with huggingface_hub

scripts/retrain_corrected_cache.sh

@@ -0,0 +1,108 @@
+#!/usr/bin/env bash
+# =============================================================================
+# retrain_corrected_cache.sh -- rebuild the whole checkpoint chain under the
+# CORRECTED YaRN parameterisation (YARN_ORIG_MAXPOS=4096).
+#
+# WHY THE FULL CHAIN, not just the arms: run7 itself was trained with the buggy
+# orig_maxpos=1024 default, so arms trained from it would start from weights
+# already adapted to the wrong rotation. run8 (the 4k base) is unaffected -- it
+# predates any YaRN scaling -- so the chain restarts there.
+#
+# WHY THIS IS THE REAL TEST: switching the cache at EVAL time is not a fix.
+# Measured on this box, corrected-vs-as-trained at N=100: 1k 90 vs 84, ~4k
+# 14-33 vs 75, 16k 17 vs 0. The corrected ramp (low=14 high=32) is a DIFFERENT
+# function from the one the weights adapted to (low=7 high=25) -- better in some
+# regimes, much worse around 4k. Only retraining resolves that.
+#
+# Recipe recovered from run7's own log + LR trace, not from memory:
+#   B=4 grad_accum=16 global_batch=1048576, Muon+AdamW ZeRO-1 symmetric, bf16
+#   mask=15, max_lr 1e-3 -> min 1e-4, rewarmup 50 to peak 0.3x (=3e-4), 750 steps
+# =============================================================================
+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=retrain_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=16384; YARN=4; ORIG=4096; NP=4; GPUS=0,1,2,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"
+COMMON="CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
+YARN_SCALE=$YARN 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"
+
+[ -x "$EXE" ] || die "$EXE missing"
+FREE=$(df -BG --output=avail /root | tail -1 | tr -dc '0-9')
+say "free disk ${FREE}G (need ~70G: 4 checkpoint sets at ~17G)"
+[ "${FREE:-0}" -ge 75 ] || die "only ${FREE}G free"
+
+# run a training job, then VERIFY the checkpoint exists -- a zero exit code is
+# not proof: bluscriptCP catches a failed save, prints [WARN] and trains on.
+train(){   # $1=label $2=logfile $3...=extra env
+  local label="$1" lf="$2"; shift 2
+  say "training $label"
+  env $ARCH $COMMON "$@" CP_DATA_ROOT="$PWD/$DATA_ROOT" \
+      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
+      > "$lf" 2>&1 || die "$label training failed -- see $lf"
+  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$lf" | head -1 | grep -oE '[0-9]+')
+  [ -n "$RN" ] || die "$label: no run number in $lf"
+  CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
+  [ -n "$CKF" ] || die "$label finished but wrote NO checkpoint (see $lf)"
+  grep -qa 'checkpoint save.*failed' "$lf" && die "$label: a checkpoint save FAILED"
+  say "  $label -> run $RN, $(basename "$CKF"), final loss $(grep -aE '^step ' "$lf" | tail -1 | grep -oE 'loss: [0-9.]+')"
+  case "$CKF" in *$'\n'*) die "ckptpath sanity: multi-line checkpoint path";; esac
+  echo "$CKF"
+}
+
+# ---- 1. arm training data (deterministic at --seed 0) -----------------------
+FILL=$(ls Data_Loader/Data/flux_train_*.bin 2>/dev/null | head -12 | tr '\n' ' ')
+[ -n "$FILL" ] || die "no flux_train shards for in2_gen filler"
+for A in a b c; do
+  if [ ! -f "in2_data/arm_$A/.done" ]; then
+    say "generating arm $A data"
+    $PY Tests/bluscriptcp/in2_gen.py --arm $A --T $T --examples 12816 --K 64 --shards 3 \
+        --seed 0 --d-min 64 --filler-shards $FILL --out-dir in2_data/arm_$A >>"$M" 2>&1 \
+        || die "arm $A train data generation failed"
+    $PY Tests/bluscriptcp/in2_gen.py --arm $A --T $T --examples 64 --split val --K 64 --shards 1 \
+        --seed 0 --d-min 64 --filler-shards $(ls Data_Loader/Data/flux_val*.bin | head -1) \
+        --out-dir in2_data/arm_$A >>"$M" 2>&1 || die "arm $A val data generation failed"
+    touch "in2_data/arm_$A/.done"
+  fi
+done
+say "arm data ready"
+
+# ---- 2. run7' : the YaRN 16k extension, corrected cache ---------------------
+if [ -f "$LOG/run7c.ckptpath" ]; then
+  R7=$(cat "$LOG/run7c.ckptpath"); say "run7' already done ($R7)"
+else
+  DATA_ROOT=Data_Loader/Data
+  R7=$(train "run7' (YaRN 16k extension, corrected cache)" "$LOG/run7c.log" \
+        CP_INIT_FROM="$PWD/$CK/blumodelcp_run8_step_1000.ckpt" \
+        CP_B=4 CP_GLOBAL_BATCH=1048576 CP_MAX_STEPS=750 \
+        CP_REWARMUP=50 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=250) || exit 1
+  echo "$R7" > "$LOG/run7c.ckptpath"
+fi
+
+# ---- 3. arms A/B/C from run7' ----------------------------------------------
+for A in a b c; do
+  if [ -f "$LOG/arm_$A.ckptpath" ]; then say "arm $A already done"; continue; fi
+  DATA_ROOT=in2_data/arm_$A
+  P=$(train "arm $A (corrected cache)" "$LOG/arm_$A.log" \
+        CP_INIT_FROM="$PWD/$R7" \
+        CP_B=4 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=400 \
+        CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=200) || exit 1
+  echo "$P" > "$LOG/arm_$A.ckptpath"
+done
+
+say "ALL DONE"
+say "  run7' : $(cat $LOG/run7c.ckptpath)"
+for A in a b c; do say "  arm $A : $(cat $LOG/arm_$A.ckptpath 2>/dev/null)"; done