ghcr.io/huggingface/text-generation-inference:sha-efb94e0-neuron

[!IMPORTANT] text-generation-inference 现已进入维护模式。今后，我们将仅接受小 bug 修复、文档改进和轻量级维护任务的拉取请求。

TGI 已发起运动，推动优化推理引擎依赖 transformers 模型架构。这种方法现已被下游推理引擎采用，我们对此做出了贡献并建议今后使用：https://github.com/vllm-project/vllm%E3%80%81https://github.com/sgl-project/sglang%EF%BC%8C%E4%BB%A5%E5%8F%8A%E5%85%B7%E6%9C%89%E4%BA%92%E5%85%BC%E5%AE%B9%E6%80%A7%E7%9A%84%E6%9C%AC%E5%9C%B0%E5%BC%95%E6%93%8E%EF%BC%88%E5%A6%82 llama.cpp 或 MLX）。

文本生成推理

一个用于文本生成推理的 Rust、Python 和 gRPC 服务器。已在 https://huggingface.co 的生产环境中使用，为 Hugging Chat、推理 API 和推理端点提供支持。

目录

• 快速开始
• Docker
• API 文档
• 使用私有或 gated 模型
• 关于共享内存（shm）的说明
• 分布式追踪
• 架构
• 本地安装
• 本地安装（Nix）
• 优化架构
• 本地运行
• 运行
• 量化
• 开发
• 测试

文本生成推理（TGI）是一个用于部署和服务大型语言模型（LLMs）的工具包。TGI 支持为最流行的开源 LLM 提供高性能文本生成，包括 Llama、Falcon、StarCoder、BLOOM、GPT-NeoX 等，更多模型参见 https://huggingface.co/docs/text-generation-inference/supported_models%E3%80%82TGI 实现了许多功能，例如：

• 简单的启动器，可服务大多数流行的 LLM
• 生产就绪（支持 Open Telemetry 分布式追踪、Prometheus 指标）
• 张量并行（Tensor Parallelism），可在多 GPU 上实现更快推理
• 使用服务器发送事件（SSE）的令牌流
• 对传入请求进行连续批处理，以提高总吞吐量
• https://huggingface.co/docs/text-generation-inference/en/messages_api%EF%BC%8C%E4%B8%8E Open AI 聊天补全 API 兼容
• 针对最流行架构优化的 transformers 推理代码，使用 https://github.com/HazyResearch/flash-attention 和 https://github.com/vllm-project/vllm
• 支持以下量化方式：
  • https://github.com/TimDettmers/bitsandbytes
  • https://arxiv.org/abs/2210.***
  • https://github.com/NetEase-FuXi/EETQ
  • https://github.com/casper-hansen/AutoAWQ
  • https://github.com/IST-DASLab/marlin
  • https://developer.nvidia.com/blog/nvidia-arm-and-intel-publish-fp8-specification-for-standardization-as-an-interchange-format-for-ai/
• https://github.com/huggingface/safetensors 权重加载
• 使用 https://arxiv.org/abs/2301.*** 的水印功能
• Logits 调整器（温度缩放、top-p、top-k、重复惩罚，更多详情参见 https://huggingface.co/docs/transformers/internal/generation_utils#transformers.LogitsProcessor%EF%BC%89
• 停止序列
• 对数概率
• https://huggingface.co/docs/text-generation-inference/conceptual/speculation%EF%BC%8C%E5%BB%B6%E8%BF%9F%E9%99%8D%E4%BD%8E%E7%BA%A6 2 倍
• https://huggingface.co/docs/text-generation-inference/conceptual/guidance%EF%BC%9A%E6%8C%87%E5%AE%9A%E8%BE%93%E5%87%BA%E6%A0%BC%E5%BC%8F%E4%BB%A5%E5%8A%A0%E9%80%9F%E6%8E%A8%E7%90%86%EF%BC%8C%E5%B9%B6%E7%A1%AE%E4%BF%9D%E8%BE%93%E5%87%BA%E7%AC%A6%E5%90%88%E7%89%B9%E5%AE%9A%E8%A7%84%E8%8C%83%E3%80%82
• 自定义提示生成：通过提供自定义提示来引导模型输出，轻松生成文本
• 微调支持：利用针对特定任务的微调模型，以实现更高的准确性和性能

硬件支持

• https://github.com/huggingface/text-generation-inference/pkgs/container/text-generation-inference
• https://github.com/huggingface/text-generation-inference/pkgs/container/text-generation-inference%EF%BC%88-rocm%EF%BC%89
• https://github.com/huggingface/optimum-neuron/tree/main/text-generation-inference
• https://github.com/huggingface/text-generation-inference/pull/1475
• https://github.com/huggingface/tgi-gaudi
• https://huggingface.co/docs/optimum-tpu/howto/serving

