gh200 setup for pytorch training (october 2026)

python3 -m pip install --user uv     # uv 0.12.x, aarch64-unknown-linux-gnu
export PATH=$HOME/.local/bin:/usr/local/cuda/bin:$PATH
uv venv --python 3.12 ~/venvs/train  # uv downloads its own cpython 3.12

PY=~/venvs/train/bin/python
uv pip install --python "$PY" --torch-backend=auto "torch==2.13.0"

uv pip install --python "$PY" --torch-backend=auto numpy tokenizers pyarrow "transformers==5.17.0" pyyaml

uv pip install --python "$PY" setuptools wheel ninja
git clone https://github.com/tgale96/grouped_gemm.git ~/src/grouped_gemm
cd ~/src/grouped_gemm && git checkout f1429a3 && git submodule update --init --recursive
GROUPED_GEMM_CUTLASS=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=24 \
  uv pip install --python "$PY" --no-build-isolation --no-deps ~/src/grouped_gemm

"""gh200 env checks: device, bf16 matmul, torch.compile, varlen attention, grouped gemm."""
import sys, time, torch, torch.nn.functional as F

ok = True
def rep(name, good, msg=""):
    global ok
    ok &= bool(good)
    print(f"{'PASS' if good else 'FAIL'} {name} {msg}", flush=True)

if not torch.cuda.is_available():
    rep("device", False, f"torch {torch.__version__}: no cuda")
    sys.exit(1)
p = torch.cuda.get_device_properties(0)
rep("device", True, f"{p.name} sm{p.major}{p.minor} {p.total_memory / 2**30:.1f} GiB "
                    f"torch {torch.__version__} cuda {torch.version.cuda}")
rep("bf16", torch.cuda.is_bf16_supported())

# bf16 matmul throughput
a = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16); b = torch.randn_like(a)
for _ in range(3): a @ b
torch.cuda.synchronize(); t = time.time()
for _ in range(20): a @ b
torch.cuda.synchronize()
rep("bf16 matmul", True, f"{2 * 8192**3 * 20 / (time.time() - t) / 1e12:.0f} TFLOPS")

# torch.compile (inductor + triton on aarch64)
f = torch.compile(lambda x: F.gelu(x) * 2 + x.sin())
x = torch.randn(4096, 4096, device="cuda")
rep("torch.compile", torch.allclose(f(x), F.gelu(x) * 2 + x.sin(), atol=1e-5))

# varlen attention (document masking) vs per-document sdpa
from torch.nn.attention.varlen import varlen_attn
q = torch.randn(4096, 9, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True)
cu = torch.tensor([0, 1000, 4096], device="cuda", dtype=torch.int32)
o = varlen_attn(q, q, q, cu, cu, 4096, 4096, window_size=(-1, 0))
def doc(s, e):
    t = q[s:e].transpose(0, 1)[None]
    return F.scaled_dot_product_attention(t, t, t, is_causal=True)[0].transpose(0, 1)
ref = torch.cat([doc(0, 1000), doc(1000, 4096)])
o.float().pow(2).mean().backward()
diff = (o - ref).abs().max().item()
rep("varlen_attn", diff < 2e-2, f"max|diff| vs per-doc sdpa {diff:.2e}")

# moe grouped gemm vs per-expert matmuls
E, d, h = 6, 576, 1536
sizes = torch.tensor([3000, 0, 5000, 1234, 7000, 150])
offs = torch.cumsum(sizes, 0).cuda().to(torch.int32)
xx = torch.randn(int(sizes.sum()), d, device="cuda", dtype=torch.bfloat16, requires_grad=True)
w = torch.randn(E, h, d, device="cuda", dtype=torch.bfloat16, requires_grad=True)
r = torch.cat([c @ w[i].t() for i, c in enumerate(torch.split(xx, sizes.tolist()))])
y = F.grouped_mm(xx, w.transpose(-2, -1), offs=offs)
y.float().pow(2).mean().backward()
rep("F.grouped_mm", (y - r).abs().max().item() == 0)
try:
    import grouped_gemm
    y2 = grouped_gemm.ops.gmm(xx, w, sizes.cuda(), trans_b=True)
    y2.float().pow(2).mean().backward()
    rep("grouped_gemm", (y2 - r).abs().max().item() == 0, grouped_gemm.__file__)
except ImportError:
    print("SKIP grouped_gemm (not installed)")

sys.exit(0 if ok else 1)

cd /tmp && ~/venvs/train/bin/python check_gh200.py

#!/usr/bin/env bash
# usage: bash setup_gh200.sh [VENV=~/venvs/train]   (GROUPED_GEMM=0 to skip the moe kernel)
set -euo pipefail
VENV=${1:-$HOME/venvs/train}; SRC=${SRC:-$HOME/src}
export PATH=$HOME/.local/bin:/usr/local/cuda/bin:$PATH CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
command -v uv >/dev/null || python3 -m pip install --user -q uv
[ -x "$VENV/bin/python" ] || uv venv -q --python 3.12 "$VENV"
PY=$VENV/bin/python
uv pip install -q --python "$PY" --torch-backend=auto "torch==2.13.0"
uv pip install -q --python "$PY" --torch-backend=auto numpy tokenizers pyarrow "transformers==5.17.0" pyyaml
if [ "${GROUPED_GEMM:-1}" = 1 ] && ! (cd /tmp && "$PY" -c "import grouped_gemm" 2>/dev/null); then
  uv pip install -q --python "$PY" setuptools wheel ninja
  [ -d "$SRC/grouped_gemm" ] || git clone -q https://github.com/tgale96/grouped_gemm.git "$SRC/grouped_gemm"
  (cd "$SRC/grouped_gemm" && git checkout -q f1429a3 && git submodule update -q --init --recursive)
  GROUPED_GEMM_CUTLASS=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=$(( $(nproc) / 2 )) \
    uv pip install --python "$PY" --no-build-isolation --no-deps "$SRC/grouped_gemm"
fi
uv pip freeze --python "$PY" > "$VENV.freeze.txt"
echo "done: $VENV ($(wc -l < "$VENV.freeze.txt") packages)"