AlphaFold2/multimer这个项目我最早是在一次结构生物学的组会上听说的当时同事讲激动了直接拿一个蛋白复合物的预测结构投到屏幕上配位关系、界面残基都标得清清楚楚。真正让我头疼的反而是后来的环境配置。官方GitHub给的是Docker方案一行命令启动确实优雅可实验室集群上没root权限、没有公共Docker我不得不转用Conda一点点搭。前后折腾了快两天把GPU驱动、CUDA、JAX、OpenMM、数据库下载这些环节全踩了一遍。后来陆续帮几个课题组搭过同样的环境也排过不少怪问题于是把这份沉淀写成指南。这篇内容适合所有准备复现AlphaFold2/multimer的人尤其是想在GPU上跑结构预测、又不想被各种环境报错劝退的研究生和工程师。我默认你对Linux有一定基础至少能打开终端、敲命令。如果你连Conda是什么都不太熟也没关系我会把每一步的“为什么”一起讲清楚。这样就算版本号变了、依赖更新了你也能顺着思路自己排雷。1. 项目整体思路与背景1.1 Alphafold2/multimer解决什么问题AlphaFold2是DeepMind在2021年发布的蛋白质结构预测工具当年在CASP14上以接近实验精度的表现震了整个结构生物学界。输入一条氨基酸序列它能输出一个高置信度的三维坐标模型。那几年大家开玩笑说“进实验室先跑AlphaFold”因为很多蛋白的晶体结构一时解不出来先用它预测一个模型做参考是很普遍的工作流。multimer是它的多链扩展专门处理蛋白复合物预测。两条链怎么结合、同源二聚体还是异源多聚体、界面残基大概在哪这类问题是传统对接方法的难点AlphaFold2/multimer在这块的效果也相当能打。很多做蛋白设计、药物靶点分析、抗原抗体研究的人如今都把multimer当成标准工具在用。这套系统的核心结构其实是“MSA搜索 深度学习推理”的组合。先用jackhmmer等工具在大型序列数据库里找同源序列构建多序列比对再把这些信息喂给Transformer架构的模型。GPU加速主要作用于模型推理阶段而MSA搜索阶段基本只吃CPU。这个定位关系很重要后面优化时会反复提到。1.2 为什么环境配置才是头号难点AlphaFold2本身的环境依赖并不复杂Python、JAX、TensorFlow、OpenMM、Biopython等。但麻烦在于三件事叠加在一起。第一模型训练和推理高度依赖JAX生态。JAX的GPU版本必须和CUDA、cuDNN版本精确匹配稍有不一致就会出现“nvidia-smi能看到卡但JAX说没有GPU”的诡异现象。这种问题最容易让新手心态爆炸。第二除了Python库还需要一堆外部二进制工具比如HHsuite里的hhblits、hhsearchHMMER里的jackhmmer、hmmbuild以及kalign多序列比对程序。这些工具如果自己编译费时费力如果直接apt安装发行版版本往往老旧。Conda在管理这类生物信息工具上有天然优势bioconda频道基本都能覆盖。第三数据库体积巨大。完整版数据大概需要2TB以上磁盘空间其中BFD库、MGnify这些都是“下载五分钟拷贝两小时”的巨型文件。如果网络环境一般下载数据库本身比装环境还折磨。所以把这套东西装好本质上是一场“版本对齐”的工程而Conda正是做版本隔离和复现的利器。1.3 Conda vs Docker怎么选官方推荐Docker方案好处是镜像里什么都有只要装好NVIDIA Container Toolkitrun一下就能用。但Docker在集群环境里经常受限尤其在高性能计算中心管理员一般不会给普通用户开Docker权限。Conda方案不需要root权限所有内容都装在自己目录的虚拟环境里天然适合无权限场景。它还能通过conda-pack一键打包迁移环境对多台机器部署非常友好这点在做计算资源调度时价值很大。我的建议是如果只是自己单机玩Docker确实省心如果是课题组共用的GPU服务器或高性能计算集群Conda几乎是唯一选择。这篇文章的核心就是讲透彻Conda路线下的所有细节包括GPU加速所以下面默认你已经准备好一台带NVIDIA显卡的Linux服务器。2. 动手之前的GPU基础环境排查2.1 驱动、CUDA到底由谁负责很多人第一次配置GPU环境时会被“CUDA”这个概念绕晕。我习惯用一个类比显卡驱动是操作系统和GPU硬件之间的翻译官装上驱动之后nvidia-smi才能看到显卡信息。CUDA是NVIDIA提供的并行计算平台程序通过它调用GPU计算能力而cuDNN是在CUDA之上专门为深度学习的卷积、注意力等运算做加速的库。关键在于“系统级CUDA”和“Conda环境级CUDA”并不是一回事。你完全不需要自己安装完整版CUDA Toolkitnvidia-smi显示的是驱动支持的CUDA版本不直接决定Python程序能不能用GPU。真正决定JAX能不能跑起来的是jaxlib内置的CUDA运行时以及它需要的cuDNN版本。Conda环境里的cudatoolkit和cudnn包可以提供这套运行时这也是后面配置的精髓。我先说结论驱动尽量装到足够新至少支持CUDA 11以上然后让Conda环境自己带一份匹配的CUDA运行时和系统全局不冲突。这样做的好处是你可以同时保留多个环境一个环境跑AlphaFold2另一个环境跑PyTorch各自的CUDA版本互不干扰。2.2 用Conda隔离CUDA运行时省掉一堆系统级烦恼在纯pip安装的世界里经常会遇到“系统CUDA版本太旧新库跑不起来”或“升级系统CUDA又怕影响其他项目”的两难境地。Conda方法可以直接绕开这个坑。在创建环境时或者之后你可以用conda install cudatoolkit11.3 cudnn8.2这样的方式把CUDA运行库放进虚拟环境内部。程序运行时优先加载环境内的库就不会去碰系统全局版本。这相当于给每个项目一个独立的“Windows系统目录”里面放着项目自己需要的运行库。好处是明显的但也要注意不是所有库都能被Conda很好地虚拟化。比如TensorFlow的旧版本对cuDNN的加载路径偶尔会写死因此你在环境里装好cudnn之后有时候还需要设置LD_LIBRARY_PATH让它明确去Conda环境目录下找.so动态库。给一个判断标准如果你的显卡是这几年的中高端型号驱动版本通常足够新直接用Conda环境的CUDA 11.x完全没问题如果你的显卡特别老比如还在用GTX 9系列那可能要留意算力是否满足要求AlphaFold2对老架构的支持不好。2.3 三条命令确认GPU在计算栈里可见在开始搭建环境前先用下面三条命令确认硬件和驱动状态。nvidia-smi这条命令应该能看到显卡型号、驱动版本、显存占用和温度。如果这条命令都报错说明驱动没装好后面一切免谈。驱动安装我不展开但建议用官方Runfile或发行版源来装不要在安装过程中同时开着图形界面会话那是另一个大坑。装好Conda基础环境后先别急着克隆项目先创建一个小环境测试JAX是否认卡。这一步能尽早暴露问题。conda create -n gpu_test python3.9 -y conda activate gpu_test pip install jax0.4 python -c import jax; print(jax.devices())如果输出里有GpuDevice或以cuda标记的设备说明驱动和JAX匹配正常如果只有CpuDevice问题基本确定在jaxlib安装的版本没带GPU支持。这是最常见的坑后面我会提怎么精准修复。另外可以用python -c import tensorflow as tf; print(tf.config.list_physical_devices(GPU))再验证一次TensorFlow不过AF2的主推理不依赖TensorFlow这里验证一下是为了排查上下游问题时少走弯路。3. Conda环境创建与依赖安装全流程3.1 创建独立虚拟环境为什么不用base很多人图省事直接在conda的base环境里装AlphaFold2。短期看着没问题但base环境往往还躺着别的项目依赖时间一长必炸。我见过最典型的场景是base环境里装了一个特定版本的numpyAlphaFold2依赖要升级它另一个项目又要求锁死旧版本升级完那个项目直接罢工。Conda虚拟环境的本质就是“项目隔离”。每条命令自己玩自己的互不污染。所以我强烈建议为AF2单独创建一个环境名字随意比如af2。conda create -n af2 python3.9 -y conda activate af2这里选Python 3.9而非3.11是因为官方的依赖测试和很多二进制包在3.9下最稳定。AlphaFold2主体的逻辑对Python小版本不算敏感但JAX、OpenMM、Biopython这些库在旧版本Python上兼容性更成熟。等官方明确支持Python 3.11以上再升级也不迟。如果Conda下载速度慢我建议先配置国内镜像源。这里以清华源为例conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/bioconda/ conda config --set show_channel_urls yes顺便说一句conda频道有优先级顺序写反了可能导致包从冷门频道安装引发冲突。把conda-forge放前bioconda放后基本够用。3.2 用Conda精装“外部程序”这一层进入af2环境后先用Conda安装那些独立二进制工具。这一层是所有MSA搜索的基础千万别跳过。conda install -c conda-forge -c bioconda hmmer hhsuite kalign aria2 -yhmmer提供jackhmmer和hmmbuild作用是搜索同源序列hhsuite提供hhblits和hhsearch用来做更灵敏的profile搜索和模板检测kalign做多重序列比对aria2是高速下载工具后面拉数据库全靠它。为什么不用apt直接装因为bioconda版本通常更新而且会连带解决依赖省得自己编译。好比你做饭与其自己磨面粉不如买质量有保障的成品面粉。装完后顺手确认一下which jackhmmer which hhblits which kalign which aria2c如果哪条命令没找到说明Conda环境PATH没生效。重新激活环境或者检查Conda是否装了但没有正确初始化。后面常见问题里我会讲conda init那一类报错。3.3 pip安装JAX/TensorFlow/OpenMM那套Python依赖外部工具到位后接下来处理Python层的依赖。AlphaFold2的官方源码仓库里有一个requirements.txt里面列了全部依赖包括numpy、scipy、absl-py、biopython、dm-tree、ml-collections、pandas、matplotlib、jax、jaxlib、tensorflow、openmm、pdbfixer、tqdm等。最稳的安装方式是克隆源码后直接用pip安装。git clone https://github.com/deepmind/alphafold.git cd alphafold pip install -r requirements.txt但这里有一个非常重要的细节官方requirements里的jax和jaxlib可能不一定带GPU支持或者版本过于保守。AF2在设计上依赖JAX这套自动微分框架来跑模型推理如果没有GPU版jaxlib就算显卡再强也只能用CPU慢慢算。所以装完requirements后建议再单独检查一次jax的设备列表如果发现没有GPU就手动安装匹配GPU版本的jaxlib。以我当时配置时的组合为例pip install jax0.3.25 jaxlib0.3.25cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html注意这里的cuda11.cudnn82后缀指的是jaxlib针对CUDA 11和cuDNN 8.2的预编译版本。你不需要系统里有这些库因为可以靠conda环境内的cudatoolkit和cudnn补上。conda install -c conda-forge cudatoolkit11.3 cudnn8.2 -y这个组合我实测下来非常稳。千万别在没确认需求的情况下随便升级jax版本AF2源码有时对新版本jax不兼容比如某些API被移除后直接崩。先跑通再谈升级。3.4 安装后的快速自检环境装完做一次快速自检能省掉后面无数排查时间。python -c import jax; print(jax.devices()) python -c import tensorflow as tf; print(tf.config.list_physical_devices(GPU)) python -c import openmm; print(openmm.version.version)再验证一下外部程序jackhmmer -h | head -n 5 hhblits -h | head -n 5如果这些命令都正常且JAX能识别GPU恭喜环境本体已经通了。接下来才是真正的“重头戏”AF2源码组织、数据库准备和multimer多链预测。4. 源码、数据库与multimer多链预测4.1 源码获取和模型权重源码直接克隆官方仓库git clone https://github.com/deepmind/alphafold.git里面有几个关键目录和文件run_alphafold.py是主入口pipeline_jax.py是JAX推理管线data/存放数据库相关逻辑。模型权重需要单独下载。权重参数是一个压缩包解压后是一堆.npz文件包括model_1到model_5这些单链模型以及model_1_multimer_v2/v3这些多聚体模型。下载地址官方GitHub的README里有用wget或aria2拉下来即可。解压后把路径记住运行脚本时通过--model_dir参数指向它。这里提个醒权重的年份版本会影响模型参数比如2022-12-14这版是针对multimer增强后的官方参数。建议直接用官方最新参数包不要用网上乱七八糟的旧包。4.2 数据库准备全量、reduced_dbs与断点续传数据库是AlphaFold2最容易劝退新手的地方。完整版由好几组数据库组成总共需要2TB以上的磁盘空间。我整理了一张常见数据库及其用途的表数据库体积参考用途BFD约1.7TB大库搜索同源序列非常丰富但网络差时可跳过MGnify约120GB宏基因组数据集的同源搜索Uniclust30约27GB快速MSA库Uniref90约60GB非冗余蛋白库PDB70约56GB模板搜索PDB mmCIF约250GB模板结构文件PDB seqres约40MBmultimer模板序列Uniprot约100MBmultimer全库序列如果你只是想验证环境或者跑一个小蛋白完全没有必要一上来就下载所有库。官方提供了reduced_dbs体积只有十几GB专门用来做流程验证和简单预测。我第一次测试就是用的reduced_dbs从下载到跑通只花了一个下午。下载时强烈推荐用aria2c多线程加断点续传比wget靠谱太多。aria2c -x 16 -s 16 -d ./alphafold/data/ https://storage.googleapis.com/alphafold-databases/reduced_dbs/reduced_dbs.tar.gz下载后解压注意目录结构必须符合run_alphafold.py的预期。它主要通过--data_dir指定根目录然后在里面找对应子目录。如果目录结构不对运行时会报“找不到数据库文件”之类的错误。如果你在校园网或科研网可以优先找校内共享盘或同事之间的镜像拷贝省时省力。数据库版本一旦固定建议整个目录当成“只读资产”保存不要反复删改。4.3 multimer和单链模型的内核差异单链预测主要用--model_presetmonomermultimer预测改成--model_presetmultimer它会调用独立的multimer模型权重。multimer与单链的本质区别在于输入从单独一条序列变成了多条链的序列信息MSA构建时既要考虑每条链自己的同源序列又要考虑链间共有的同源信息模型内部还会对链的排列组合做specialized处理。在操作层面输入FASTA文件时多条链的序列放在同一个文件里。每条链可以有不同的名字如果你想预测同源二聚体两条序列写相同名字即可异源二聚体就写不同名字。这个细节很多人第一次接触时会踩坑因为单链预测根本不需要考虑这个。一个很重要的注意点multimer对显存的要求比单链高不少。同样长度的序列两条链合起来推理不仅仅是一倍显存那么简单注意力交互的计算量会明显增加。我的经验是如果跑两条长度为400残基的链24GB显存的显卡基本够但再多几条链或者链更长就该考虑32GB及以上显存或者用多卡并行和BF16来腾空间。5. GPU加速优化从OOM、BF16到五模型并行5.1 合理理解pipeline里的“GPU瓶颈”AlphaFold2的完整流程可以拆成四段MSA搜索、特征构建、模型推理、结构弛豫。其中MSA搜索主要跑在CPU上因为jackhmmer和hhblits这类工具根本不支持GPU特征构建同样在CPU上完成模型推理是真正吃GPU的环节最后的弛豫用的是OpenMM通常也是CPU负载。这就引出很多人的第一个困惑我明明有GPU为什么看监控发现GPU使用率有时候很低答案是整个流程大部分时间在做MSA搜索和特征构建GPU只在推理阶段火力全开。所以如果你只跑一个小蛋白可能感觉不到GPU加速明显但跑大蛋白或多聚体时推理时间会显著拉长GPU的作用才会突显。因此GPU优化的目标不是让每个环节都吃满GPU而是尽量压缩推理阶段的时间同时避免显存成为瓶颈。另外如果机器有多个CPU核心在MSA阶段把线程数调上去往往会比纠结GPU参数更能提升整体感受。很多“CPU、GPU、内存占用都不高但卡”的现象其实就是MSA阶段的进程数没给够或者磁盘IO在拖后腿。5.2 显存预分配和OOM实战JAX在GPU上有一个和PyTorch不太一样的行为默认情况下XLA会在第一次运行时就为自己预分配几乎全部显存。这本来是好事因为避免了频繁分配释放带来的碎片但代价是多进程共卡时容易互相抢显存。尤其在同一个实验室里别人也在跑别的模型你的AF2直接把显存占满对方就OOM了。解决方法是设置环境变量来限制预分配比例export XLA_PYTHON_CLIENT_MEM_FRACTION0.9这个值可以根据实际情况调整。如果你确定这张卡只给AF2用可以设成0.95把计算资源吃透如果还要跟别人共享建议设成0.5甚至更低。注意这个环境变量要在运行run_alphafold.py之前export而且它影响的是JAX/XLA的分配策略。另外还有一个容易忽视的参数JAX_ENABLE_X64。如果你手贱设成了1JAX会强制用64位浮点与整数显存占用直接翻倍很多OOM就是这么来的。默认情况下不要开它。如果仍然OOM就先从输入序列下手把蛋白拆短或者换成小一点的复合物验证流程。硬件该升级就升级有些问题不是软件能解决的。5.3 用BF16等技巧换速度AlphaFold2源码默认用float32精度推理。理论上用更低精度能明显提速特别是在Ampere架构及以后的GPU上。你可以用JAX的环境变量XLA_USE_BF161来强制XLA把部分计算转到bfloat16。export XLA_USE_BF161我在A100和RTX 4090上都实测过开启BF16之后推理速度通常能提升20%到40%pLDDT分数几乎没有肉眼可见的差异。但必须承认对于某些边界情况低精度可能让结构局部区域的质量略微下降。所以我的建议是先跑一个已知蛋白做对照看看预测结构差异是否在可接受范围内再决定生产环境是否开启。另一个思路是减少不必要的模板搜索和MSA迭代次数。AF2本身把MSA视为限制步骤如果可用的同源序列很少也可以考虑调整迭代参数。但这是生物学层面的权衡别为了快而牺牲预测质量。5.4 多卡并行跑五个模型ensembleAlphaFold2每次预测会跑5个模型model_1到model_5最终从结果里挑一个置信度最高的结构输出。官方脚本是串行执行这5个模型的单卡耗时就是5倍模型推理时间。如果你手头有几张GPU完全可以把这5个模型并行跑。一种很实用的做法是手动拆成多个进程每个进程绑定不同的CUDA_VISIBLE_DEVICES。举个例子for model in model_1 model_2 model_3 model_4 model_5; do CUDA_VISIBLE_DEVICES0 python run_alphafold.py \ --fasta_pathstarget.fasta \ --data_dir./data \ --model_dir./params \ --model_names$model \ --output_dir./output_$model done wait这个脚本会让5个模型同时跑在不同卡的显存切片上如果只有两张卡可以把0和1重复分配让每张卡处理2到3个模型。跑完后从每个output_$model里找到对应的.pkl文件读取里面的ranking_confidence分数分数最高的那个就是最终推荐结构。官方没有提供一键合并脚本但自己写一个小工具读取pkl文件很简单用Python的pickle加glob就能搞定。这样优化下来五模型总耗时几乎等于单个模型的耗时加速效果非常可观。这里顺便提一下JAX底层会以CTA线程块为单位在GPU上调度kernel你不需要手动控制这些细节但理解这一点能帮你判断为什么某些GPU型号在AF2上特别快核心多、显存带宽高、Tensor Core支持好效果自然好。老卡就算显存够也可能因为Tensor Core能力弱而慢不少。6. 常见问题与排查速查6.1 一张速查表把我在实践中遇到过的典型问题整理成一张表方便你遇到时第一时间对照。现象可能原因排查与解决conda activate报错提示要先跑conda initShell没初始化conda执行source ~/.bashrc或按提示运行conda init bash后重开终端nvidia-smi正常但JAX提示只有CPUpip安装的jaxlib不是GPU版手动安装带cuda标签的jaxlib并保证conda环境内cudatoolkit/cudnn版本匹配TensorFlow/OpenMM找不到GPU环境内缺少CUDA动态库conda install cudatoolkit cudnn必要时在运行前设置LD_LIBRARY_PATH指向环境lib目录一运行就OOMJAX默认预分配显卡全显存export XLA_PYTHON_CLIENT_MEM_FRACTION0.8或降低并行进程数数据库下载一半断掉网络不稳定用aria2c多线程加断点续传必要时重新执行同一条下载命令提示找不到模型权重文件--model_dir路径不对或未解压检查参数包解压目录里是否存在model_1.npz等文件跑到jackhmmer阶段长时间不动数据库文件损坏或路径不对CPU线程不给力确认--data_dir下目录结构对比官方目录规范把MSA线程数调大import openmm报找不到libstdc系统和conda的GCC库版本冲突尝试conda install -c conda-forge libstdcxx-ng环境和另一个项目互相污染没有做环境隔离严格使用独立conda环境必要时候用conda-pack迁移完整环境显存占用不高但整个机器卡顿可能卡在MSA的CPU计算或磁盘IO用top查看CPU占用用iostat查看IO避免在机械硬盘上跑大数据库6.2 一个典型“内存不高但卡住”的排查实录一次同事跑AF2nvidia-smi看GPU利用率只有个位数内存也不高但程序就是卡住不出结果。他怀疑是GPU驱动问题。我登上去先看进程状态发现一个jackhmmer的CPU占用接近100%另一个hhblits的IO等待时间特别长。原因很快清楚了他的数据库放在了挂载的NFS共享盘上跨网络读小文件的速度极慢MSA搜索阶段大量随机读写小文件把IO活活拖死。解决办法是把数据库复制到本地固态硬盘上再跑同时给jackhmmer更多的CPU线程。换完磁盘后同样的任务从“跑了一晚上没出结果”变成了“40分钟跑完”。这个案例说明GPU优化只是整个系统的冰山一角CPU和磁盘IO往往才是更隐蔽的瓶颈。再分享一个conda相关的细节。有次运行conda install时提示某个包冲突我试了半天才意识到是之前把defaults和conda-forge频道混用导致的。后来我干脆为这个项目固定频道的优先级使用strict的channel priority冲突少了很多。做法是conda config --set channel_priority strict如果在配置过程中实在解决不了问题还可以把环境导出重来用conda env export environment.yml记录下当前状态再在干净的机器上重建。这比手动一个一个排查命令要高效。最后再分享一个我常用的部署技巧如果你需要在很多台机器上复现同一个环境用conda-pack把它打包成tar.gz拷到目标机器解压后激活就行根本不用重新联网装依赖。这在没有外网的高性能计算节点上尤其好用。pip install conda-pack conda pack -n af2 -o af2_env.tar.gz解压到目标机器后执行:mkdir -p ~/af2_env tar -xzf af2_env.tar.gz -C ~/af2_env source ~/af2_env/bin/activate我自己实际跑下来最深的感触是AlphaFold2/multimer本身对环境的要求不算夸张但如果你把conda、驱动、数据库和JAX都一次性看成系统工程去理解后面几乎所有报错都能靠“版本对齐”四个字解决。第一次上手先用reduced_dbs把一个短蛋白完整跑通不要一上来就上两个TB的大库。跑通了流程再加数据、调GPU、上集群都会顺很多。祝各位一次复现成功。