Hugging Face 嵌入容器於 Amazon SageMaker 發布

透過 Text Embedding Inference (TEI) 的高效能推論

Hugging Face Embedding Container 由 Text Embedding Inference (TEI) 提供動力，這是一種記憶體效能高的解決方案，專為高效能提供嵌入模型服務而設計。TEI 針對包括 E5、GTE、Ember 與 FlagEmbedding 在內的熱門模型族群優化抽取流程。

TEI 驅動容器的主要技術特點包括：

  • 最佳化推論： 使用 Flash Attention、Candle 與 cuBLASLt 以最佳化 transformer 程式碼。
  • 高效資源管理： 實作基於 token 的動態批次處理與 safetensors 權重載入。
  • 快速部署： 具備小型 Docker 映像檔與快速啟動時間，且不需要模型圖形編譯步驟。
  • 生產就緒度： 包含 Prometheus 指標與透過 Open Telemetry 的分散式追蹤。

支援的模型架構

此容器支援廣泛的嵌入模型架構，包括：

  • BERT/CamemBERT： 例子包括 BAAI/bge-large-en-v1.5 與 Snowflake/snowflake-arctic-embed-m-v1.5。
  • RoBERTa： 例子包括 sentence-transformers/all-roberta-large-v1。
  • XLM-RoBERTa： 例子包括 sentence-transformers/paraphrase-xlm-r-multilingual-v1。
  • NomicBert： 例子包括 jinaai/jina-embeddings-v2-base-en。
  • JinaBert： 例子包括 nomic-ai/nomic-embed-text-v1.5。

部署與效能基準

部署透過 sagemaker Python SDK 管理。使用者可使用 get_huggingface_llm_image_uri 方法取得容器 URI，且針對 CPU 與 GPU 實例提供不同的映像檔。

效能會因所選的實例類型而有顯著差異。以下測試以嵌入 100 萬個 token（共 3,900 個每次 256 token 的請求）且使用 10 個同時執行緒為基礎：

GPU 效能 (ml.g5.xlarge)

  • 硬體： 1 個 NVIDIA A10G GPU。
  • 吞吐量： 約每秒 130 個請求。
  • 延遲： 10 個同時請求時為 4 毫秒。
  • 總處理時間： 約 30 秒處理 100 萬個 token。

CPU 效能 (ml.c6i.2xlarge)

  • 硬體： 4 核心 Intel Ice-Lake vCPU，8GB 記憶體。
  • 吞吐量： 約每秒 5 個請求（包含從歐洲到 us-east-1 的網路延遲）。
  • 延遲： 10 個同時請求時為 2 秒（由 CloudWatch 測量）。
  • 總處理時間： 約 841 秒處理 100 萬個 token。

Sources