Blame for scripts/ext32k_b4.sh

#!/usr/bin/env bash
# =============================================================================
# ext32k_b4.sh -- YaRN extension 16k -> 32k on b4, then evaluate 1k..32k.
#
# WHY b4 AND NOT b3, at a 32,768 window:
#   b3 documents are 16,384-32,768 tokens -- mean 21,224 = 65% of the window, and
#   the MINIMUM is 16,384. So every b3 sample carries an EOT boundary plus part of
#   a second document, and no dependency can span 32k. That is the same defect as
#   the original flux corpus (median 826 tok), one scale up.
#   b4 documents are 32,768-65,536 -- the minimum EQUALS the window, so every
#   sample is one contiguous document and a 32k-spanning dependency is possible.
#   Truncating a 44k document to 32k is fine; a straddled window is not.
#
# YaRN INVARIANT (the bug that cost 27pp at 16k):
#   SCALE * ORIG_MAXPOS must equal seq_len EXACTLY, and ORIG_MAXPOS must be the
#   length the parent was actually TRAINED at.  2 * 16384 = 32768. The default
#   ORIG_MAXPOS is 1024 (YARNOps.cpp:57) and would silently be wrong.
#
# PARENT: l2 = the best 16k arm (b2 docs, which fit the 16k window). Extending
#   from the strongest 16k model, not from run7'.
#
# GLOBAL BATCH held at 524,288 (B=1, T=32768, grad_accum=16) so the optimizer
#   trajectory matches every 16k arm. Only the window changed.
#
# MEMORY: 2x H200 = 143 GB. cp=2 puts 16,384 tokens on each rank -> ~130.8 GB,
#   91% utilisation. That is close enough that fragmentation could still OOM, so
#   PHASE 1 is a 3-step probe and the full run only starts if it survives.
# =============================================================================
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/envv/bin/python3
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=2; GPUS=0,1
ORIG=16384; SCALE=2                       # 2 * 16384 == 32768 exactly
STEPS=${STEPS:-750}
STABLE=${STABLE:-$(( STEPS - 200 ))}      # WSD: flat LR, minimal legal decay tail
FREQ=${FREQ:-50}
KEEP=${KEEP:-1}                           # disk is tight on this box (~45 GB free)
DATA=lendata/b4_32768_65536
PARENT="$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt"
HAY="$PWD/lendata/haystack/b2val_haystack.bin"
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"

$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
[ -x "$EXE" ]    || die "$EXE not built"
[ -f "$PARENT" ] || die "parent l2 missing: $PARENT"
[ -f "$HAY" ]    || die "haystack missing"
[ -d "$DATA" ]   || die "$DATA missing"
ntr=$(ls "$DATA"/*train*.bin 2>/dev/null | wc -l); nva=$(ls "$DATA"/*val*.bin 2>/dev/null | wc -l)
[ "$ntr" -ge 1 ] && [ "$nva" -ge 1 ] || die "need >=1 train and >=1 val shard in $DATA (have $ntr/$nva)"
have=$(( $(du -cb "$DATA"/*train*.bin | tail -1 | cut -f1) / 2 ))
need=$(( STEPS*524288 ))
say "=== 16k -> 32k YaRN extension on b4 ==="
say "$ntr train shards = $(echo "scale=3; $have/1000000000" | bc) B tokens; need $(echo "scale=3; $need/1000000000" | bc) B for $STEPS steps"
[ "$have" -lt "$need" ] && say "WARNING: data repeats $(echo "scale=2; $need/$have" | bc)x"
say "YaRN: SCALE=$SCALE ORIG_MAXPOS=$ORIG -> $((SCALE*ORIG)) (must equal CP_T=$T)"
[ $((SCALE*ORIG)) -eq $T ] || die "YaRN invariant violated"

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 CP_WARMUP=100 CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3"

# ---- PHASE 1: does 32k fit at cp=2 on 143 GB? -----------------------------
if [ ! -f "$LOG/probe.ok" ]; then
  say "PHASE 1: 3-step memory probe (projected 130.8 GB of 143 GB = 91%)"
  env $ARCH $COMMON CP_STABLE=0 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 -- 32k does not fit at cp=2 on this box."
    say "  32k needs cp=4 (67.5 GB/GPU). Options: a 4-GPU box, or stop here."
    grep -aoE 'cudaMallocAsync\(size=[0-9]+\) failed' "$LOG/probe.log" | head -2 | sed 's/^/    /' >&2
    exit 1
  fi
  grep -aE '^step +[0-9]+ \|' "$LOG/probe.log" | tail -1 | sed 's/^/  /' >&2
  say "PROBE OK: peak $(grep -aoE 'mem: [0-9]+MB' "$LOG/probe.log" | tail -1)"
  touch "$LOG/probe.ok"
fi

# ---- PHASE 2: the extension ------------------------------------------------
if [ -f "$LOG/ext.ckpt" ] && [ -f "$(cat "$LOG/ext.ckpt")" ]; then
  say "extension already trained -> $(cat "$LOG/ext.ckpt")"
else
  say "PHASE 2: $STEPS steps, stable=$STABLE, ckpt every $FREQ keep $KEEP"
  env $ARCH $COMMON CP_STABLE=$STABLE \
      CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP CP_CKPT_NEW_RUN=1 \
      CP_INIT_FROM="$PWD/$PARENT" CP_DATA_ROOT="$PWD/$DATA" \
      CP_B=1 CP_GLOBAL_BATCH=524288 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 the full run"
  RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/ext.log" | head -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 "run $RN wrote NO checkpoint"
  n=$(ls checkpoints_bluscriptcp/rank_*/"$(basename "$CKF")" 2>/dev/null | wc -l)
  [ "$n" -ge "$NP" ] || die "only $n/$NP rank shards"
  echo "$CKF" > "$LOG/ext.ckpt"
  say "ext32k -> run $RN, $(basename "$CKF"), loss $(grep -aE '^step ' "$LOG/ext.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
  say "  peak mem $(grep -aoE 'mem: [0-9]+MB' "$LOG/ext.log" | tail -1)"
fi

# ---- PHASE 3: evaluate, including the new 32k point ------------------------
EXT=$(basename "$(cat "$LOG/ext.ckpt")")
EXTRA="ext32k=$EXT"
[ -f "$CK/lenbucket_l2_blumodelcp_run2_step_400.ckpt" ] && \
  EXTRA="$EXTRA l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt"
MODELS=$(echo $EXTRA | tr ' ' '\n' | cut -d= -f1 | tr '\n' ' ')
say "PHASE 3: eval models $MODELS"
# 32744 + MAX_NEW 24 = 32768 exactly; 32760 would pad to 33024 and abort.
for LEN in 4096 8192 16360 32744; do
  say "eval PROMPT_LEN=$LEN"
  GEN_DATA_ROOT=$DATA HAY="$HAY" PY=$PY OUT=ext32k_len_$LEN TASKS="niah_single" \
  CKPTS="$MODELS" N=100 PROMPT_LEN=$LEN MAX_NEW=24 T=$T \
  DEPTHS="0.0,0.25,0.5,0.75,1.0" YARN_ORIG=$ORIG YARN=$SCALE \
  GPUS=0,1 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"
  [ -s "ext32k_len_$LEN/summary.csv" ] && say "  -> $(sed -n 1p ext32k_len_$LEN/summary.csv) / $(sed -n 2p ext32k_len_$LEN/summary.csv)"
done
say "=== EXT32K DONE ==="