setup_arm.sh

#!/bin/bash
# Build the aarch64 (GH200 / Grace-Hopper) virtualenv: .venv-arm
#
# GH200 nodes are ARM (aarch64) + CUDA, so the x86_64 .venv and the pinned
# x86 requirements.txt do NOT work there. This installs a MINIMAL, arch-native
# set sufficient for DCScore training (dcscore_rl/train.py). It deliberately does
# NOT install vllm / unsloth — DCScore doesn't use them, and they're the packages
# most likely to lack working aarch64 builds.
#
# RUN THIS ON A GH200 NODE (not the login node), e.g.:
#   srun -p <arm-gpu-partition> --constraint=gh200 --gres=gpu:1 --nodes=1 --mem=64G --time=2:00:00 --pty bash
#   cd /path/to/<repo> && ./setup_arm.sh
#
# NOTE: ARM ML wheels are finicky. If a pin fails to resolve, drop the version or
# adjust the torch CUDA index (cu126 below) to match the node's CUDA module.
set -euo pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
VENV="${VENV:-$HERE/.venv-arm}"

if [ "$(uname -m)" != "aarch64" ]; then
    echo "ERROR: this is for aarch64 (GH200). You're on $(uname -m)."
    echo "       srun onto a gh200 node first (see header), then re-run."
    exit 1
fi

module load cuda/12.6 2>/dev/null || module load cuda 2>/dev/null || true

python3 -m venv "$VENV"
# shellcheck disable=SC1091
source "$VENV/bin/activate"
pip install --upgrade pip

# aarch64 CUDA build of torch (NVIDIA sbsa wheels). Match cu12x to the node's CUDA.
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126

# Pin the HF / ML libs to the SAME versions as the x86 env (requirements.txt) so
# the code behaves identically — e.g. transformers 4.57.6 still has
# AutoModelForVision2Seq (newer versions removed it -> ImportError). These all
# ship aarch64 wheels (noarch or manylinux_aarch64).
pip install \
    transformers==4.57.6 tokenizers==0.22.2 safetensors==0.7.0 \
    accelerate==1.12.0 peft==0.18.1 trl==0.26.2 datasets==4.3.0 \
    qwen-vl-utils==0.0.14 pycocoevalcap==1.2 \
    pillow==12.1.1 numpy==2.2.6 python-dotenv==1.2.1 groq \
    openai httpx

# bitsandbytes: arch-native (aarch64 CUDA), latest compatible build.
pip install bitsandbytes

python - <<'PY'
import torch
print("torch", torch.__version__, "| cuda available:", torch.cuda.is_available(),
      "|", torch.cuda.get_device_name(0) if torch.cuda.is_available() else "no gpu")
PY
echo ""
echo "ARM venv ready: $VENV"
echo "config/env.sh auto-selects it on aarch64, so just: sbatch dcscore_rl/run_train.sh"