[DimaMolod]

Reported by a user whose large predictions ran out of memory with unified memory on and RAM to spare.

With unified memory on, XLA's memory limit is VRAM × max(1, XLA_CLIENT_MEM_FRACTION) (xla/pjrt/gpu/gpu_helpers.cc, the same in the XLA versions pinned by jax 0.5.3 and 0.9.1). That limit holds VRAM and host spill together. structure_inference_xla_mem_fraction: auto set the fraction to host RAM / VRAM, so the whole job was capped at its host RAM, and about one GPU's worth of the requested RAM could never be used. Measured on an H100 80 GB with --mem=128G: auto (1.607) ran out at 127 GiB, only about 48 GiB of it spilled. R/V + 1 (2.607) reached 206 GiB.

Changes

• auto is now 1 + (host RAM − structure_inference_ram_bytes) / VRAM, never below 1. The margin is the backend's inference base RAM (8000 MB for AF3, 16000 MB for AF2), which stays free for the process itself. It is logged in the [unified-memory] line.
• On a MIG slice that nvidia-smi -L does not list, the job now falls back to --query-gpu instead of failing. Snakemake runs job shells with set -euo pipefail, so the pipeline's grep exit 1 used to end the job before inference. That bug is already on main.
• README and config: the fraction is described as a multiple of VRAM covering VRAM and spill together. The README no longer says XLA_PYTHON_CLIENT_MEM_FRACTION is exported; the shell unsets it, because with both names set JAX's CUDA plugin fails and JAX runs on CPU without raising.

Testing

• New dry-run tests run the rendered unified-memory block in strict mode against a fake nvidia-smi. They cover both backends, several GPUs, MIG slices listed and unlisted, failing and [N/A] readings, host RAM at or below the margin, a pinned fraction, and unified memory off. Putting back the old formula, or removing the MIG || true, makes them fail.
• The same block in the AF2 and AF3 2.9.0 images, across 4 configs × 8 nvidia-smi behaviours × 3 host-RAM values: 276/276 as expected. A de_DE locale still prints a decimal point.
• On an H100 inside both images, AlphaPulldown's import order gave XLA a limit of 198.66 GiB = 2.509 × 79.18 GiB. So the exported value is what XLA uses, even though ColabFold sets XLA_PYTHON_CLIENT_MEM_FRACTION=4.0 on import (see notes).
• Full suite: 173 passed, 1 skipped, 2 xfailed.

Notes, not changed here

• Spilled pages are not counted against the job's memory cgroup on EL8 with cgroup v1: about 126 GiB spilled while the cgroup peaked at 7.4 GiB. So --mem does not limit spill. The margin keeps the node from being overcommitted, not the job from being killed.
• ColabFold's batch.py sets TF_FORCE_UNIFIED_MEMORY=1 and XLA_PYTHON_CLIENT_MEM_FRACTION=4.0 on import if they are unset. AlphaPulldown avoids it only because run_structure_prediction.py starts the GPU backend before anything imports ColabFold. With the imports reversed, unified-memory runs fell back to CPU without an error, and AF2 with unified memory off got a 4.0 fraction. The ordering has no comment or test in AlphaPulldown.
• The AF3 image sets TF_FORCE_UNIFIED_MEMORY=1 and XLA_CLIENT_MEM_FRACTION=0.95 in its own environment, so structure_inference_unified_memory: false still uses managed memory there, with the limit rounded up to plain VRAM. It still fails fast instead of spilling, as before.
• Params and shell changes rerun finished predictions only for users without rerun-triggers: mtime, which both bundled profiles set.

Throughput review (sacct, ~145k structure_inference jobs since March 2026)

• Memory model still holds. The largest host RSS per 1000-token bin (MaxRSS of the srun step, which excludes spill) is at most half the default request for both backends. A new test pins these anchors. With the new formula XLA's limit only grows, so no job is worse off.
• Walltime: no change. Completed folds ran 12 min at the median and 84 min at p99. In July, when both kinds of request were running, jobs asking ≤180 min did not wait less than jobs asking >720 min. GPU queue time here follows availability and fair-share, not the walltime asked, and a timeout wastes the whole run, so the defaults stay.
• Retries: no change. In a sample of 400 failed logs, 37% were JAX finding no GPU (131 of 149 on one node, gpu39), which a retry on another node fixes. 29% were unsupported-flag config errors, which the parse-time flag guard now catches. 9% were XLA out of memory, which this PR targets. Of the folds killed by a SLURM memory limit, 75% finished after the ×1.1 escalation.
• Spill is slow. On an H100 a matmul loop ran at the same speed with or without unified memory while it fitted in VRAM, but about 17× slower once ~20 GiB spilled. So the AF3 image's own unified-memory default costs nothing, and the README now advises keeping AlphaFold 2 off 16 GB MIG slices, where 21 of the 35 sampled out-of-memory errors happened.

🤖 Generated with Claude Code