Skip to content

AlphaFold3

赛题介绍

AlphaFold3简介

蛋白质结构预测是一项基础性且具有挑战性的任务，半个世纪以来一直是生物学领域的重大难题。革命性的人工智能模型 AlphaFold 以卓越的性能成功解决了这一困扰科学界数十年的挑战，并因此，其核心研究者荣获 2024 年诺贝尔化学奖的一半奖项。最新版本 AlphaFold3 配备了新颖的特性，能够准确预测配体、蛋白质和核酸复合物的结构，有望为药物发现、疾病治疗及基础生物学过程的理解带来革命性突破。 AlphaFold3 采用基于扩散的架构，输入可以是单个氨基酸序列或多个生物分子的序列。该过程包含两个阶段：数据流水线和模型推理。运行于CPU 的数据流水线通过 Jackhmmer/Nhmmer 在遗传数据库中进行搜索，为蛋白质和 RNA 实体构建多序列比对（MSA）。运行于 GPU 的模型推理阶段则利用 MSA、模板和原始序列，将其输入 Pairformer 和扩散模块以生成预测结构。原始模型推理阶段通常运行于 GPU，但本次比赛提供了可在 CPU 上运行的版本。 在初赛环节，我们仅关注模型推理阶段。我们提供十二个单蛋白质序列样本及预生成的 MSA 和模板作为任务输入，无需执行数据流水线阶段。同时，提供了 CPU 版本的推理代码。参赛选手需在保证正确性的前提下，尽可能提升模型在 CPU 上的运行性能。 该题目的比赛平台是探索1000集群上的CPU节点，限制为280核。也就是说每个节点核数都用满的话(56核)，只能使用5个节点进行推理(我们提供的代码是单机版本，多机需要自己探索)。当然你也可以选择每个节点使用28核，共10个节点，不过这样的通信开销会变大，理论上不会有更好的效果

相关论文

Accurate structure prediction of biomolecular interactions with AlphaFold 3

仓库结构

代码获取

为方便选手使用，我们提供了以下 GitHub 仓库，这些仓库基于原作者版本进行了少量修改，以便更稳定地运行

git clone https://github.com/lvjx04/alphafold3.git
git clone https://github.com/lvjx04/xfold.git

原作者仓库地址 - AlphaFold3 - xfold

关于WORK目录

每个账号下都有自己WORK目录，所有作业必须在WORK目录下提交。详细见探索1000集群使用文档

alphafold3

这是AlphaFold3的官方代码仓库，使用jax实现，可以在GPU和TPU上高效运行。 但经测试发现，其对 CPU 的支持较差，因此本仓库主要用于环境配置参考。 若选手希望基于 JAX 进行后续优化，也可自行配置环境并解决潜在兼容性问题。

xfold

推荐使用的 xfold 仓库已将模型从 JAX 迁移至 PyTorch，对 CPU 运行环境支持更好。 我们提供了运行该仓库代码的完整环境配置及示例脚本。

主要关注的文件包括但不限于以下内容：

  • run_alphafold.py：主执行脚本,负责整个AlphaFold3预测流程的编排,包括数据管道处理、模型初始化、推理执行和结果输出
  • xfold/alphafold.py：AlphaFold3主模型类,集成Evoformer、扩散头、预测头等所有子系统,实现完整的前向传播流程

  • xfold/nn/：AlphaFold3模型的所有核心神经网络模块,实现蛋白质结构预测的各个子系统

    • diffusion_head.py: 扩散头模块，通过去噪扩散过程迭代生成3D原子坐标，包含DiffusionHead类和噪声调度函数
    • diffusion_transformer.py: 扩散Transformer架构，包含DiffusionTransformer(24层自注意力块)和DiffusionCrossAttTransformer(原子级交叉注意力)，用于结构去噪过程中的特征变换
    • atom_cross_attention.py: 原子级交叉注意力模块，包含AtomCrossAttEncoder(将token特征编码为原子特征)和AtomCrossAttDecoder(解码为原子位置更新)，是连接token级和原子级表示的桥梁
    • pairformer.py: PairformerBlock模块，通过三角乘法、网格注意力和过渡层处理成对表示，是Evoformer的核心迭代单元(48层)
    • attention.py: 注意力机制实现，包含GridSelfAttention(行列注意力)和MSAAttention(MSA注意力)，用于特征交互
    • head.py: 预测头模块，包含DistogramHead(预测残基间距离分布)和ConfidenceHead(预测pLDDT、PAE等置信度指标)
    • primitives.py: 基础神经网络构建块，包含Transition(门控前馈网络)和OuterProductMean(MSA外积均值)，为上层模块提供可复用的基础操作
    • triangle_multiplication.py: TriangleMultiplication模块，实现三角更新算法，用于更新成对表示

  • xfold/fast_nn/: 提供基于 Triton 的高性能算子实现，可通过配置在 PyTorch 原生实现和 Triton 优化实现间切换

    • layer_norm.py: LayerNorm 的 Triton 优化实现，通过融合 kernel 减少内存访问，提供LayerNorm类可替代torch.nn.LayerNorm
    • attention.py: 注意力机制的Triton优化实现，包含融合的点积注意力 kernel，支持可选的mask和bias
    • gated_linear_unit.py: 门控线性单元的Triton优化实现，融合矩阵乘法和SiLU激活
    • config.py: 配置模块，通过layer_norm_implementation、dot_product_attention_implementation、gated_linear_unit_implementation等参数控制使用 PyTorch 或 Triton 后端

