Using TEI on AMD Instinct GPUs (ROCm)

Text Embeddings Inference supports AMD Instinct GPUs (MI200, MI300 series) using ROCm.

Prerequisites

• AMD Instinct GPU (MI200, MI300 series) with ROCm drivers on the host

Option A: Docker (recommended)

The easiest way to run TEI on AMD GPUs is with the pre-built Docker image:

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

Then test it:

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

Option B: Build from source inside the TEI ROCm image

If you want to modify TEI and rebuild it, use the official TEI ROCm image (the pinned image below) as your base environment instead of AMD’s rocm/pytorch image. It already ships the exact ROCm build of PyTorch, flash-attn, the compute kernels, and the installed Python backend that TEI expects, so you skip the PyTorch/dependency setup entirely and avoid version-mismatch issues — you only rebuild the parts you change.

Step 1: Start the container

From the repository root of your checkout, mount it and override the entrypoint to get a shell. The image digest pins the rocm-latest build used here; choose a matching image when working with a different TEI revision:

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, the compute kernels, and the text-embeddings-server Python backend are already installed in this image, and a prebuilt text-embeddings-router is on the PATH. The steps below only cover rebuilding after you edit the source.

Step 2: Install Rust (only needed to rebuild the router)

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

Step 3: Rebuild the Rust router from your source

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

If you also changed the Python backend, generate protobuf stubs in the mounted checkout before reinstalling it in editable mode. The image’s installed package does not populate a fresh checkout’s ignored text_embeddings_server/pb directory. The subshell keeps the remaining commands at the repository root.

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

Step 4: Launch your build

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

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

Once the server is ready, you can test it with a simple embed request:

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

Verifying GPU detection

After launch you should see a log line confirming ROCm was detected:

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

You can also verify from Python:

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

Notes

This is a work in progress — more model support and optimized operations for AMD GPUs are coming soon.