unparallelled commited on
Commit
9305db1
·
verified ·
1 Parent(s): 12735da

Upload scripts/abl_phase1.sh with huggingface_hub

Browse files
Files changed (1) hide show
  1. scripts/abl_phase1.sh +99 -0
scripts/abl_phase1.sh ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # =============================================================================
3
+ # ABLATION phase 1: 4k -> 16k. Two arms, one per H200, run CONCURRENTLY.
4
+ #
5
+ # ARM=const constant 8e-5 (GPU 0)
6
+ # ARM=cosine cosine 8e-5 -> 8e-6 (GPU 1)
7
+ #
8
+ # Both use Muon + AdamW, mirroring bluscript_zero_ckpt.cpp's BLU_OPTIMIZER=muon:
9
+ # Muon drives the >=2-D weight matrices, an aux AdamW path takes everything else,
10
+ # and ONE schedule drives both because mcfg.lr == mcfg.adamw_lr (ratio 1:1, which
11
+ # set_lr preserves). Pretraining used 8e-4/6e-4 (4:3); here both arms share one
12
+ # value per the ablation spec.
13
+ #
14
+ # WHAT HAD TO CHANGE IN THE TRAINER (2026-09-18):
15
+ # CP_OPTIMIZER, CP_MAX_LR and CP_MIN_LR were read NOWHERE. max_lr/min_lr were
16
+ # hardcoded 6e-4/6e-5 and every prior CP run that passed CP_OPTIMIZER=muon
17
+ # silently trained with AdamW. All three are now wired.
18
+ #
19
+ # Warmup is 1% of the phase's steps: 12 of 1200.
20
+ # Everything else is identical to p1final so the comparison is clean.
21
+ # =============================================================================
22
+ set -uo pipefail
23
+ cd "$(dirname "$0")"
24
+ ROOT="$(cd ../.. && pwd)"
25
+ export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
26
+ export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
27
+ export PATH=/usr/local/cuda/bin:$PATH
28
+ PY=/root/Adhi/BluTrain/envv/bin/python
29
+ EXE=./build/bluscriptCP_exec
30
+
31
+ ARM=${ARM:?set ARM=const or ARM=cosine}
32
+ GPU=${GPU:?set GPU=0 or 1}
33
+ T=16384; SCALE=4; ORIG=4096
34
+ STEPS=${STEPS:-1200}
35
+ WARM=$(( STEPS / 100 )) # 1% of this phase's steps
36
+ MAXLR=8e-5
37
+ case "$ARM" in
38
+ const) MINLR=8e-5; STABLE=$(( STEPS - WARM - 2 )) ;; # flat: tail is the 2-step minimum
39
+ cosine) MINLR=8e-6; STABLE=0 ;; # warmup -> cosine all the way down
40
+ *) echo "ARM must be const or cosine"; exit 1 ;;
41
+ esac
42
+ GB=524288
43
+ DATA=Data_Loader/Data/ablmix16
44
+ CK=checkpoints_bluscriptcp/rank_0
45
+ PARENT="$CK/varybatch.ckpt"
46
+ LOG=p1_${ARM}_logs; mkdir -p "$LOG"; M="$LOG/master.log"
47
+ say(){ echo "[$(date -u +%H:%M:%S)] [$ARM] $*" | tee -a "$M" >&2; }
48
+ die(){ say "FATAL: $*"; exit 1; }
49
+
50
+ [ -x "$EXE" ] || die "$EXE not built"
51
+ S=build/objects/.cp_rope_flags
52
+ grep -q 'CP_FUSED_ROPE=1' "$S" && grep -q 'CP_ATTN_SPLIT=1' "$S" || die "binary lacks split kernels"
53
+ [ "$(stat -Lc%s "$PARENT")" = 1362674564 ] || die "parent wrong size"
54
+ [ -d "$DATA" ] || die "$DATA missing"
55
+ [ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant broken"
56
+ [ "$STABLE" -eq 0 ] || [ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "WSD guard would fire"
57
+ $PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
58
+
59
+ ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
60
+ BASE="CP_SIZE=1 CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
61
+ YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG CP_ROPE_THETA=500000 \
62
+ CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 \
63
+ CP_OPTIMIZER=muon CP_MAX_LR=$MAXLR CP_MIN_LR=$MINLR \
64
+ CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=1.0 \
65
+ CP_STABLE=$STABLE CP_CKPT=1 CP_CKPT_FREQ=200 CP_CKPT_KEEP=3 \
66
+ CP_CKPT_DIR=checkpoints_${ARM}"
67
+
68
+ say "=== phase 1: 4k -> 16k, arm=$ARM on GPU $GPU ==="
69
+ say "LR : max $MAXLR min $MINLR warmup $WARM (1% of $STEPS) stable $STABLE"
70
+ say "optim : Muon (>=2-D matrices) + aux AdamW, both at the SAME lr"
71
+ say "parent : varybatch data: $DATA steps: $STEPS"
72
+
73
+ if [ ! -f "$LOG/probe.ok" ]; then
74
+ say "PROBE 3 steps"
75
+ env $ARCH $BASE CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$PWD/$PARENT" \
76
+ CP_DATA_ROOT="$PWD/$DATA" CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=3 \
77
+ CUDA_VISIBLE_DEVICES=$GPU mpirun --allow-run-as-root -np 1 "$EXE" > "$LOG/probe.log" 2>&1
78
+ RC=$?
79
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log" && die "PROBE OOM"
80
+ N=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe.log")
81
+ P=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log"|tail -1)
82
+ [ "$RC" -eq 0 ] && [ "$N" -ge 1 ] || { say "PROBE FAILED rc=$RC steps=$N"; tail -14 "$LOG/probe.log" >&2; exit 1; }
83
+ # ASSERT the trainer actually took the new env vars -- all three were no-ops
84
+ # until today, and a silent fallback would look like a healthy run.
85
+ grep -qai 'muon' "$LOG/probe.log" || die "no Muon in probe log -- CP_OPTIMIZER did not take"
86
+ LR0=$(grep -aoE '^step +0 \|.*lr [0-9.e+-]+' "$LOG/probe.log" | grep -oE 'lr [0-9.e+-]+' | awk '{print $2}')
87
+ say "PROBE OK: $N steps, peak $P, step-0 lr=$LR0 (expect ~8.0e-05 at warmup 1/$WARM)"
88
+ touch "$LOG/probe.ok"
89
+ fi
90
+
91
+ say "TRAINING $STEPS steps"
92
+ env $ARCH $BASE CP_INIT_FROM="$PWD/$PARENT" CP_CKPT_NEW_RUN=1 CP_DATA_ROOT="$PWD/$DATA" \
93
+ CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=$STEPS \
94
+ CUDA_VISIBLE_DEVICES=$GPU mpirun --allow-run-as-root -np 1 "$EXE" >> "$LOG/train.log" 2>&1
95
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM mid-run"
96
+ RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/train.log" | head -1 | grep -oE '[0-9]+')
97
+ [ -n "$RN" ] && echo "$RN" > "$LOG/runnum"
98
+ say "done, run ${RN:-?}, $(grep -aE '^step ' "$LOG/train.log"|tail -1|grep -oE 'loss: [0-9.]+')"
99
+ say "=== PHASE 1 $ARM COMPLETE ==="