在运行时配置 NIM#

AlphaFold2 NIM 允许在运行时进行多项自定义，从而实现

  • 调整 NIM 的性能

  • 利用您独特的硬件

  • 修改运行时变量（例如 NIM 缓存的位置）

下面，我们将详细介绍各种 AlphaFold2 NIM 特定的运行时配置选项。

简要回顾：启动 AlphaFold2 NIM#

让我们首先回顾一下如何启动 NIM

export LOCAL_NIM_CACHE=~/.cache/nim

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:2.0.0

记住

  • 选项 -p 设置 NIM 的端口

  • -e 选项定义环境变量，这些变量在运行时传递到 NIM 的容器中。

  • --rm 在容器存在时将其移除

  • -it 允许直接在 CLI 与容器交互

配置 NIM 缓存的位置#

NIM 缓存的位置可以在本地文件系统和 NIM 内部进行更改。当需要更大的驱动器来存储模型权重时，这非常有用。

要在主机文件系统上更改 NIM 缓存位置，请更改上述命令中的源挂载，我们在环境变量 LOCAL_NIM_CACHE 中设置了该挂载

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:2.0.0

这只会影响容器外部的位置。要更改容器内部缓存的位置，您必须

  1. 确保挂载新的目标路径

  2. 在 NIM_CACHE_PATH 环境变量中提供新路径。

这是一个示例，我们将容器内的缓存移动到 /workdir/cache/af2/

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_API_KEY \
    -e NIM_CACHE_PATH=/workdir/cache/af2 \ ## Note: we must set this environment variable...
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \ ## and provide the destination mount.
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:2.0.0

为 NIM 请求使用备用端口#

如果您有其他 HTTP 或 gRPC 服务器正在运行（例如，其他 NIM），您可能需要通过为您的 NIM 使用另一个端口来使 8000 端口可用。这可以通过以下方式完成

  1. 通过设置 -p 选项来更改公开的端口。

  2. 将 NIM_HTTP_API_PORT 环境变量设置为新端口。

这是一个示例，我们将 NIM 设置为在端口 7979 上运行

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_API_KEY \
    -e NIM_HTTP_API_PORT=7979 \ ## We must set the NIM_HTTP_API_PORT environment variable...
    -e NIM_CACHE_PATH=/workdir/cache/af2 \
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \
    -p 7979:7979 \ ## as well as forward the port to host.
    nvcr.io/nim/deepmind/alphafold2:2.0.0

配置并行 MSA 运行器的数量和每个运行器可用的线程数#

AlphaFold2 NIM 在许多方面都得到了优化。其中一项优化包括能够根据您的机器上可用的 CPU 核心数、可用的 GPU 数量以及可用 SSD 的速度来调整 MSA 流程以获得更好的性能（注意：我们不建议在使用硬盘驱动器进行存储的机器上运行 AlphaFold2 NIM，因为它会显着降低 NIM 的性能）。

AlphaFold2 NIM 内 MSA 有两个可调并行层，由两个不同的环境变量设置

  1. NIM_PARALLEL_MSA_RUNNERS：此变量控制 MSA 流程的粗粒度并行性。

  2. NIM_PARALLEL_THREADS_PER_MSA：此变量控制每个 MSA 流程使用的线程数。

提示

NIM_PARALLEL_MSA_RUNNERS 和 NIM_PARALLEL_THREADS_PER_MSA 的乘积应 >= NIM 可用的 CPU 核心数。

例如，在具有 48 个 CPU 核心的机器上，一个合适的设置是

NIM_PARALLEL_MSA_RUNNERS=4
NIM_PARALLEL_THREADS_PER_MSA=12

这将分配 4 个运行器，每个运行器有 12 个线程。

注意

我们不建议使用超过 5 个运行器，并且每个 MSA 的线程数通常应在 8 到 16 之间。

注意

NIM_PARALLEL_MSA_RUNNERS=3 和 NIM_PARALLEL_THREADS_PER_MSA=8 的默认值适用于更多机器，并且大多数用户不需要修改这些变量。

要在运行时修改这些变量，请在 NIM 启动命令中将它们作为环境变量传递。例如，在 64 核机器上，当运行许多 MSA 数据库时，我们可能会将 NIM_PARALLEL_MSA_RUNNERS 增加到 6，并将 NIM_PARALLEL_THREADS_PER_MSA 减少到 10，以便我们使工作负载更好地适应我们的机器

docker run --rm --name alphafold2 --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_API_KEY \
    -e NIM_PARALLEL_MSA_RUNNERS=6 \ ## We set the number of parallel runners to 6, perhaps because we have used many MSA databases...
    -e NIM_PARALLEL_THREADS_PER_MSA=10 \ ## and we set the number of threads per MSA to 12, to 
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2:2.0.0
在本页