2023-02-03 04:43:37 -07:00
< div align = "center" >
2022-11-02 10:29:56 -06:00
# Text Generation Inference
2022-10-08 04:30:12 -06:00
2023-02-03 04:43:37 -07:00
< a href = "https://github.com/huggingface/text-generation-inference" >
< img alt = "GitHub Repo stars" src = "https://img.shields.io/github/stars/huggingface/text-generation-inference?style=social" >
< / a >
< a href = "https://github.com/huggingface/text-generation-inference/blob/main/LICENSE" >
< img alt = "License" src = "https://img.shields.io/github/license/huggingface/text-generation-inference" >
< / a >
< a href = "https://huggingface.github.io/text-generation-inference" >
< img alt = "Swagger API documentation" src = "https://img.shields.io/badge/API-Swagger-informational" >
< / a >
2022-10-11 08:50:54 -06:00
2022-10-18 07:19:03 -06:00
![architecture ](assets/architecture.jpg )
< / div >
2023-02-03 04:43:37 -07:00
A Rust, Python and gRPC server for text generation inference. Used in production at [HuggingFace ](https://huggingface.co )
to power LLMs api-inference widgets.
## Table of contents
- [Features ](#features )
2023-04-17 10:43:24 -06:00
- [Optimized Architectures ](#optimized-architectures )
2023-02-03 04:43:37 -07:00
- [Get Started ](#get-started )
- [Docker ](#docker )
2023-02-08 09:53:33 -07:00
- [API Documentation ](#api-documentation )
- [A note on Shared Memory ](#a-note-on-shared-memory-shm )
2023-02-13 05:02:45 -07:00
- [Distributed Tracing ](#distributed-tracing )
2023-02-03 04:43:37 -07:00
- [Local Install ](#local-install )
- [CUDA Kernels ](#cuda-kernels )
- [Run BLOOM ](#run-bloom )
- [Download ](#download )
- [Run ](#run )
- [Quantization ](#quantization )
- [Develop ](#develop )
- [Testing ](#testing )
2023-04-27 11:16:35 -06:00
2022-10-27 06:25:29 -06:00
## Features
2022-10-11 08:50:54 -06:00
2023-03-03 10:42:20 -07:00
- Serve the most popular Large Language Models with a simple launcher
- Tensor Parallelism for faster inference on multiple GPUs
2023-02-08 14:30:11 -07:00
- Token streaming using Server-Sent Events (SSE)
2023-05-01 06:16:50 -06:00
- [Continous batching of incoming requests ](https://github.com/huggingface/text-generation-inference/tree/main/router ) for increased total throughput
- Optimized transformers code for inference using [flash-attention ](https://github.com/HazyResearch/flash-attention ) on the most popular architectures
2022-11-02 10:29:56 -06:00
- Quantization with [bitsandbytes ](https://github.com/TimDettmers/bitsandbytes )
2022-10-27 06:25:29 -06:00
- [Safetensors ](https://github.com/huggingface/safetensors ) weight loading
2023-03-03 10:42:20 -07:00
- Watermarking with [A Watermark for Large Language Models ](https://arxiv.org/abs/2301.10226 )
2023-04-04 05:27:46 -06:00
- Logits warper (temperature scaling, top-p, top-k, repetition penalty, more details see [transformers.LogitsProcessor ](https://huggingface.co/docs/transformers/internal/generation_utils#transformers.LogitsProcessor ))
2022-12-12 10:25:22 -07:00
- Stop sequences
2022-12-15 09:03:56 -07:00
- Log probabilities
2023-03-03 10:42:20 -07:00
- Production ready (distributed tracing with Open Telemetry, Prometheus metrics)
2022-10-27 06:25:29 -06:00
2023-04-11 08:38:22 -06:00
## Optimized architectures
2022-10-11 08:50:54 -06:00
2022-11-07 04:53:56 -07:00
- [BLOOM ](https://huggingface.co/bigscience/bloom )
2023-04-11 11:16:41 -06:00
- [FLAN-T5 ](https://huggingface.co/google/flan-t5-xxl )
2023-02-14 05:02:16 -07:00
- [Galactica ](https://huggingface.co/facebook/galactica-120b )
2023-04-11 08:38:22 -06:00
- [GPT-Neox ](https://huggingface.co/EleutherAI/gpt-neox-20b )
- [Llama ](https://github.com/facebookresearch/llama )
2023-04-11 11:16:41 -06:00
- [OPT ](https://huggingface.co/facebook/opt-66b )
- [SantaCoder ](https://huggingface.co/bigcode/santacoder )
2022-10-27 06:25:29 -06:00
2023-02-14 05:02:16 -07:00
Other architectures are supported on a best effort basis using:
2022-11-04 11:03:04 -06:00
`AutoModelForCausalLM.from_pretrained(<model>, device_map="auto")`
or
`AutoModelForSeq2SeqLM.from_pretrained(<model>, device_map="auto")`
2023-02-03 04:43:37 -07:00
## Get started
### Docker
2022-10-11 08:50:54 -06:00
2023-02-03 04:43:37 -07:00
The easiest way of getting started is using the official Docker container:
```shell
model=bigscience/bloom-560m
num_shard=2
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run
2023-02-08 09:53:33 -07:00
docker run --gpus all --shm-size 1g -p 8080:80 -v $volume:/data ghcr.io/huggingface/text-generation-inference:latest --model-id $model --num-shard $num_shard
2023-02-03 04:43:37 -07:00
```
2023-03-23 11:03:45 -06:00
**Note:** To use GPUs, you need to install the [NVIDIA Container Toolkit ](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html ). We also recommend using NVIDIA drivers with CUDA version 11.8 or higher.
2022-10-11 08:50:54 -06:00
2023-04-29 03:53:42 -06:00
To see all options to serve your models (in the [code ](https://github.com/huggingface/text-generation-inference/blob/main/launcher/src/main.rs ) or in the cli:
```
text-generation-launcher --help
```
2023-02-03 04:43:37 -07:00
You can then query the model using either the `/generate` or `/generate_stream` routes:
2022-10-08 04:30:12 -06:00
2023-02-03 04:43:37 -07:00
```shell
curl 127.0.0.1:8080/generate \
-X POST \
2023-03-07 10:52:22 -07:00
-d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":17}}' \
2023-02-03 04:43:37 -07:00
-H 'Content-Type: application/json'
```
2022-10-08 04:30:12 -06:00
```shell
2023-02-03 04:43:37 -07:00
curl 127.0.0.1:8080/generate_stream \
-X POST \
2023-03-07 10:52:22 -07:00
-d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":17}}' \
2023-02-03 04:43:37 -07:00
-H 'Content-Type: application/json'
2022-10-08 04:30:12 -06:00
```
2023-03-03 10:42:20 -07:00
or from Python:
```shell
2023-03-07 10:52:22 -07:00
pip install text-generation
2023-03-03 10:42:20 -07:00
```
2023-03-07 10:52:22 -07:00
```python
from text_generation import Client
2023-03-03 10:42:20 -07:00
2023-03-07 10:52:22 -07:00
client = Client("http://127.0.0.1:8080")
print(client.generate("What is Deep Learning?", max_new_tokens=17).generated_text)
2023-03-03 10:42:20 -07:00
2023-03-07 10:52:22 -07:00
text = ""
for response in client.generate_stream("What is Deep Learning?", max_new_tokens=17):
if not response.token.special:
text += response.token.text
print(text)
```
2023-03-03 10:42:20 -07:00
2023-02-03 04:43:37 -07:00
### API documentation
2022-10-08 04:30:12 -06:00
2023-02-03 04:43:37 -07:00
You can consult the OpenAPI documentation of the `text-generation-inference` REST API using the `/docs` route.
The Swagger UI is also available at: [https://huggingface.github.io/text-generation-inference ](https://huggingface.github.io/text-generation-inference ).
2023-02-13 05:02:45 -07:00
### Distributed Tracing
`text-generation-inference` is instrumented with distributed tracing using OpenTelemetry. You can use this feature
by setting the address to an OTLP collector with the `--otlp-endpoint` argument.
2023-02-08 09:53:33 -07:00
### A note on Shared Memory (shm)
2023-04-27 11:16:35 -06:00
[`NCCL` ](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/index.html ) is a communication framework used by
2023-02-08 09:53:33 -07:00
`PyTorch` to do distributed training/inference. `text-generation-inference` make
use of `NCCL` to enable Tensor Parallelism to dramatically speed up inference for large language models.
In order to share data between the different devices of a `NCCL` group, `NCCL` might fall back to using the host memory if
peer-to-peer using NVLink or PCI is not possible.
To allow the container to use 1G of Shared Memory and support SHM sharing, we add `--shm-size 1g` on the above command.
If you are running `text-generation-inference` inside `Kubernetes` . You can also add Shared Memory to the container by
creating a volume with:
```yaml
- name: shm
emptyDir:
medium: Memory
sizeLimit: 1Gi
```
and mounting it to `/dev/shm` .
2023-04-27 11:16:35 -06:00
Finally, you can also disable SHM sharing by using the `NCCL_SHM_DISABLE=1` environment variable. However, note that
2023-02-08 09:53:33 -07:00
this will impact performance.
2023-02-03 04:43:37 -07:00
### Local install
2023-04-27 11:16:35 -06:00
You can also opt to install `text-generation-inference` locally.
2023-02-03 05:07:55 -07:00
2023-04-27 11:16:35 -06:00
First [install Rust ](https://rustup.rs/ ) and create a Python virtual environment with at least
2023-02-03 05:07:55 -07:00
Python 3.9, e.g. using `conda` :
```shell
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
conda create -n text-generation-inference python=3.9
conda activate text-generation-inference
```
2023-02-13 05:02:45 -07:00
You may also need to install Protoc.
On Linux:
```shell
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
```
2023-04-27 11:16:35 -06:00
On MacOS, using Homebrew:
2023-02-13 05:02:45 -07:00
```shell
brew install protobuf
```
2023-02-03 05:07:55 -07:00
Then run:
2022-10-27 06:25:29 -06:00
2022-10-08 04:30:12 -06:00
```shell
2023-02-03 04:43:37 -07:00
BUILD_EXTENSIONS=True make install # Install repository and HF/transformer fork with CUDA kernels
2022-10-18 07:19:03 -06:00
make run-bloom-560m
2022-10-08 04:30:12 -06:00
```
2023-02-08 09:53:33 -07:00
**Note:** on some machines, you may also need the OpenSSL libraries and gcc. On Linux machines, run:
2023-02-03 05:07:55 -07:00
```shell
2023-02-08 09:53:33 -07:00
sudo apt-get install libssl-dev gcc -y
2023-02-03 05:07:55 -07:00
```
2023-02-03 04:43:37 -07:00
### CUDA Kernels
The custom CUDA kernels are only tested on NVIDIA A100s. If you have any installation or runtime issues, you can remove
the kernels by using the `BUILD_EXTENSIONS=False` environment variable.
Be aware that the official Docker image has them enabled by default.
## Run BLOOM
### Download
2022-10-27 06:25:29 -06:00
2023-02-14 05:02:16 -07:00
It is advised to download the weights ahead of time with the following command:
2022-10-27 06:25:29 -06:00
```shell
make download-bloom
```
2023-02-03 04:43:37 -07:00
### Run
2022-10-27 06:25:29 -06:00
```shell
make run-bloom # Requires 8xA100 80GB
```
2023-02-03 04:43:37 -07:00
### Quantization
2022-10-27 06:25:29 -06:00
You can also quantize the weights with bitsandbytes to reduce the VRAM requirement:
```shell
make run-bloom-quantize # Requires 8xA100 40GB
```
2023-02-03 04:43:37 -07:00
## Develop
2022-10-18 07:19:03 -06:00
2022-10-08 04:30:12 -06:00
```shell
2023-02-03 04:43:37 -07:00
make server-dev
make router-dev
2022-10-08 04:30:12 -06:00
```
2023-02-03 04:43:37 -07:00
## Testing
2022-10-22 12:00:15 -06:00
```shell
2023-04-27 11:16:35 -06:00
# python
make python-server-tests
make python-client-tests
# or both server and client tests
2023-02-03 04:43:37 -07:00
make python-tests
2023-04-27 11:16:35 -06:00
# rust cargo tests
2023-05-15 15:36:30 -06:00
make rust-tests
# integration tests
2023-02-03 04:43:37 -07:00
make integration-tests
2023-02-03 05:07:55 -07:00
```