用於 AWS Trainium 和 Inferentia 的 Neuron 後端

Neuron 後端允許在 AWS Trainium 和 Inferentia 系列晶片上部署 TGI。

支援以下硬體目標：

Trainium 1，
Inferentia 2。

功能特色

支援基本的 TGI 功能：

連續批次處理 (continuous batching)，
Token 串流 (token streaming)，
使用 transformers 進行貪婪搜尋 (greedy search) 和多項式取樣 (multinomial sampling)。

從 Hugging Face Hub 部署服務

針對特定模型部署 NeuronX TGI 服務最簡單的方法是遵循模型頁面 (model card) 上的部署說明：

點擊右側的「部署 (Deploy)」按鈕，
選擇您的部署服務（支援「Inference Endpoints」和「SageMaker」），
選擇「AWS Trainium & Inferentia」，
按照說明進行操作。

在專用主機上部署服務

只需運行帶有兩組參數的 text-generation-inference 容器，即可啟動該服務：

docker run <system_parameters> ghcr.io/huggingface/text-generation-inference:3.3.5-neuron <service_parameters>

系統參數：用於映射主機與服務之間的埠、磁碟區和裝置，
服務參數：會轉發給 text-generation-launcher。

部署服務時，您需要一個預先編譯的 Neuron 模型。Neuron TGI 後端支援兩種主要操作模式：

您可以部署已匯出至 Neuron 的模型，
或者，您可以利用 Neuron 模型快取 (Neuron Model Cache) 來匯出您自己的模型。

常見系統參數

每當啟動 TGI 服務時，我們強烈建議您掛載一個作為容器內 /data 目錄的共享磁碟區：這是模型將被快取的地方，以加速服務後續的執行個體化。

另請注意，應使足夠多的 Neuron 裝置對容器可見；需了解每個 Neuron 裝置有兩個核心（因此在兩個核心上部署時，需要至少曝露一個裝置）。在生產環境中曝露裝置的建議方式是明確使用 --device 選項（例如 --device /dev/neuron0），並根據需要曝露的裝置數量重複此選項。

注意：或者，對於快速的本地測試，也可以在 privileged（特權）模式下啟動服務，以取得對所有 Neuron 裝置的存取權。

最後，如果您需要存取受限儲存庫 (gated repositories)，您可能需要匯出 HF_TOKEN。

以下是僅曝露第一個裝置的服務實例化範例：

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       <service_parameters>

使用來自 🤗 Hugging Face Hub 的標準模型（建議）

我們維護了一個 Neuron 模型快取，其中包含最熱門的架構和部署參數，網址為 aws-neuron/optimum-neuron-cache。

如果您只是想在不事先匯出至 Neuron 的情況下快速嘗試使用模型來運行該服務，在滿足某些條件下也是可以的：

啟動服務時必須指定匯出參數（或使用預設參數），
模型配置必須已被快取。

下方的程式碼片段展示了如何從 Hub 標準模型部署服務：

export HF_TOKEN=<YOUR_TOKEN>
docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       --device=/dev/neuron2 \
       --device=/dev/neuron3 \
       -e HF_TOKEN=${HF_TOKEN} \
       -e HF_AUTO_CAST_TYPE="fp16" \
       -e HF_NUM_CORES=8 \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id meta-llama/Meta-Llama-3-8B \
       --max-batch-size 1 \
       --max-input-length 3164 \
       --max-total-tokens 4096

使用匯出至本地路徑的模型

或者，您可以先在本地將模型匯出為 Neuron 格式。

然後，您可以在共享磁碟區內部署該服務：

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id /data/<neuron_model_path>

注意：您不需要指定任何服務參數，因為它們將會從模型匯出配置中自動推導出來。但是，您必須曝露足夠的裝置，以符合匯出階段指定的核心數量。

使用來自 🤗 Hugging Face Hub 的 Neuron 模型

在組織內共享 Neuron 模型最簡單的方法是將其推送到 Hugging Face Hub，這樣就可以直接部署，而無需再次進行匯出。

下方的程式碼片段展示了如何從 Hub Neuron 模型部署服務：

docker run -p 8080:80 \
       -v $(pwd)/data:/data \
       --device=/dev/neuron0 \
       --device=/dev/neuron1 \
       -e HF_TOKEN=${HF_TOKEN} \
       ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron \
       --model-id <organization>/<neuron-model>

選擇服務參數

使用以下指令列出可用的服務參數：

docker run ghcr.io/huggingface/text-generation-inference:<VERSION>-neuron --help

推理端點的配置總是在吞吐量與延遲之間做取捨：平行處理更多的請求會帶來更高的吞吐量，但也會增加延遲。

Neuron 模型具有靜態輸入維度 [batch_size, max_length]。

這對以下參數增加了幾項限制：

--max-batch-size 必須設定為 batch size，
--max-input-length 必須小於 max_length，
--max-total-tokens 必須設定為 max_length（這是針對每個請求的）。

雖然並非嚴格必要，但對於高效的預填充 (prefilling) 很重要：

--max-batch-prefill-tokens 應該設定為 batch_size * max-input-length。

選擇正確的 Batch Size

如前一段所述，Neuron 模型的靜態 Batch Size 會直接影響端點的延遲和吞吐量。

請參閱 text-generation-inference 以獲取最佳化提示。

請注意，主要限制是確保模型能夠在執行個體可用的總裝置記憶體內（每個 Neuron 核心 16GB，每個裝置 2 個核心）容納指定的 batch_size。

查詢服務

您可以使用 /generate 或 /generate_stream 路由來查詢模型：

curl 127.0.0.1:8080/generate \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'

curl 127.0.0.1:8080/generate_stream \
    -X POST \
    -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":20}}' \
    -H 'Content-Type: application/json'

注意：請將 127.0.0.1:8080 替換為您的實際 IP 位址與埠號。

在 GitHub 上更新

Text Generation Inference