unparallelled commited on
Commit
24f5146
·
verified ·
1 Parent(s): 9fcfe82

Upload scripts/t128_train.sh with huggingface_hub

Browse files
Files changed (1) hide show
  1. scripts/t128_train.sh +105 -0
scripts/t128_train.sh ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env bash
2
+ # =============================================================================
3
+ # 64k extension, 8x H100, with the CORRECT configuration.
4
+ #
5
+ # THE MISTAKE THIS AVOIDS: the 32k run trained at CP_T=32,768 with SCALE=32,
6
+ # which squeezed the whole context into virtual positions 0..1,023 -- 32 tokens
7
+ # per virtual unit. The model re-conditioned to that compressed regime and lost
8
+ # the positional sensitivity that made compression useful: 90 -> 49 by step 800.
9
+ #
10
+ # HERE THE INVARIANT HOLDS: SCALE x ORIG == CP_T -> 16 x 4,096 == 65,536
11
+ # -> positions span the FULL virtual range 0..4,095
12
+ # -> 16 tokens per virtual unit, the minimum compression 64k allows
13
+ # -> this is what p1final was trained under (4 x 4,096 == 16,384) and p1final
14
+ # is the best model this project has produced.
15
+ # Evaluate afterwards at x64; compression helps at INFERENCE, not at TRAINING.
16
+ #
17
+ # LR: CP_MAX_LR / CP_MIN_LR are NO-OPS (0 occurrences in bluscriptCP.cpp).
18
+ # max_lr 6e-4 hardcoded :151, x REWARMUP_PEAK 0.3 = 1.8e-4 peak, min 6e-5.
19
+ # TAIL=200 gives a REAL cosine decay. Every prior run used TAIL=2 and never
20
+ # annealed -- the 2026-09-09 log flagged this ("all of it UNDECAYED").
21
+ #
22
+ # DATA: mix, never pure b5. b5 50 / b4 25 / b3 15 / b1 10. b5 documents are
23
+ # 65,536-131,072 tokens, at or above the window, so every window is contiguous
24
+ # single-document text. b5 ends at 36% of its epoch; both prior collapses were
25
+ # at ~80%.
26
+ # =============================================================================
27
+ set -uo pipefail
28
+ cd "$(dirname "$0")"
29
+ ROOT="$(cd ../.. && pwd)"
30
+ export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
31
+ export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
32
+ export PATH=/usr/local/cuda/bin:$PATH
33
+ PY=/root/Adhi/BluTrain/envv/bin/python
34
+ EXE=./build/bluscriptCP_exec
35
+ LOG=t128_logs; mkdir -p "$LOG"; M="$LOG/master.log"
36
+ T=131072; NP=8; SCALE=32; ORIG=4096
37
+ CPSZ=${CPSZ:-8}; DP=$(( NP / CPSZ ))
38
+ STEPS=${STEPS:-1200}; WARM=100; TAIL=200; FREQ=200; KEEP=2
39
+ STABLE=$(( STEPS - WARM - TAIL ))
40
+ GB=524288
41
+ DATA=Data_Loader/Data/mix128
42
+ CK=checkpoints_bluscriptcp/rank_0
43
+ PARENT="$CK/blumodelcp_run2_step_1200.ckpt" # t64, the 64k model
44
+ say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
45
+ die(){ say "FATAL: $*"; exit 1; }
46
+
47
+ [ -x "$EXE" ] || die "$EXE not built"
48
+ S=build/objects/.cp_rope_flags
49
+ grep -q 'CP_FUSED_ROPE=1' "$S" && grep -q 'CP_ATTN_SPLIT=1' "$S" || die "binary lacks split kernels ($(cat $S))"
50
+ [ -f "$PARENT" ] || die "parent missing"
51
+ [ "$(stat -c%s "$PARENT")" = 2423335522 ] || die "parent wrong size (expect t64 2,423,335,522)"
52
+ [ -d "$DATA" ] || die "$DATA missing"
53
+ [ $((SCALE*ORIG)) -eq $T ] || die "YaRN INVARIANT BROKEN: $SCALE*$ORIG != $T -- this is the 32k mistake"
54
+ [ $(( WARM + STABLE )) -lt $(( STEPS - 1 )) ] || die "illegal WSD"
55
+ $PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
56
+ U=$($PY - "$DATA" "$T" "$GB" <<'PY'
57
+ import glob,os,sys
58
+ d,BT,G=sys.argv[1],int(sys.argv[2]),int(sys.argv[3]); GA=G//BT; t=0
59
+ for p in sorted(glob.glob(f"{d}/train_shard_*.bin")):
60
+ n=os.path.getsize(p)//2; pos=0; b=0
61
+ while pos+BT+1<=n:
62
+ b+=1; pos+=BT
63
+ if pos+(BT+1)>n: break
64
+ t+=b
65
+ print(t//GA)
66
+ PY
67
+ )
68
+ [ "$STEPS" -le "$U" ] || die "STEPS=$STEPS > usable=$U (loader would WRAP silently)"
69
+
70
+ ARCH="CP_N_EMBD=1536 CP_N_LAYER=24 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=1"
71
+ BASE="CP_SIZE=$CPSZ CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
72
+ YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG CP_ROPE_THETA=500000 \
73
+ CP_ZERO=1 CP_ZERO_V2=1 CP_BF16=1 CP_OPTIMIZER=muon \
74
+ CP_WARMUP=$WARM CP_REWARMUP=$WARM CP_REWARMUP_PEAK=0.3 \
75
+ CP_STABLE=$STABLE CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"
76
+
77
+ say "=== 680M 131k extension, 8x H100, cp=$CPSZ dp=$DP ==="
78
+ say "parent : t64 (98 @65,512 -- the 64k model, staged ladder)"
79
+ say "YaRN : SCALE=$SCALE x ORIG=$ORIG = $((SCALE*ORIG)) == CP_T INVARIANT HOLDS"
80
+ say " virtual range 0..$(( (T-1)/SCALE )) (full), $SCALE tokens/virtual unit"
81
+ say "LR : peak 1.8e-4 (6e-4 x 0.3), warmup $WARM / stable $STABLE / DECAY $TAIL -> 6e-5"
82
+ say "data : $DATA b6=8.9 b5=39 b4=25 b3=15 b1=10, usable $U steps, running $STEPS"
83
+
84
+ if [ ! -f "$LOG/probe_cp${CPSZ}.ok" ]; then
85
+ say "PROBE cp=$CPSZ ($((T/CPSZ)) tok/rank, grad_accum $((GB/(T*DP))))"
86
+ env $ARCH $BASE CP_STABLE=0 CP_CKPT=0 CP_INIT_FROM="$PWD/$PARENT" \
87
+ CP_DATA_ROOT="$PWD/$DATA" CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=3 \
88
+ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 mpirun --allow-run-as-root -np $NP "$EXE" > "$LOG/probe_cp${CPSZ}.log" 2>&1
89
+ RC=$?
90
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/probe_cp${CPSZ}.log" && die "PROBE OOM at cp=$CPSZ"
91
+ N=$(grep -acE '^step +[0-9]+ \|' "$LOG/probe_cp${CPSZ}.log"); P=$(grep -aoE 'mem: [0-9]+MB' "$LOG/probe_cp${CPSZ}.log"|tail -1)
92
+ [ "$RC" -eq 0 ] && [ "$N" -ge 1 ] && [ -n "$P" ] || { say "PROBE FAILED rc=$RC steps=$N"; tail -14 "$LOG/probe_cp${CPSZ}.log" >&2; exit 1; }
93
+ say "PROBE OK cp=$CPSZ: peak $P of 81559 MiB, dt $(grep -aoE 'dt: [0-9.]+ms' "$LOG/probe_cp${CPSZ}.log"|tail -1)"
94
+ touch "$LOG/probe_cp${CPSZ}.ok"
95
+ fi
96
+
97
+ say "TRAINING $STEPS steps"
98
+ env $ARCH $BASE CP_INIT_FROM="$PWD/$PARENT" CP_CKPT_NEW_RUN=1 CP_DATA_ROOT="$PWD/$DATA" \
99
+ CP_B=1 CP_GLOBAL_BATCH=$GB CP_MAX_STEPS=$STEPS \
100
+ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 mpirun --allow-run-as-root -np $NP "$EXE" >> "$LOG/train.log" 2>&1
101
+ grep -qaiE 'out of memory|bad_alloc' "$LOG/train.log" && die "OOM mid-run"
102
+ RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/train.log" | head -1 | grep -oE '[0-9]+')
103
+ [ -n "$RN" ] || die "no run number"; echo "$RN" > "$LOG/runnum"
104
+ say "done, run $RN, $(grep -aE '^step ' "$LOG/train.log"|tail -1|grep -oE 'loss: [0-9.]+')"
105
+ say "=== 131K TRAINING COMPLETE ==="