#!/usr/bin/env bash
# =============================================================================
# b3_long_run.sh -- long b3 arm at 16k: does MORE long-context training flatten
# the retrieval-vs-length slope, or only lift the intercept?
#
#   ./b3_long_run.sh            first launch (branches from run7')
#   ./b3_long_run.sh resume <N> continue run N after a stop
#
# WHY: l1/l2/l3 all landed on an 8k->16k slope of -18 to -29 pp/doubling and
# document length only moved the intercept. l3 (b3, truncated) tied l2 exactly
# at 16k (+0.0pp), so length is exhausted as a lever at this window. Step count
# is the remaining untested one at 16k.
#
# NEAR-ZERO DECAY -- constant LR for 97.5% of the run.
#   A literally decay-free schedule is REFUSED by bluscriptCP.cpp:1006:
#     warmup + stable >= max_steps-1 -> die ("leaves no cosine decay tail")
#   The comment there gives the reason: it stops you running "right up to the
#   point the checkpoint is unusable". So the closest legal setting is a minimal
#   tail: stable = steps-200 -> flat LR for steps 100..3899, decay over the last
#   100 (2.5%). Every checkpoint you would stop at sits in the flat region, which
#   is the stop/resume property you wanted, and the run still ends decayed.
#   get_lr() computes decay_start = warmup_start + warmup + stable, and returns
#   max_lr for every step below it. Setting CP_STABLE = CP_MAX_STEPS pushes
#   decay_start to 4100 > 4000, so the cosine tail is NEVER entered: warmup 100
#   steps to 3.0e-4, then flat 3.0e-4 to the end.
#   THE TRADE: WSD's decay phase is where a large part of the final loss drop
#   happens. A checkpoint taken mid-stable UNDERSTATES what the model would do
#   after a decay. When you decide to stop for real, run a short decay from the
#   last checkpoint (CP_STABLE=0 CP_MAX_STEPS=<stop+10%>) before the headline
#   eval. Constant LR is what makes stop/resume free -- that is the point.
#
# CHECKPOINTS: CP_CKPT_FREQ=25 as requested. Each save is 17.28 GB (8.18 GB
# cp=2 shard + 0.46 GB zero_g1, x2 ranks) and 88.5 GB is free, so
# CP_CKPT_KEEP=2 -> 34.6 GB, leaving 54 GB. KEEP=1 would work but leaves no
# fallback if a save is interrupted mid-write; KEEP=4+ does not fit.
#
# RESUME CORRECTNESS: CP_REWARMUP_ANCHOR MUST be pinned to 0 on every resume.
# Its default (-1) falls back to the *current* start_step, which re-anchors a
# fresh 100-step warmup ramp at every resume -- the LR would sawtooth instead of
# staying flat. bluscriptCP.cpp:2105 documents exactly this trap.
# =============================================================================
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=b3long_logs; mkdir -p "$LOG"
M="$LOG/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

STEPS=${STEPS:-4000}
STABLE=${STABLE:-$(( STEPS - 200 ))}   # 100-step tail: the minimum the guard allows
FREQ=${FREQ:-25}
KEEP=${KEEP:-2}
NP=2; GPUS=0,1; T=16384
PARENT="$CK/blumodelcp_run6_step_750.ckpt"      # run7', same parent as l3@400
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
BASE="CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
YARN_SCALE=4 YARN_ORIG_MAXPOS=4096 \
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_STABLE=$STABLE CP_WARMUP=100 CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 \
CP_CKPT=1 CP_CKPT_FREQ=$FREQ CP_CKPT_KEEP=$KEEP"

$PY -c "import tiktoken" 2>/dev/null || die "tiktoken missing"
[ -x "$EXE" ] || die "$EXE not built"
[ -d lendata/b3_16384_32768 ] || die "b3 data missing"

MODE=${1:-start}
if [ "$MODE" = resume ]; then
  RUN=${2:-}; [ -n "$RUN" ] || die "usage: $0 resume <run_number>"
  ls $CK/blumodelcp_run${RUN}_step_*.ckpt >/dev/null 2>&1 || die "no checkpoint for run $RUN"
  say "RESUMING run $RUN (anchor pinned to 0 so warmup does not restart)"
  EXTRA="CP_CKPT_RESUME=$RUN CP_REWARMUP_ANCHOR=0"
else
  [ -f "$PARENT" ] || die "parent run7' missing: $PARENT"
  say "=== long b3 arm: START ==="
  say "parent run7' (same as l3@400, so l3@400 vs this isolates STEP COUNT)"
  EXTRA="CP_INIT_FROM=$PWD/$PARENT CP_CKPT_NEW_RUN=1"
fi

say "steps=$STEPS stable=$STABLE -> decay_start=$((100+STABLE)); flat 3.0e-4 for steps 100..$((99+STABLE)), then $((STEPS-100-STABLE))-step tail"
say "ckpt every $FREQ steps, keep $KEEP (= $(echo "scale=1; $KEEP*17.28" | bc) GB), free $(df -BG --output=avail /root | tail -1 | tr -dc '0-9')G"
say "tokens $(( STEPS*524288 )) = $(echo "scale=1; $STEPS*524288/2930000000*100" | bc)% of b3's 2.93 B"

env $ARCH $BASE $EXTRA CP_DATA_ROOT="$PWD/lendata/b3_16384_32768" \
    CP_B=1 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=$STEPS \
    CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
    >> "$LOG/l3long.log" 2>&1
rc=$?
say "trainer exited rc=$rc"
grep -qaiE 'out of memory|bad_alloc' "$LOG/l3long.log" && die "OOM"
RN=$(grep -aoE 'CP_Training_log[0-9]+' "$LOG/l3long.log" | tail -1 | grep -oE '[0-9]+')
say "run number: ${RN:-unknown}"
# prove the schedule stayed FLAT -- a cosine tail would show LR falling here
say "LR trace: $(grep -aE '^step +100 ' "$LOG/l3long.log" | grep -oE 'lr [0-9.e-]+' | head -1) at step 100"
say "          $(grep -aE '^step +2000 ' "$LOG/l3long.log" | grep -oE 'lr [0-9.e-]+' | head -1) at step 2000  (must equal step 100)"
say "          $(grep -aE '^step ' "$LOG/l3long.log" | tail -1 | grep -oE 'lr [0-9.e-]+') at the last logged step"
CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
[ -n "$CKF" ] && { echo "$CKF" > "$LOG/l3long.ckpt"; say "latest ckpt: $(basename "$CKF")"; }
say "to continue:  ./b3_long_run.sh resume ${RN:-<run>}"