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个节点，不过这样的通信开销会变大，理论上不会有更好的效果

仓库结构

代码获取

为方便选手使用，我们提供了以下 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

示例结果

终端输出

高亮那一行的时间就是我们关注的性能指标


Skipping running the data pipeline.
Found local device: cpu
Building model from scratch...
loading the model parameters...
Loading from /home/weights/af3.bin
Processing fold input 37aa_2JO9
Skipping data pipeline...
Output directory: output/37aa_2jo9
Writing model input JSON to output/37aa_2jo9
Predicting 3D structure for 37aa_2JO9 for seed(s) (1,)...
Featurising data for seeds (1,)...
Featurising 37aa_2JO9 with rng_seed 1.
I1008 21:58:45.525593 22623381854016 pipeline.py:160] processing 37aa_2JO9, random_seed=1
Featurising 37aa_2JO9 with rng_seed 1 took 0.58 seconds.
Featurising data for seeds (1,) took  5.10 seconds.
Running model inference for seed 1 took  69.47 seconds.
Extracting output structures (one per sample) for seed 1...
Extracting output structures (one per sample) for seed 1 took  0.16 seconds.
Running model inference and extracting output structures for seed 1 took  69.63 seconds.
Running model inference and extracting output structures for seeds (1,) took  69.63 seconds.
Writing outputs for 37aa_2JO9 for seed(s) (1,)...
Done processing fold input 37aa_2JO9.

优化提示

NUMA亲和性优化，更换openmp实现，更换内存分配器。可以参考PyTorch Tuning Guide
attention的实现性能较差，可以尝试优化
不同迭代之间的某些数据是相同的，可以在第一次得到之后存储下来，避免重复计算
更高效的推理框架，OpenVINO，TensorRT，ONNX Runtime
多机推理，本次比赛我们限制核数为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

评分规则

该赛题占比20%，前8个测例，每个占1.5%，后4个每个占2%
成功运行并得到合法结果：30%
性能分：70% 计算方法为\(\frac{T_{fastest}}{T_{team}}\)

优化报告与代码

初赛结束时，需要交一份压缩包，命名格式为队伍编号，来帮助我们判断你们这道题的优化的合法性。要求如下：

一份 report.pdf 详细说明优化方法
一份最终版本的整个所使用到的代码(比如你是按照我们的指导运行的，需要包含alphafold3和xfold这两个文件夹)
包含提交密码的文本文件

Alphafold3提交链接

决赛说明

决赛共有9个测例，每个测例的提交文件和初赛一致：决赛测例下载链接

或可直接使用wget下载：

wget https://cloud.tsinghua.edu.cn/f/5b695e129cea40d38486/?dl=1 -O thuscc25_af3_final_input.zip

序号	测例	基准分数
1	8w4w_a_358_358	0.9702
2	8c4z_a_112_b_282_394	0.9251
3	8fso_a_411_411	0.6482
4	8uc3_a_196_d_105_c_105_b_196_602	0.9660
5	9jcv_a_606_606	0.9070
6	8x0w_a_470_b_464_934	0.7902
7	8rw5_b_991_991	0.8812
8	8qft_a_1174_1174	0.8665
9	8j7f_a_289_b_289_c_289_d_289_e_20_1176	0.7649

提交方式依旧为

./submit.sh AF3-7 alphafold3-7.txt

基准分数要求同初赛，不能低于90%

注意事项

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