运行程序

环境配置

module load compilers/gcc/v12.2.0
export CXX=/apps/compilers/gcc/v12.2.0/bin/g++
export CC=/apps/compilers/gcc/v12.2.0/bin/gcc

module load soft/anaconda3/config
conda init --all
bash # or zsh
conda create -n af3 python=3.11 -y
conda activate af3

cd alphafold3
git checkout v3.0.0-thuscc25
pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
pip3 install --no-deps . -vvv
build_data

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip3 install einops

模型参数和输入测例下载

模型参数下载 

mkdir weights && cd weights
wget https://cloud.tsinghua.edu.cn/f/a7129732456a4521a882/?dl=1 -O af3.bin
wget https://cloud.tsinghua.edu.cn/f/32e1c913ac1948dcaf23/?dl=1 -O fourier_bias.npy
wget https://cloud.tsinghua.edu.cn/f/f196faf302d24fe6af78/?dl=1 -O fourier_weight.npy

输入测例下载 

cd xfold
wget https://cloud.tsinghua.edu.cn/f/417438c75bce4d42bf16/?dl=1 -O input.zip
unzip input.zip

运行

我们提供了示例脚本，选手只需要更改其中的相关路径为自己的

cd xfold
sbatch example.sh

优化提示

  1. NUMA亲和性优化，更换openmp实现，更换内存分配器。可以参考PyTorch Tuning Guide
  2. attention的实现性能较差，可以尝试优化
  3. 不同迭代之间的某些数据是相同的，可以在第一次得到之后存储下来，避免重复计算
  4. 更高效的推理框架，OpenVINO，TensorRT，ONNX Runtime
  5. 多机推理，本次比赛我们限制核数为280核，这里需要你自己写多机推理代码

正确性检验

我们将以提供的 baseline 代码为标准，使用其生成的 ranking_scores.csv 文件中的 5 个分数的平均值作为基准分数。参赛队伍的结果不得低于该基准分数的 90%

各测例对应的基准分数如下表所示：

测例 基准分数
37aa_2jo9 0.76297803
107aa_1tce 0.83502734
301aa_3db6 0.95459096
436aa_1muu 0.90951743
583aa_1cf3 0.97798391
740aa_4a5s 0.96837966
959aa_7lzh 0.68904242
1024aa_5l6h 0.96428085
1284aa_4xwk 0.81414889
1290aa_1g9a 0.70253949
1303aa_1t3t 0.97658630
1491aa_5kis 0.95324605

评分规则

  • 每个测例占总分数的\(\frac{1}{12}\)
  • 成功运行并得到合法结果：30%
  • 性能分：70%。计算方法为\(0.7 \times \frac{T_{fastest}}{T_{team}}\)

注意事项

  1. 不允许使用低于 16 位的精度
  2. 不得修改 AlphaFold3 模型参数
  3. 不得修改model_config.json文件，包括但不限于循环次数和扩散次数等
  4. 不得修改提供的输入 JSON 文件
  5. 不得更改原始 AlphaFold3 的输出文件，包括结构文件和置信度文件
  6. 提交的推理日志文件必须包含预测开始时间、各阶段耗时（包括推理、输出结构提取等）及推理过程中的其他关键内容