GPU And Clusters

g executes statistical kernels through JAX. Choose the target device in TOML:

[compute]
device = "gpu" # or "cpu"

CPU support is installed by the base runtime dependencies. GPU support requires the GPU installation flow and a scheduler allocation on a GPU node.

GPU acceleration is workload-dependent. Single-trait runs can be limited by BGEN decode, host-device transfer, or output writing rather than JAX compute.

Cold, Warm, And Steady-State Runs

The first process on a node may pay for Python import, JAX backend initialization, and JAX compilation. A later run with the same shapes and cache policy can be faster. Treat these as different measurements:

The persistent cache is always enabled and defaults to <platform temporary directory>/<user>/g-jax-cache. [compute].jax_cache_dir overrides that location. Put overrides on local or fast user-writable storage, and do not share CPU cache artifacts across nodes with different CPU features.

When GPU May Not Help

GPU execution is not automatically faster. CPU can match or beat GPU for small or I/O-bound runs when:

• the scan has one phenotype and few variants;
• BGEN decode or sample alignment dominates;
• host-to-device transfer dominates compute;
• Parquet writing dominates runtime;
• approximate-Firth candidate density is low enough that GPU work is sparse;
• the command repeatedly changes shapes and recompiles.

What To Check First

Run these checks inside the same environment and allocation where the scan will run:

hostname
uv run python -c "import jax; print(jax.devices())"
uv run g regenie --help

If JAX does not list the expected GPU, fix the node allocation, driver, CUDA, or JAX install before tuning g.

Probe the JAX Runtime

Run this check in the same environment and on the same kind of node where the scan will run:

uv run python -c "import jax; print(jax.devices())"

If JAX cannot see the expected accelerator, fix the driver, CUDA, or JAX wheel environment before measuring g performance. On a login node without NVIDIA devices, a CPU-only result can be expected even when the GPU environment is otherwise valid.

Generic SLURM GPU Job

Install and sync the checkout before submitting the job. Inside the batch script, use uv run --no-sync so the job uses the already-created .venv/ instead of trying to change the environment while the scan is running:

#!/usr/bin/env bash
#SBATCH --job-name=g-regenie2
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=04:00:00

set -euo pipefail

cd /path/to/g
export UV_CACHE_DIR="${SCRATCH:-$HOME}/.cache/uv"
export UV_LINK_MODE=copy

uv run --no-sync g regenie \
  --config /path/to/gpu.toml \
  --qt \
  --bgen /path/to/genotypes.bgen \
  --sample /path/to/genotypes.sample \
  --phenoFile /path/to/phenotypes.tsv \
  --phenoCol phenotype_continuous \
  --pred /path/to/regenie_step1_qt_pred.list \
  --out /path/to/output/g_gpu_regenie2

Adjust #SBATCH options for your site's partitions, accounts, GPU resource syntax, and memory policy.

Generic SLURM CPU Job

Large CPU scans should also run on a compute node rather than a login node:

#!/usr/bin/env bash
#SBATCH --job-name=g-regenie2-cpu
#SBATCH --cpus-per-task=16
#SBATCH --mem=64G
#SBATCH --time=04:00:00

set -euo pipefail

cd /path/to/g
export UV_CACHE_DIR="${SCRATCH:-$HOME}/.cache/uv"
export UV_LINK_MODE=copy

uv run --no-sync g regenie \
  --config /path/to/cpu.toml \
  --qt \
  --bgen /path/to/genotypes.bgen \
  --sample /path/to/genotypes.sample \
  --phenoFile /path/to/phenotypes.tsv \
  --phenoCol phenotype_continuous \
  --pred /path/to/regenie_step1_qt_pred.list \
  --out /path/to/output/g_cpu_regenie2

Cluster Notes

• Do not run GPU scans or large CPU scans on a login node.
• Keep UV_CACHE_DIR and run output on user-writable storage with enough quota.
• Use UV_LINK_MODE=copy on shared filesystems where hardlinks or reflinks are unreliable.
• Run uv sync before submitting production jobs, then use uv run --no-sync inside batch jobs.
• Use upstream regenie to create Step 1 predictions before running g Step 2.

The gauss development-server recipes and benchmark wrappers are documented in Server Gauss SLURM.

Runtime Knobs

Important runtime knobs include:

Scheduler queue depths, packed8 BGEN compatibility validation, decode tiling, packed8 selection, and Parquet grouping/compression are internal genotype/engine/output policies.

Fair performance comparisons require equivalent statistical modes. Compare score-only to score-only, and compare approximate Firth only when both tools use approximate Firth with the same fallback threshold.

For broader tuning and measurement guidance, see Performance Guide.