接入 NVIDIA A100、吞吐量提高 10 倍！Milvus GPU 版本使用指南

作为为数不多的支持 GPU 的向量数据库产品，Milvus 2.3 在吞吐量和低延迟方面都带来了显著的变化，尤其是与此前的 CPU 版本相比，不仅吞吐量提高了 10 倍，还能将延迟控制在极低的水准。

不过，正如我前面提到的，鲜有向量数据库支持 GPU，这其中除了有技术门槛较高的因素外，还涉及诸多不确定性的问题。那么，Milvus 为什么要做一件充满挑战的事情？

回想 Milvus 决定支持 GPU 的场景，很多细节仍旧历历在目。当时，随着 LLM 的兴起，用户对于向量数据库的性能提出了更高的要求，尤其是在一些对性能、延迟有着极高要求的场景，只通过 CPU 索引来支撑的难度越来越高，而 GPU 有着非常强大的并行处理能力。

因此，我们决定在 Milvus 2.3 版本中支持 GPU。幸运的是，来自 NVIDIA 的小伙伴给予了我们诸多支持，他们主动给我们提供了Rapid Raft GPU 索引接入 Milvus 的支持代码。在 NVIDIA 和 Milvus 团队的共同努力下，Milvus GPU 版本如约而至。

接下来，我们就来看看如何使用 Milvus GPU 版本。

01.

CUDA 驱动安装

首先，在我们的宿主机环境中，需要检查系统中是否已经正确的识别 NVIDIA 显卡，在命令行中输入：

lspci

在输出的设备中，看到 NVIDIA 字段，则说明该系统中已经安装了 NVIDIA 显卡。

00:00.0 Host bridge: Intel Corporation 440FX - 82441FX PMC [Natoma]
00:01.0 ISA bridge: Intel Corporation 82371SB PIIX3 ISA [Natoma/Triton II]
00:01.3 Non-VGA unclassified device: Intel Corporation 82371AB/EB/MB PIIX4 ACPI (rev 08)
00:03.0 VGA compatible controller: Amazon.com, Inc. Device 1111
00:04.0 Non-Volatile memory controller: Amazon.com, Inc. Device 8061
00:05.0 Ethernet controller: Amazon.com, Inc. Elastic Network Adapter (ENA)
00:1e.0 3D controller: NVIDIA Corporation TU104GL [Tesla T4] (rev a1)
00:1f.0 Non-Volatile memory controller: Amazon.com, Inc. NVMe SSD Controller

以上是我的环境中的输出，可以看到其中识别到一张 NVIDIA T4 显卡。

接下来，可以去 NVIDIA 官方网站，添加 repo installer。

https://developer.nvidia.com/cuda-downloads

以 ubuntu 20.04 为例：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update

这里需要注意，如果宿主机没有 cuda 的需要，可以不安装 cuda；如果是 server 版本的操作系统，可以安装 headless 版本驱动；如果是 desktop 版本的系统，可以安装普通版本。

• server

sudo apt install nvidia-headless-535
sudo apt install nvidia-utils-535

• desktop

sudo apt install nvidia-driver-535
sudo apt install nvidia-utils-535

安装完成之后，需要重启一下系统，使得驱动生效。重启完成，可以输入：

nvidia-smi

如果可以看到详细的显卡状态，即表示驱动安装 OK。

Milvus GPU 版本镜像使用 cuda 11.8 打包，如果是 NVIDIA Tesla 系列专业显卡，需要的最小驱动版本>=450.80.02；如果是游戏显卡，需要驱动版本>=520.61.05。

Milvus GPU 镜像支持 Compute Capability 为 6.1、7.0、7.5、8.0 的 NVIDIA 显卡，查看显卡型号对应的Compute Capability，请参阅 https://developer.nvidia.com/cuda-gpus。NVIDIA Container Toolkit 安装则参考 https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html

02.

Milvus GPU 可配环境变量

Milvus GPU 版本目前仅支持单个 Milvus 进程单个显卡，Milvus GPU 版本默认使用 device 0 显卡。如果需要使用多卡，可以通过启动多个 Milvus 进程或者容器，然后配合 CUDA_VISIBLE_DEVICES 环境变量来实现多卡部署。

容器环境下，可以通过 -e 来设置该环境变量:

sudo docker run --rm -e NVIDIA_VISIBLE_DEVICES=3 milvusdb/milvus:v2.3.0-gpu-beta

在 docker-compose 环境中，可以通过 device_ids 字段来设置，参考 https://docs.docker.com/compose/gpu-support/

注意，为单个 Milvus 进程或者单个容器配置了多张卡可见，Milvus 也只能使用其中的一张卡。

KNOWHERE_STREAMS_PER_GPU 环境变量可以用来设置 cuda stream 的并发数。适当的调大此参数，有可能获得更好的性能，但是，也会带来更多的显存开销。

KNOWHERE_GPU_MEM_POOL_SIZE 环境变量可以用来设置显存池大小。如果不设置改环境变量，Milvus 会自动分配当前 GPU 的一半内存作为显存池，如果在服务的过程中，出现显存池容量不足，那么 Milvus 会自动尝试再次增加显存池大小，默认上限是整个显存大小。

export KNOWHERE_GPU_MEM_POOL_SIZE=2048;4096

以上设置显存池初始大小为 2048 MB，最大显存池大小为 4096 MB。

如果在一张卡上部署 2 个 Milvus进程，那么，此环境变量一定需要合理的分配，否则 Milvus 会出现显存竞争崩溃的情况。

03.

Milvus GPU 编译

本地编译 Milvus GPU 版本需要依赖英伟达提供的 cuda-toolkit，在安装 cuda-toolkit 请先完成 NVIDIA 驱动的安装：

sudo apt install --no-install-recommends cuda-toolkit 

本地编译 Milvus GPU 之前，我们需要先安装部分依赖软件以及工具：

sudo apt install python3-pip libopenblas-dev libtbb-dev pkg-config

安装 conan：

pip3 install conan==1.59.0 --user
expoprt PATH=$PATH:~/.local/bin

安装较新版本的 cmake>=3.23，参考 https://apt.kitware.com。golang 的安装可以参考 https://go.dev/doc/install。

make milvus-gpu

启动 milvus standalone 模式：

cd bin
sudo ./milvus run standalone

在配置好 nvidia-docker 或者 docker 替代品的情况下，可以很方便地使用 Milvus 提供的 docker-compose.yml 文件来完成容器化部署。

用户可以从 milvus repo 中获取 docker-compose.yml 文件，地址 https://github.com/milvus-io/milvus/blob/master/deployments/docker/gpu/standalone/docker-compose.yml

 docker-compose up -d

至此，便可完成 standalone 模式的 Milvus 部署。

如果用户宿主机有多张显卡，可以通过修改 docker-compose.yml 中的 device_ids 字段来修改映射到 Milvus 的显卡。

以上就是关于 Milvus GPU 版本使用说明的全部内容，大家有任何疑问都可以与我们交流。下一篇文章，我们将详解 Knowhere 2.0，欢迎阅读！