第三章 训练环境搭建：硬件需求与软件配置

说实话，搞AlphaFold训练，环境搭不好后面全是坑。我见过太多人花几周配环境，结果一跑就崩。今天咱们把这块彻底理清楚——从硬件怎么选，到软件怎么装，再到用Docker一步到位。走起。

3.1 硬件选型：GPU还是TPU？

你第一个要回答的问题是：你到底用GPU还是TPU？我个人的观点很明确——如果你只做推理或微调，一块高显存GPU就够了。但要是你想从头训练完整的AlphaFold2模型，那TPU几乎是刚需。为什么？显存和算力差距摆在那儿。

GPU：显存是绝对瓶颈

AlphaFold的模型参数加上中间激活，一次前向就要吃掉大几十GB。我记得第一次拿RTX 3090（24GB显存）跑小批量测试，batch size设为1都差点炸了。如果你只有单卡，建议至少40GB显存——A100 40GB或80GB、H100、或者A6000。V100 32GB勉强能跑，但得把模型精度降到bf16，配合梯度累积。

下面是推荐配置：

GPU型号	显存	适合场景	备注
A100 80GB	80GB	完整训练 / 微调	最佳选择
A100 40GB	40GB	微调 / 中等规模训练	需用梯度累积
H100 80GB	80GB	新架构，算力更强	驱动/库需更新
V100 32GB	32GB	小批量调试	显存紧张，建议降精度
RTX 4090	24GB	仅推理 / 轻量微调	训练？别想了

TPU：云上大力飞砖

TPU是谷歌的亲儿子，DeepMind训练AlphaFold用的就是TPU Pod。TPU v3-32或v4-64，显存共享到几百GB。但你要注意：TPU只能用JAX，不能跑PyTorch。而且TPU的云端计费很贵，我当年用过一次TPU v3-8跑测试，几个小时就花了上千块。

内存和存储也不能忽视

除了GPU，CPU内存至少128GB起步——因为数据预处理时，你要把整个PDB结构数据加载到内存。曾经我搭了一台64GB内存的机器，结果预处理阶段直接OOM，swap都打满了。存储方面，建议上NVMe SSD，至少2TB。AlphaFold的训练数据和中间特征文件加起来轻松超过1TB。

# 我的内存监控脚本（片段）
import psutil
while True:
    mem = psutil.virtual_memory()
    if mem.percent > 90:
        print("警告：内存使用超过90%，请检查进程！")
    time.sleep(10)

3.2 软件依赖：从Python到JAX

硬件到位后，软件环境就是最容易踩坑的地方。简单来说，整个依赖链长这样：

1. 系统层：Ubuntu 20.04或22.04，不要用CentOS（CUDA驱动兼容性差）
2. GPU驱动：NVIDIA驱动 ≥ 525，用官方runfile装最稳
3. CUDA工具包：CUDA 11.8或12.1，看JAX版本要求
4. cuDNN：8.6以上，要和CUDA匹配
5. Python：3.10或3.11，别用3.12（很多库还没适配）
6. JAX：硬性要求，AlphaFold的加速核心
7. Flax + Haiku + Optax：模型框架和优化器库
8. 其他依赖：numpy、tensorflow、pandas、h5py等

我一般用Miniconda管理环境，下面是完整的安装命令（以CUDA 12.1为例）：

# 创建环境
conda create -n alphafold python=3.10
conda activate alphafold

# 安装CUDA相关（通过pip/cuda可以免装独立CUDA）
pip install "jax[cuda12_local]==0.4.35" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# 安装AlphaFold依赖
pip install dm-haiku optax flax tensorflow pandas h5py dm-tree

# 验证JAX是否识别GPU
python -c "import jax; print(jax.devices())"

版本兼容性速查表

JAX版本	CUDA版本	cuDNN版本	推荐Python
0.4.35	11.8 / 12.1	8.6+	3.10
0.4.24	11.8	8.6	3.9~3.11
0.4.10	11.4~11.8	8.3~8.6	3.8~3.10

为什么我强调版本？因为JAX对CUDA小版本很敏感，差了0.1都可能报"cuda driver too old"。我的经验是：先确定JAX版本，再去NVIDIA官网下载对应的CUDA和cuDNN，不要盲目升级。

3.3 使用Docker镜像快速搭建

最省心的方式——用Docker。AlphaFold官方已经提供了现成镜像，你只需要一个GPU驱动和NVIDIA Container Toolkit，其他全在容器里。我个人认为，Docker是避免环境配置打架的唯一正道。尤其是当你需要同时跑不同版本的AlphaFold时，一个容器一个环境，互不干扰。

具体步骤：

1. 安装Docker和nvidia-docker2
2. 拉取官方镜像（以v2.3.2为例）
3. 启动容器并挂载数据目录

# 1. 安装NVIDIA Container Toolkit（假设已安装Docker）
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update && sudo apt-get install -y nvidia-container-toolkit
sudo systemctl restart docker

# 2. 拉取AlphaFold镜像
docker pull alphafold2:v2.3.2  # 如果官方有更新，换成最新tag

# 3. 运行容器（注意挂载数据和配置）
docker run --gpus all -it \
    -v /path/to/data:/data \
    -v /path/to/output:/output \
    -v /path/to/config:/config \
    alphafold2:v2.3.2

自定义你的Docker镜像

官方镜像不一定完全满足你的需求——比如你想集成自己的预处理脚本，或者换一个JAX版本。这时候最好自己写Dockerfile。下面是我常用的模板：

FROM nvidia/cuda:12.1.0-runtime-ubuntu20.04

# 安装基础工具
RUN apt-get update && apt-get install -y wget git python3-pip

# 设置Python
RUN pip install --upgrade pip
RUN pip install jax[cuda12_local]==0.4.35
RUN pip install dm-haiku optax flax tensorflow

# 拷贝AlphaFold代码
COPY ./alphafold /app/alphafold
WORKDIR /app

# 默认入口
CMD ["python", "-m", "alphafold.run_alphafold"]

构建命令：

docker build -t my-alphafold:v1 .

这样你就有完全可控的训练环境了。我强烈建议每个项目都单独build一个镜像，记录Dockerfile到Git里——这样复现结果的时候不会抓狂。

好了，环境搭好只是第一步。后面咱们要聊数据怎么预处理、配置文件怎么写——别急，慢慢来。