Hugging Face's logo

unparallelled
/
BluScriptCP

Safetensors
Model card Files
xet
 Community
BluScriptCP
File size: 6,712 Bytes
14ea6ec
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
#!/usr/bin/env bash
# =============================================================================
# ext_32k.sh -- YaRN-extend the best 16k model to 32k, then evaluate.
#
# WHY THIS RUN EXISTS
#   The length ladder ended in a tie: l3 (b3, 16-32k docs TRUNCATED to the 16k
#   window) matched l2 (b2, 8-16k docs that FIT) to the decimal at 16,360 --
#   35.0 vs 35.0, and -6.0/+2.0/+0.0 pp across 4k/8k/16k, none resolvable.
#   Conclusion: document length helps only while documents fit inside the window.
#   So widen the window first, THEN feed b3, whose documents now fit.
#
# PARENT: l2, the best 16k arm, not run7'. l2 is where the 16k ladder topped out.
#
# THE YaRN INVARIANT that cost 27pp when it was wrong:
#   YARN_SCALE * YARN_ORIG_MAXPOS must equal CP_T EXACTLY, and ORIG_MAXPOS must
#   be the length the weights were ACTUALLY trained at (16384 here, not the
#   1024 default at YARNOps.cpp:57).
#     2 * 16384 = 32768 == CP_T   OK
#
# MEMORY -- 2 GPUS IS TIGHT. cp cannot exceed the GPU count, so cp=2 puts
# T/2 = 16,384 tokens on each rank:
#     projected 130.8 GB/GPU = 91.4% of an H200's 143 GB
#   That is close enough that fragmentation may still OOM, so STEP 1 is a 3-step
#   probe and the script aborts rather than wasting hours. This does NOT fit an
#   80 GB H100 (measured: it OOMs).
#
# TOKEN BUDGET: 750 steps x 1,048,576 = 786 M tokens, matching run7's original
# 4k->16k extension exactly, so the two extension stages are comparable.
#
# GATE ON SHORT CONTEXT, NOT PERPLEXITY. The 2026-09-02 mismatch condition
# IMPROVED 16k (0 -> 14-17%) while COLLAPSING 4k from 74% to 16%. Perplexity at
# the new length would have looked fine. The eval therefore spans 1k..32k and a
# regression at 1k/4k is a failure even if 32k improves.
# =============================================================================
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=ext32k_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=${NP:-2}; GPUS=${GPUS:-0,1}
ORIG=16384; SCALE=2
STEPS=${STEPS:-750}
PARENT="$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt"
DATA=lendata/b3_16384_32768
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"

[ $(( SCALE * ORIG )) -eq $T ] || die "YaRN invariant broken: $SCALE * $ORIG != $T"
$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing: $PY -m pip install tiktoken"
[ -x "$EXE" ]    || die "$EXE not built (make libtensor SM_ARCH=90 && make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp)"
[ -f "$PARENT" ] || die "parent l2 missing: $PARENT"
[ -d "$DATA" ]   || die "b3 data missing: $DATA"

say "=== 16k -> 32k YaRN extension ==="
say "CP_T=$T YARN_SCALE=$SCALE YARN_ORIG_MAXPOS=$ORIG (product $(( SCALE*ORIG )) == $T)"
say "parent $(basename $PARENT), data $DATA, cp=$NP, $STEPS steps x 1,048,576 = 786 M tokens"

COMMON="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_MAX_LR=1e-3 CP_MIN_LR=1e-4"

# ---- STEP 1: 3-step memory probe -------------------------------------------
if [ ! -f "$LOG/probe.ok" ]; then
  say "probe: 3 steps at T=$T to confirm it fits before committing hours"
  env $ARCH $COMMON CP_REWARMUP=1 CP_REWARMUP_PEAK=0.3 CP_CKPT=0 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=65536 CP_MAX_STEPS=3 \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$LOG/probe.log" 2>&1
  if grep -qaiE 'out of memory|bad_alloc' "$LOG/probe.log"; then
    say "PROBE OOM at T=$T with cp=$NP. Projected 130.8 GB/GPU."
    say "  This needs more ranks: cp=4 -> 67.5 GB/GPU. Not runnable on $NP GPUs."
    exit 1
  fi
  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | tail -1 | sed 's/^/  probe: /' | tee -a "$M"
  say "probe peak mem: $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
  touch "$LOG/probe.ok"
else
  say "probe already passed"
fi

# ---- STEP 2: the extension --------------------------------------------------
if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
  say "extension already done -> $(cat "$LOG/ext.ckpt")"
else
  say "training $STEPS steps"
  env $ARCH $COMMON CP_REWARMUP=50 CP_REWARMUP_PEAK=0.3 \
      CP_CKPT=1 CP_CKPT_KEEP=2 CP_CKPT_FREQ=50 CP_CKPT_NEW_RUN=1 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=1048576 CP_MAX_STEPS=$STEPS \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      >> "$LOG/ext.log" 2>&1
  grep -qaiE 'out of memory|bad_alloc' "$LOG/ext.log" && die "OOM during training"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/ext.log" | tail -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || die "no run number"
  CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || die "no checkpoint written"
  bn=$(basename "$CKF"); n=$(ls checkpoints_bluscriptcp/rank_*/"$bn" 2>/dev/null | wc -l)
  [ "$n" -ge "$NP" ] || die "only $n/$NP rank shards for $bn"
  echo "$CKF" > "$LOG/ext.ckpt"
  say "ext32k -> run $RN, $bn, final loss $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
fi

# ---- STEP 3: eval 1k..32k, matched cache -----------------------------------
EXT=$(basename "$(cat "$LOG/ext.ckpt")")
EXTRA="l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt ext32k=$EXT"
say "eval: l2 (16k parent) vs ext32k, matched YARN_ORIG=$ORIG SCALE=$SCALE"
# PROMPT_LEN + MAX_NEW must pad up to <= CP_T (PAD=256): 32744+24 = 32768 exactly
for LEN in 1024 4096 16360 32744; do
  say "PROMPT_LEN=$LEN"
  GEN_DATA_ROOT=$DATA HAY="$PWD/lendata/haystack/b2val_haystack.bin" \
  PY=$PY OUT=ext32k_len_$LEN TASKS="niah_single" CKPTS="l2 ext32k" N=100 \
  PROMPT_LEN=$LEN MAX_NEW=24 T=$T DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN=$SCALE YARN_ORIG=$ORIG GPUS=$GPUS JOBS_PER_GPU=1 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN nonzero at $LEN"
  say "  -> $(sed -n 1p ext32k_len_$LEN/summary.csv 2>/dev/null) / $(sed -n 2p ext32k_len_$LEN/summary.csv 2>/dev/null)"
done
say "=== 32k EXTENSION DONE ==="
say "REGRESSION CHECK: compare 1024 and 4096 against l2. A 32k gain bought with a"
say "  short-context collapse is the 2026-09-02 mismatch failure repeating."