ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | # One-time setup of a remote GPU box for batched ATOM training (docs/BATCHED_TRAINING.md §Runbook). | |
| # Run from the repo root on the GPU box: bash scripts/setup_gpu_box.sh | |
| # Env: PYTHON=python3 (>= 3.10) VENV=.venv TORCH_INDEX="" (e.g. https://download.pytorch.org/whl/cu128 to | |
| # force a CUDA build; empty = PyPI default wheel, which is CUDA-enabled on Linux x86_64) | |
| # SKIP_TESTS=0 BENCH=1 | |
| # No HF / transformers / tokenizers are installed or needed: torch + numpy (+ pytest for the checks). | |
| set -euo pipefail | |
| cd "$(dirname "$0")/.." | |
| PYTHON=${PYTHON:-python3}; VENV=${VENV:-.venv}; TORCH_INDEX=${TORCH_INDEX:-}; SKIP_TESTS=${SKIP_TESTS:-0}; BENCH=${BENCH:-1} | |
| echo "== host: $(hostname) $(date -Is)" | |
| echo "== GPU"; if command -v nvidia-smi >/dev/null; then | |
| nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader | |
| nvidia-smi | sed -n '1,4p' | grep -i "cuda version" || true | |
| else echo "nvidia-smi NOT found: no NVIDIA driver -> CUDA training impossible"; exit 1; fi | |
| echo "== disk (repo fs)"; df -h . | tail -1 | |
| echo "== CPU/RAM"; nproc; free -g | sed -n '1,2p' | |
| $PYTHON - <<'P' | |
| import sys | |
| assert sys.version_info >= (3, 10), f"python >= 3.10 needed, got {sys.version}" | |
| print("python", sys.version.split()[0]) | |
| P | |
| [ -x "$VENV/bin/python" ] || $PYTHON -m venv "$VENV" | |
| "$VENV/bin/python" -m pip install -q --upgrade pip | |
| if [ -n "$TORCH_INDEX" ]; then "$VENV/bin/pip" install -q torch --index-url "$TORCH_INDEX"; else "$VENV/bin/pip" install -q torch; fi | |
| "$VENV/bin/pip" install -q numpy pytest | |
| "$VENV/bin/python" - <<'P' | |
| import torch | |
| print("torch", torch.__version__, "cuda build", torch.version.cuda, "cuda available", torch.cuda.is_available()) | |
| assert torch.cuda.is_available(), "torch cannot see the GPU: set TORCH_INDEX to a CUDA wheel matching the driver" | |
| for i in range(torch.cuda.device_count()): | |
| p = torch.cuda.get_device_properties(i) | |
| print(f" gpu{i}: {p.name} {p.total_memory/2**30:.1f} GiB sm_{p.major}{p.minor}") | |
| x = torch.randn(1024, 1024, device="cuda"); print(" matmul ok", float((x @ x).sum()) == float((x @ x).sum())) | |
| P | |
| if [ "$SKIP_TESTS" != "1" ]; then | |
| echo "== batched equivalence tests (CPU + CUDA)" | |
| PYTHONPATH=. "$VENV/bin/python" -m pytest -q test/test_batched_train.py | |
| fi | |
| if [ "$BENCH" = "1" ]; then | |
| echo "== short GPU benchmark (bytes/s vs B)" | |
| mkdir -p docs/artifacts/batched_training | |
| PYTHONPATH=. "$VENV/bin/python" tools/bench_batched.py --device cuda --threads 4 --nice 0 --k 128 --windows 2 \ | |
| --batch 64 256 1024 --out docs/artifacts/batched_training/bench_gpu_$(hostname).json | |
| fi | |
| echo "== setup OK" | |