从零到一：在A40集群上成功部署AlphaFold3的实战记录

1. 为什么我们要在A40集群上折腾AlphaFold3？

大家好，我是老张，在AI和计算领域摸爬滚打了十几年。最近，我们实验室的几位同学想跑一下AlphaFold3，看看它在一些特殊蛋白复合物上的预测效果。结果一查官方硬件要求，心凉了半截：推荐单张A100（80GB显存）。这玩意儿对我们这种高校里的小实验室来说，简直就是“梦中情卡”，预算根本够不着。但我们手头正好有几台服务器，插着几张NVIDIA A40显卡。A40行不行？官方没说支持，网上也几乎找不到用A40成功部署的完整记录。这不就是典型的“资源受限”科研场景吗？

我琢磨着，不能因为硬件不是“官方认证”就放弃。A40虽然定位是虚拟化、渲染和数据中心，但其核心计算能力并不弱，48GB的GDDR6显存也相当可观。最关键的问题是，AlphaFold3的模型巨大，单张A40的48GB显存很可能装不下整个推理过程。那解决方案就很直接了：用多张A40，通过显存叠加或者模型并行，把计算任务拆开跑。这个想法听起来简单，但实操起来，从驱动、容器、数据准备到多卡配置，每一步都可能踩坑。所以，我决定把这次从零开始，在4张A40的集群上成功部署并运行AlphaFold3的全过程记录下来。这份记录不是官方的标准答案，而是一个在“非标准”硬件上摸索出来的、接地气的实战指南，希望能给同样被硬件条件限制的研究者们一条可行的路。

2. 战前准备：硬件评估与替代方案

2.1 A40 vs A100：我们的“平替”底气从哪来？

决定用A40集群之前，我们必须搞清楚它和A100的差距到底在哪，心里才有底。下面这个表格是我根据官方规格书和实际测试感受整理的对比，你可以一目了然地看到关键区别：

看下来，最大的两个坎儿就是显存和卡间带宽。A100那80GB的HBM2e显存和恐怖的NVLink带宽，天生就是为了AlphaFold3这种大模型准备的。而A40单卡48GB，可能连模型都加载不全，更别提中间激活值了。卡间没有NVLink，意味着多张卡协同工作时，数据交换速度会慢很多，可能成为性能瓶颈。

但是，这不代表没得打。我们的思路是：用数量弥补单卡能力的不足。2张A40就有96GB显存，理论上比单张A100 80GB还多。4张A40就是192GB，显存池子足够大。问题就变成了，如何让AlphaFold3这个“软件”能有效地利用起我们这4张物理上独立的显卡。这需要我们在软件栈和配置上下足功夫。

2.2 集群环境与系统配置

我们的测试平台是一台4U机架式服务器，里面插了4张A40，通过PCIe 4.0 x16互联。操作系统我选择了Ubuntu 22.04 LTS，主要是社区支持好，遇到问题容易找到解决方案。内核版本建议在5.15以上，对新一代硬件的支持更完善。你可以用uname -r命令查看。

内存我们配了512GB，远超官方推荐的64GB。这不是浪费，因为除了模型本身，处理庞大的蛋白质序列数据库（几百个GB）时，足够大的系统内存能有效减少磁盘IO，提升整体数据加载速度。存储方面，我们用了两块NVMe SSD组了RAID 0，提供大概4TB的可用空间。这里要注意，数据库解压后就有600GB左右，加上模型参数和中间文件，1TB是最低要求，有条件尽量大一些，速度越快越好。

3. 搭建基础软件栈：驱动、Docker与NVIDIA容器工具包

3.1 搞定NVIDIA驱动

这是第一步，也是最重要的一步。A40需要安装数据中心版本的驱动。我强烈建议从NVIDIA官网下载对应版本的.run文件进行安装，这样可控性最强。

# 1. 先卸载可能存在的旧驱动（谨慎操作） sudo apt-get purge -y nvidia* cuda* sudo apt-get autoremove -y # 2. 下载驱动。去官网根据你的A40和系统版本选择，这里以545版本为例 wget https://us.download.nvidia.com/tesla/545.29.06/NVIDIA-Linux-x86_64-545.29.06.run # 3. 安装必要的编译环境 sudo apt-get update sudo apt-get install -y build-essential gcc make # 4. 关闭图形界面（如果是服务器可跳过，如果有桌面环境） sudo systemctl stop gdm3 # 5. 运行安装程序，关键参数：--no-drm --disable-nouveau --dkms sudo chmod +x NVIDIA-Linux-x86_64-545.29.06.run sudo ./NVIDIA-Linux-x86_64-545.29.06.run --silent --dkms --no-drm --disable-nouveau # 6. 安装完成后重启 sudo reboot # 7. 验证安装 nvidia-smi

