BGE-M3 embedding + rerank server

Self-hosted server cho BAAI/bge-m3 (dense + sparse + ColBERT multi-vector) và BAAI/bge-reranker-v2-m3 (cross-encoder rerank). Một image chạy được cả trên máy CPU lẫn máy GPU — thiết bị được chọn lúc khởi động.

Chạy nhanh

cd tmp/embedding-server
cp .env.example .env
sed -i "s/^EMBEDDING_API_KEY=.*/EMBEDDING_API_KEY=$(openssl rand -hex 32)/" .env
grep EMBEDDING_API_KEY .env          # nhớ key này để đưa cho dxai2

docker compose up -d --build         # mặc định dùng wheel torch CPU
docker compose logs -f               # lần đầu tải ~4.5GB weights

Trên máy có GPU NVIDIA (cần nvidia-container-toolkit) — overlay tự đổi sang wheel cu128:

docker compose -f docker-compose.yml -f docker-compose.gpu.yml up -d --build

torch bị ghim ở 2.8.0 chứ không thấp hơn: BAAI/bge-m3 chỉ publish pytorch_model.bin (không có safetensors), mà transformers từ chối torch.load với torch < 2.6 vì CVE-2025-32434. torch 2.5.x sẽ crash lúc load model.

Không dùng Docker:

./run.sh                             # tự tạo .venv, tự chọn wheel torch cpu/cu124

Kiểm tra:

curl -s localhost:8080/health | python3 -m json.tool
./smoke_test.sh

API

Mọi endpoint trừ /health cần header Authorization: Bearer $EMBEDDING_API_KEY.

POST /embed

{
  "texts": ["văn bản 1", "văn bản 2"],
  "return_dense": true,      // vector 1024 chiều, đã L2-normalize
  "return_sparse": true,     // {token_id: weight} — lexical, dùng cho hybrid search
  "return_colbert": false,   // [n_token][1024] — nặng, chỉ bật khi cần late interaction
  "max_length": 8192
}

Trả về:

{
  "model": "BAAI/bge-m3",
  "dense":   [[0.01, ...], [...]],
  "sparse":  [{"6006": 0.24, "39093": 0.31}, {...}],
  "colbert": null,
  "usage": {"texts": 2}
}

sparse là dict token_id → weight, nạp thẳng được vào sparse vector của Qdrant / Milvus / Elasticsearch. Điểm lexical giữa query và doc = tổng w_query[t] * w_doc[t] trên các token chung.

POST /rerank

{
  "query": "cách cài cloud gaming",
  "documents": ["doc a", "doc b", "doc c"],
  "top_n": 5,
  "return_documents": false,
  "normalize": false        // true => bọc sigmoid, đưa score về (0,1)
}

Trả results đã sắp giảm dần theo relevance_score, index trỏ về vị trí trong mảng documents gốc.

GET /health

Không cần auth — trả device, dtype, cấu hình batching, độ sâu hàng đợi.

Thiết kế

Dynamic micro-batching (app/batcher.py) là phần quan trọng nhất về hiệu năng. Request vào một asyncio.Queue; một background task gom mọi request tương thích đã tới trong lúc lượt inference trước đang chạy, cộng thêm cửa sổ chờ MAX_WAIT_MS (mặc định 10ms), rồi chạy một lượt forward duy nhất. Nhiều client bắn request lẻ vẫn được gộp thành batch lớn.

Trong mỗi batch, text được sắp theo độ dài trước khi chia micro-batch, nên mỗi micro-batch chỉ pad tới độ dài của văn bản dài nhất trong nhóm đó thay vì tới max_length. Với dữ liệu độ dài lệch nhau, đây là khác biệt lớn.

Micro-batch bị chặn bởi cả số lượng sequence (MAX_BATCH_SIZE) lẫn tổng token sau khi pad (MAX_BATCH_TOKENS). Chặn theo số lượng thôi là không đủ: 16 sequence × 8192 token cấp phát vài GB activation, nhân với 4 worker là chạm trần RAM. Ngưỡng token giữ cho batch dài tự động co lại.

