BluScriptCP / scripts /b4_continue.sh
unparallelled's picture
Upload scripts/b4_continue.sh with huggingface_hub
7fc2e13 verified
Raw History Blame Contribute Delete
7.06 kB
#!/usr/bin/env bash
# =============================================================================
# b4_continue.sh -- CONTINUE ext32k through the rest of b4 at a 32k window.
#
# ext32k already trained 750 steps = 0.393 B tokens = 34% of b4's train split.
# This runs steps 750 -> 2227, the remaining 1,477 steps = 0.775 B tokens,
# finishing exactly one pass over b4 (1.168 B of 1.171 B available, no wrap).
#
# WHY CP_CKPT_RESUME AND NOT CP_INIT_FROM:
# CP_INIT_FROM loads WEIGHTS ONLY and starts at step 0, so the DataLoader would
# restart at token 0 and replay the 0.393 B ext32k already saw. CP_CKPT_RESUME
# restores start_step=750, and bluscriptCP.cpp:2084 then calls
# skip_batches(750 * 16 = 12,000) to fast-forward the stream to 0.393 B. That
# is what makes this a true continuation of the epoch rather than a replay.
#
# KNOWN DEGRADATION -- the ZeRO optimizer sidecar is missing.
# Yesterday's teardown uploaded the two .ckpt files but NOT the .zero_g1
# sidecar, and that box is gone. Per bluscriptCP.cpp:2032, "without the sidecar
# a bf16 run resumes with group 1's Adam m/v at zero -- it trains on, so the
# loss curve only looks slightly wrong." Weights and group-0 optimizer state
# restore normally; group 1's momentum rebuilds over tens of steps. Expect a
# small loss transient right after step 750 that is an artifact, not a signal.
#
# LEARNING RATE -- CP_MAX_LR / CP_MIN_LR DO NOT EXIST.
# bluscriptCP.cpp hardcodes max_lr=6e-4 (line 131) and min_lr=6e-5 (line 132),
# and the env-parsing block never reads CP_MAX_LR or CP_MIN_LR -- there is not
# even a float env helper. Every script in this project passed them as a no-op.
# So the real schedule is 6e-4 * CP_REWARMUP_PEAK(0.3) = 1.8e-4 peak, 6e-5 floor.
# ext32k's first 750 steps ran at exactly that (verified from its log: 1.8000e-04
# at steps 99..649, 6.0030e-05 at 749), so this continuation MUST keep the same
# values or it would not be a continuation. Do not "fix" the LR here.
#
# SCHEDULE: re-warm 100 steps from 750 back to 1.8e-4, flat to 2027, then a
# 200-step decay tail. decay_start = 750+100+1177 = 2027. The guard at
# bluscriptCP.cpp:1006 needs CP_WARMUP+CP_STABLE < CP_MAX_STEPS-1:
# 100+1177 = 1277 < 2226. This is a SECOND schedule on top of ext32k's
# completed one, so "ext32k@750 vs this" is not a steps-only contrast.
# =============================================================================
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/Adhi/BluTrain/envv/bin/python
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=b4cont_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=32768; NP=2; GPUS=0,1; ORIG=16384; SCALE=2
RUN=2; START=750; STEPS=2227; WARM=100; STABLE=1177; FREQ=50; KEEP=2
EVALS="1100 1650 2227" # 750 is already evaluated (the ext32k result)
DATA=lendata/b4_32768_65536
HAY="$PWD/lendata/haystack/b2val_haystack.bin"
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="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 CP_REWARMUP_ANCHOR=$START CP_STABLE=$STABLE \
CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"
$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
[ -x "$EXE" ] || die "not built"
for r in 0 1; do
[ -f "checkpoints_bluscriptcp/rank_$r/blumodelcp_run${RUN}_step_${START}.ckpt" ] || \
die "rank_$r ext32k step-$START checkpoint missing"
done
[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant"
say "=== CONTINUE ext32k through the rest of b4 ==="
say "resume run $RUN at step $START -> $STEPS ($((STEPS-START)) new steps, $(echo "scale=3; ($STEPS-$START)*524288/1000000000" | bc) B tokens)"
say "loader fast-forwards $((START*16)) batches = $(echo "scale=3; $START*524288/1000000000" | bc) B; reads to $(echo "scale=3; $STEPS*524288/1000000000" | bc) B of 1.171 B -- no wrap"
say "WARNING: .zero_g1 sidecar absent -- group 1 Adam moments start at zero (small transient)"
say "eval points: $EVALS"
run_seg(){
local TGT="$1"
[ -f "$LOG/seg_${TGT}.done" ] && { say "segment to $TGT done"; return 0; }
say "training to step $TGT"
( for i in $(seq 1 6000); do
n=$(ls checkpoints_bluscriptcp/rank_*/blumodelcp_run${RUN}_step_${TGT}.ckpt 2>/dev/null | wc -l)
if [ "$n" -ge "$NP" ]; then
sleep 5
for p in $(pgrep -f 'mpirun.*bluscriptCP_exec') $(pgrep -f 'build/bluscriptCP_exec'); do kill "$p" 2>/dev/null; done
exit 0
fi
sleep 20
done ) &
local WPID=$!
env $ARCH $BASE CP_CKPT_RESUME=$RUN CP_DATA_ROOT="$PWD/$DATA" \
CP_B=1 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=$STEPS \
CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
>> "$LOG/train.log" 2>&1
kill $WPID 2>/dev/null; wait $WPID 2>/dev/null
grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM"
local SRC="blumodelcp_run${RUN}_step_${TGT}.ckpt"
local n; n=$(ls checkpoints_bluscriptcp/rank_*/"$SRC" 2>/dev/null | wc -l)
[ "$n" -ge "$NP" ] || die "step $TGT checkpoint missing ($n/$NP)"
local SAFE="b4c${TGT}_${SRC}"
for r in 0 1; do
[ -f "checkpoints_bluscriptcp/rank_$r/$SAFE" ] || \
cp "checkpoints_bluscriptcp/rank_$r/$SRC" "checkpoints_bluscriptcp/rank_$r/$SAFE"
done
echo "$SAFE" > "$LOG/seg_${TGT}.ckpt"; touch "$LOG/seg_${TGT}.done"
say "step $TGT reached, preserved $SAFE, loss $(grep -aE '^step ' "$LOG/train.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
}
eval_seg(){
local TGT="$1"
[ -f "$LOG/eval_${TGT}.done" ] && return 0
local SAFE; SAFE=$(cat "$LOG/seg_${TGT}.ckpt")
say "EVAL @ step $TGT"
for LEN in 8192 16360 32744; do
GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=b4c_${TGT}_len_$LEN TASKS="niah_single" \
CKPTS="b4c$TGT" 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=0,1 JOBS_PER_GPU=1 BF16=1 GREEDY=1 \
CKPT_EXTRA="b4c$TGT=$SAFE" ./run_longeval.sh >> "$LOG/eval_${TGT}_${LEN}.log" 2>&1 \
|| say " WARN eval @$TGT len $LEN"
[ -s "b4c_${TGT}_len_$LEN/summary.csv" ] && \
say " step $TGT len $LEN -> $(sed -n 2p b4c_${TGT}_len_$LEN/summary.csv)"
done
touch "$LOG/eval_${TGT}.done"
}
for TGT in $EVALS; do run_seg "$TGT" || die "segment $TGT failed"; eval_seg "$TGT"; done
say "=== B4 CONTINUATION DONE ==="
say "progression vs ext32k@750 (niah_single):"
for LEN in 8192 16360 32744; do
line=" len $LEN: 750=(see ext32k)"
for TGT in $EVALS; do
v=$(sed -n 2p "b4c_${TGT}_len_$LEN/summary.csv" 2>/dev/null | cut -d, -f2)
line="$line ${TGT}=${v:-?}"
done
say "$line"
done