File size: 2,940 Bytes
567a9bd | 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 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | #!/bin/bash
# B1 — pi0.5 LoRA SFT on turning_on_radio, PyTorch + peft, with live HF uploads.
#
# Validated by job 2135: LoRA applies (trainable 758.3M / 3.682B = 20.60%, the bulk
# being the fully-trainable action expert — RLinf's recipe wraps only the VLM), 200
# steps ran clean, loss trended down, TRAIN_PT_rc=0.
#
# save_interval is 2,000 steps (~2h at the measured 3.4 s/it), NOT the config's
# 10,000: training is moving to a B200 host in ~8h, and at 10,000 the first
# checkpoint would land after ~13h — i.e. nothing to carry over. 2,000 gives four
# checkpoints inside that window.
#
# Uploads run in a *separate* process watching the output dir, so a Hub outage or
# rate limit cannot kill the training job.
#SBATCH --job-name=b1_sft
#SBATCH --partition=baram
#SBATCH --gres=gpu:baram:2,gpu:debug:2
#SBATCH --cpus-per-task=32
#SBATCH --mem=110G
#SBATCH --time=4-00:00:00
#SBATCH --output=/data2/hyeongjinkim/behavior-challenge/logs/b1_sft_%j.log
#SBATCH --error=/data2/hyeongjinkim/behavior-challenge/logs/b1_sft_%j.log
set -uo pipefail
# Invalid HF_TOKEN in the login env shadows ~/.cache/huggingface/token.
unset HF_TOKEN
ROOT=/data2/hyeongjinkim/behavior-challenge
CFG=${CFG:-pi05_b1k_pytorch_local}
EXP_NAME=${EXP_NAME:-b1_sft_$(date +%Y%m%d)}
BATCH=${BATCH:-16}
STEPS=${STEPS:-50000}
SAVE_INTERVAL=${SAVE_INTERVAL:-2000}
LOG_INTERVAL=${LOG_INTERVAL:-100}
HF_REPO=${HF_REPO:-madokalif/pi05-b1k-turning-on-radio-lora}
NGPU=$(echo "$CUDA_VISIBLE_DEVICES" | tr ',' '\n' | wc -l)
export OPENPI_LORA=${OPENPI_LORA:-1}
export OPENPI_LORA_RANK=${OPENPI_LORA_RANK:-32}
CKPT_DIR=$ROOT/repos/openpi-sft/outputs/checkpoints/$CFG/$EXP_NAME
echo "=== job $SLURM_JOB_ID $EXP_NAME ==="
echo "cfg=$CFG gpus=$NGPU batch=$BATCH steps=$STEPS save_every=$SAVE_INTERVAL"
echo "lora=$OPENPI_LORA rank=$OPENPI_LORA_RANK"
echo "hf_repo=$HF_REPO"
echo "ckpt_dir=$CKPT_DIR"
nvidia-smi --query-gpu=index,name,memory.total --format=csv
cd "$ROOT/repos/openpi-sft"
# Uploader first, so it is already watching when the first checkpoint lands.
.venv/bin/python "$ROOT/upload_ckpt_hf.py" \
--watch-dir "$CKPT_DIR" \
--repo-id "$HF_REPO" \
--interval 300 > "$ROOT/logs/hf_upload_${SLURM_JOB_ID}.log" 2>&1 &
UPLOADER_PID=$!
trap 'echo "stopping uploader $UPLOADER_PID"; kill $UPLOADER_PID 2>/dev/null' EXIT
echo "uploader pid=$UPLOADER_PID -> logs/hf_upload_${SLURM_JOB_ID}.log"
.venv/bin/torchrun --standalone --nnodes=1 --nproc_per_node="$NGPU" \
scripts/train_pytorch.py "$CFG" \
--exp_name="$EXP_NAME" \
--overwrite \
--batch-size="$BATCH" \
--num-train-steps="$STEPS" \
--save-interval="$SAVE_INTERVAL" \
--log-interval="$LOG_INTERVAL"
rc=$?
echo "TRAIN_B1_SFT_rc=$rc"
# Catch anything written after the uploader's last scan before the job exits.
echo "--- final upload sweep ---"
.venv/bin/python "$ROOT/upload_ckpt_hf.py" --watch-dir "$CKPT_DIR" --repo-id "$HF_REPO" --once 2>&1 | tail -5
exit $rc
|