diff --git a/README.md b/README.md index 9c66b13..a3a1571 100644 --- a/README.md +++ b/README.md @@ -270,10 +270,12 @@ you hit these. # Example for EMBL gpu-el8 — replace nodes with your cluster's (nothing is hard-coded): structure_inference_gpu_vram_headroom: 1.0 # <1.0 tolerates that fraction of host spill structure_inference_gpu_tiers: + - {min_vram_gb: 16, nodes: "gpu60,gpu61,gpu62,gpu63,gpu64,gpu65,gpu66,gpu67,gpu68"} # RTX PRO 4500, 16GB MIG - {min_vram_gb: 24, nodes: "gpu21,gpu22,gpu29,gpu30,gpu31,gpu32,gpu33,gpu34,gpu35,gpu36,gpu37"} - {min_vram_gb: 40, nodes: "gpu25,gpu26,gpu27,gpu28"} - {min_vram_gb: 48, nodes: "gpu40,gpu41,gpu42,gpu43,gpu44,gpu45,gpu46,gpu47,gpu48"} - {min_vram_gb: 80, nodes: "gpu38,gpu39"} + - {min_vram_gb: 96, nodes: "gpu50,gpu51,gpu52,gpu53"} # RTX PRO 6000 Blackwell ``` When set this drives `--exclude` per job and **overrides** `structure_inference_gpu_model` (the two @@ -285,9 +287,10 @@ you hit these. - **Exclude specific nodes** with `slurm_exclude_nodes` → passed verbatim to `sbatch --exclude` (e.g. `"gpu50,gpu51"`). Use it as a fallback for nodes whose GPU the container can't use — e.g. a CUDA compute capability newer than the container's bundled `ptxas` (fails `ptxas too old` / - `UNIMPLEMENTED`). The RTX PRO 6000 / Blackwell failure mode seen on EMBL `gpu50-53` was an - old/pre-Tokamax AlphaFold 3 image issue; updated AF3 v3.0.2/Tokamax images should run on those - cards, so excluding them is not proof of RTX compatibility. + `UNIMPLEMENTED`). For Blackwell (sm_120) that is purely an image-age problem: AlphaPulldown + **2.5.0 containers are verified working** on RTX PRO 6000 (`gpu50-53`) and on the RTX PRO 4500 + 16 GB MIG slices (`gpu60-68`) for both AF3 and AF2 — see [Blackwell GPUs](#blackwell-gpus). Only + exclude those nodes while you are still on a pre-2.5.0 image. `--exclude` is allowed in `slurm_extra` whereas `--constraint`/`--gres`/`--gpus` are not, so it is the supported way to drop a few nodes while keeping the rest of the partition. - **`structure_inference_max_runtime`** caps per-job wall time (minutes). Wall time scales as @@ -297,6 +300,50 @@ you hit these. +
+Blackwell GPUs (sm_120) and MIG slices + +#### Blackwell GPUs + +Blackwell cards (compute capability 12.0 / sm_120 — RTX PRO 4500 and RTX PRO 6000) need an +AlphaPulldown **2.5.0 or newer** container. Support comes from the image, not the driver: the +containers ship their own CUDA runtime as pip `nvidia-*` wheels, and pre-2.5.0 AlphaFold 3 images +bundle jaxlib 0.4.34 on CUDA 12.6, whose `ptxas` cannot target sm_120. Those images die at the very +first kernel compilation with `ptxas does not support CC 12.0` / `UNIMPLEMENTED: ptxas too old`, +before any inference runs. You cannot patch around it from outside the container — jaxlib calls its +own bundled `ptxas`, so `XLA_FLAGS=--xla_gpu_cuda_data_dir` and `PATH` have no effect, and swapping +in only a newer `ptxas` still leaves the CUDA runtime and cuDNN too old for the real kernels. + +Verified with real inference on 2.5.0 containers, with confidence scores matching the older cards: + +| GPU | Compute capability | AlphaFold 3 | AlphaFold 2 | +|-----|--------------------|-------------|-------------| +| RTX PRO 4500 Blackwell (16 GB MIG slice) | 12.0 | ✅ | ✅ | +| RTX PRO 6000 Blackwell (96 GB) | 12.0 | ✅ | ✅ | +| H100 PCIe | 9.0 | ✅ | — | +| A100 40 GB | 8.0 | ✅ | ✅ | +| RTX 3090 | 8.6 | ✅ | ✅ | + +For AF3 all three attention implementations (`triton`/Tokamax, `cudnn`, `xla`) work on Blackwell, +so no `--flash_attention_implementation` override is needed. + +#### MIG slices + +Nodes sliced with MIG (at EMBL, `gpu60-68` are RTX PRO 4500 cards split into 16 GB `1g.16gb` +instances) need no special `slurm_gres`: a plain `gpu:1` request lands on one slice and SLURM sets +`CUDA_VISIBLE_DEVICES=MIG-`. Route work to them by size with `structure_inference_gpu_tiers` +(a `min_vram_gb: 16` tier) — they suit monomers and small complexes, while larger jobs should land +on the 96 GB RTX PRO 6000 tier. + +One MIG caveat this workflow already handles: `nvidia-smi --query-gpu=memory.total` reports the +**parent card** (e.g. 32623 MiB) rather than the slice (~16 GB). Since `structure_inference_xla_mem_fraction: +auto` is `host RAM / GPU VRAM`, taking that number at face value would roughly halve the fraction and +effectively switch off host spill exactly where it is most needed. The workflow therefore reads the +slice's own profile from `nvidia-smi -L` when `CUDA_VISIBLE_DEVICES` is a MIG UUID, and falls back to +`--query-gpu` on whole cards. + +
+
Unified memory for large complexes (structure_inference_unified_memory) diff --git a/config/config.yaml b/config/config.yaml index b7be0bf..a615e7f 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -172,19 +172,28 @@ structure_inference_gpu_model: "" # gpu-el8 partition - replace nodes with your cluster's; nothing is hard-coded. # structure_inference_gpu_vram_headroom: 1.0 # <1.0 tolerates that fraction of host spill # structure_inference_gpu_tiers: +# - {min_vram_gb: 16, nodes: "gpu60,gpu61,gpu62,gpu63,gpu64,gpu65,gpu66,gpu67,gpu68"} # RTX PRO 4500 Blackwell, 16GB MIG slices # - {min_vram_gb: 24, nodes: "gpu21,gpu22,gpu29,gpu30,gpu31,gpu32,gpu33,gpu34,gpu35,gpu36,gpu37"} # RTX 3090 # - {min_vram_gb: 40, nodes: "gpu25,gpu26,gpu27,gpu28"} # A100 40GB # - {min_vram_gb: 48, nodes: "gpu40,gpu41,gpu42,gpu43,gpu44,gpu45,gpu46,gpu47,gpu48"} # L40s/A40 48GB # - {min_vram_gb: 80, nodes: "gpu38,gpu39"} # H100 PCIe 80GB -# Note: RTX PRO 6000 / Blackwell nodes (gpu50-53, 96GB at EMBL) were incompatible -# with pre-Tokamax AlphaFold 3 containers that bundled an old JAX/JAX-Triton/ptxas -# stack. Updated AF3 v3.0.2/Tokamax containers should run there; keep those nodes -# in slurm_exclude_nodes only when using old images or locally built containers that -# still fail with "ptxas too old" / UNIMPLEMENTED. +# - {min_vram_gb: 96, nodes: "gpu50,gpu51,gpu52,gpu53"} # RTX PRO 6000 Blackwell 96GB +# Blackwell (compute capability 12.0 / sm_120): RTX PRO 6000 (gpu50-53, 96GB) and +# RTX PRO 4500 (gpu60-68, sliced into 16GB MIG instances) are verified working with +# AlphaPulldown 2.5.0 containers - AF3 (jax 0.9.1/ptxas 12.9/cuDNN 9.17/Tokamax, all +# three attention implementations) and AF2 (jax 0.5.3 on ptxas 12.9/cuDNN 9.2x), with +# confidence scores matching A100/H100/3090. Only PRE-2.5.0 images fail there, at the +# first kernel compilation with "ptxas too old" / UNIMPLEMENTED; that is an image +# problem, so exclude those nodes only while you are still on an old image. +# The 16GB MIG slices suit monomers and small complexes; send bigger jobs to the +# 96GB RTX PRO 6000 cards via the tiers above. A plain "gpu:1" gres lands on a MIG +# slice correctly (SLURM sets CUDA_VISIBLE_DEVICES=MIG-...), so no gres change is +# needed to use them. # H100-SXM/H200/B200 live on the separate gpu-training partition (not routed here). # Optional: comma-separated nodes to keep structure_inference OFF, passed to sbatch # as --exclude. Useful as a fallback for GPUs the prediction container cannot use -# (e.g. a CUDA compute capability the bundled ptxas is too old for). Example: +# (e.g. a CUDA compute capability the bundled ptxas is too old for - which for +# Blackwell means a pre-2.5.0 image; see the note above). Example: # slurm_exclude_nodes: "gpu50,gpu51,gpu52,gpu53" # slurm_exclude_nodes: "" # Cap structure_inference wall time (minutes) so retry scaling (1440 * attempt) cannot diff --git a/workflow/Snakefile b/workflow/Snakefile index 29e5a7a..5ae1f41 100644 --- a/workflow/Snakefile +++ b/workflow/Snakefile @@ -679,9 +679,28 @@ rule structure_inference: unset XLA_PYTHON_CLIENT_MEM_FRACTION _aj_frac="{params.xla_mem_fraction}" if [ "$_aj_frac" = "auto" ]; then - _aj_gpu_query=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null || true) - _aj_gpu_mb="${{_aj_gpu_query%%$'\n'*}}" - _aj_gpu_mb="${{_aj_gpu_mb//[!0-9]/}}" + # On a MIG slice --query-gpu reports the *parent* card (e.g. 32623 MiB + # for a 16 GB 1g.16gb slice), which would roughly halve the fraction and + # so switch off host spill exactly where it is most needed. Read the + # slice's own profile out of `nvidia-smi -L` first, and fall back to + # --query-gpu on whole (non-MIG) cards. + _aj_gpu_mb="" + case "${{CUDA_VISIBLE_DEVICES:-}}" in + MIG-*) + _aj_mig_gb=$(nvidia-smi -L 2>/dev/null \ + | grep -F "${{CUDA_VISIBLE_DEVICES}}" \ + | sed -n 's/.*MIG [0-9]\+g\.\([0-9]\+\)gb.*/\1/p' \ + | head -1) + if [ -n "$_aj_mig_gb" ]; then + _aj_gpu_mb=$(( _aj_mig_gb * 1024 )) + fi + ;; + esac + if [ -z "$_aj_gpu_mb" ]; then + _aj_gpu_query=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null || true) + _aj_gpu_mb="${{_aj_gpu_query%%$'\n'*}}" + _aj_gpu_mb="${{_aj_gpu_mb//[!0-9]/}}" + fi if [ -n "$_aj_gpu_mb" ] && [ "$_aj_gpu_mb" -gt 0 ]; then _aj_frac=$(awk -v r={resources.mem_mb} -v g="$_aj_gpu_mb" 'BEGIN{{printf "%.3f", r/g}}') else