AlphaFold 3 性能优化实战指南：数据管线、GPU 推理、编译桶与内存配置全解析

• 人工智能
• 基础模型
• 深度学习
• 生物信息学
• 科学计算

导读

AlphaFold 3 的推理管线分为数据管线（遗传序列搜索与模板搜索，纯 CPU）与模型推理（特征化 + 扩散模型前向传播，依赖 GPU）两大阶段，两者的性能瓶颈、硬件需求与调优手段完全不同。本文以仓库 docs/performance.md 为核心，结合 run_alphafold.py、docker/Dockerfile 与模型配置源码，系统讲解如何通过磁盘提速、CPU 并行、分阶段运行、编译桶（compilation buckets）、pair_transition_shard_spec分片、Flash Attention 选择与环境变量调优，让 AlphaFold 3 在单卡 A100/H100 上以高吞吐运行，并掌握在 40 GB 显存 A100、V100、P100 等受限硬件上完成折叠的配置方案。

整体性能画像：数据管线与模型推理

数据管线（Data Pipeline）的耗时因素

数据管线负责遗传序列搜索（Jackhmmer / Nhmmer）与模板搜索，其运行时间会因以下因素产生显著波动：

• 输入规模：链数量、序列长度直接影响搜索量；
• 同源序列数量：找到的同源序列越多，MSA 构建越耗时；
• 可用硬件：其中磁盘速度对遗传搜索的影响尤为突出，因为 HMMER 工具需要反复读取大型序列数据库文件。

因此文档给出的性能改进建议依次是：

1. 提升磁盘速度，例如借助 RAM-backed 文件系统（如 tmpfs）存放数据库；
2. 增加可用 CPU 核数并加强并行化；
3. 为深 MSA 预留足够内存：序列具有深度 MSA 时，Jackhmmer 或 Nhmmer 可能需要远超推荐 64 GB RAM的内存，部署时需按输入规模评估内存上限。

仓库在 src/alphafold3/data/pipeline.py 中为 Jackhmmer 与 Nhmmer 分别提供了jackhmmer_n_cpu与nhmmer_n_cpu配置项（默认均为 8），并在 run_alphafold.py 中暴露为命令行 flag：默认取min(cpu_count, 8)，且注释明确指出"超过 8 核几乎不会再带来额外加速"，可作为 CPU 资源分配的直接依据。

模型推理的基准定位

文档指出：AlphaFold 3 论文（Nature，表 8）给出的 inference timing 基于16 张 40 GB 的 NVIDIA A100；而本仓库支持在单张 80 GB NVIDIA A100上以面向高吞吐优化的配置运行。下表使用 GPU 秒数（即 16 卡场景乘以 16）对比两种配置的免编译（compile-free）推理耗时：

从源码看，单卡高吞吐的支撑点包括：默认使用 Triton Flash Attention（src/alphafold3/model/model_config.py 中flash_attention_implementation = 'triton'，run_alphafold.py 提供--flash_attention_implementation切换）、bfloat16 精度与编译桶机制。这些内容将在下文逐一展开。

分阶段运行管线：解耦 CPU 数据管线与 GPU 推理

run_alphafold.py支持分阶段执行，用以优化资源利用，典型收益场景：

1. 成本/资源优化：把纯 CPU 的数据管线与需要 GPU 的模型推理拆分到不同机器或不同时段执行；
2. MSA/模板结果复用：缓存数据管线产物（带 MSA 与模板的增强 JSON），供**不同 seed 或不同特征变体（如更换配体）**的多次推理复用，避免重复搜索。

入口脚本在 run_alphafold.py 定义了--run_data_pipeline与--run_inference两个开关（默认均为 True），并限制两者不能同时为 False。

仅运行数据管线：--norun_inference

python3 run_alphafold.py \ --json_path input.json \ --output_dir output/ \ --norun_inference

该阶段只生成 MSA 与模板，不做特征化和模型推理，但在运行时长、CPU 与内存开销上可能相当昂贵。产物是已用 MSA 和模板增强的 JSON 文件（对应源码write_fold_input_json写出的*_data.json），可直接作为后续推理阶段的输入。

仅运行特征化与模型推理：--norun_data_pipeline

python3 run_alphafold.py \ --json_path augmented.json \ --output_dir output/ \ --norun_data_pipeline

该阶段跳过数据管线，只做特征化与推理，要求输入 JSON 已包含预计算的 MSA 和模板。特征化环节由 src/alphafold3/data/featurisation.py 的validate_fold_input严格校验：蛋白链必须携带 unpaired MSA、paired MSA 与 Templates，RNA 链必须携带 unpaired MSA，缺失即抛错——这与文档的"必须包含预计算 MSA/模板"要求一一对应。

加速器硬件要求与实测对照

官方支持的两种配置

仓库官方支持并经过数值精度与吞吐效率全面测试的配置为：

