How to avoid GPU OOM when predicting large multimer complexes (>2500 residues) in AlphaFold 3?
2 hours agoWe are attempting to fold a 4-chain ternary complex (~2600 residues total) using the local AF3 pipeline on a single NVIDIA A100 80GB SXM. The pairformer attention matrix explodes during the diffusion head evaluation, throwing CUDA out of memory. Has anyone found working chunking or CPU-offloading flags?
Answers (2)
You can enable chunked pairformer cross-attention and memory checkpointing inside the model runner config. In your run script or wrapper, set `pairformer_chunk_size=128` and offload the MSA sequence representations after the initial embedding pass:
# Modify runner config:
model_config = dict(
pairformer_chunk_size=128,
gradient_checkpointing=True,
offload_msa_to_cpu=True,
diffusion_batch_size=1
)
predictions = model.predict(multimer_features)Note that official AlphaFold 3 from DeepMind is implemented in JAX/Haiku using JAX sharding (shard_map / pjit) across multi-GPU nodes. If running community PyTorch implementations, sequence parallelism is supported via PyTorch FSDP or DeepSpeed ZeRO-3.