快速开始

Docker

有关详细的入门指南，请参见 https://huggingface.co/docs/text-generation-inference/quicktour%E3%80%82%E6%9C%80%E7%AE%80%E5%8D%95%E7%9A%84%E5%85%A5%E9%97%A8%E6%96%B9%E5%BC%8F%E6%98%AF%E4%BD%BF%E7%94%A8%E5%AE%98%E6%96%B9 Docker 容器：

model=HuggingFaceH4/zephyr-7b-beta
# 与 Docker 容器共享卷以避免每次运行都下载权重
volume=$PWD/data

docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data \
ghcr.io/huggingface/text-generation-inference:3.3.5 --model-id $model

然后您可以发送请求，例如：

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'

您还可以使用 https://huggingface.co/docs/text-generation-inference/en/messages_api 来获取与 Open AI 聊天补全 API 兼容的响应。

curl localhost:8080/v1/chat/completions \
-X POST \
-d '{
"model": "tgi",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": "What is deep learning?"
}
],
"stream": true,
"max_tokens": 20
}' \
-H 'Content-Type: application/json'

[!NOTE] 要使用 NVIDIA GPU，您需要安装 https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html%E3%80%82%E6%88%91%E4%BB%AC%E8%BF%98%E5%BB%BA%E8%AE%AE%E4%BD%BF%E7%94%A8 CUDA 版本 12.2 或更高的 NVIDIA 驱动程序。在没有 GPU 或 CUDA 支持的机器上运行 Docker 容器时，只需移除 --gpus all 标志并添加 --disable-custom-kernels 即可。请注意，CPU 并非本项目的目标平台，因此性能可能欠佳。

[!NOTE] TGI 支持 AMD Instinct MI210 和 MI250 GPU。详情可参见 https://huggingface.co/docs/text-generation-inference/installation_amd#using-tgi-with-amd-gpus%E3%80%82%E8%A6%81%E4%BD%BF%E7%94%A8 AMD GPU，请使用 docker run --device /dev/kfd --device /dev/dri --shm-size 1g -p 8080:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:3.3.5-rocm --model-id $model 替换上述命令。

要查看服务模型的所有选项（在 https://github.com/huggingface/text-generation-inference/blob/main/launcher/src/main.rs 或命令行中）：

text-generation-launcher --help

API 文档

您可以通过 /docs 路由查阅 text-generation-inference REST API 的 OpenAPI 文档。Swagger UI 也可在以下地址获取：https://huggingface.github.io/text-generation-inference%E3%80%82

使用私有或 gated 模型

你可以选择使用 HF_TOKEN 环境变量来配置 text-generation-inference 所使用的令牌。这使你能够访问受保护的资源。

例如，如果你想提供 gated 的 Llama V2 模型变体：

1. 访问 https://huggingface.co/settings/tokens
2. 复制你的 CLI READ 令牌
3. 导出 HF_TOKEN=

或使用 Docker：

model=meta-llama/Meta-Llama-3.1-8B-Instruct
volume=$PWD/data # 与 Docker 容器共享卷以避免每次运行都下载权重
token=

docker run --gpus all --shm-size 1g -e HF_TOKEN=$token -p 8080:80 -v $volume:/data \
ghcr.io/huggingface/text-generation-inference:3.3.5 --model-id $model

关于共享内存（shm）的说明

https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/index.html 是 PyTorch 用于分布式训练/推理的通信框架。text-generation-inference 利用 NCCL 实现张量并行（Tensor Parallelism），以显著加速大型语言模型的推理速度。

为了在 NCCL 组的不同设备之间共享数据，如果无法通过 NVLink 或 PCI 进行对等通信，NCCL 可能会回退到使用主机内存。

为允许容器使用 1G 共享内存并支持 SHM 共享，我们在上述命令中添加了 --shm-size 1g。

如果你在 Kubernetes 中运行 text-generation-inference，也可以通过创建以下卷为容器添加共享内存：

- name: shm
  emptyDir:
    medium: Memory
    sizeLimit: 1Gi

并将其挂载到 /dev/shm。

最后，你也可以通过设置 NCCL_SHM_DISABLE=1 环境变量来禁用 SHM 共享。但请注意，这会影响性能。

分布式追踪

text-generation-inference 通过 OpenTelemetry 实现了分布式追踪功能。你可以通过 --otlp-endpoint 参数设置 OTLP 收集器的地址来使用此功能。默认服务名称可以通过 --otlp-service-name 参数覆盖。

架构