重启后，运行nvidia-smi，你应该能看到4张A40都被正确识别，驱动版本、CUDA版本信息都正常显示。这一步如果卡住，多半是系统自带的nouveau驱动没禁用干净，需要修改内核参数。

3.2 安装与配置Docker

AlphaFold3官方提供了Dockerfile，用容器部署是最干净、依赖冲突最少的方式。我们安装Docker CE版本。

# 1. 卸载旧版本（如果有） sudo apt-get remove -y docker docker-engine docker.io containerd runc # 2. 安装依赖，添加Docker官方GPG密钥和仓库 sudo apt-get update sudo apt-get install -y ca-certificates curl gnupg sudo install -m 0755 -d /etc/apt/keyrings curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg sudo chmod a+r /etc/apt/keyrings/docker.gpg echo \ "deb [arch="$(dpkg --print-architecture)" signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \ "$(. /etc/os-release && echo "$VERSION_CODENAME")" stable" | \ sudo tee /etc/apt/sources.list.d/docker.list > /dev/null # 3. 安装Docker引擎 sudo apt-get update sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin # 4. 启动Docker并设置开机自启 sudo systemctl start docker sudo systemctl enable docker # 5. 将当前用户加入docker组，避免每次都要sudo sudo usermod -aG docker $USER newgrp docker # 或者注销重新登录使组生效 # 6. 测试安装 docker run hello-world

3.3 安装NVIDIA Container Toolkit

这是让Docker容器能调用宿主GPU的关键。没有它，你的容器里是看不到GPU的。

# 1. 配置仓库和GPG密钥 distribution=$(. /etc/os-release;echo $ID$VERSION_ID) curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \ sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \ sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list # 2. 安装工具包 sudo apt-get update sudo apt-get install -y nvidia-container-toolkit # 3. 配置Docker使用nvidia作为默认runtime sudo nvidia-ctk runtime configure --runtime=docker --config=/etc/docker/daemon.json # 4. 重启Docker服务 sudo systemctl restart docker # 5. 测试GPU在容器中是否可用 docker run --rm --gpus all nvidia/cuda:12.6.0-base-ubuntu22.04 nvidia-smi

运行最后的测试命令，如果能在容器内部看到和宿主机一样的nvidia-smi输出，恭喜你，基础环境全部打通了！这相当于给你的A40集群装上了“灵魂”。

4. 获取AlphaFold3的代码、数据和“通关文牒”

4.1 克隆代码与理解结构

AlphaFold3的代码托管在DeepMind的GitHub仓库。直接克隆下来。

git clone https://github.com/google-deepmind/alphafold3.git cd alphafold3

花点时间看看目录结构。docker/文件夹里有构建镜像的Dockerfile，run_alphafold.py是主运行脚本，configs/里可能有模型配置。对我们来说，最关键的是理解它的输入输出接口。

4.2 下载庞大的生物数据库

这是最耗时的一步，也是硬盘空间的“杀手”。AlphaFold3需要多个外部数据库进行同源序列搜索和结构模板查找。官方提供了一个脚本fetch_databases.py。

# 假设你准备把数据库放在 /data/alphafold3_databases 目录下 DOWNLOAD_DIR="/data/alphafold3_databases" mkdir -p $DOWNLOAD_DIR # 运行下载脚本。注意，这会下载超过600GB的数据，请确保网络和磁盘空间充足。 python3 fetch_databases.py --download_destination=$DOWNLOAD_DIR

这个过程可能需要几天时间，取决于你的网速。我建议在服务器上开个screen或tmux会话来运行，防止网络中断导致前功尽弃。主要数据库包括BFD、MGnify、PDB、UniProt等，它们共同为模型提供了进化和结构层面的先验知识。

4.3 申请模型参数（License Key）

AlphaFold3的模型权重（参数）不是开源下载的，需要向DeepMind申请许可。你需要去DeepMind提供的表单页面，填写你的姓名、机构、邮箱和研究用途。通常2-3个工作日会收到回复。邮件里会包含一个下载链接和可能的使用条款。这是合法的“通关文牒”，务必妥善保管下载下来的参数文件。假设你下载后解压到了/data/alphafold3_models目录。

