OA0 = Omni AI 0
OA0 是一个探索 AI 的论坛
现在注册
已注册用户请  登录
OA0  ›  代码  ›  Text Generation Inference (TGI) — HuggingFace 推理服务

Text Generation Inference (TGI) — HuggingFace 推理服务

 
  api ·  2026-02-28 00:56:11 · 3 次点击  · 0 条评论  

[!CAUTION]
text-generation-inference 现已进入维护模式。未来,我们将仅接受针对次要错误修复、文档改进和轻量级维护任务的拉取请求。

TGI 开创了依赖 transformers 模型架构的优化推理引擎运动。这一方法现已被下游推理引擎采纳,我们参与贡献并推荐未来使用:vllmSGLang,以及具有互操作性的本地引擎,如 llama.cpp 或 MLX。

优化 TGI 部署 # Text Generation Inference GitHub Repo stars Swagger API documentation 一个用于文本生成推理的 Rust、Python 和 gRPC 服务器。在 [Hugging Face](https://huggingface.co) 的生产环境中用于驱动 Hugging Chat、推理 API 和推理端点。

目录

Text Generation Inference (TGI) 是一个用于部署和服务大型语言模型 (LLM) 的工具包。TGI 为最流行的开源 LLM 提供高性能文本生成,包括 Llama、Falcon、StarCoder、BLOOM、GPT-NeoX 以及更多模型。TGI 实现了许多功能,例如:

  • 简单的启动器,用于服务最流行的 LLM
  • 生产就绪(使用 Open Telemetry 进行分布式追踪,Prometheus 指标)
  • 张量并行,用于在多个 GPU 上加速推理
  • 使用服务器发送事件 (SSE) 进行令牌流式传输
  • 对传入请求进行连续批处理以提高总吞吐量
  • 与 Open AI Chat Completion API 兼容的 Messages API
  • 针对推理优化的 transformers 代码,在最流行的架构上使用 Flash AttentionPaged Attention
  • 支持以下量化方法:
  • bitsandbytes
  • GPT-Q
  • EETQ
  • AWQ
  • Marlin
  • fp8
  • 加载 Safetensors 权重
  • 使用 A Watermark for Large Language Models 进行水印处理
  • Logits 处理器(温度缩放、top-p、top-k、重复惩罚等,更多细节参见 transformers.LogitsProcessor
  • 停止序列
  • 对数概率
  • 推测解码 ~2 倍延迟降低
  • 引导/JSON 输出。指定输出格式以加速推理并确保输出符合某些规范。
  • 自定义提示生成:通过提供自定义提示轻松生成文本,以引导模型的输出
  • 微调支持:利用针对特定任务微调的模型,以获得更高的准确性和性能

硬件支持

快速开始

Docker

详细入门指南,请参阅快速导览。最简单的入门方式是使用官方 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'

您也可以使用 TGI 的 Messages API 来获得与 Open AI Chat Completion 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'

注意: 要使用 NVIDIA GPU,您需要安装 NVIDIA Container Toolkit。我们还建议使用 CUDA 版本 12.2 或更高的 NVIDIA 驱动程序。在没有 GPU 或 CUDA 支持的机器上运行 Docker 容器时,只需移除 --gpus all 标志并添加 --disable-custom-kernels 即可,请注意 CPU 并非本项目预期平台,因此性能可能不理想。

注意: TGI 支持 AMD Instinct MI210 和 MI250 GPU。详细信息可在支持的硬件文档中找到。要使用 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 替换上述命令。

要查看服务模型的所有选项(在代码中或在 CLI 中):

text-generation-launcher --help

API 文档

您可以使用 /docs 路由查阅 text-generation-inference REST API 的 OpenAPI 文档。
Swagger UI 也可在以下地址访问:https://huggingface.github.io/text-generation-inference

使用私有或受控模型

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

例如,如果您想服务受控的 Llama V2 模型变体:

  1. 前往 https://huggingface.co/settings/tokens
  2. 复制您的 CLI READ token
  3. 导出 HF_TOKEN=<您的 CLI READ token>

或使用 Docker:

model=meta-llama/Meta-Llama-3.1-8B-Instruct
volume=$PWD/data # 与 Docker 容器共享一个卷,避免每次运行都下载权重
token=<您的 cli READ 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) 的说明

NCCLPyTorch 用于进行分布式训练/推理的通信框架。text-generation-inference 利用 NCCL 实现张量并行,以显著加速大型语言模型的推理。

为了在 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 参数覆盖。

架构

TGI 架构

Adyen 关于 TGI 内部工作原理的详细博文:使用 TGI 进行大规模 LLM 推理 (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 虚拟环境,例如使用 condapython 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

注意: 在某些机器上,您可能还需要 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

注意: 当您在非 NixOS 系统上使用 Nix 时,必须创建一些符号链接,以使 CUDA 驱动程序库对 Nix 包可见。

对于 TGI 开发,您可以使用 impure 开发环境:

nix develop .#impure

# 仅在第一次启动开发环境或更新 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
)

所有开发依赖项(cargo、Python、Torch 等)都可以在此开发环境中使用。

优化的架构

TGI 开箱即用,可以为所有现代模型服务优化模型。它们可以在此列表中找到。

其他架构在尽力而为的基础上支持,使用:

AutoModelForCausalLM.from_pretrained(<model>, device_map="auto")

AutoModelForSeq2SeqLM.from_pretrained(<model>, 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

4 位量化可使用 bitsandbytes 的 NF4 和 FP4 数据类型。可以通过向 text-generation-launcher 提供 --quantize bitsandbytes-nf4--quantize bitsandbytes-fp4 作为命令行参数来启用。

有关量化的更多信息,请参阅量化文档

开发

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
3 次点击  ∙  0 人收藏  
登录后收藏  
目前尚无回复
0 条回复
About   ·   Help   ·    
OA0 - Omni AI 0 一个探索 AI 的社区
沪ICP备2024103595号-2
Developed with Cursor