AlphaFold3 的性能表现

跟我一起来探索 AlphaFold3 的世界吧！

数据管道

数据管道（即基因序列搜索和模板搜索）的运行时间可能会因输入大小、发现的同源序列数量以及可用硬件（特别是磁盘速度会影响基因搜索速度）而有显著变化。如果您想提高性能，建议提高磁盘速度（例如，通过使用基于RAM的文件系统），或增加可用CPU核心并增加更多并行化。此外，请注意，对于具有深度MSA的序列，Jackhmmer或Nhmmer可能需要超过推荐的64 GB RAM的大量RAM。

模型推理

AlphaFold 3论文的补充信息中的表8提供了在 16卡NVIDIA A100 上配置运行AlphaFold 3的编译无关推理时间，每个设备有40 GB的内存。相比之下，这个仓库支持在单卡NVIDIA A100上运行AlphaFold 3，具有80 GB的内存，并优化了配置以最大化吞吐量。

意思就是：单卡80G A100比16卡 40G A100优化的更好。

我们在下面的表格中比较了这两种设置的编译无关推理时间，使用GPU秒（即使用16个A100时乘以16）。这个仓库中的设置在所有令牌大小上都更高效（至少快2倍），表明它适用于高通量应用。

分阶段运行管道

run_alphafold.py脚本可以分阶段执行，以优化资源利用率。这可能很有用：

1. 将CPU专用数据管道与模型推理（需要GPU）分开，以优化成本和资源使用。
2. 缓存MSA/模板搜索的结果，然后在不同的种子或不同特征变化（例如，配体）的多次推理中重用增强的JSON文件。

仅数据管道

使用–norun_inference启动run_alphafold.py，生成多重序列比对（MSAs）和模板，而不运行特征化和模型推理。这个阶段在运行时间、CPU和RAM使用方面可能相当昂贵。输出将是增强了MSAs和模板的JSON文件，然后可以直接用作运行推理的输入。

仅特征化和模型推理

使用–norun_data_pipeline启动run_alphafold.py，跳过数据管道，只运行特征化和模型推理。这个阶段需要输入JSON文件包含预先计算的MSAs和模板。

加速器硬件要求

我们官方支持以下配置，并对其进行了广泛的测试，以确保数值精度和吞吐量效率：

• 单卡 NVIDIA A100 (80 GB)
• 单卡 NVIDIA H100 (80 GB)

1.  NVIDIA A100 (40 GB)

AlphaFold 3可以在单个NVIDIA A100 (40 GB)上运行，需要进行以下配置更改：

启用统一内存。

在model_config.py中调整pair_transition_shard_spec：

pair_transition_shard_spec: Sequence[_Shape2DType] = (

(2048, None),

(3072, 1024),

(None, 512),

)

虽然数值准确，但这种配置的吞吐量将低于NVIDIA A100 (80 GB)的设置，因为可用内存较少。

2. NVIDIA V100 (16 GB)

虽然您可以在单个NVIDIA V100上使用–flash_attention_implementation=xla标志在run_alphafold.py上运行AlphaFold 3，处理多达1,280个令牌的序列，但这种配置尚未经过数值精度或吞吐量效率的测试，因此请谨慎操作。

其他标志

1. 编译时间问题的XLA标志解决方案

为了解决已知的XLA问题，导致编译时间大大增加，必须设置以下环境变量（默认在提供的Dockerfile中设置）。

ENV XLA_FLAGS=”–xla_gpu_enable_triton_gemm=false”

2. GPU内存

以下环境变量（默认在Dockerfile中设置）允许折叠单个大小高达5,120个 token 的输入，在单个具有80 GB内存的A100上：

ENV XLA_PYTHON_CLIENT_PREALLOCATE=true

ENV XLA_CLIENT_MEM_FRACTION=0.95

统一内存

如果您想在内存较少的GPU上运行AlphaFold 3（例如，具有40 GB内存的A100），我们建议启用统一内存。启用统一内存允许程序在没有足够空间时将GPU内存溢出到主机内存。这可以防止OOM，但代价是程序变慢，因为访问主机内存而不是设备内存。

您可以通过在Dockerfile中设置以下环境变量来启用统一内存：

ENV XLA_PYTHON_CLIENT_PREALLOCATE=false

ENV TF_FORCE_UNIFIED_MEMORY=true

ENV XLA_CLIENT_MEM_FRACTION=3.2

3. JAX持久化编译缓存

您可能还想使用JAX持久化编译缓存，以避免在运行之间不必要的模型重新编译。您可以在run_alphafold.py中使用–jax_compilation_cache_dir <YOUR_DIRECTORY>标志启用编译缓存。

更详细的说明可在JAX文档

https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html#persistent-compilation-cache

中找到，更具体地说是关于在Google Cloud

https://jax.readthedocs.io/en/latest/persistent_compilation_cache.html#persistent-compilation-cache

上使用的说明。特别是，请注意，如果您想使用非本地文件系统，如Google Cloud Storage，您将需要安装etilshttps://github.com/google/etils

（这在AlphaFold 3 Docker容器中默认不包括）。