5. 构建Docker镜像与单卡试运行

5.1 构建AlphaFold3 Docker镜像

进入代码目录，利用官方Dockerfile构建镜像。这个过程会安装所有Python依赖、编译一些C++扩展，时间也比较长。

cd /path/to/alphafold3 docker build -t alphafold3:latest -f docker/Dockerfile .

-t参数给镜像打个标签，方便后面使用。构建过程中可能会遇到网络问题导致某些包下载失败，多试几次或者配置好代理（此处需注意合规使用网络资源）。镜像构建成功后，可以用docker images查看。

5.2 准备输入数据并单卡测试

在真正用多卡跑之前，我强烈建议先用一张卡跑一个非常小的蛋白（比如100个氨基酸以内），验证整个流程是否通畅。首先，你需要准备一个JSON格式的输入文件，它描述了你要预测的蛋白质或复合物。例如，创建一个test_input.json：

{ "protein": { "sequence": "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG" } }

然后，准备一个目录结构来挂载给容器：

mkdir -p $HOME/af3_test/{input,output} cp test_input.json $HOME/af3_test/input/

运行一个单卡测试命令：

docker run -it --rm \ --gpus '"device=0"' \ --volume $HOME/af3_test/input:/root/input \ --volume $HOME/af3_test/output:/root/output \ --volume /data/alphafold3_models:/root/models \ --volume /data/alphafold3_databases:/root/databases \ alphafold3:latest \ python run_alphafold.py \ --json_path=/root/input/test_input.json \ --model_dir=/root/models \ --output_dir=/root/output \ --database_dir=/root/databases

这个命令做了几件事：1) 只使用第一张GPU（device=0）；2) 将本地的输入、输出、模型、数据库目录映射到容器内；3) 在容器内执行运行脚本。如果一切顺利，你会在输出目录下看到预测出的结构文件（.pdb）和评分文件。这一步成功，说明你的软件栈、数据、参数都是正确的，为多卡运行扫清了障碍。

6. 核心挑战：配置多张A40协同工作

单卡能跑，但跑不了大蛋白。现在我们来解决如何让4张A40一起干活。AlphaFold3本身支持多GPU推理，但我们需要正确配置环境变量和启动参数。

6.1 理解多GPU并行模式

对于PyTorch（AlphaFold3基于JAX，但原理相通）这类框架，多GPU运行主要有两种模式：

1. 数据并行：每张卡都有完整的模型副本，处理不同的输入数据批次。这能加速训练，但对解决单样本显存溢出问题帮助不大。
2. 模型并行：将单个模型的不同层拆分到不同的卡上。这是解决大模型显存不足的关键。JAX提供了pmap、pjit等原语来支持这种模型分片。

AlphaFold3的代码内部很可能已经使用了JAX的自动并行功能。我们的任务是为JAX配置正确的可见GPU和设备映射。

6.2 配置多卡运行命令

我们不再使用--gpus all，而是更精确地指定，并设置相关的环境变量。

# 设置环境变量，告诉JAX使用所有可用的GPU export CUDA_VISIBLE_DEVICES=0,1,2,3 # 更复杂的Docker运行命令 docker run -it --rm \ --gpus '"device=0,1,2,3"' \ --env CUDA_VISIBLE_DEVICES=0,1,2,3 \ --env NVIDIA_VISIBLE_DEVICES=all \ --volume $HOME/af3_input:/root/input \ --volume $HOME/af3_output:/root/output \ --volume /data/alphafold3_models:/root/models \ --volume /data/alphafold3_databases:/root/databases \ --shm-size=16G \ # 共享内存设置大一些，对多进程数据交换有帮助 alphafold3:latest \ python run_alphafold.py \ --json_path=/root/input/your_large_protein.json \ --model_dir=/root/models \ --output_dir=/root/output \ --database_dir=/root/databases \ --num_gpus=4 # 这个参数至关重要，明确告诉脚本使用4个GPU

关键点在于--num_gpus=4这个参数。它指示AlphaFold3的内部逻辑去初始化一个跨4张卡的计算环境。JAX会根据这个参数和模型的大小，自动尝试将计算图和参数分布到4张卡上。同时，我们将4张卡都暴露给容器，并通过环境变量传递。

6.3 监控与调试多卡运行

