File size: 1,164 Bytes
74d1994 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | #!/bin/bash
# RL GRPO for the gemma-4-26B-A4B NLA (self-contained trainer, bf16, GPU0).
# reward = -log reconstruction MSE; --train-critic REQUIRED (else advantages collapse).
set -u
set -a; source /root/.nla_env; set +a
export PYTHONUNBUFFERED=1
export WANDB_ENTITY=MATS10-CS-JB
export WANDB_RUN_GROUP=g4a
export TOKENIZERS_PARALLELISM=false
DATA=${DATA:-/data/gemma4_nla/or}
CK=/ckpts/gemma4_nla
BASE=google/gemma-4-26B-A4B-it
AV=${AV:-$CK/av/iter_0001000}
AR=${AR:-$CK/ar/iter_0001000}
CUDA_VISIBLE_DEVICES=${GPU:-0} python3 -m nla.train_rl_self_contained \
--av-ckpt "$AV" --ar-ckpt "$AR" \
--base-ckpt "$BASE" --quant none \
--rl-parquet $DATA/rl_shuf.parquet --sidecar $DATA/rl_shuf.parquet \
--save-dir $CK/rl \
--num-steps 500 --batch-prompts 16 --group-size 16 \
--max-new-tokens 150 --temperature 1.0 \
--lr 1e-5 --kl-beta 0.01 --clip-eps 0.2 \
--train-critic --critic-lr 5e-5 \
--logp-micro-batch 2 --max-rows 30000 \
--save-every 20 --eval-every 20 --eval-n-prompts 20 --eval-skip-rows 35000 \
--eval-n-hallucination 0 --eval-n-karvonen 0 \
--max-grad-norm 1.0 \
--wandb-project cot-oracle --wandb-name g4a_rl_grpo --seed 0
|