Request bị huỷ (client timeout / ngắt kết nối) được loại khỏi batch trước khi chạy — nếu không, một đợt timeout sẽ khiến server nghiền vài phút cho những request không còn ai nhận. Hàng đợi có trần (QUEUE_MAXSIZE, trả 503 khi đầy) và mỗi request có REQUEST_TIMEOUT (trả 504).

Chỉ một forward pass chạy tại một thời điểm trong mỗi process (một ThreadPoolExecutor(max_workers=1) và một asyncio.Lock dùng chung cho cả embed lẫn rerank) — process đã được cấp sẵn một phần cứng cố định, chạy chồng chỉ làm thrash.

Layout tiến trình do entrypoint.sh / run.sh quyết định:

	GPU	CPU
workers	= số GPU nhìn thấy được	4
dtype	fp16	fp32 (fp16 trên CPU bị emulate, chậm hơn)
threads/worker	4	cores / workers
micro-batch	64	16
bộ nhớ	~2.5GB VRAM mỗi GPU	~5GB RSS mỗi worker

Multi-GPU: mỗi worker chiếm một card riêng. uvicorn không cho worker biết nó là worker thứ mấy, nên mỗi process tự giành một ordinal từ một file đếm có khoá (WORKER_INDEX_FILE, xoá lúc boot) rồi bind cuda:{ordinal % n_gpu}. Không có sharding model — mỗi GPU giữ một bản đầy đủ và phục vụ độc lập, nên throughput tăng gần tuyến tính theo số card. Giới hạn số card bằng NVIDIA_VISIBLE_DEVICES=2,3 trong .env nếu các GPU khác đang bận việc khác.

Trên máy 80 core hiện tại: 4 worker × 20 thread, ~20GB RAM. Tăng WORKERS trong .env nếu còn RAM và tải chủ yếu là request nhỏ đồng thời.

Model được cài trực tiếp trên transformers (app/models.py) chứ không qua FlagEmbedding: server tự kiểm soát dtype, threading, batching, và tránh multi-process encode pool của FlagEmbedding vốn tranh core với worker pool của uvicorn. Ba head tái hiện đúng bản gốc — dense = CLS đã normalize, sparse = relu(sparse_linear(h)) max-pool theo token id (bỏ special token), colbert = colbert_linear(h[:, 1:]) masked rồi normalize.

Expose ra internet

Server tự nó chỉ có Bearer key. Đặt sau một lớp TLS:

# Cloudflare Tunnel — không cần mở port, có TLS sẵn
cloudflared tunnel --url http://localhost:8080

Nếu đi qua reverse proxy trên cùng máy, đặt BIND_ADDR=127.0.0.1 trong .env để container không nghe trực tiếp ra ngoài. Đổi key: sửa .env rồi docker compose up -d.

MAX_TEXTS_PER_REQUEST và MAX_DOCUMENTS_PER_RERANK (mặc định 512) chặn một request đơn lẻ nuốt hết CPU. Muốn siết thêm thì hạ hai giá trị này.

Tuning

Env	Mặc định	Ghi chú
WORKERS	1 (GPU) / 4 (CPU)	Mỗi worker một bản model đầy đủ
MAX_BATCH_SIZE	0 = tự chọn	Sequence mỗi forward (64 GPU / 16 CPU)
MAX_BATCH_TOKENS	0 = tự chọn	Token đã pad mỗi forward (65536 GPU / 16384 CPU) — cái thực sự chặn RAM
MAX_WAIT_MS	10	Tăng → throughput cao hơn, latency cao hơn
QUEUE_MAXSIZE	256	Hàng đợi đầy → 503 thay vì phình vô hạn
REQUEST_TIMEOUT	300	Quá hạn → 504, và job bị loại khỏi batch
MAX_LENGTH	8192	Trần token của BGE-M3; hạ xuống 1024 nếu doc ngắn để nhanh hơn nhiều
RERANK_MAX_LENGTH	1024	Áp cho cặp (query, doc)
GPU_DTYPE	fp16	bf16 trên Ampere trở lên nếu gặp overflow
ENABLE_RERANKER	1	Tắt để tiết kiệm ~2.3GB mỗi worker

Cách chỉnh có tác động lớn nhất trên CPU: hạ MAX_LENGTH xuống đúng độ dài thực tế của dữ liệu. Chi phí attention tăng theo bình phương độ dài chuỗi.