#!/usr/bin/env bash
# Reach probe for the 1.06B proxy (36L, d1536, KV4, hd128, untied).
#
# The IN2 test needs the base to be SUFFICIENT, not good:
#   (1) denom = nll_absent - nll_floor comfortably above needle_score's 0.5-nat
#       withholding gate, and
#   (2) a reach boundary INSIDE the window, so a failing region remains to improve.
# Reference points: run100 (120M) denom 8.84 / reach 3045 @ T=4096; the 600M at
# step 500 gave denom 10.73 / reach >=3061.
#
# Loads via CP_INIT_FROM (weights only): the base was trained with ZeRO-1 + Muon on
# 4 GPUs, and this probe is a forward-only eval on 1 GPU with no ZeRO -- a normal
# resume would try to parse ShardedMuon's optimizer state as nn::AdamW's.
set -euo 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

CKPT=${CKPT:-$PWD/checkpoints_bluscriptcp/blumodelcp_run8_step_1000.ckpt}
OUT=${OUT:-probe_run8}
A1B="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
DATA=${DATA:-$PWD/Data_Loader/Data}

MODEL_KEY=1b_run8 INIT_FROM="$CKPT" \
MODEL_ARCH="$A1B CP_ATTN_MODE=ring CP_ATTN_FUSION=split" \
HEAD_DIM=128 FACTORS="${FACTORS:-1}" CACHES="${CACHES:-ntk}" YARN_SCALE="${YARN_SCALE:-}" \
BLOCKS=${BLOCKS:-3} NEEDLE_LEN=16 SEED=0 GPUS=${GPUS:-0} \
FILLER=${FILLER:-Data_Loader/Data/flux_val_000000.bin} DATA_ROOT="$DATA" \
NEEDLE_DIR=$OUT/needles CSV_DIR=$OUT/csv STATE_DIR=$OUT/state RESULTS=$OUT/reach.csv \
  bash Tests/bluscriptcp/needle_sweep.sh 2>&1 | tail -14