unparallelled commited on
Commit
8289a8d
·
verified ·
1 Parent(s): 0405416

Upload scripts/t32_train.sh with huggingface_hub

Browse files
Files changed (1) hide show
  1. scripts/t32_train.sh +100 -0
scripts/t32_train.sh ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # =============================================================================
3
+ # 32k training at the CORRECTED scale -- the cell that has never been run.
4
+ #
5
+ # WHY SCALE=32 AND NOT 8:
6
+ # The invariant SCALE*ORIG==CP_T would force SCALE=8 at T=32,768, putting the
7
+ # far end at virtual 4,096 -- the badly-conditioned edge, and exactly the
8
+ # phase-1 mistake that made 1,200 steps look worthless.
9
+ # YARNOps.cpp treats seq_len != scale*orig_maxpos as a WARNING, not an error
10
+ # ("(R2) ... warn, don't hard-fail"), so CP_T=32,768 with SCALE=32 builds a
11
+ # cache SIZED 32,768 but SCALED for 131,072: positions land at virtual
12
+ # 0..1,023. The angle at a position is (pos/scale)*omega, independent of cache
13
+ # LENGTH, so this MATCHES the inference config we will evaluate at:
14
+ # train CP_T=32,768 x32 -> virtual 0..1,023
15
+ # infer 65,512 tok x64 -> virtual 0..1,023
16
+ # The R2 warning is EXPECTED here and is not a misconfiguration.
17
+ #
18
+ # PARENT: p1final, not base680. At x64 p1final beats base680 by +8/+11/+7pp,
19
+ # so it is the better starting point despite having been trained at x4.
20
+ #
21
+ # MEMORY: the 3-step probe peaked at 77,397 of 81,559 MB (95%). Tight. Saves
22
+ # every 200 steps so an OOM late in the run does not lose everything.
23
+ # =============================================================================
24
+ set -uo pipefail
25
+ cd "$(dirname "$0")"
26
+ ROOT="$(cd ../.. && pwd)"
27
+ export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
28
+ export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
29
+ export PATH=/usr/local/cuda/bin:$PATH
30
+ PY=/root/Adhi/BluTrain/envv/bin/python
31
+ EXE=./build/bluscriptCP_exec
32
+ LOG=t32_logs; mkdir -p "$LOG"; M="$LOG/master.log"
33
+ T=32768; NP=2; CPSZ=2; SCALE=32; ORIG=4096
34
+ STEPS=${STEPS:-1200}; WARM=100; TAIL=2; FREQ=200; KEEP=3
35
+ STABLE=$(( STEPS - WARM - TAIL ))
36
+ GB=524288
37
+ DATA=Data_Loader/Data/t32mix
38
+ CK=checkpoints_bluscriptcp/rank_0
39
+ PARENT="$CK/blumodelcp_run12_step_1200.ckpt"
40
+ say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
41
+ die(){ say "FATAL: $*"; exit 1; }
42
+
43
+ [ -x "$EXE" ] || die "$EXE missing"
44
+ S=build/objects/.cp_rope_flags
45
+ grep -q 'CP_FUSED_ROPE=1' "$S" && grep -q 'CP_ATTN_SPLIT=1' "$S" || die "binary lacks split kernels"
46
+ [ -f "$PARENT" ] || die "parent p1final missing"
47
+ [ "$(stat -c%s "$PARENT")" = 5141244514 ] || die "parent wrong size"
48
+ [ -d "$DATA" ] || die "$DATA missing"
49
+ U=$($PY - "$DATA" "$T" "$GB" <<'PY'
50
+ import glob,os,sys
51
+ d,BT,GBt=sys.argv[1],int(sys.argv[2]),int(sys.argv[3]); GA=GBt//BT; t=0
52
+ for p in sorted(glob.glob(f"{d}/train_shard_*.bin")):
53
+ n=os.path.getsize(p)//2; pos=0; b=0
54
+ while pos+BT+1<=n:
55
+ b+=1; pos+=BT
56
+ if pos+(BT+1)>n: break
57
+ t+=b
58
+ print(t//GA)
59
+ PY
60
+ )
61
+ [ "$STEPS" -le "$U" ] || die "STEPS=$STEPS > usable=$U (loader would WRAP silently)"
62
+ [ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "illegal WSD"
63
+
64
+ ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
65
+ BASE="CP_SIZE=$CPSZ CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
66
+ YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG CP_ROPE_THETA=500000 \
67
+ CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 CP_OPTIMIZER=muon \
68
+ CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=0.3 \
69
+ CP_STABLE=$STABLE CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"
70
+
71
+ say "=== 680M 32k training at the CORRECTED scale ==="
72
+ say "parent : p1final (run12 step 1200)"
73
+ say "data : $DATA b4=50 b3=30 b1=20, usable $U steps, running $STEPS"
74
+ say "YaRN : CP_T=$T SCALE=$SCALE ORIG=$ORIG -> virtual 0..$(( (T-1)/SCALE ))"
75
+ say " (R2 seq_len != scale*orig warning is EXPECTED and intended)"
76
+ say "mem : probe peaked 77,397/81,559 MB -- 95%, tight"
77
+
78
+ if [ ! -f "$LOG/probe.ok" ]; then
79
+ say "PROBE: 3 steps at the EXACT training config (grad_accum $((GB/T)))"
80
+ env $ARCH $BASE CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$PWD/$PARENT" \
81
+ CP_DATA_ROOT="$PWD/$DATA" CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=3 \
82
+ CUDA_VISIBLE_DEVICES=0,1 mpirun --allow-run-as-root -np $NP "$EXE" > "$LOG/probe.log" 2>&1
83
+ RC=$?
84
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log" && die "PROBE OOM at the real batch -- 32k will not fit"
85
+ N=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe.log"); P=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log"|tail -1)
86
+ [ "$RC" -eq 0 ] && [ "$N" -ge 1 ] && [ -n "$P" ] || { say "PROBE FAILED rc=$RC steps=$N"; tail -12 "$LOG/probe.log" >&2; exit 1; }
87
+ say "PROBE OK: $N steps, peak $P, dt $(grep -aoE 'dt: [0-9.]+ms' "$LOG/probe.log"|tail -1)"
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=0,1 mpirun --allow-run-as-root -np $NP "$EXE" >> "$LOG/train.log" 2>&1
95
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM mid-run -- checkpoints up to the last save survive"
96
+ RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/train.log" | head -1 | grep -oE '[0-9]+')
97
+ [ -n "$RN" ] || die "no run number"; echo "$RN" > "$LOG/runnum"
98
+ say "done, run $RN, $(grep -aE '^step ' "$LOG/train.log"|tail -1|grep -oE 'loss: [0-9.]+')"
99
+ say "peak mem $(grep -aoE 'mem: [0-9]+MB' "$LOG/train.log"|tail -1)"
100
+ say "=== 32K TRAINING COMPLETE ==="