• 1 张 NVIDIA A100（80 GB）
• 1 张 NVIDIA H100（80 GB）

两种配置下的免编译推理耗时对比如下：

H100 相比 A100 在各 token 规模上约快 1.8 倍。需注意这些是免编译（compile-free）的纯推理耗时，首次运行仍会包含模型编译开销（见"编译桶"与"JAX 持久化编译缓存"两节）。

其他硬件配置

NVIDIA A100（40 GB）：上限 4,352 tokens

单张 40 GB A100 可处理至多 4,352 tokens的输入，需要两项配置修改：

1. 启用 unified memory（见下文环境变量小节）；
2. 调整pair_transition_shard_spec（位于 src/alphafold3/model/model_config.py，默认值为((2048, None), (None, 1024))），改为：

pair_transition_shard_spec: Sequence[_Shape2DType] = ( (2048, None), (3072, 1024), (None, 512), )

该配置条目的格式为(num_tokens_upper_bound, shard_size)，shard_size=None表示无上限。结合 src/alphafold3/model/network/modules.py 中get_shard_size的实现（按顺序匹配首个满足上界的分片规格）与 modules.py 中sharded_apply的应用，逐条解读如下：

• (2048, None)：序列不超过 2,048 tokens 时不进行分片；
• (3072, 1024)：序列不超过 3,072 tokens 时按 1,024 的块大小分片；
• (None, 512)：其余更长的序列统一按 512 的块大小分片。

分片作用于 Evoformer 与 Diffusion Transformer 中的 pair transition 块（TransitionBlock），通过对 pair 表示按残基维度切块计算来降低瞬时显存峰值。该配置数值上依然精确，但由于可用显存更小，吞吐会低于 80 GB A100 的方案。

NVIDIA V100（CUDA Capability 7.x）

CUDA Capability 7.x 设备存在已知数值问题，需设置环境变量XLA_FLAGS包含--xla_disable_hlo_passes=custom-kernel-fusion-rewriter。设置后，配合 unified memory，单张 V100 可处理至多 1,280 tokens。此外，run_alphafold.py 在启动推理时会主动校验：7.x 卡上若XLA_FLAGS缺少该 flag，或--flash_attention_implementation未设为"xla"，会直接抛错终止，避免数值错误静默发生。

NVIDIA P100

单张 P100 可处理至多 1,024 tokens，且无需任何配置修改。同理，入口脚本要求 GPU compute capability 不低于 6.0。

其他设备

未在其他设备上执行大规模数值测试，但从实现上推断其数值结果是准确的；不过，对 7.x 设备的已知数值问题与对应的XLA_FLAGS规避手段仍然适用。

编译桶（Compilation Buckets）机制

XLA 编译模型的开销显著，AlphaFold 3 通过编译桶避免过度重编译：即用单个模型编译覆盖一段输入规模区间。特征化时，系统为输入选择能容纳它的最小桶，并填充（padding）至该桶的 token 数；若后续输入落在同一桶内，则可复用之前的编译结果。

桶数量的配置存在权衡：桶越多 → 重编译次数越多，但 padding 越少；桶越少则反之。

默认桶与超桶输入

默认最大桶为5,120 tokens（定义于 run_alphafold.py 的--buckets默认值，序列为256,512,768,1024,1280,1536,2048,2560,3072,3584,4096,4608,5120）。处理超过最大桶的输入会触发为该输入尺寸新建桶并重编译模型——对应 src/alphafold3/model/pipeline/pipeline.py 中calculate_bucket_size的行为：当num_tokens超过最大桶时打印 warning，并直接返回输入自身的 token 数作为新桶。

此时应通过--buckets追加更大的桶尺寸。文档给出的示例：假设三个输入 token 数分别为5132, 5280, 5342，使用默认桶会触发三次独立编译（每个尺寸各一次）；改为传入：

--buckets 256,512,768,1024,1280,1536,2048,2560,3072,3584,4096,4608,5120,5376

三个输入都会落到 5,376 桶，模型只编译一次。文档特别注明：对该示例而言--buckets 5376已足够；给出多桶示例是为覆盖更广泛的输入规模分布。

桶的边界约束

源码要求桶列表必须严格递增（calculate_bucket_size会校验prev < curr，否则抛ValueError）；若输入 token 数超过最大桶且未自定义更大的桶，系统仍会以输入实际尺寸建新桶运行，但会伴随一次重编译。

环境变量与附加 Flag 调优

XLA 编译时间规避（默认开启）

为规避已知的、会大幅拉长编译时间的 XLA 问题，必须设置以下环境变量（docker/Dockerfile 已默认设置）：

ENV XLA_FLAGS="--xla_gpu_enable_triton_gemm=false"

即禁用 Triton GEMM。注意：该设置与 7.x 卡的规避 flag互斥，二者只能取其一。

CUDA Capability 7.x 显卡

