spb/wpai0002_uqo

1#!/usr/bin/env bash2# HTR-Large (1.07 B) chain on one H100 (GPU given as $1): 12 288-token batches OOM at 71 GB, so 8 192 / 6 144 tokens.3set -uo pipefail4cd ~/wpai00025export PATH="$HOME/.local/bin:$PATH" CUDA_VISIBLE_DEVICES=${1:-0} PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True6P=data/pairs/medium_v1_haiku/validated/pairs.jsonl; S=data/splits/medium_v1/splits.jsonl; T=artifacts/tokenizers/medium_v07COR=data/parsed/medium_v0; SP0=data/splits/medium_v0/splits.jsonl8SZ=large; SUF=lambda9# Stage 1 (lr 2e-4) diverged at step ≈ 11 500 (grad norm 150–300, val CE 1.74 → 3.6); stage 2 restarts from its best10# checkpoint with a lower peak lr (INIT/LR/STEPS overridable: INIT=... LR=... STEPS=... bash scripts/lambda_chain_large.sh 0).11INIT=${INIT:-}; LR=${LR:-2e-4}; STEPS=${STEPS:-40000}; INITARG=""; [ -n "$INIT" ] && INITARG="--init-from $INIT"12.venv/bin/python scripts/pretrain_denoise.py --corpus $COR --splits $SP0 --tokenizer $T --run denoise_m1_${SZ}_${SUF} --arch m1 --size $SZ --seed 1 --max-steps $STEPS --eval-every 1000 --warmup 1000 --lr $LR --variants-per-passage 2 --max-tokens 8192 --max-batch 128 --patience 6 --device cuda --amp bf16 --num-workers 16 $INITARG > train_denoise_${SUF}.log 2>&113.venv/bin/python scripts/train.py --pairs $P --splits $S --tokenizer $T --run medium_m1_${SZ}_${SUF} --preset m1 --size $SZ --seed 1 --init-from artifacts/runs/denoise_m1_${SZ}_${SUF}/best --lr 6e-5 --max-steps 30000 --eval-every 1000 --warmup 500 --max-tokens 6144 --max-batch 128 --patience 6 --device cuda --amp bf16 --num-workers 16 > train_medium_${SUF}.log 2>&114.venv/bin/python scripts/predict.py --pairs $P --splits $S --split val --checkpoint artifacts/runs/medium_m1_${SZ}_${SUF}/best --system htr-m1-${SZ}-${SUF} --out artifacts/predictions/medium_val --limit 3000 --device cuda --batch-size 64 > predict_${SUF}_val.log 2>&115for SPL in test heldout_domain heldout_period; do16  .venv/bin/python scripts/predict.py --pairs $P --splits $S --split $SPL --checkpoint artifacts/runs/medium_m1_${SZ}_${SUF}/best --system htr-m1-${SZ}-${SUF} --out artifacts/predictions/medium_${SPL} --device cuda --batch-size 64 > predict_${SUF}_${SPL}.log 2>&117done18echo "CHAIN ${SUF} DONE" >> chain_done.txt19