Использование TEI на GPU AMD Instinct (ROCm)

Text Embeddings Inference поддерживает AMD Instinct GPU (серии MI200, MI300) с использованием ROCm.

Предварительные требования

• AMD Instinct GPU (серии MI200, MI300) с драйверами ROCm на хосте

Вариант A: Docker (рекомендуется)

Самый простой способ запустить TEI на GPU AMD — использовать готовый образ Docker:

model=BAAI/bge-base-en-v1.5
volume=$PWD/data  # share a volume to avoid re-downloading weights

docker run \
  --device /dev/kfd --device /dev/dri/renderD128 \
  --group-add video \
  --ipc=host \
  -p 8080:80 \
  -v $volume:/data \
  --pull always \
  ghcr.io/huggingface/text-embeddings-inference:rocm-latest \
  --model-id $model --dtype bfloat16

Затем протестируйте его:

curl http://localhost:8080/v1/embeddings \
    -H 'Content-Type: application/json' \
    -d '{
      "input": "What is Deep Learning?",
      "model": "text-embeddings-inference",
      "encoding_format": "float"
    }'

Вариант B: сборка из исходного кода внутри образа TEI ROCm

Если вы хотите модифицировать TEI и пересобрать его, используйте официальный образ TEI ROCm (закреплённый образ ниже) в качестве базовой среды вместо образа rocm/pytorch от AMD. Он уже содержит точную сборку PyTorch для ROCm, flash-attn, вычислительные ядра и установленный Python-бэкенд, который ожидает TEI, поэтому вы полностью пропускаете настройку PyTorch и зависимостей и избегаете проблем с несовпадением версий — вы пересобираете только те части, которые меняете.

Шаг 1: Запустите контейнер

Из корня репозитория вашего клонирования смонтируйте его и переопределите entrypoint, чтобы получить оболочку. Дайджест образа фиксирует сборку rocm-latest, используемую здесь; при работе с другой ревизией TEI выбирайте соответствующий образ:

docker run -it --device=/dev/kfd --device=/dev/dri/renderD128 \
  --group-add video --ipc=host --shm-size 8g \
  -v "$PWD:/workspace" -w /workspace \
  --entrypoint bash \
  ghcr.io/huggingface/text-embeddings-inference@sha256:9b3f24f8740d4c3fa656666e07a383a50a9e4f3273d3856980841cd8d45e727f

PyTorch (ROCm), flash-attn, вычислительные ядра и Python-бэкенд text-embeddings-server уже установлены в этом образе, а предварительно собранный text-embeddings-router находится в PATH. Приведённые ниже шаги описывают только пересборку после изменения исходного кода.

Шаг 2: установка Rust (требуется только для пересборки роутера)

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
source "$HOME/.cargo/env"

# Required by the router's protobuf build step.
apt-get update && apt-get install -y protobuf-compiler pkg-config libssl-dev

Шаг 3: Пересоберите Rust-роутер из исходного кода

cargo build --release \
    --no-default-features \
    --features python,http \
    --bin text-embeddings-router

Если вы также изменили Python-бэкенд, сгенерируйте protobuf-заглушки в смонтированном каталоге репозитория перед повторной установкой пакета в режиме редактирования. Установленный в образе пакет не заполняет игнорируемый каталог text_embeddings_server/pb в новом каталоге репозитория. Подоболонка сохраняет оставшиеся команды в корне репозитория.

(
  cd backends/python/server
  make gen-server
  pip install --no-deps -e .
)

Шаг 4: Запуск сборки

model=BAAI/bge-base-en-v1.5

./target/release/text-embeddings-router --model-id $model --dtype bfloat16 --port 8080

Когда сервер будет готов, вы можете протестировать его с помощью простого запроса на встраивание:

curl http://localhost:8080/v1/embeddings \
    -H 'Content-Type: application/json' \
    -d '{
      "input": "What is Deep Learning?",
      "model": "text-embeddings-inference",
      "encoding_format": "float"
    }'

Проверка обнаружения GPU

После запуска вы должны увидеть строку в журнале, подтверждающую обнаружение ROCm:

INFO text_embeddings_server::utils::device: ROCm / HIP version: X.Y.Z

Вы также можете проверить это из Python:

import torch
print(torch.cuda.is_available())  # True
print(torch.version.hip)          # e.g. 6.2.12345-...

Примечания

Это работа в процессе — поддержка дополнительных моделей и оптимизированные операции для GPU AMD скоро будут добавлены.