適用於 AWS Trainium 和 Inferentia 的 Neuron 後端

Neuron 後端支援在 AWS Trainium 和 Inferentia 系列晶片上部署 TGI 服務。

支援以下硬體目標：

Trainium 1，
Inferentia 2。

功能

支援基本的 TGI 功能：

連續批處理，
令牌流式傳輸，
使用 transformers 進行貪婪搜尋和多項式取樣。

從 Hugging Face Hub 部署服務

部署特定模型的 NeuronX TGI 服務的簡單方法是按照模型卡中的部署說明操作：

點選右側的“部署”按鈕，
選擇您的部署服務（支援“推理端點”和“SageMaker”），
選擇“AWS Trainium & Inferentia”，
按照說明操作。

在專用主機上部署服務

只需執行 text-generation-inference 容器並提供兩組引數即可啟動服務：

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

系統引數用於對映主機和服務之間的埠、卷和裝置，
服務引數轉發給 `text-generation-launcher`。

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

您可以將服務部署在已經匯出為 Neuron 格式的模型上，
或者，您也可以利用 Neuron 模型快取匯出自己的模型。

通用系統引數

無論何時啟動 TGI 服務，我們都強烈建議您將一個共享卷掛載為容器中的 `/data`：這將是模型快取的位置，以加快服務的進一步例項化。

另請注意，應使足夠的 Neuron 裝置對容器可見，因為每個 Neuron 裝置都有兩個核心（因此在兩個核心上部署時，您需要至少暴露一個裝置）。在生產環境中暴露裝置的推薦方法是顯式使用 `--device` 選項（例如 `--device /dev/neuron0`），並根據要暴露的裝置數量重複多次。

注意：或者，為了快速進行本地測試，也可以在 `privileged` 模式下啟動服務以訪問所有 Neuron 裝置。

最後，如果您想訪問受限儲存庫，您可能需要匯出 `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>

使用 🤗 HuggingFace Hub 上的標準模型（推薦）

我們在 aws-neuron/optimum-neuron-cache 中維護了一個最受歡迎架構和部署引數的 Neuron 模型快取。

如果您只想在不先將模型匯出到 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>

注意：您無需指定任何服務引數，因為它們都將從模型匯出配置中推匯出來。但是，您必須暴露足夠的裝置以匹配匯出階段指定的核心數量。

使用 🤗 HuggingFace 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`（它是每個請求的）。

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

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

選擇正確的批處理大小

如前一段所述，Neuron 模型的靜態批處理大小直接影響端點延遲和吞吐量。

有關最佳化提示，請參閱 text-generation-inference。

請注意，主要限制是能夠將指定 `batch_size` 的模型適配到例項上可用的總裝置記憶體中（每個 Neuron 核心 16GB，每個裝置有 2 個核心）。

查詢服務

您可以使用 `/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 上更新

文字生成推理