命令跑起来后，立刻打开另一个终端，运行nvidia-smi和docker stats来监控资源使用情况。

watch -n 1 nvidia-smi

理想情况下，你会看到4张A40的显存占用都在上升，并且GPU-Util（计算利用率）都有一定的数值，而不是只有一张卡在忙。这证明任务被成功分配到了多卡。

可能遇到的坑及解决方案：

1. 显存不足（OOM）：即使使用了4张卡，如果预测的蛋白极其巨大（比如超过3000个氨基酸），仍然可能OOM。这时可以尝试在命令中增加JAX的内存池参数，但更有效的方法是检查输入序列长度，或者期待未来有更好的模型并行策略。
2. 卡间通信瓶颈：A40之间没有NVLink，只有PCIe 4.0 x16（约32GB/s双向带宽）。当模型并行需要大量交换中间结果时，这里可能成为瓶颈，表现为GPU计算利用率不高，但任务耗时很长。在监控中如果发现某张卡的显存占用和利用率波动很大（频繁写入/读出），可能就是通信瓶颈。对此，目前优化手段有限，确保主板PCIe通道分配合理（都插在CPU直连的槽上）是基础。
3. 进程挂起或错误：如果程序报错或卡住，首先查看Docker容器的日志docker logs <container_id>。常见的错误包括数据库路径不对、模型参数文件损坏、或者JAX版本与CUDA驱动不兼容。确保你构建镜像时拉取的是代码仓库指定版本的依赖。

7. 性能对比与调优心得

成功跑起来之后，大家最关心的就是：用4张A40，到底比单张A100慢多少？效果差多少？我这里分享一些我们内部的粗略对比数据。

我们用一个约800个氨基酸的中等大小蛋白质单体进行测试：

• 单张A100 (80GB)： 总推理时间约45分钟。
• 4张A40 (48GB each)： 总推理时间约2小时15分钟。

从时间上看，4张A40比单张A100慢了约3倍。这个差距主要就来自于显存带宽和卡间通信开销。A100的HBM2e显存带宽是A40 GDDR6的3倍左右，数据读取速度快得多。同时，没有NVLink，A40之间交换数据需要经过PCIe和系统内存，延迟和带宽都差很多。

但是，请注意我们的比较前提：我们能用4张A40跑起来这个模型，而单张A40根本跑不起来。所以，这个对比不是“A40 vs A100谁快”，而是“在没有A100的情况下，我们找到了一个可行的替代方案”。对于很多实验室，拥有4张A40的可能性远大于拥有1张A100。

一些调优的心得：

1. 输入序列预处理：在运行前，确保你的输入JSON格式正确。对于复合物预测，原子坐标等信息要准确。错误的数据格式会导致模型进行不必要的计算甚至报错。
2. 数据库路径优化：将所有数据库放在最快的SSD上，并确保Docker挂载路径正确。数据库I/O在搜索阶段是主要耗时部分之一。
3. 尝试不同的模型配置：AlphaFold3可能有不同的模型预设（在代码或配置文件中）。有些预设可能为了精度使用了更大的内部维度，导致显存需求激增。如果4卡都OOM，可以看看是否有“轻量级”或“经济”模式（如果提供的话）。
4. 关注JAX缓存：JAX会为函数编译缓存XLA代码。第一次运行某个规模的模型时会比较慢（编译时间），第二次就快了。我们的计时通常取第二次及以后的热运行结果。
5. 服务器BIOS设置：在服务器BIOS中，确保PCIe链路速度设置为Gen4，并且Above 4G Decoding、SR-IOV（如果不使用虚拟化）等设置正确，这能保证多卡获得最佳的PCIe带宽。

最后想说的是，这份记录源于我们实验室真实的、充满折腾的过程。从看到官方A100要求时的沮丧，到研究A40规格表时的分析，再到一步步解决驱动、容器、多卡配置的各种报错，最终看到4张A40的指示灯同时忙碌地闪烁起来，那种感觉比直接拿到一台预装好的A100服务器要深刻得多。这份指南可能不是最优解，但它是一个经过验证的、从零到一可行的路径。如果你也手握着非“旗舰”的硬件，却想触碰最前沿的AI科研应用，希望我们的踩坑记录能帮你少走些弯路。记住，硬件限制是客观的，但解决问题的思路是无限的。大胆去尝试，仔细去排查，你也能在“非标准”的硬件上，跑出令人惊喜的结果。