AlphaFold3 的性能表现
跟我一起来探索 AlphaFold3 的世界吧！
原文
https://github.com/google-deepmind/alphafold3/blob/main/docs/performance.md
数据管道
数据管道（即基因序列搜索和模板搜索）的运行时间可能会因输入大小、发现的同源序列数量以及可用硬件（特别是磁盘速度会影响基因搜索速度）而有显著变化。如果您想提高性能，建议提高磁盘速度（例如，通过使用基于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倍），表明它适用于高通量应用。
| Token数 | 单卡A100 80GB (GPU秒) | 16卡A100 40GB (GPU秒) | 提高倍数 |
|---|---|---|---|
| 1024 | 62 | 352 | 5.7倍 |
| 2048 | 275 | 1136 | 4.1倍 |
| 3072 | 703 | 2016 | 2.9倍 |
| 4096 | 1434 | 3648 | 2.5倍 |
| 5120 | 2547 | 5552 | 2.2倍 |
分阶段运行管道
run_alphafold.py脚本可以分阶段执行，以优化资源利用率。这可能很有用：
将CPU专用数据管道与模型推理（需要GPU）分开，以优化成本和资源使用。 缓存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，但代价是程序变慢，因为访问主机内存而不是设备内存。要了解更多信息，请查看NVIDIA博客文章。
https://developer.nvidia.com/blog/unified-memory-cuda-beginners/
您可以通过在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容器中默认不包括）。
