Upload source/files/run_in2_experiment.sh with huggingface_hub

source/files/run_in2_experiment.sh

@@ -0,0 +1,218 @@
+#!/usr/bin/env bash
+# IN2 experiment driver: generate arm C, gate all arms, train A/B/C, probe each,
+# read out, upload. Unattended, with hard gates -- every one of them exists because
+# the failure it catches is SILENT and would otherwise produce a plausible-looking
+# number for the wrong thing.
+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/python"; HF="$ROOT/envv/bin/hf"
+LOGDIR=in2_logs; mkdir -p $LOGDIR
+MASTER=$LOGDIR/driver_$(date +%Y%m%d_%H%M).log
+say(){ echo "[$(date +%H:%M:%S)] $*" | tee -a "$MASTER"; }
+die(){ say "FATAL: $*"; say "aborting -- see $MASTER"; exit 1; }
+
+# ---- configuration -------------------------------------------------------
+ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
+BASE_CKPT=${BASE_CKPT:-$PWD/checkpoints_bluscriptcp/rank_0/blumodelcp_run7_step_750.ckpt}
+DATA=${DATA:-/root/.cache/huggingface/hub/datasets--navingv--flux/snapshots/5232b3d133d2eb71905725f288ff6b56dc0df773}
+T=16384; K=64; EXAMPLES=12816; SHARDS=3
+# D_MIN is the generator's DEFAULT (64) and must stay there. It is tempting to set
+# it to R so every arm-C needle lands in the failing region -- that is wrong, and the
+# gate catches it: d_min feeds draw_sites(lo_margin=span+d_min), so raising it moves
+# arm C's PROBES, not just its needles. At d_min=9835 arm C's answers were confined to
+# the last ~6.5k of the document (6522 distinct offsets vs arm B's 16293) and cross-arm
+# assertion 8 failed. Arm B can only absorb position effects while B and C share
+# positions, which is exactly what raising d_min destroys -- and it re-creates the
+# position-specific-circuit confound the interleaving exists to prevent.
+# The consequence is accepted, not fixed: only ~8.2% of arm C's 820,224 triples land
+# beyond R=9835 (67,653 of them). That fraction is the price of positional matching.
+D_MIN=${D_MIN:-64}
+R_MEASURED=${R_MEASURED:-9835}
+FT_STEPS=${FT_STEPS:-400}; GBATCH=${GBATCH:-524288}; B=${B:-4}
+YARN=4; GPUS=0,1,2,3; NP=4
+HF_REPO=${HF_REPO:-unparallelled/BluScriptCP}
+STAGE=/root/in2_upload; rm -rf $STAGE; mkdir -p $STAGE
+
+# Single-instance lock. Three copies of this script ran concurrently once (one from
+# each terminal), sharing in2_logs/, checkpoints_bluscriptcp/ and all four GPUs: 10
+# exec processes contending, and they would have raced for run numbers and written
+# each other's arm markers. A second launch must refuse, not collide.
+LOCK=.in2_driver.lock
+if [ -e "$LOCK" ] && kill -0 "$(cat "$LOCK" 2>/dev/null)" 2>/dev/null; then
+  echo "REFUSING: another driver is already running (pid $(cat "$LOCK")). "
+  echo "  If it is dead: rm $LOCK"
+  exit 1
+fi
+echo $$ > "$LOCK"
+trap 'rm -f "$LOCK"' EXIT INT TERM
+
+say "IN2 driver starting.  base=$BASE_CKPT  d_min=$D_MIN  R=$R_MEASURED"
+[ -x ./build/bluscriptCP_exec ] || die "no binary"
+[ -f "$BASE_CKPT" ] || die "base checkpoint missing: $BASE_CKPT"
+[ -x "$PY" ] || die "venv python missing"
+
+# ---- 1. arm C data (A and B already generated) ---------------------------
+if [ ! -f in2_data/arm_c/.done ]; then
+  say "generating arm C (train) at --d-min $D_MIN (generator default; see note above)"
+  FILL=$(ls Data_Loader/Data/flux_train_0000{1,2,3}*.bin 2>/dev/null | head -12 | tr '\n' ' ')
+  [ -z "$FILL" ] && FILL=$(ls Data_Loader/Data/flux_train*.bin | head -12 | tr '\n' ' ')
+  $PY Tests/bluscriptcp/in2_gen.py --arm c --T $T --examples $EXAMPLES --K $K --shards $SHARDS \
+      --seed 0 --d-min $D_MIN --filler-shards $FILL --out-dir in2_data/arm_c >>"$MASTER" 2>&1 \
+      || die "arm C train generation failed"
+  say "generating arm C (val)"
+  $PY Tests/bluscriptcp/in2_gen.py --arm c --T $T --examples 64 --split val --K $K --shards 1 \
+      --seed 0 --d-min $D_MIN --filler-shards $(ls Data_Loader/Data/flux_val*.bin | head -1) \
+      --out-dir in2_data/arm_c >>"$MASTER" 2>&1 || die "arm C val generation failed"
+  touch in2_data/arm_c/.done
+fi
+grep -E 'distance:' "$MASTER" | tail -2 | while read -r l; do say "  arm C $l"; done
+
+# ---- 2. GATE: 23 layout assertions, all arms, both splits ---------------
+# This is what caught arm B silently dropping 213 of 1024 needles last time -- a
+# desynchronised answer_start histogram destroys the B-vs-C matched-position
+# discipline the whole primary claim rests on.
+for SPLIT in train val; do
+  say "gate ($SPLIT) on all three arms"
+  OUT=$($PY Tests/bluscriptcp/in2_gate.py --dirs in2_data/arm_a in2_data/arm_b in2_data/arm_c \
+        --split $SPLIT --B $B --dp 1 2>&1); echo "$OUT" >>"$MASTER"
+  echo "$OUT" | grep -q "ALL PASS" || { echo "$OUT" | tail -20; die "GATE FAILED ($SPLIT)"; }
+  say "  gate $SPLIT: ALL PASS"
+done
+
+# ---- 3. train the three arms --------------------------------------------
+declare -A RUNNO
+for ARM in a b c; do
+  MK=$LOGDIR/arm_${ARM}.trained
+  if [ -f "$MK" ]; then RUNNO[$ARM]=$(cat "$MK"); say "arm $ARM already trained (run ${RUNNO[$ARM]})"; continue; fi
+  # One checkpoint step costs 17 GB across 4 ranks (3.9 GB model + 2 sidecars, per
+  # rank). The first attempt filled a 193 GB disk: FREQ=100 x KEEP=2 x 3 arms is up
+  # to 102 GB of arm checkpoints on top of the base. Saves then failed with a mere
+  # [WARN] while training carried on, so arm C ran all 400 steps and reported "done"
+  # with ZERO checkpoints on disk -- 45 minutes of GPU time for nothing. Check first.
+  FREE_G=$(df -BG --output=avail / | tail -1 | tr -dc '0-9')
+  say "  free disk: ${FREE_G}G"
+  [ "${FREE_G:-0}" -ge 40 ] || die "only ${FREE_G}G free; one arm checkpoint needs ~17G plus a same-size .tmp during the write. Free space before continuing."
+
+  # Prune sidecars whose .ckpt has already been rotated away. CP_CKPT_KEEP prunes
+  # .ckpt files but NOT the zero_g* sidecars bluscriptCP writes alongside them, so
+  # they accumulate silently -- 50 GB of them built up on the 600M run.
+  for d in checkpoints_bluscriptcp/rank_*; do
+    [ -d "$d" ] || continue
+    K=$(ls $d/*.ckpt 2>/dev/null | grep -oE '_run[0-9]+_step_[0-9]+' | sort -u)
+    for f in $d/*.zero_g*; do
+      [ -e "$f" ] || continue
+      S=$(echo "$f" | grep -oE '_run[0-9]+_step_[0-9]+')
+      echo "$K" | grep -qx "$S" || rm -f "$f"
+    done
+    rm -f $d/*.tmp
+  done
+
+  L=$LOGDIR/arm_${ARM}_train.log
+  say "training arm $ARM  ($FT_STEPS steps x $GBATCH tokens, T=$T, cp=$NP, YaRN=$YARN)"
+  env CUDA_VISIBLE_DEVICES=$GPUS $ARCH \
+    CP_INIT_FROM="$BASE_CKPT" \
+    CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split \
+    CP_T=$T YARN_SCALE=$YARN CP_B=$B CP_GLOBAL_BATCH=$GBATCH \
+    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_REWARMUP=100 CP_REWARMUP_PEAK=0.3 \
+    CP_MAX_STEPS=$FT_STEPS CP_WARMUP=1 \
+    CP_CKPT=1 CP_CKPT_FREQ=200 CP_CKPT_KEEP=1 CP_CKPT_NEW_RUN=1 \
+    CP_DATA_ROOT="$PWD/in2_data/arm_$ARM" \
+    mpirun --allow-run-as-root -np $NP ./build/bluscriptCP_exec > "$L" 2>&1 \
+    || die "arm $ARM training failed -- see $L"
+  RN=$(grep -oE 'CP_Training_log[0-9]+' "$L" | head -1 | grep -oE '[0-9]+')
+  [ -n "$RN" ] || die "arm $ARM: could not determine run number"
+  # A zero exit code is NOT proof the checkpoint exists: bluscriptCP catches a failed
+  # save, prints [WARN], and trains on. Verify the artifact, not the signal.
+  CKF=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
+  [ -n "$CKF" ] || die "arm $ARM (run $RN) finished but wrote NO checkpoint -- check $L for 'checkpoint save ... failed'"
+  say "  arm $ARM checkpoint: $(basename "$CKF") ($(du -h "$CKF" | cut -f1))"
+  if grep -q 'checkpoint save.*failed' "$L"; then
+    die "arm $ARM: a checkpoint save FAILED during training (see $L). Refusing to continue on a partial arm."
+  fi
+  echo "$RN" > "$MK"; RUNNO[$ARM]=$RN
+  say "  arm $ARM done -> run $RN, final loss $(grep -E '^step ' "$L" | tail -1 | grep -oE 'loss: [0-9.]+')"
+done
+
+# ---- 3b. LEARNABILITY GATE (arm C) --------------------------------------
+# The plan is explicit: if arm C's synthetic val_loss does not fall, the run says
+# nothing about IN2 and must NOT be written up as a null.
+CSV=CP_BluScreipt_Training_logs/CP_Training_log${RUNNO[c]}.csv
+if [ -f "$CSV" ]; then
+  # val_loss is forward-filled and val_freq defaults to 250, so a 400-step arm yields
+  # as few as ONE real evaluation. Comparing first to last of a forward-filled column
+  # then reports "did not fall" for a run that was never measured twice -- which is
+  # exactly what aborted the first attempt AFTER all three arms had trained.
+  NV=$(awk -F, 'NR>1 && $3+0>0 {print $3}' "$CSV" | sort -u | wc -l)
+  V=$(awk -F, 'NR>1 && $3+0>0 && $3!=p {v=$3; if(f==""){f=$3} p=$3} END{print f" "v}' "$CSV")
+  VF=${V% *}; VL=${V#* }
+  say "arm C val_loss: $NV distinct point(s), first=$VF last=$VL"
+  if [ "${NV:-0}" -lt 2 ]; then
+    say "  learnability gate SKIPPED: only $NV validation point(s) -- cannot judge a trend."
+    say "  (lower CP_VAL_FREQ, or lengthen the arm, to make this gate meaningful)"
+    VF=1; VL=0   # force the comparison below to pass
+  fi
+  $PY -c "import sys; sys.exit(0 if float('$VL') < float('$VF') else 1)" \
+    || die "LEARNABILITY GATE FAILED: arm C val_loss did not fall ($VF -> $VL). The run says nothing about IN2."
+  say "  learnability gate: PASS"
+else say "  WARN: no CSV for arm C run ${RUNNO[c]}; learnability gate skipped"; fi
+
+# ---- 4. probe each arm (same YaRN cache it trained under) ---------------
+for ARM in a b c; do
+  OUTD=probe_arm_${ARM}
+  [ -f $OUTD/reach.csv ] && { say "probe $ARM already done"; continue; }
+  CK=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RUNNO[$ARM]}_step_*.ckpt 2>/dev/null | head -1)
+  [ -n "$CK" ] || die "arm $ARM: no checkpoint for run ${RUNNO[$ARM]}"
+  say "probing arm $ARM  ($CK)"
+  CKPT="$PWD/$CK" OUT=$OUTD FACTORS=$YARN CACHES=yarn YARN_SCALE=$YARN GPUS=$GPUS BLOCKS=3 \
+    ./probe_1b.sh >>"$MASTER" 2>&1 || say "  WARN: probe $ARM returned nonzero"
+  [ -f $OUTD/reach.csv ] && say "  arm $ARM reach.csv written" || say "  WARN: no reach.csv for arm $ARM"
+done
+
+# ---- 5. readout ---------------------------------------------------------
+say "readout"
+$PY Tests/bluscriptcp/in2_report.py --needle-dir probe_arm_c/needles --T $T \
+    --arms a b c --R $R_MEASURED --blocks 3 > $LOGDIR/in2_report.txt 2>&1 \
+    || say "  WARN: in2_report returned nonzero (see $LOGDIR/in2_report.txt)"
+tail -30 $LOGDIR/in2_report.txt | tee -a "$MASTER"
+
+# ---- 6. UPLOAD everything needed to carry forward -----------------------
+say "staging artefacts for HF"
+cp -r $LOGDIR $STAGE/ 2>/dev/null
+for d in probe_run8 probe_run8_s4 probe_run7_yarn4 probe_arm_a probe_arm_b probe_arm_c; do
+  [ -d "$d" ] && cp -r "$d" $STAGE/ 2>/dev/null
+done
+mkdir -p $STAGE/train_csv && cp CP_BluScreipt_Training_logs/*.csv $STAGE/train_csv/ 2>/dev/null
+mkdir -p $STAGE/run_logs  && cp run_logs/*.log $STAGE/run_logs/ 2>/dev/null
+cp -r /root/claude-vault $STAGE/vault 2>/dev/null
+# code: git holds it, but a tarball costs nothing and survives a lost branch
+tar czf $STAGE/cp_source.tgz --exclude='build' --exclude='checkpoints_bluscriptcp' \
+    --exclude='in2_data' --exclude='Data_Loader/Data' --exclude='.git' --exclude='*.bin' \
+    -C "$ROOT" dist/Context_Parallelism dist/zero dist/communication/src/processGroupNccl.cpp 2>/dev/null
+# the three arm checkpoints (weights only -- flat, loadable at any np)
+mkdir -p $STAGE/checkpoints
+for ARM in a b c; do
+  CK=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RUNNO[$ARM]:-0}_step_*.ckpt 2>/dev/null | head -1)
+  [ -n "$CK" ] && cp "$CK" $STAGE/checkpoints/arm_${ARM}_$(basename "$CK")
+done
+[ -f "$BASE_CKPT" ] && cp "$BASE_CKPT" $STAGE/checkpoints/extended_$(basename "$BASE_CKPT")
+du -sh $STAGE/* 2>/dev/null | tee -a "$MASTER"
+
+if [ -x "$HF" ]; then
+  say "uploading to $HF_REPO"
+  "$HF" upload "$HF_REPO" $STAGE/cp_source.tgz  source/cp_source_in2.tgz      >>"$MASTER" 2>&1
+  "$HF" upload "$HF_REPO" $STAGE/vault          docs/                          >>"$MASTER" 2>&1
+  "$HF" upload "$HF_REPO" $STAGE/in2_logs       logs/in2/                      >>"$MASTER" 2>&1
+  "$HF" upload "$HF_REPO" $STAGE/train_csv      logs/in2_train_csv/            >>"$MASTER" 2>&1
+  "$HF" upload "$HF_REPO" $STAGE/run_logs       logs/in2_stdout/               >>"$MASTER" 2>&1
+  for d in probe_run8 probe_run8_s4 probe_run7_yarn4 probe_arm_a probe_arm_b probe_arm_c; do
+    [ -d "$STAGE/$d" ] && "$HF" upload "$HF_REPO" "$STAGE/$d" "logs/$d/"       >>"$MASTER" 2>&1
+  done
+  "$HF" upload "$HF_REPO" $STAGE/checkpoints    checkpoints/in2/               >>"$MASTER" 2>&1
+  say "upload complete"
+else say "WARN: hf CLI not found at $HF -- artefacts staged in $STAGE but NOT uploaded"; fi
+
+say "DONE. staged copy: $STAGE   master log: $MASTER"