Adyen 关于 TGI 内部工作原理的详细博客文章：LLM inference at scale with TGI (Martin Iglesias Goyanes - Adyen, 2024)

本地安装

你也可以选择在本地安装 text-generation-inference。

首先克隆仓库并进入目录：

git clone https://github.com/huggingface/text-generation-inference
cd text-generation-inference

然后安装 Rust 并创建至少 Python 3.9 的 Python 虚拟环境，例如使用 conda 或 python venv：

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

# 使用 conda
conda create -n text-generation-inference python=3.11
conda activate text-generation-inference

# 使用 python venv
python3 -m venv .venv
source .venv/bin/activate

你可能还需要安装 Protoc。

在 Linux 上：

PROTOC_ZIP=protoc-21.12-linux-x86_64.zip
curl -OL https://github.com/protocolbuffers/protobuf/releases/download/v21.12/$PROTOC_ZIP
sudo unzip -o $PROTOC_ZIP -d /usr/local bin/protoc
sudo unzip -o $PROTOC_ZIP -d /usr/local 'include/*'
rm -f $PROTOC_ZIP

在 macOS 上，使用 Homebrew：

brew install protobuf

然后运行：

BUILD_EXTENSIONS=True make install # 安装仓库及包含 CUDA 内核的 HF/transformer 分支
text-generation-launcher --model-id mistralai/Mistral-7B-Instruct-v0.2

[!NOTE] 在某些机器上，你可能还需要 OpenSSL 库和 gcc。在 Linux 机器上，运行：

> sudo apt-get install libssl-dev gcc -y
>

本地安装（Nix）

另一种选择是使用 Nix 在本地安装 text-generation-inference。目前，我们仅支持在带有 CUDA GPU 的 x86_64 Linux 上使用 Nix。使用 Nix 时，所有依赖项都可以从二进制缓存中获取，无需在本地构建。

首先按照说明安装 Cachix 并启用 Hugging Face 缓存。设置缓存很重要，否则 Nix 会在本地构建许多依赖项，这可能需要数小时。

之后，你可以使用 nix run 运行 TGI：

cd text-generation-inference
nix run --extra-experimental-features nix-command --extra-experimental-features flakes . -- --model-id meta-llama/Llama-3.1-8B-Instruct

[!NOTE] 当你在非 NixOS 系统上使用 Nix 时，必须创建一些符号链接，使 CUDA 驱动库对 Nix 包可见。

对于 TGI 开发，你可以使用 impure 开发 shell：

nix develop .#impure

# 仅在首次启动开发 shell 或更新 protobuf 后需要执行
(
cd server
mkdir text_generation_server/pb || true
python -m grpc_tools.protoc -I../proto/v3 --python_out=text_generation_server/pb \
--grpc_python_out=text_generation_server/pb --mypy_out=text_generation_server/pb ../proto/v3/generate.proto
find text_generation_server/pb/ -type f -name "*.py" -print0 -exec sed -i -e 's/^\(import.*pb2\)/from . \1/g' {} \;
touch text_generation_server/pb/__init__.py
)

此开发 shell 中包含所有开发依赖项（cargo、Python、Torch 等）。

优化架构

TGI 开箱即可支持所有现代模型的优化版本。这些模型可在https://huggingface.co/docs/text-generation-inference/supported_models%E4%B8%AD%E6%89%BE%E5%88%B0%E3%80%82

其他架构通过以下方式提供尽力支持：

AutoModelForCausalLM.from_pretrained( , device_map="auto")

或

AutoModelForSeq2SeqLM.from_pretrained( , device_map="auto")

本地运行

运行

text-generation-launcher --model-id mistralai/Mistral-7B-Instruct-v0.2

量化

你也可以运行预量化权重（AWQ、GPTQ、Marlin）或使用 bitsandbytes、EETQ、fp8 进行动态量化，以减少 VRAM 需求：

text-generation-launcher --model-id mistralai/Mistral-7B-Instruct-v0.2 --quantize

4bit 量化可通过 https://arxiv.org/pdf/2305.***.pdf%E5%AE%9E%E7%8E%B0%E3%80%82%E5%8F%AF%E9%80%9A%E8%BF%87%E5%90%91 text-generation-launcher 提供 --quantize bitsandbytes-nf4 或 --quantize bitsandbytes-fp4 命令行参数来启用。

有关量化的更多信息，请参阅https://huggingface.co/docs/text-generation-inference/en/conceptual/quantization%E3%80%82

开发

make server-dev
make router-dev

测试

# python 测试
make python-server-tests
make python-client-tests
# 或同时运行服务端和客户端测试
make python-tests
# rust cargo 测试
make rust-tests
# 集成测试
make integration-tests