Hugging Face 的 LOGO

使用 Infinity 加速 AMD 上的嵌入和重排序模型

社区文章 发布日期：2024年12月3日
赞
6
Michael's avatar
Michael
michaelfeil

这篇客座文章由 michaelfeil/infinity 的维护者撰写，Infinity 是一个流行的开源库，用于在 AMD 上使用高吞吐量引擎推理文本嵌入、重排序、视觉嵌入 (clip, colqwen) 和音频嵌入模型。本文“包含一篇教程，介绍如何通过 Pytorch 和 ONNX 快速部署基于 AMD 的嵌入解决方案在 ROCm 上，并将其与在 Nvidia 上实现相同目标进行比较，以及一份针对 AMD 的快速优化指南。特别感谢 embeddedllm.com 的贡献者支持构建适用于 ROCm 的 Optimum-amd/Onnx！”

让 AMD 更受欢迎？！

为什么选择 AMD GPU——在与 Hackathon 参加者、社区聚会的人交流以及关注 Localllama subreddit 上的讨论时，我常常觉得 AMD 被遗忘，其能力被低估，甚至有人说需要数年才能赶上。与此同时，你可以访问你喜欢的网站（https://pytorch.ac.cn/get-started/locally/），并通过一行安装命令获取带有 rocm 的 torch。社区在某种程度上是正确的——根据 2024 年 12 月 Infinity 匿名用户提交的数据，目前几乎没有 AMD GPU 用于运行 Infinity，上周，仅有 0.7% 的 Infinity 使用量在 AMD 上。

image/png 图片：Infinity 使用情况，2024年11月26日-12月3日，不包括 MPS 和 CPU 目标。

作为一个真正的开源项目，保持厂商中立性很有趣，只要软件易于维护——Infinity 已经添加了 Apple MPS 支持，拥有同类最佳的 CPU 支持，并且还计划在不久的将来兼容 AWS Inferentia——是时候解决 AMD 了！！

教程

回顾：Infinity 如何通过 Docker 在 CUDA 和 CPU 上运行

要在 Docker 上使用加速镜像 + 加速器特定安装。加速器特定安装有助于从 Docker 主机共享驱动程序特定详细信息并启用设备共享。

对于 Nvidia，您只需安装 nvidia-container-toolkit nvidia-container-toolkit。

完成此操作后，Infinity 指令已在几个流行仓库中更新，例如 snowflake-arctic-embed-m。

# remove # to run example
docker run \
 --gpus all \ # --gpus mounts the NVIDIA GPUS - and requires the nvidia-docker toolkit.
 -p 7997:7997 \ # port forwarding
 michaelf34/infinity:0.0.70 \ # selecting the Dockerfile.nvidia image from infinity
 v2 \
 --model-id Snowflake/snowflake-arctic-embed-m \
 --engine torch # tell infinity to run the model using the pytorch backend engine.

