Upload scripts/abl_phase1.sh with huggingface_hub

5bacf6f verified about 2 months ago

6.05 kB

#!/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}
# OPT=muon (default, reproduces the 2026-09-18 arms) or OPT=adamw.
# adamw is the CONTROL that separates the two variables yesterday's run confounded:
#   p1final    1.8e-4 WSD  + AdamW
#   const/muon 8e-5 const  + Muon      <- differs from p1final in BOTH
#   const/adamw 8e-5 const + AdamW     <- differs from each in exactly ONE
OPT=${OPT:-muon}
case "$OPT" in muon|adamw) ;; *) echo "OPT must be muon or adamw"; exit 1 ;; esac
T=16384; SCALE=4; ORIG=4096
STEPS=${STEPS:-1200}
WARM=$(( STEPS / 100 ))                 # 1% of this phase's steps
MAXLR=8e-5
case "$ARM" in
  const)  MINLR=8e-5; STABLE=$(( STEPS - WARM - 2 )) ;;   # flat: tail is the 2-step minimum
  cosine) MINLR=8e-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/ablmix16
CK=checkpoints_bluscriptcp/rank_0
PARENT="$CK/varybatch.ckpt"
LOG=p1_${ARM}_${OPT}_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"
[ "$(stat -Lc%s "$PARENT")" = 1362674564 ] || die "parent wrong size"
[ -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"

OPTENV=""; [ "$OPT" = muon ] && OPTENV="CP_OPTIMIZER=muon"
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=1 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 \
${OPTENV} 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}_${OPT}"

say "=== phase 1: 4k -> 16k, arm=$ARM opt=$OPT on GPU $GPU ==="
say "LR      : max $MAXLR  min $MINLR  warmup $WARM (1% of $STEPS)  stable $STABLE"
say "optim   : $([ "$OPT" = muon ] && echo "Muon (>=2-D matrices) + aux AdamW, same lr" || echo "AdamW only (ZeRO-1 v2)")"
say "parent  : varybatch   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=$GPU mpirun --allow-run-as-root -np 1 "$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.
  # Assert the CONSTRUCTED optimizer, in both directions. The [ZeRO-1] line
  # used to print configured state, so a Muon run and an AdamW run were
  # indistinguishable in the log -- which is how CP_OPTIMIZER stayed a no-op
  # across every run before 2026-09-18.
  WANT=$([ "$OPT" = muon ] && echo "optimizer=MUON+adamw-aux" || echo "optimizer=adamw")
  grep -qa "$WANT" "$LOG/probe.log" || die "probe built the WRONG optimizer -- wanted '$WANT', log says '$(grep -ao "optimizer=[A-Za-z+-]*" "$LOG/probe.log" | head -1)'"
  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=$GPU mpirun --allow-run-as-root -np 1 "$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 1 $ARM COMPLETE ==="

	#!/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}
	# OPT=muon (default, reproduces the 2026-09-18 arms) or OPT=adamw.
	# adamw is the CONTROL that separates the two variables yesterday's run confounded:
	# p1final 1.8e-4 WSD + AdamW
	# const/muon 8e-5 const + Muon <- differs from p1final in BOTH
	# const/adamw 8e-5 const + AdamW <- differs from each in exactly ONE
	OPT=${OPT:-muon}
	case "$OPT" in muon\|adamw) ;; *) echo "OPT must be muon or adamw"; exit 1 ;; esac
	T=16384; SCALE=4; ORIG=4096
	STEPS=${STEPS:-1200}
	WARM=$(( STEPS / 100 )) # 1% of this phase's steps
	MAXLR=8e-5
	case "$ARM" in
	const) MINLR=8e-5; STABLE=$(( STEPS - WARM - 2 )) ;; # flat: tail is the 2-step minimum
	cosine) MINLR=8e-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/ablmix16
	CK=checkpoints_bluscriptcp/rank_0
	PARENT="$CK/varybatch.ckpt"
	LOG=p1_${ARM}_${OPT}_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"
	[ "$(stat -Lc%s "$PARENT")" = 1362674564 ] \|\| die "parent wrong size"
	[ -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"

	OPTENV=""; [ "$OPT" = muon ] && OPTENV="CP_OPTIMIZER=muon"
	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=1 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 \
	${OPTENV} 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}_${OPT}"

	say "=== phase 1: 4k -> 16k, arm=$ARM opt=$OPT on GPU $GPU ==="
	say "LR : max $MAXLR min $MINLR warmup $WARM (1% of $STEPS) stable $STABLE"
	say "optim : $([ "$OPT" = muon ] && echo "Muon (>=2-D matrices) + aux AdamW, same lr" \|\| echo "AdamW only (ZeRO-1 v2)")"
	say "parent : varybatch 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=$GPU mpirun --allow-run-as-root -np 1 "$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.
	# Assert the CONSTRUCTED optimizer, in both directions. The [ZeRO-1] line
	# used to print configured state, so a Muon run and an AdamW run were
	# indistinguishable in the log -- which is how CP_OPTIMIZER stayed a no-op
	# across every run before 2026-09-18.
	WANT=$([ "$OPT" = muon ] && echo "optimizer=MUON+adamw-aux" \|\| echo "optimizer=adamw")
	grep -qa "$WANT" "$LOG/probe.log" \|\| die "probe built the WRONG optimizer -- wanted '$WANT', log says '$(grep -ao "optimizer=[A-Za-z+-]*" "$LOG/probe.log" \| head -1)'"
	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=$GPU mpirun --allow-run-as-root -np 1 "$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 1 $ARM COMPLETE ==="