Upload phase3_9b_131k.sh with huggingface_hub

phase3_9b_131k.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+# Phase 3: Train 9B indexer at 131k context on MI300X
+# CRITICAL: must train at 131k — 32k training produces indexers that degrade past 64k
+# Recipe: large subsample (32768) + reuse-k=5, SVD init, 600 steps
+# Gate: recall @ top_p=0.9 ≈ paper Table 12 (extrapolated ~50-55k at 131k)
+set -ex
+
+source ~/rtpurbo-venv/bin/activate
+export HSA_OVERRIDE_GFX_VERSION=9.4.2
+export HF_TOKEN=$(cat ~/.hf_token)
+export PYTHONPATH=~/work
+export HF_HOME=~/hf_cache
+export PYTORCH_HIP_ALLOC_CONF=expandable_segments:True
+
+cd ~/work
+
+# First: run the 9B probe (all-retrieval, needed for stage-1)
+echo "=== Running 9B all-retrieval probe ==="
+python3 -m train.surgeries._rtpurbo_probe_chunked \
+    --model Qwen/Qwen3.5-9B \
+    --all-retrieval \
+    --seq-len 8192 \
+    --calibration-size 4 \
+    --out probe_results/mi300x_9b_allretr.json \
+    --device cuda \
+    --dtype bfloat16 \
+    2>&1
+
+echo "=== PROBE DONE, starting stage-1 at 131k ==="
+
+# Stage-1: 131k context, large subsample, reuse-k=5
+python3 train/rtpurbo_stage1.py \
+    --model Qwen/Qwen3.5-9B \
+    --probe probe_results/mi300x_9b_allretr.json \
+    --out checkpoints/mi300x_9b_131k \
+    --seq-len 131072 \
+    --chunked-teacher \
+    --cq 4096 \
+    --teacher-query-sample 32768 \
+    --reuse-k 5 \
+    --teacher-query-min-pos 8192 \
+    --no-pack \
+    --svd-init \
+    --svd-batches 4 \
+    --steps 600 \
+    --lr 1e-3 \
+    --warmup 100 \
+    --weight-decay 0.01 \
+    --batch-size 1 \
+    --device cuda \
+    --dtype bfloat16 \
+    --data-config configs/rtpurbo_pg19_public.json \
+    --log-every 5 \
+    --save-every 25 \
+    --resume \
+    2>&1
+
+echo "=== STAGE-1 9B 131k DONE ==="