unparallelled commited on
Commit
bcbc864
·
verified ·
1 Parent(s): 3c22efa

Upload scripts/p2_680m.sh with huggingface_hub

Browse files
Files changed (1) hide show
  1. scripts/p2_680m.sh +97 -0
scripts/p2_680m.sh ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # =============================================================================
3
+ # Phase 2 of the 680M ladder: 16k -> 32k. Parent is phase 1's final checkpoint.
4
+ #
5
+ # YaRN: SCALE is target / ORIGINAL PRETRAIN length, always 4096 -- not
6
+ # target / previous phase. varybatch was pretrained at 4096 (config.md,
7
+ # bluscript_zero_ckpt.cpp), so phase 1 used 4x4096 and phase 2 uses 8x4096.
8
+ # Re-basing ORIG to 16384 with SCALE=2 would also satisfy SCALE*ORIG==CP_T but
9
+ # moves YaRN's ramp boundaries (low/high) and is NOT what the method prescribes.
10
+ #
11
+ # DATA: src2 holds shards disjoint from phase 1's (b3 from 4, b1 from 2), so
12
+ # this phase never re-trains on tokens phase 1 already saw.
13
+ #
14
+ # MEMORY: phase 1 measured 50,117 MB at T=16,384. Activations scale ~linearly
15
+ # in T, so this is expected to be tight on 80 GB cards. The probe decides -- if
16
+ # it OOMs, the fix is CP_SIZE=2 -> grad_accum up, or wait for a 4-GPU box.
17
+ # =============================================================================
18
+ set -uo pipefail
19
+ cd "$(dirname "$0")"
20
+ ROOT="$(cd ../.. && pwd)"
21
+ export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
22
+ export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
23
+ export PATH=/usr/local/cuda/bin:$PATH
24
+ PY=/root/Adhi/BluTrain/envv/bin/python
25
+ EXE=./build/bluscriptCP_exec
26
+ LOG=p2_logs; mkdir -p "$LOG"; M="$LOG/master.log"
27
+ T=32768; NP=2; CPSZ=2; GPUS=0,1
28
+ ORIG=4096; SCALE=8
29
+ STEPS=${STEPS:-1200}; WARM=100; TAIL=2; FREQ=200; KEEP=3
30
+ STABLE=$(( STEPS - WARM - TAIL ))
31
+ GB=524288
32
+ DATA=Data_Loader/Data/p2mix
33
+ say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
34
+ die(){ say "FATAL: $*"; exit 1; }
35
+
36
+ P1RUN=$(cat p1_logs/runnum 2>/dev/null) || die "no phase-1 runnum"
37
+ PARENT=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${P1RUN}_step_*.ckpt 2>/dev/null | head -1)
38
+ [ -n "$PARENT" ] && [ -f "$PARENT" ] || die "no phase-1 checkpoint found for run $P1RUN"
39
+ grep -qa 'PHASE 1 TRAINING COMPLETE' p1_logs/master.log 2>/dev/null \
40
+ || die "phase 1 has not reported COMPLETE -- refusing to branch off a mid-run checkpoint"
41
+
42
+ $PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
43
+ [ -x "$EXE" ] || die "$EXE not built"
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
+ [ -d "$DATA" ] || die "$DATA missing"
47
+ [ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant: $SCALE*$ORIG != $T"
48
+ U=$($PY - "$DATA" "$T" <<'PY'
49
+ import glob,os,sys
50
+ d,BT=sys.argv[1],int(sys.argv[2]); GA=524288//BT; t=0
51
+ for p in sorted(glob.glob(f"{d}/train_shard_*.bin")):
52
+ n=os.path.getsize(p)//2; pos=0; b=0
53
+ while pos+BT+1<=n:
54
+ b+=1; pos+=BT
55
+ if pos+(BT+1)>n: break
56
+ t+=b
57
+ print(t//GA)
58
+ PY
59
+ )
60
+ [ "$STEPS" -le "$U" ] || die "STEPS=$STEPS exceeds usable=$U (loader would WRAP silently)"
61
+ [ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "illegal WSD"
62
+
63
+ say "=== 680M phase 2: 16k -> 32k (2x H100, cp=$CPSZ) ==="
64
+ say "parent: $PARENT"
65
+ say "data: $DATA b4=50% b3=30% b1=20% usable $U steps, running $STEPS"
66
+ say "YaRN: SCALE=$SCALE ORIG=$ORIG -> $((SCALE*ORIG)) == CP_T=$T OK"
67
+
68
+ ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
69
+ BASE="CP_SIZE=$CPSZ CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
70
+ YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG CP_ROPE_THETA=500000 \
71
+ CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 CP_OPTIMIZER=muon \
72
+ CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=0.3 \
73
+ CP_STABLE=$STABLE CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"
74
+
75
+ if [ ! -f "$LOG/probe.ok" ]; then
76
+ say "PHASE 2a: 3-step probe ($((T/CPSZ)) tok/rank) -- memory is the open question here"
77
+ env $ARCH $BASE CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$PWD/$PARENT" \
78
+ CP_DATA_ROOT="$PWD/$DATA" CP_B=1 CP_GLOBAL_BATCH=$(( T * NP / CPSZ )) CP_MAX_STEPS=3 \
79
+ CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" > "$LOG/probe.log" 2>&1
80
+ RC=$?
81
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log" && die "PROBE OOM at T=$T -- needs >2 GPUs"
82
+ grep -qaE 'terminate called|Signal: Aborted' "$LOG/probe.log" && {
83
+ say "PROBE CRASHED:"; grep -aA1 'terminate called' "$LOG/probe.log"|grep -a 'what()'|head -1 >&2; exit 1; }
84
+ N=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe.log"); P=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log"|tail -1)
85
+ [ "$RC" -eq 0 ] && [ "$N" -ge 1 ] && [ -n "$P" ] || { say "PROBE FAILED rc=$RC steps=$N peak='$P'"; tail -12 "$LOG/probe.log" >&2; exit 1; }
86
+ say "PROBE OK: $N steps, peak $P of 81559 MiB/card"; touch "$LOG/probe.ok"
87
+ fi
88
+
89
+ say "PHASE 2b: training $STEPS steps"
90
+ env $ARCH $BASE CP_INIT_FROM="$PWD/$PARENT" CP_CKPT_NEW_RUN=1 CP_DATA_ROOT="$PWD/$DATA" \
91
+ CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=$STEPS \
92
+ CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" >> "$LOG/train.log" 2>&1
93
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM"
94
+ RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/train.log" | head -1 | grep -oE '[0-9]+')
95
+ [ -n "$RN" ] || die "no run number"; echo "$RN" > "$LOG/runnum"
96
+ say "training done, run $RN, $(grep -aE '^step ' "$LOG/train.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
97
+ say "=== PHASE 2 TRAINING COMPLETE ==="