对所有 7.x GPU（如 V100），XLA_FLAGS必须改为包含--xla_disable_hlo_passes=custom-kernel-fusion-rewriter；由于此类 GPU 不支持 Triton GEMM 内核，无需（也不应）同时禁用 Triton GEMM：

ENV XLA_FLAGS="--xla_disable_hlo_passes=custom-kernel-fusion-rewriter"

GPU 内存预分配（A100/H100 80 GB 默认配置）

以下环境变量（docker/Dockerfile 默认设置）支撑在单张 80 GB A100 或 H100上折叠至多 5,120 tokens的单个输入：

ENV XLA_PYTHON_CLIENT_PREALLOCATE=true ENV XLA_CLIENT_MEM_FRACTION=0.95

PREALLOCATE=true让 JAX 启动时预占显存，MEM_FRACTION=0.95允许使用 95% 的设备显存，从而把最大可用显存留给模型。

Unified Memory：突破 5,120 tokens / 小显存 GPU

若需处理超过 5,120 tokens 的输入，或在显存更小的 GPU（如 40 GB A100）上运行，推荐启用 unified memory。其原理是当显存不足时将 GPU 内存换出到主机内存，避免 OOM；代价是访问主机内存比设备内存慢，程序变慢。启用方式：

ENV XLA_PYTHON_CLIENT_PREALLOCATE=false ENV TF_FORCE_UNIFIED_MEMORY=true ENV XLA_CLIENT_MEM_FRACTION=3.2

注意此配置与上文的 80 GB 默认配置互斥，应按显存规格选用。XLA_CLIENT_MEM_FRACTION=3.2表示允许 JAX 申请超过物理显存的内存比例（配合 TF_FORCE_UNIFIED_MEMORY 交由 unified memory 兜底）。

Flash Attention 实现选择

run_alphafold.py 提供--flash_attention_implementation，可选triton、cudnn、xla：

• triton（默认）：最快，测试最充分，需 Ampere 及更新架构（含 A100/H100）；
• cudnn：cuDNN 实现的 Flash Attention，同样需 Ampere 及以上；
• xla：XLA 注意力实现（无 Flash Attention），跨设备可移植，也是 7.x 卡的唯一选择（入口脚本会强制校验）。

JAX 持久化编译缓存

为避免多次运行间的模型重复编译，可在 run_alphafold.py 中使用--jax_compilation_cache_dir指定缓存目录：

python3 run_alphafold.py \ --json_path input.json \ --output_dir output/ \ --jax_compilation_cache_dir /path/to/cache

入口脚本在main开头将其写入jax.config.update('jax_compilation_cache_dir', ...)（run_alphafold.py）。若要使用非本地文件系统（如 GCS）作为缓存，需额外安装etils（默认 Docker 容器未包含）。

其他值得关注的推理参数

除性能专项配置外，run_alphafold.py 中还有几个直接影响推理开销的参数，组合使用时需一并评估：

• --num_recycles（默认 10）：扩散前 Trunk 重循环次数，越大越耗时；
• --num_diffusion_samples（默认 5）：每 seed 生成的扩散样本数，线性放大推理成本；
• --num_seeds：批量生成连续 seed，便于复用数据管线结果做多次采样；
• --gpu_device：多卡系统上将每次运行固定到指定 GPU；
• --save_embeddings：输出 Trunk 的 single/pair embeddings，会额外写盘；
• --conformer_max_iterations：RDKit 构象搜索迭代上限，影响配体相关输入的特征化时长。

实践建议汇总

1. 优先提升磁盘与 CPU：数据库放 RAM-backed 文件系统，CPU 核数按需分配（Jackhmmer/Nhmmer 超过 8 核收益有限）；
2. 大输入先跑数据管线：用--norun_inference产出增强 JSON 并缓存，再对多个 seed/配体变体用--norun_data_pipeline复用；
3. 按 token 规模配桶：常规输入使用默认桶；接近或超过 5,120 tokens 时用--buckets补齐更大的桶，减少重编译次数；
4. 按显存选配置：80 GB 卡用默认PREALLOCATE=true+MEM_FRACTION=0.95；40 GB A100 或超 5,120 tokens 改用 unified memory 三件套，并按文档调整pair_transition_shard_spec；
5. 7.x 卡务必设XLA_FLAGS并搭配--flash_attention_implementation xla，入口脚本会校验，未设置将无法启动推理；
6. 持久化编译缓存：高频推理场景开启--jax_compilation_cache_dir，跨机器共享时另装etils。

以上所有配置均以本仓库 docs/performance.md、run_alphafold.py、docker/Dockerfile 及 src/alphafold3/model/model_config.py 为准；性能数值为文档所载的官方测试结果，实际收益请结合自身输入规模与硬件实测。

• 人工智能
• 基础模型
• 深度学习
• 生物信息学
• 科学计算