Upload scripts/abl_phase2.sh with huggingface_hub

scripts/abl_phase2.sh

@@ -0,0 +1,101 @@
+#!/usr/bin/env bash
+# =============================================================================
+# ABLATION phase 1: 4k -> 16k. Two arms, one per H200, run CONCURRENTLY.
+#
+#   ARM=const   constant 8e-5              (GPU 0)
+#   ARM=cosine  cosine 8e-5 -> 8e-6        (GPU 1)
+#
+# Both use Muon + AdamW, mirroring bluscript_zero_ckpt.cpp's BLU_OPTIMIZER=muon:
+# Muon drives the >=2-D weight matrices, an aux AdamW path takes everything else,
+# and ONE schedule drives both because mcfg.lr == mcfg.adamw_lr (ratio 1:1, which
+# set_lr preserves). Pretraining used 8e-4/6e-4 (4:3); here both arms share one
+# value per the ablation spec.
+#
+# WHAT HAD TO CHANGE IN THE TRAINER (2026-09-18):
+#   CP_OPTIMIZER, CP_MAX_LR and CP_MIN_LR were read NOWHERE. max_lr/min_lr were
+#   hardcoded 6e-4/6e-5 and every prior CP run that passed CP_OPTIMIZER=muon
+#   silently trained with AdamW. All three are now wired.
+#
+# Warmup is 1% of the phase's steps: 12 of 1200.
+# Everything else is identical to p1final so the comparison is clean.
+# =============================================================================
+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
+export PATH=/usr/local/cuda/bin:$PATH
+PY=/root/Adhi/BluTrain/envv/bin/python
+EXE=./build/bluscriptCP_exec
+
+ARM=${ARM:?set ARM=const or ARM=cosine}
+GPU=${GPU:?set GPU=0 or 1}
+T=65536; SCALE=16; ORIG=4096
+STEPS=${STEPS:-1200}
+WARM=$(( STEPS / 100 ))                 # 1% of this phase's steps
+MAXLR=4e-5
+case "$ARM" in
+  const)  MINLR=4e-5; STABLE=$(( STEPS - WARM - 2 )) ;;   # flat: tail is the 2-step minimum
+  cosine) MINLR=4e-6; STABLE=0 ;;                          # warmup -> cosine all the way down
+  *) echo "ARM must be const or cosine"; exit 1 ;;
+esac
+GB=524288
+DATA=Data_Loader/Data/ablmix64
+CK=checkpoints_bluscriptcp/rank_0
+PARENT=$(ls -t checkpoints_${ARM}/rank_0/blumodelcp_run*_step_*.ckpt 2>/dev/null | head -1)
+[ -n "$PARENT" ] || { echo "FATAL: phase 1 ($ARM) produced no checkpoint"; exit 1; }
+LOG=p2_${ARM}_logs; mkdir -p "$LOG"; M="$LOG/master.log"
+say(){ echo "[$(date -u +%H:%M:%S)] [$ARM] $*" | tee -a "$M" >&2; }
+die(){ say "FATAL: $*"; exit 1; }
+
+[ -x "$EXE" ] || die "$EXE not built"
+S=build/objects/.cp_rope_flags
+grep -q 'CP_FUSED_ROPE=1' "$S" && grep -q 'CP_ATTN_SPLIT=1' "$S" || die "binary lacks split kernels"
+[ -s "$PARENT" ] || die "phase-1 parent missing"
+grep -qa "PHASE 1 $ARM COMPLETE" p1_${ARM}_logs/master.log 2>/dev/null || die "phase 1 ($ARM) did not report COMPLETE -- refusing to branch off a mid-run checkpoint"
+[ -d "$DATA" ] || die "$DATA missing"
+[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant broken"
+[ "$STABLE" -eq 0 ] || [ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "WSD guard would fire"
+$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
+
+ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
+BASE="CP_SIZE=2 CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
+YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG CP_ROPE_THETA=500000 \
+CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 \
+CP_OPTIMIZER=muon CP_MAX_LR=$MAXLR CP_MIN_LR=$MINLR \
+CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=1.0 \
+CP_STABLE=$STABLE CP_CKPT=1 CP_CKPT_FREQ=200 CP_CKPT_KEEP=3 \
+CP_CKPT_DIR=checkpoints_${ARM}"
+
+say "=== phase 2: 16k -> 64k, arm=$ARM (cp=2, BOTH GPUs) ==="
+say "LR      : max $MAXLR  min $MINLR  warmup $WARM (1% of $STEPS)  stable $STABLE"
+say "optim   : Muon (>=2-D matrices) + aux AdamW, both at the SAME lr"
+say "parent  : $(basename "$PARENT")   data: $DATA   steps: $STEPS"
+
+if [ ! -f "$LOG/probe.ok" ]; then
+  say "PROBE 3 steps"
+  env $ARCH $BASE CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$PWD/$PARENT" \
+      CP_DATA_ROOT="$PWD/$DATA" CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=3 \
+      CUDA_VISIBLE_DEVICES=0,1 mpirun --allow-run-as-root -np 2 "$EXE" > "$LOG/probe.log" 2>&1
+  RC=$?
+  grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log" && die "PROBE OOM"
+  N=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe.log")
+  P=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log"|tail -1)
+  [ "$RC" -eq 0 ] && [ "$N" -ge 1 ] || { say "PROBE FAILED rc=$RC steps=$N"; tail -14 "$LOG/probe.log" >&2; exit 1; }
+  # ASSERT the trainer actually took the new env vars -- all three were no-ops
+  # until today, and a silent fallback would look like a healthy run.
+  grep -qai 'muon' "$LOG/probe.log" || die "no Muon in probe log -- CP_OPTIMIZER did not take"
+  LR0=$(grep -aoE '^step +0 \|.*lr [0-9.e+-]+' "$LOG/probe.log" | grep -oE 'lr [0-9.e+-]+' | awk '{print $2}')
+  say "PROBE OK: $N steps, peak $P, step-0 lr=$LR0 (expect ~8.0e-05 at warmup 1/$WARM)"
+  touch "$LOG/probe.ok"
+fi
+
+say "TRAINING $STEPS steps"
+env $ARCH $BASE CP_INIT_FROM="$PWD/$PARENT" CP_CKPT_NEW_RUN=1 CP_DATA_ROOT="$PWD/$DATA" \
+    CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=$STEPS \
+    CUDA_VISIBLE_DEVICES=0,1 mpirun --allow-run-as-root -np 2 "$EXE" >> "$LOG/train.log" 2>&1
+grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM mid-run"
+RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/train.log" | head -1 | grep -oE '[0-9]+')
+[ -n "$RN" ] && echo "$RN" > "$LOG/runnum"
+say "done, run ${RN:-?}, $(grep -aE '^step ' "$LOG/train.log"|tail -1|grep -oE 'loss: [0-9.]+')"
+say "=== PHASE 2 $ARM COMPLETE ==="