去年我被拉去给一个结构生物学课题组搭 AlphaFold2/multimer 的预测环境。当时总觉得官方 GitHub 写得够清楚了照着 Docker 跑就行。结果课题组那边的机器根本没有 Docker 权限只能走 Conda 裸装。这一走就是两个星期的坑驱动、CUDA、cuDNN、JAX、OpenMM 全都在版本矩阵里互相咬着。这篇文章就是把当初踩出来的路线固化下来给那些和我一样不能在服务器上开 Docker、必须用 Conda 手工配置 AlphaFold2/multimer 的人。文中会用尽量直白的方式讲清楚三件事环境里到底有哪些模块、它们之间如何匹配、以及 GPU 加速调到什么程度才算真正“跑得快”。无论你是刚接触结构预测的湿实验同学还是被分配了运维任务的生信工程师照着这套流程走一遍应该能避开大部分我踩过的雷。1. 动手前先回答三个问题显卡、数据库与运行场景1.1 单体模型和multimer的压力差在哪个环节AlphaFold2 的单体模型处理的是单条氨基酸序列输入一个 FASTA 就能预测这条链的三维坐标。而 multimer 模型要处理的是复合物也就是两条或更多的链在三维空间中如何组装。很多人以为 multimer 只是把序列简单拼接实际完全不是这样。multimer 在 MSA多序列比对阶段有一个专门的“配对”逻辑会尝试把不同链放到同一个物种/基因组背景下去找共进化信号这会让 MSA 特征矩阵的规模明显变大。到了 Evoformer 和 attention 部分计算复杂度又是和序列总长度近似平方的关系总长翻一倍显存和算力需求远不止翻一倍。所以我给课题组建议时第一件事不是让他们马上下载代码而是先想明白日常预测的蛋白大概多长是 300 个残基的单体还是 3000 个残基的五聚体。这两个场景选卡和选数据库的策略完全不一样。1.2 显存、驱动和CUDA的匹配边界AlphaFold2 官方只支持 NVIDIA GPUAMD 卡别指望官方镜像能直接跑。模型推理的主力是 JAX底下的矩阵乘靠 cuBLAS/cuDNN所以驱动和 CUDA 版本是第一个分水岭。先看显卡能选什么档。我大概整理了这样一张经验表显卡显存适合场景RTX 3060/406012GB单体小蛋白、两条短链的小二聚体RTX 3090/409024GB单体 2000 残基左右、中小复合物A500024GB接近 3090适合长时间跑批A600048GB三聚体/四聚体部分大单体A100/A80040/80GB大复合物、多链 MSA 的重负载注意这只是参考不是硬指标。真正决定显存爆不爆的是模型内部的 attention 计算以及你跑推理时同时承载的模型副本数量。驱动这块有个非常容易误解的点nvidia-smi最上面显示的 “CUDA Version” 并不是你实际用的 CUDA 版本它只是当前驱动能支持的上限。比如驱动 525 显示 “CUDA Version: 12.0”不代表你的环境里就有 CUDA 12.0只代表驱动允许你装到那么新的 toolkit。AlphaFold2 v2.3.x 官方环境锁定的是 CUDA 11.8 和 cuDNN 8.2Linux 驱动最好在 520 以上。如果你手头的机器是那种只能用 470 老驱动的旧平台那就得回头找 JAX 0.4 时代更早的 CUDA 11.2 兼容轮子晚点我会在排错部分再说。1.3 磁盘和内存预算最容易算漏的两笔账AlphaFold2 运行起来不只是“装个软件”那么简单它的核心运行还需要下载一批巨大的序列数据库。完整版的数据库体积大概是 2.3TB 到 2.6TB其中大头是 BFD 数据库。如果只用 reduced_dbs 模式跳过 BFD总体可以压到 500GB 左右但搜索范围小了预测精度会打折扣。数据库大致是这么分布的数据库作用磁盘占用约UniRef90MSA 搜索主库59GBMGnify宏基因组序列库64GBBFD大而全的蛋白库1.7TBUniProtmultimer 配对辅助库12GBPDB70模板搜索19GBPDB mmCIF / seqres模板结构200GB内存也是个隐形杀手。很多人只盯着显卡显存结果在 MSA 阶段直接看到进程被系统 Kill。jackhmmer 搜索时每个进程的物理内存消耗虽然不算夸张但 AlphaFold2 会同时维护 MSA 的临时文件和内存对象遇到多链蛋白时32GB 内存会比较紧张64GB 以上才舒服。所以配置前先确认三件事磁盘够不够 2.5TB内存凑不凑得到 64GB以及这张卡在不需要跑 Docker 的前提下能不能被系统识别。2. Conda环境搭建把一条Dockerfile拆成十步慢慢盘2.1 换源并创建python3.9环境Conda 裸装 AlphaFold2第一步是从创建干净环境开始。我一般会先把 conda 源切到国内可用的镜像这样后面装依赖会快很多。conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/ conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/ conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ conda config --set show_channel_urls yes接着创建环境conda create -n alphafold python3.9 -y conda activate alphafold为什么锁 python3.9而不是现在风头正盛的 3.11因为 AlphaFold2 v2.3.x 官方只声明支持 3.9/3.10OpenMM 8.0 对 3.9 的轮子最成熟。3.11 虽然也能硬装但我不建议在排错成本最高的时刻去踩超出官方测试范围的坑。如果你是在脚本里非交互式地 activate经常遇到 “CommandNotFoundError: Your shell has not been properly configured to use conda activate” 之类的提示。这个报错本质上不是 conda 没装好而是 shell 没有执行过 conda init。解决办法conda init bash source ~/.bashrc在系统管理脚本里更稳妥的写法是source /opt/miniconda3/etc/profile.d/conda.sh conda activate alphafold2.2 补齐hmmer/hhsuite/kalign这些“非pip”依赖AlphaFold2 的 MSA 搜索不是用 Python 库完成的而是直接调用外部的命令行工具jackhmmer来自 HMMER 套件、hhblits/hhsearch来自 HH-suite、kalign。这些工具二进制不在 PyPI 上所以不能用 pip 装。最省事的办法是靠 Conda 的 bioconda 通道。conda install -c conda-forge -c bioconda hmmer hhsuite kalign -y在服务器上没有 sudo 权限的时候这一条命令能省掉非常多手工编译的麻烦。手工编译 HMMER 本身不复杂但 HH-suite 依赖 3.3.0 以上的 CMake 配置装错版本会导致 hhblits 运行时报一堆隐晦的段错误。用 conda 装完以后建议先验证一下which jackhmmer which hhblits which hhsearch which kalign只要这四个路径都能打印出来AlphaFold2 最基础的外部依赖就到位了。2.3 安装JAX/TensorFlow/OpenMM版本锁死的背后逻辑克隆官方仓库之后安装 Python 依赖的顺序很重要。先看看仓库里有什么git clone https://github.com/google-deepmind/alphafold.git cd alphafold然后安装 requirements.txt再强制覆盖成 GPU 版 JAXpip install -r requirements.txt pip install --upgrade jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html这里有个大多数人第一次都会搞错的地方requirements.txt 里锁定的是 CPU 版 JAX因为官方在纯 CPU 环境也能跑推理只是慢到怀疑人生。你要做的不是删掉 requirements.txt而是在它之后再用 CUDA 轮子覆盖一遍 JAX。等安装完成建议再精确锁一次版本防止后续 pip 操作把 jaxlib 降级pip install jax0.4.14 jaxlib0.4.14cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlTensorFlow 那块同样有讲究。官方 requirements.txt 用的是tensorflow-cpu。AlphaFold2 的模型推理发生在 JAX 里TensorFlow 只承担一些数据准备、指标计算之类的辅助工作所以没必要装 tensorflow-gpu。很多人手滑装了 GPU 版 TensorFlow等于在同一进程里引进了另一套 cuDNN/CUDA 依赖轻则报一堆找不到符号的错误重则直接把 JAX 的 GPU 上下文搞崩。OpenMM 负责最后的能量优化 relax 阶段。用 conda 安装会比 pip 更省心conda install -c conda-forge openmm8.0.0 -y装完验证导入import openmm print(openmm.__version__)如果是 v2.2 及更早的 AlphaFold2OpenMM 的补丁脚本是必须的到了 v2.3 改用 OpenMM 8.0这个历史包袱基本卸掉了。但前提是别在同一个环境里混装 pip 版 OpenMM 和 conda 版 OpenMM我遇到过的 platform 加载问题九成源于这个。2.4 运行前做一次快速的GPU可见性检查环境装完先不要急着跑完整预测。用 Python 确认 JAX 真的能看到 GPUimport jax print(jax.devices())如果能打印出类似[StreamExecutorGpuDevice(0, Tesla V100-SXM2-16GB)]的输出说明 JAX 的 CUDA 路径已经通了。如果只打印出 CPU 设备就去检查 jaxlib 的版本和 LD_LIBRARY_PATH别继续往下跑。接在 VSCode 或 PyCharm 里调试时也记得把解释器切到alphafold这个 conda 环境否则你在终端里明明能 import一进 IDE 就报 ModuleNotFoundError。这个问题看起来低级但我见过不少人在这个环节浪费了一晚上。3. 数据库与multimer参数该下载的从来不只是模型权重3.1 全套数据库的体积清单与下载策略官方仓库提供了数据库下载脚本。假设数据目录是/data/alphafold直接执行bash scripts/download_all_data.sh /data/alphafold这个脚本会把 UniRef90、MGnify、BFD、UniProt、PDB70、PDB mmCIF、PDB seqres、obsolete 映射等一次性拉下来。如果你只用单体小蛋白理论上可以用 reduced_dbs 省下 BFD 那 1.7TB但下载脚本仍然会把全套拉下来除非你手动挑库。下载阶段最常见的痛苦是“下到一半断网又重新开始”。脚本底层用的是 wget稳定性一般。如果服务器在公网带宽不稳定的环境我建议用 screen 或 tmux 把脚本挂到后台大文件改用 aria2c 配合多线程手动拉conda install -c conda-forge aria2 -y aria2c -x 16 -s 16 某个数据库的下载URL数据库目录结构建议保持官方脚本默认的规范因为 run_alphafold.py 在拼路径时会去固定的子目录里找文件。如果你自己改动目录层级后续填--uniref90_database_path这类参数的时候很容易填错。3.2 multimer模型权重与monomer权重怎么一起拿到很多人以为 multimer 是 AlphaFold2 的一个分支版本要单独下载整套软件。其实不是。同一个仓库、同一套依赖同时支持单体和多体区别只在模型参数文件和运行参数。模型参数下载脚本是在download_all_data.sh里自动执行的但如果你数据库是之前单独拉的可以显式补一次bash scripts/download_alphafold_params.sh /data/alphafold下载完之后数据目录里会出现类似params_model_1.npz一直到params_model_5.npz以及对应的params_model_1_multimer.npz到params_model_5_multimer.npz。每个 npz 文件都对应一个模型副本AlphaFold2 会跑 5 个模型的 ensemble最后按置信度排序输出。monomer 和 multimer 的参数是完全独立的。你只在跑 multimer 时才会加载*_multimer.npz跑单体时加载的是不带 multimer 后缀的那份。3.3 FASTA多链序列写法与第一个multimer任务multimer 模式对 FASTA 的格式有额外要求同一个 fasta 里多链不是用多个分开而是在同一个序列行里用冒号:把链序列连起来。比如一个二聚体my_complex MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQ:GSKGKLGCSDLKDLKSDLKSDLFASDLKASDKF这是官方支持的写法。你用多个头写成两条 record反而会被当作多个独立单体任务处理跑出来的结果不是复合物。运行命令大致是这个样子python run_alphafold.py \ --fasta_paths/data/jobs/my_complex.fasta \ --data_dir/data/alphafold \ --output_dir/data/output \ --model_presetmultimer \ --db_presetfull_dbs \ --max_template_date2024-01-01 \ --enable_gpu_relaxTrue第一次跑通 multimer 时观察两件事一是日志里有没有 “Running model_1_multimer” 的字样二是输出目录的 JSON 结果里有没有iptm这个字段。iptm 是复合物界面预测质量的参考指标。没有iptm的输出说明你运行命令里根本没有用 multimer 模型多半是--model_preset漏了或者拼错了。4. GPU加速优化把等待时间从小时压到分钟4.1 JAX的显存预分配它会一口气吞掉大半张卡JAX 在首次执行 GPU 计算时会按照 XLA 的策略预分配显存。默认情况下 JAX 会预分配 75% 甚至更高的显存这个设计和很多人的直觉相反——不是“用到再申请”而是“先占住一大块”。如果这张卡是专用跑 AlphaFold2 的那预分配是好事可以减少反复申请显存的损耗。但如果一张卡还要跑别的任务预分配就会导致别的进程分配不到显存。按需调整export XLA_PYTHON_CLIENT_MEM_FRACTION0.9或者完全关闭预分配export XLA_PYTHON_CLIENT_PREALLOCATEfalse我的经验是专用卡用MEM_FRACTION0.9共享卡用PREALLOCATEfalse。但注意关闭预分配之后显存碎片化可能更严重长序列蛋白反而更容易爆显存所以共享卡上跑大蛋白时要多留几个心眼。4.2 MSA搜索阶段才是真正的“吃CPU大户”如果你盯着 nvidia-smi看到 GPU 利用率不高、显存占用也不高但整个任务就是卡住不动不要立刻怀疑 GPU 出问题了。AlphaFold2 的流程分两大块前面是 MSA 搜索和模板搜索几乎纯 CPU后面才是 GPU 上的模型推理。对于单体蛋白jackhmmer 在 UniRef90 里搜一次可能就要十几分钟到半小时。对于 multimer每个链都要做搜索还要额外跑 UniProt 辅助配对时间更长。所以如果你跑批的是几十个蛋白真正值得优化的不是 GPU 推理而是把多个蛋白拆成多个进程并行让 CPU 和 GPU 同时别闲着。reduced_dbs 模式在这个阶段很有价值。想快速验证整个环境通不通、不想等半小时 MSA可以先跑 reduced_dbs--db_presetreduced_dbs等验证通过再切回 full_dbs。4.3 relax阶段用GPU还是CPU得看显存余量AlphaFold2 最后会把预测结构丢给 OpenMM 做一轮能量优化官方叫 relax。这个阶段支持 GPU 加速参数是--enable_gpu_relaxTrue。我个人的实测感觉是小单体蛋白500 残基以下GPU relax 和 CPU relax 的差距不明显CPU 跑也就几分钟但到了大三聚体、四聚体GPU relax 可以把几十分钟压缩到几分钟。前提是 GPU 推理阶段结束之后显存还有余量给 OpenMM。如果显存已经贴着上限跑relax 阶段再去申请显存反而可能触发 OOM这时候老老实实保留 CPU relax 更稳。另外提醒一点OpenMM 的 GPU 加速依赖 CUDA platform。如果你在 conda 里装了 OpenMM 但没装 cudatoolkit或者 LD_LIBRARY_PATH 里混进了别的 CUDA 版本relax 阶段会出现 “CUDA platform is not available” 的报错。可以先跑一个最小测试from openmm import Platform print([Platform.getPlatform(i).getName() for i in range(Platform.getNumPlatforms())])输出里要有 CUDAGPU relax 才能正常工作。4.4 多卡利用率单任务一张卡多任务排队跑AlphaFold2 官方代码没有在单任务内部做数据并行或多卡拆分。你在一张 A100 上跑的单体蛋白不会因为插了 8 张卡就变快 8 倍。所以想让整个集群满负荷正确思路是把不同蛋白的任务分别塞到不同的 GPU 上。一个朴素又可靠的跑批方案for i in 0 1 2 3; do CUDA_VISIBLE_DEVICES$i nohup python run_alphafold.py \ --fasta_paths/data/jobs/protein_$i.fasta \ --data_dir/data/alphafold \ --output_dir/data/outputs/run_$i \ --model_presetmonomer \ --db_presetfull_dbs \ --max_template_date2024-01-01 \ /data/logs/run_$i.log 21 done这样四个任务各自绑一张卡。如果你用的是 Slurm也可以做成 array job逻辑一样。看 nvidia-smi 时不要期待单张卡一直 100% 满载很多时间其实花在 CPU 的 MSA 搜索上只要 GPU 在“有活干”的窗口里能把算力跑满就已经是健康的利用率了。5. 踩坑实录跑不通时我都是这么定位的5.1 conda activate报错与shell初始化Conda 环境建好之后最常见的第一个报错就是这个CommandNotFoundError: Your shell has not been properly configured to use conda activate.原因很直白你的 shell 没有初始化 conda 的 hook。新装的 Miniconda 不会自动修改当前终端的配置。解决conda init bash source ~/.bashrc如果是通过 nohup 后台脚本执行任务脚本里也要先 source conda.sh 再 activate否则 crontab 或 systemd 之类的环境里根本找不到 conda 命令。5.2 JAX找不到libcudnn.so.8不是没装是装乱了JAX 的 CUDA 轮子对 cuDNN 版本非常敏感。v0.4.14 对应的 CUDA 轮子期望的是 cuDNN 8.2如果你的环境里因为别的包把 cuDNN 顶到了 8.6 或 9.x运行时就会报Could not load dynamic library libcudnn.so.8排查步骤执行python -c import jax; print(jax.devices())看是不是已经悄悄回退到 CPU。nvidia-smi确认显卡能被系统看到。echo $LD_LIBRARY_PATH看有没有旧版 CUDA toolkit 的路径混在后面。在用户目录下找一下有没有被 pip 塞进去的旧 jaxlib有的话先卸载干净。最粗暴但有效的修复还是那两行命令pip uninstall jax jaxlib -y pip install jax0.4.14 jaxlib0.4.14cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlJAX 的 CUDA 轮子大部分时候是自带 CUDA/cuDNN 库的真正导致报错的原因往往不是“缺”而是“乱”——同一个环境里存在两个互相冲突的版本。所以我一再建议AlphaFold2 只和它自己的依赖住在一个 conda 环境里不要顺手把做深度学习用的 PyTorch 也装进去。PyTorch 的 CUDA 依赖和 JAX 的依赖并不是总能和平共处。5.3 OpenMM的Platform加载失败与relax阶段报错relax 阶段如果出现Platform not found: CUDA先检查 OpenMM 是不是 conda-forge 版再看LD_LIBRARY_PATH里有没有把 conda 的 lib 目录排到前面export LD_LIBRARY_PATH$CONDA_PREFIX/lib:$LD_LIBRARY_PATH还有一点OpenMM 8.0 的 Python API 已经是import openmm不是早期版本的from simtk.openmm import *。如果网上搜到老教程抄代码时得多留个心眼。检查方式python -c import openmm; print(openmm.Platform.getNumPlatforms())能输出大于等于 2CPU、CUDA就是正常的。5.4 MSA阶段进程消失先看是不是内存被杀有时候任务跑着跑着终端没有任何报错就是进程没了。这种时候别只顾着看 GPU先执行dmesg | tail -n 50如果里面有Out of memory: Killed process基本可以确定是物理内存不够被内核 OOM Killer 干掉了。解决方案不是调显存而是要么给机器加内存或 swap要么拆小任务/使用 reduced_dbs 减小 MSA 搜索规模。如果错误信息是 XLA 的RESOURCE_EXHAUSTED那才是显存不够。应对方式在前面讲过调低XLA_PYTHON_CLIENT_MEM_FRACTION或者换更大显存的卡或者换更短的序列。在这个问题上没有动态魔法参数能彻底解决attention 的平方复杂度就摆在那里。5.5 GPU占用不高但很“卡”的错觉“GPU 没满、显存没满、但任务就是慢”的场景大概率发生在第一阶段 MSA 搜索或 XLA 编译。XLA 在第一次执行某个形状的模型时要做编译这段时间 GPU 可能完全没有负载纯 CPU 在顶。不要因为 nvidia-smi 看起来“闲着”就反复杀进程重启那只会让编译一遍遍重新来。判断方法很简单看日志进度。如果日志停在 MSA 相关步骤它就是正常的 CPU 阶段如果日志停在 “Running model” 且 GPU 又不工作再检查 JAX 是否真的用了 GPU。6. 配置完成后的验证与日常使用建议6.1 用gpu-burn给显卡做个体检我不太建议一装完环境就直接跑十几个小时的预测。新卡、二手卡、租的云 GPU最好先跑几分钟压力测试确认没有硬件层面的隐性故障。开源工具 gpu-burn 很方便git clone https://github.com/wilicc/gpu-burn.git cd gpu-burn make ./gpu_burn 300300 秒跑完日志里会给出实际计算速度。如果和官方标称的浮点算力差得离谱或者中途出现显存报错尽早换卡别让 AlphaFold2 给你“试错”。6.2 从输出JSON确认模型真的跑对了跑完一个任务--output_dir下会生成 PDB 文件和 pickle/json 结果。单体看 pLDDT90 以上是高置信度70-90 是可以参考的置信区间50 以下基本不能用于真实结构判断。multimer 还要多看一个iptm。它衡量的是链间界面的预测一致度0.8 以上说明复合物的整体拓扑大概率可信0.5-0.8 属于有待实验验证的范围。AlphaFold2 最终给出的ranking_confidence是把 pLDDT 和 iptm 按一定权重合成出来的源码里在 confidence 模块里数值越高越优先看这个模型。另外一个验证点确认输出结构数量。默认 5 个模型都会出结果如果输出目录里只有 1 个模型说明运行参数里可能限制了模型数量或者运行中断了一半。重新看刚才的日志重点找Running model_1_multimer到Running model_5_multimer这几行。6.3 长期跑批的三条经验跑批多了之后最影响效率的反而不是依赖而是文件管理和日志。我总结三条最实用的经验第一数据库目录设为只读固定在某个共享存储路径下。每个任务都用同一份数据库不要重复下载。第二为每个任务写独立日志记录开始时间、每个阶段耗时、最终耗时。AlphaFold2 的 MSA 阶段和推理阶段耗时差异非常大有日志之后才能判断哪些蛋白适合优先跑哪些可以先在 reduced_dbs 下试一轮。第三重跑同一序列前先备份旧输出。AlphaFold2 的输出目录不会自动跳过去重命名直接覆盖会导致之前的中间结果和结构文件丢失。多肽段的实验数据本来就很贵别让输出管理拖后腿。如果你和我一样是在一个 “有卡、没 Docker 权限、又要多人共用” 的集群里运维 AlphaFold2把这套 Conda 流程里的版本关系理顺之后后面加卡、加人、切换模型预设都会很顺手。我现在的日常习惯是AlphaFold2 用独立的alphafold环境PyTorch 训练用另一个环境物理隔离各管各的 CUDA。两个星期踩坑换来的就这几个字——版本隔离比聪明更重要。