它正在运行。如果我们要在 CPU 上完成同样的事情，使用 ONNX 可能会更好。幸好，arctic 模型有一组额外的 ONNX 权重，位于 [`./onnx/](https://huggingface.co/Snowflake/snowflake-arctic-embed-m/tree/main/onnx)。

docker run \
 \ # no gpus mounted
 -p 7997:7997 \ # port forwarding
 michaelf34/infinity:0.0.70-cpu \ 
 v2 \
 --model-id Snowflake/snowflake-arctic-embed-m \
 --engine optimum # tell infinity to run the model using the onnx/optimum backend engine.

教程 1：在 AMD ROCm 上运行 AMD 嵌入模型

以下是如何在 AMD 上重现该示例。

首先，我们需要兼容的 GPU 并安装 AMD-container-toolkit https://rocm.docs.amd.com/projects/install-on-linux/en/latest/how-to/docker.html。对于 AMD，该软件被称为 ROCm 内核模式驱动程序，amdgpu-dkms。

安装了 `amdgpu-dkms` 和 AMD 驱动程序后，任何 AMD Pytorch 镜像都可以启动。

在这种情况下，该镜像基于 [rocm/pytorch:rocm6.2.3_ubuntu22.04_py3.10_pytorch_release_2.3.0] (https://github.com/michaelfeil/infinity/blob/main/libs/infinity_emb/Dockerfile.amd_auto)，需要在主机上安装兼容 rocm6.2.3 的版本。您可以通过 michaelf34/infinity:0.0.70-amd 从 Docker Hub 拉取镜像，无需登录。

我们不是通过 `--gpus` 挂载 GPU，而是将以下内容添加到 docker 运行命令中。

--security-opt seccomp=unconfined \ # mount the devices for 
--device=/dev/kfd \
--device=/dev/dri \

将其整合在一起

docker run -it \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \ 
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  -p 7997:7997 \
  michaelf34/infinity:0.0.70-amd \
  v2 \
 --model-id Snowflake/snowflake-arctic-embed-m \
 --engine torch # tell infinity to run the model using the onnx/optimum backend engine.

教程 2：通过 ONNX 在 MI300x 上运行 AMD ROCm 重排序模型

要在 AMD 上通过 ONNX 运行相同的模型，我们必须打开一个额外的构建参数 (--build-arg GPU_ARCH=gfx942)，为 gfx942（仅限 MI300x！）GPU 架构构建 onnxruntime。此外，我们选择 --engine optimum 进行 ONNX 推理。对于其他构建目标，您可以使用 michaelf34/infinity:0.0.70-amd-gfx94a (MI200x) 和 michaelf34/infinity:0.0.70-amd-gfx1100 (选定的 AMD Radeon 显卡)。

docker run -it \
  --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined \ # mount the devices for 
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  -p 7997:7997 \
  michaelf34/infinity:0.0.70-amd-gfx942 \ # selecting the Dockerfile.rocm with GPU_ARCH=gfx942 build arch
  v2 \
 --model-id Snowflake/snowflake-arctic-embed-m --engine optimum --device cuda \
 --model-id mixedbread-ai/mxbai-rerank-base-v1 --engine optimum --device cuda # repeating multiple `model-id` allows for multiple model launch.

教程 3：无需 Docker 运行设置 (Runpod.io)

本博客并非由 Runpod 赞助。我使用它是因为截至 2024 年，它是在 AMD-MI300x 上运行 AMD 镜像的一种非常简单的方法。

启动配置如下：在 UI 中，我选择 MI300x 节点并修改 `port`，`image=michaelf34/infinity:0.0.70-amd-gfx942`，Infinity CLI 版本（使用 `v2`！），以及两个模型：`--model-id Snowflake/snowflake-arctic-embed-m --engine optimum --model-id mixedbread-ai/mxbai-rerank-base-v1 --engine optimum`。

image/png

一旦启动完成，您应该在控制台中看到以下日志输出：INFO: Uvicorn running on http://0.0.0.0:7997 (Press CTRL+C to quit)。

此外，日志中还包含一些在预热期间收集的基本性能统计数据：每秒 487 (len=512) - 12021 (len=3) 个嵌入和高达每秒 4836 个重排序请求。这不包括 API + 动态批处理开销，但仍然非常可观，特别是对于短请求！单个副本可能能够处理多达 1000 个并发 RAG 用户发送的短查询——一点也不差！此时，基础设施的其他部分可能成为瓶颈！

点击此处显示完整日志 
2024-12-03T08:43:23.843767712Z INFO:     Waiting for application startup.
2024-12-03T08:43:23.847049714Z INFO     2024-12-03 08:43:23,843 infinity_emb INFO:        infinity_server.py:92
2024-12-03T08:43:23.847099298Z          Creating 2engines:
2024-12-03T08:43:23.847104256Z          engines=['Snowflake/snowflake-arctic-embed-m',
2024-12-03T08:43:23.847108382Z          'mixedbread-ai/mxbai-rerank-base-v1']
2024-12-03T08:43:23.847677370Z INFO     2024-12-03 08:43:23,846 infinity_emb INFO: Anonymized   telemetry.py:30
2024-12-03T08:43:23.847687605Z          telemetry can be disabled via environment variable
2024-12-03T08:43:23.847690820Z          `DO_NOT_TRACK=1`.
2024-12-03T08:43:23.852227822Z INFO     2024-12-03 08:43:23,851 infinity_emb INFO:           select_model.py:64
2024-12-03T08:43:23.852244597Z          model=`Snowflake/snowflake-arctic-embed-m` selected,
2024-12-03T08:43:23.852248152Z          using engine=`optimum` and device=`cuda`
2024-12-03T08:43:24.307934415Z INFO     2024-12-03 08:43:24,304 infinity_emb INFO: Found 7 utils_optimum.py:244
2024-12-03T08:43:24.307985381Z          onnx files: [PosixPath('onnx/model.onnx'),
2024-12-03T08:43:24.307990809Z          PosixPath('onnx/model_bnb4.onnx'),
2024-12-03T08:43:24.307994915Z          PosixPath('onnx/model_fp16.onnx'),
2024-12-03T08:43:24.307998901Z          PosixPath('onnx/model_int8.onnx'),
2024-12-03T08:43:24.308002676Z          PosixPath('onnx/model_q4.onnx'),
2024-12-03T08:43:24.308006482Z          PosixPath('onnx/model_quantized.onnx'),
2024-12-03T08:43:24.308010748Z          PosixPath('onnx/model_uint8.onnx')]
2024-12-03T08:43:24.309548020Z INFO     2024-12-03 08:43:24,307 infinity_emb INFO: Using   utils_optimum.py:248
2024-12-03T08:43:24.309575701Z          onnx/model.onnx as the model
2024-12-03T08:43:24.707931233Z The ONNX file onnx/model.onnx is not a regular name used in optimum.onnxruntime, the ORTModel might not behave as expected.
2024-12-03T08:43:33.299022861Z 2024-12-03 08:43:33.298756964 [W:onnxruntime:, session_state.cc:1168 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.
2024-12-03T08:43:33.299078624Z 2024-12-03 08:43:33.298775843 [W:onnxruntime:, session_state.cc:1170 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.
2024-12-03T08:43:34.174298204Z INFO     2024-12-03 08:43:34,171 infinity_emb INFO: Getting   select_model.py:97
2024-12-03T08:43:34.174339565Z          timings for batch_size=32 and avg tokens per
2024-12-03T08:43:34.174344533Z          sentence=3
2024-12-03T08:43:34.174348629Z                  0.38     ms tokenization
2024-12-03T08:43:34.174352464Z                  2.22     ms inference
2024-12-03T08:43:34.174356300Z                  0.06     ms post-processing
2024-12-03T08:43:34.174360326Z                  2.66     ms total
2024-12-03T08:43:34.174364172Z          embeddings/sec: 12021.05
2024-12-03T08:43:34.326102645Z INFO     2024-12-03 08:43:34,324 infinity_emb INFO: Getting  select_model.py:103
2024-12-03T08:43:34.326140922Z          timings for batch_size=32 and avg tokens per
2024-12-03T08:43:34.326145960Z          sentence=512
2024-12-03T08:43:34.326177467Z                  7.80     ms tokenization
2024-12-03T08:43:34.326181933Z                  57.67    ms inference
2024-12-03T08:43:34.326186220Z                  0.13     ms post-processing
2024-12-03T08:43:34.326190196Z                  65.60    ms total
2024-12-03T08:43:34.326194262Z          embeddings/sec: 487.80
2024-12-03T08:43:34.326716801Z INFO     2024-12-03 08:43:34,325 infinity_emb INFO: model    select_model.py:104
2024-12-03T08:43:34.326738783Z          warmed up, between 487.80-12021.05 embeddings/sec
2024-12-03T08:43:34.326748608Z          at batch_size=32
2024-12-03T08:43:34.332365460Z INFO     2024-12-03 08:43:34,331 infinity_emb INFO:           select_model.py:64
2024-12-03T08:43:34.332386802Z          model=`mixedbread-ai/mxbai-rerank-base-v1` selected,
2024-12-03T08:43:34.332391909Z          using engine=`optimum` and device=`cuda`
2024-12-03T08:43:34.737521909Z INFO     2024-12-03 08:43:34,736 infinity_emb INFO: Found 2 utils_optimum.py:244
2024-12-03T08:43:34.737551914Z          onnx files: [PosixPath('onnx/model.onnx'),
2024-12-03T08:43:34.737554037Z          PosixPath('onnx/model_quantized.onnx')]
2024-12-03T08:43:34.738028694Z INFO     2024-12-03 08:43:34,737 infinity_emb INFO: Using   utils_optimum.py:248
2024-12-03T08:43:34.738033541Z          onnx/model.onnx as the model
2024-12-03T08:43:35.260076178Z The ONNX file onnx/model.onnx is not a regular name used in optimum.onnxruntime, the ORTModel might not behave as expected.
2024-12-03T08:43:42.332373238Z 2024-12-03 08:43:42.332105549 [W:onnxruntime:, session_state.cc:1168 VerifyEachNodeIsAssignedToAnEp] Some nodes were not assigned to the preferred execution providers which may or may not have an negative impact on performance. e.g. ORT explicitly assigns shape related ops to CPU to improve perf.
2024-12-03T08:43:42.332428290Z 2024-12-03 08:43:42.332129364 [W:onnxruntime:, session_state.cc:1170 VerifyEachNodeIsAssignedToAnEp] Rerunning with verbose output on a non-minimal build will show node assignments.
2024-12-03T08:43:43.528743277Z INFO     2024-12-03 08:43:43,526 infinity_emb INFO: Getting   select_model.py:97
2024-12-03T08:43:43.528793812Z          timings for batch_size=32 and avg tokens per
2024-12-03T08:43:43.528802886Z          sentence=4
2024-12-03T08:43:43.528809295Z                  0.49     ms tokenization
2024-12-03T08:43:43.528814814Z                  6.12     ms inference
2024-12-03T08:43:43.528820282Z                  0.01     ms post-processing
2024-12-03T08:43:43.528825580Z                  6.62     ms total
2024-12-03T08:43:43.528831038Z          embeddings/sec: 4836.28
2024-12-03T08:43:43.859858228Z INFO     2024-12-03 08:43:43,858 infinity_emb INFO: Getting  select_model.py:103
2024-12-03T08:43:43.859906941Z          timings for batch_size=32 and avg tokens per
2024-12-03T08:43:43.859912008Z          sentence=512
2024-12-03T08:43:43.859916225Z                  25.66    ms tokenization
2024-12-03T08:43:43.859920161Z                  106.47   ms inference
2024-12-03T08:43:43.859923996Z                  0.03     ms post-processing
2024-12-03T08:43:43.859927862Z                  132.16   ms total
2024-12-03T08:43:43.859952338Z          embeddings/sec: 242.13
2024-12-03T08:43:43.860544431Z INFO     2024-12-03 08:43:43,859 infinity_emb INFO: model    select_model.py:104
2024-12-03T08:43:43.860571682Z          warmed up, between 242.13-4836.28 embeddings/sec at
2024-12-03T08:43:43.860577340Z          batch_size=32
2024-12-03T08:43:43.861875996Z INFO     2024-12-03 08:43:43,861 infinity_emb INFO:         batch_handler.py:443
2024-12-03T08:43:43.861889265Z          creating batching engine
2024-12-03T08:43:43.865405636Z INFO     2024-12-03 08:43:43,862 infinity_emb INFO: ready   batch_handler.py:512
2024-12-03T08:43:43.865426647Z          to batch requests.
2024-12-03T08:43:43.866914465Z INFO     2024-12-03 08:43:43,866 infinity_emb INFO:         batch_handler.py:443
2024-12-03T08:43:43.866928997Z          creating batching engine
2024-12-03T08:43:43.869706908Z INFO     2024-12-03 08:43:43,867 infinity_emb INFO: ready   batch_handler.py:512
2024-12-03T08:43:43.869716492Z          to batch requests.
2024-12-03T08:43:43.871849030Z INFO     2024-12-03 08:43:43,870 infinity_emb INFO:       infinity_server.py:106
2024-12-03T08:43:43.871875049Z          ♾️  Infinity - Embedding Inference Server
2024-12-03T08:43:43.871881699Z          MIT License; Copyright (c) 2023-now Michael Feil
2024-12-03T08:43:43.871887428Z          Version 0.0.70
2024-12-03T08:43:43.871896681Z          Open the Docs via Swagger UI:
2024-12-03T08:43:43.871900798Z          http://0.0.0.0:7997/docs
2024-12-03T08:43:43.871908990Z          Access all deployed models via 'GET':
2024-12-03T08:43:43.871913136Z          curl http://0.0.0.0:7997/models
2024-12-03T08:43:43.871921018Z          Visit the docs for more information:
2024-12-03T08:43:43.871924893Z          https://michaelfeil.github.io/infinity
2024-12-03T08:43:43.872579980Z INFO:     Application startup complete.
2024-12-03T08:43:43.872964823Z INFO:     Uvicorn running on http://0.0.0.0:7997 (Press CTRL+C to quit)

检查 Runpod 设置，它告诉我我的应用程序在 https://3e00bu865kgyuf-7997.proxy.runpod.net 上运行。

快速向 /models 端点发送 curl 请求并向 embeddings 端点发送 POST 请求，显示模型确实在运行！

$ curl https://3e00bu865kgyuf-7997.proxy.runpod.net/models

返回

{"data":[{"id":"Snowflake/snowflake-arctic-embed-m","stats":{"queue_fraction":0.0,"queue_absolute":0,"results_pending":0,"batch_size":32},"object":"model","owned_by":"infinity","created":1733215830,"backend":"optimum","capabilities":["embed"]},{"id":"mixedbread-ai/mxbai-rerank-base-v1","stats":{"queue_fraction":0.0,"queue_absolute":0,"results_pending":0,"batch_size":32},"object":"model","owned_by":"infinity","created":1733215830,"backend":"optimum","capabilities":["rerank"]}],"object":"list"}infinity-emb-py3.10(base)

curl -X 'POST' \
  'https://3e00bu865kgyuf-7997.proxy.runpod.net/embeddings' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "model": "Snowflake/snowflake-arctic-embed-m",
  "encoding_format": "float",
  "user": "string",
  "input": [
    "This is a sample sentence, to verify embeddings run on AMD"
  ],
  "modality": "text"
}'

AMD Infinity 优化：

我个人用于最大化实例处理请求数量的一些策略

  1. 小批量大小和小模型首选 `onnx` (`--engine optimum`)，否则使用 pytorch `--engine torch`。这是因为 ONNX 在小型计算机视觉模型上具有低 CPU 开销和延迟优化。经验法则：`model_size * batch_size * max_context_length` 越大，您应该使用 `torch` 的次数越多。
  2. 使用 `—engine torch` 时尝试 torch compile `—compile`。这会延长启动时间，但可以为您带来可观的优势。
  3. 根据您的显存调整 `--batch-size`。如果您有大量可用的显存，您可能会使用更大的批处理大小。更大的批处理大小将降低*权重与中间激活*的比例，并可能将小模型推入计算受限状态（最小化显存用于获取模型权重的时间比例）。在 MI300x 上，`--batch-size 128` 值得一试！
  4. 考虑增加更多 GPU：Infinity 中的 `--device-id` 选择目标设备。在最佳情况下，您有多个 GPU，当所有 GPU 都已挂载时，通过添加 `--device-id 0,1,2,3` 到 CLI 可以将吞吐量提高约 4 倍。
  5. 不要将 CPU 用于长上下文嵌入。运行短查询通常没问题，但重新嵌入数据库可能需要很长时间，而且在大型 CPU 集群上$/token效率不高。

目前在 AMD 上构建的缺点

由于此帖子不受任何组织赞助或附属，因此目前有一些关于预期挑战和从使 Infinity 与 AMD 兼容中学到的一些个人思考。

  • 融合的 torch 内核有时不可用。Infinity 大量使用了嵌套张量（即 `torch._nested_tensor_from_mask`）和编码器模型改进。这些不属于 torch 的公共 API，目前对 AMD ROCM 和 Apple MPS 的支持有限或没有。
  • rocm-pytorch docker 镜像非常大。解压后有 80GB 那么大。在基础镜像中，我们有 3 个现有的 torch 安装，肯定有改进的潜力。这也超出了 Github CI 的自动化构建范围。
  • Poetry/pip/uv 安装生态系统更倾向并默认为 Nvidia，在 CPU/ROCm 上通过自定义 pip URL 解决问题很困难，惨痛的教训是：截至 2024 年 12 月，`pip install --extra-index-url ..` 是唯一有效的方法。
  • 维护者之路偶尔会坎坷。请准备好偶尔从头开始构建 onnxruntime，这可能需要多达 200GiB 的临时存储空间，并且在一台高端机器上构建镜像可能需要长达一小时。

总结：

在这篇博客文章中，您学习了如何在 AMD 上运行嵌入模型，例如 Snowflake/snowflake-arctic-embed-m。欢迎分享本教程，并为 Infinity 贡献代码！

社区

注册或登录发表评论

赞
6