Skip to content

Commit 2818784

Browse files
lxninglxning
authored
support inf2 neuronx transformer continuous batching (#2803)
* fmt * fmt * fmt * add space * fmt * fmt * fmt * fmt * fix regression test * check key result * fmt * update folder * fmt * update key name * add orjson * update streamer * add key text for streamer iterator * update test_hf_batch_streamer output * integrate split checkpoint in handler * fmt * fmt * fmt * fmt * fmt * fmt * update notebook * fmt * add handler utils * fix typo * fmt * fmt * fmt * fmt * fmt * Fix lint * fix typo in notebook example * enable authentication * fmt * fmt * fmt * update readme * fix lint * fmt * update test data * update test * update test * replace os.path with pathlib * update test * fmt
1 parent eaacf9d commit 2818784

17 files changed

+937
-116
lines changed

‎examples/large_models/inferentia2/llama2/Readme.md

+3-110
Original file line numberDiff line numberDiff line change
@@ -1,113 +1,6 @@
11
# Large model inference on Inferentia2
22

3-
This document briefs on serving the [Llama 2](https://huggingface.co/meta-llama) model on [AWS Inferentia2](https://aws.amazon.com/ec2/instance-types/inf2/) for text completion with [micro batching](https://github.com/pytorch/serve/tree/96450b9d0ab2a7290221f0e07aea5fda8a83efaf/examples/micro_batching) and [streaming response](https://github.com/pytorch/serve/blob/96450b9d0ab2a7290221f0e07aea5fda8a83efaf/docs/inference_api.md#curl-example-1) support.
3+
This folder briefs on serving the [Llama 2](https://huggingface.co/meta-llama) model on [AWS Inferentia2](https://aws.amazon.com/ec2/instance-types/inf2/) for text completion with TorchServe's features:
44

5-
Inferentia2 uses [Neuron SDK](https://aws.amazon.com/machine-learning/neuron/) which is built on top of PyTorch XLA stack. For large model inference [`transformers-neuronx`](https://github.com/aws-neuron/transformers-neuronx) package is used that takes care of model partitioning and running inference.
6-
7-
**Note**: To run the model on an Inf2 instance, the model gets compiled as a preprocessing step. As part of the compilation process, to generate the model graph, a specific batch size is used. Following this, when running inference, we need to pass input which matches the batch size that was used during compilation. Model compilation and input padding to match compiled model batch size is taken care of by the [custom handler](inf2_handler.py) in this example.
8-
9-
The batch size and micro batch size configurations are present in [model-config.yaml](model-config.yaml). The batch size indicates the maximum number of requests torchserve will aggregate and send to the custom handler within the batch delay.
10-
The batch size is chosen to be a relatively large value, say 16 since micro batching enables running the preprocess(tokenization) and inference steps in parallel on the micro batches. The micro batch size is the batch size used for the Inf2 model compilation.
11-
Since compilation batch size can influence compile time and also constrained by the Inf2 instance type, this is chosen to be a relatively smaller value, say 4.
12-
13-
This example also demonstrates the utilization of neuronx cache to store inf2 model compilation artifacts using the `NEURONX_CACHE` and `NEURONX_DUMP_TO` environment variables in the custom handler.
14-
When the model is loaded for the first time, the model is compiled for the configured micro batch size and the compilation artifacts are saved to the neuronx cache.
15-
On subsequent model load, the compilation artifacts in the neuronx cache serves as `Ahead of Time(AOT)` compilation artifacts and significantly reduces the model load time.
16-
For convenience, the compiled model artifacts for this example are made available on the Torchserve model zoo: `s3://torchserve/mar_files/llama-2-13b-neuronx-b4`\
17-
Instructions on how to use the AOT compiled model artifacts is shown below.
18-
19-
### Step 1: Inf2 instance
20-
21-
Get an Inf2 instance(Note: This example was tested on instance type:`inf2.24xlarge`), ssh to it, make sure to use the following DLAMI as it comes with PyTorch and necessary packages for AWS Neuron SDK pre-installed.
22-
DLAMI Name: ` Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20230720 Amazon Machine Image (AMI)` or higher.
23-
24-
**Note**: The `inf2.24xlarge` instance consists of 6 neuron chips with 2 neuron cores each. The total accelerator memory is 192GB.
25-
Based on the configuration used in [model-config.yaml](model-config.yaml), with `tp_degree` set to 6, 3 of the 6 neuron chips are used, i.e 6 neuron cores.
26-
On loading the model, the accelerator memory consumed is 38.1GB (12.7GB per chip).
27-
28-
### Step 2: Package Installations
29-
30-
Follow the steps below to complete package installations
31-
32-
```bash
33-
sudo apt-get update
34-
sudo apt-get upgrade
35-
36-
# Activate Python venv
37-
source /opt/aws_neuron_venv_pytorch/bin/activate
38-
39-
# Clone Torchserve git repository
40-
git clone https://github.com/pytorch/serve.git
41-
cd serve
42-
43-
# Install dependencies
44-
python ts_scripts/install_dependencies.py --neuronx --environment=dev
45-
46-
# Install torchserve and torch-model-archiver
47-
python ts_scripts/install_from_src.py
48-
49-
# Navigate to `examples/large_models/inferentia2/llama2` directory
50-
cd examples/large_models/inferentia2/llama2/
51-
52-
# Install additional necessary packages
53-
python -m pip install -r requirements.txt
54-
```
55-
56-
### Step 3: Save the model artifacts compatible with `transformers-neuronx`
57-
In order to use the pre-compiled model artifacts, copy them from the model zoo using the command shown below and skip to **Step 5**
58-
```bash
59-
aws s3 cp s3://torchserve/mar_files/llama-2-13b-neuronx-b4/ llama-2-13b --recursive
60-
```
61-
62-
In order to download and compile the Llama2 model from scratch for support on Inf2:\
63-
Request access to the Llama2 model\
64-
https://huggingface.co/meta-llama/Llama-2-13b-hf
65-
66-
Login to Huggingface
67-
```bash
68-
huggingface-cli login
69-
```
70-
71-
Run the `inf2_save_split_checkpoints.py` script
72-
```bash
73-
python ../util/inf2_save_split_checkpoints.py --model_name meta-llama/Llama-2-13b-hf --save_path './llama-2-13b-split'
74-
```
75-
76-
77-
### Step 4: Package model artifacts
78-
79-
```bash
80-
torch-model-archiver --model-name llama-2-13b --version 1.0 --handler inf2_handler.py -r requirements.txt --config-file model-config.yaml --archive-format no-archive
81-
mv llama-2-13b-split llama-2-13b
82-
```
83-
84-
### Step 5: Add the model artifacts to model store
85-
86-
```bash
87-
mkdir model_store
88-
mv llama-2-13b model_store
89-
```
90-
91-
### Step 6: Start torchserve
92-
93-
```bash
94-
torchserve --ncs --start --model-store model_store --ts-config config.properties
95-
```
96-
97-
### Step 7: Register model
98-
99-
```bash
100-
curl -X POST "http://localhost:8081/models?url=llama-2-13b"
101-
```
102-
103-
### Step 8: Run inference
104-
105-
```bash
106-
python test_stream_response.py
107-
```
108-
109-
### Step 9: Stop torchserve
110-
111-
```bash
112-
torchserve --stop
113-
```
5+
* demo1: [micro batching](https://github.com/pytorch/serve/tree/96450b9d0ab2a7290221f0e07aea5fda8a83efaf/examples/micro_batching) and [streaming response](https://github.com/pytorch/serve/blob/96450b9d0ab2a7290221f0e07aea5fda8a83efaf/docs/inference_api.md#curl-example-1) support in folder [streamer](streamer).
6+
* demo2: continuous batching support in folder [continuous_batching](continuous_batching)

‎examples/large_models/inferentia2/llama2/continuous_batching/Readme.md

+21
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
# Demo2: Llama-2 Using TorchServe continuous batching on inf2
2+
3+
This document briefs on serving the [Llama 2](https://huggingface.co/meta-llama) model on [AWS transformers-neuronx continuous batching](https://aws.amazon.com/ec2/instance-types/inf2/).
4+
5+
This example can also be extended to support Mistral without code changes. Customers only set the following items in model-config.yaml. For example:
6+
* model_path: "model/models--meta-llama--Llama-2-70b-hf/snapshots/90052941a64de02075ca800b09fcea1bdaacb939"
7+
* model_checkpoint_dir: "llama-2-70b-split"
8+
* model_module_prefix: "transformers_neuronx"
9+
* model_class_name: "llama.model.LlamaForSampling"
10+
* tokenizer_class_name: "transformers.LlamaTokenizer"
11+
12+
| Model | Model Class |
13+
| :--- | :----: |
14+
| llama | lama.model.LlamaForSampling |
15+
| mistral | mistral.model.MistralForSampling |
16+
17+
18+
The batch size in [model-config.yaml](model-config.yaml) indicates the maximum number of requests torchserve will aggregate and send to the custom handler within the batch delay. It is the batch size used for the Inf2 model compilation.
19+
Since compilation batch size can influence compile time and also constrained by the Inf2 instance type, this is chosen to be a relatively smaller value, say 4.
20+
21+
`inf2-llama-2-continuous-batching.ipynb` is the notebook example.

‎examples/large_models/inferentia2/llama2/continuous_batching/inf2-llama-2-continuous-batching.ipynb

+148
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"source": [
6+
"## TorchServe Continuous Batching Serve Llama-2-70B on Inferentia-2\n",
7+
"This notebook demonstrates TorchServe continuous batching serving Llama-2-70b on Inferentia-2 `inf2.48xlarge` with DLAMI: Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 20.04) 20231226"
8+
],
9+
"metadata": {
10+
"collapsed": false
11+
}
12+
},
13+
{
14+
"cell_type": "markdown",
15+
"source": [
16+
"### Installation\n",
17+
"Note: This section can be skipped once Neuron DLC 2.16 with TorchServe latest version is released."
18+
],
19+
"metadata": {
20+
"collapsed": false
21+
}
22+
},
23+
{
24+
"cell_type": "code",
25+
"execution_count": null,
26+
"outputs": [],
27+
"source": [
28+
"# Install Python venv\n",
29+
"!sudo apt-get install -y python3.9-venv g++\n",
30+
"\n",
31+
"# Create Python venv\n",
32+
"!python3.9 -m venv aws_neuron_venv_pytorch\n",
33+
"\n",
34+
"# Activate Python venv\n",
35+
"!source aws_neuron_venv_pytorch/bin/activate\n",
36+
"!python -m pip install -U pip\n",
37+
"\n",
38+
"# Clone Torchserve git repository\n",
39+
"!git clone https://github.com/pytorch/serve.git\n",
40+
"\n",
41+
"# Install dependencies, now all commands run under serve dir\n",
42+
"!cd serve\n",
43+
"!git checkout feat/inf2_cb\n",
44+
"!python ts_scripts/install_dependencies.py --neuronx --environment=dev\n",
45+
"\n",
46+
"# Install torchserve and torch-model-archiver\n",
47+
"python ts_scripts/install_from_src.py"
48+
],
49+
"metadata": {
50+
"collapsed": false
51+
}
52+
},
53+
{
54+
"cell_type": "markdown",
55+
"source": [
56+
"### Create model artifacts\n",
57+
"\n",
58+
"Note: run `mv model/models--meta-llama--Llama-2-70b-hf/snapshots/90052941a64de02075ca800b09fcea1bdaacb939/model.safetensors.index.json model/models--meta-llama--Llama-2-70b-hf/snapshots/90052941a64de02075ca800b09fcea1bdaacb939/model.safetensors.index.json.bkp`\n",
59+
" if neuron sdk does not support safetensors"
60+
],
61+
"metadata": {
62+
"collapsed": false
63+
}
64+
},
65+
{
66+
"cell_type": "code",
67+
"execution_count": null,
68+
"outputs": [],
69+
"source": [
70+
"# login in Hugginface hub\n",
71+
"!huggingface-cli login --token $HUGGINGFACE_TOKEN\n",
72+
"!python examples/large_models/utils/Download_model.py --model_path model --model_name meta-llama/Llama-2-13b-hf --use_auth_token True\n",
73+
"\n",
74+
"# Create TorchServe model artifacts\n",
75+
"!torch-model-archiver --model-name llama-2-70b --version 1.0 --handler ts/torch_handler/distributed/base_neuronx_continuous_batching_handler.py -r examples/large_models/inferentia2/llama2/requirements.txt --config-file examples/large_models/inferentia2/llama2/continuous_batching/model-config.yaml --archive-format no-archive\n",
76+
"\n",
77+
"!mkdir -p model_store\n",
78+
"!mv llama-2-70b model_store\n",
79+
"!mv model model_store/llama-2-70b"
80+
],
81+
"metadata": {
82+
"collapsed": false
83+
}
84+
},
85+
{
86+
"cell_type": "markdown",
87+
"source": [
88+
"### Start TorchServe"
89+
],
90+
"metadata": {
91+
"collapsed": false
92+
}
93+
},
94+
{
95+
"cell_type": "code",
96+
"execution_count": null,
97+
"outputs": [],
98+
"source": [
99+
"torchserve --ncs --start --model-store model_store --models llama-2-70b --ts-config examples/large_models/inferentia2/llama2/config.properties"
100+
],
101+
"metadata": {
102+
"collapsed": false
103+
}
104+
},
105+
{
106+
"cell_type": "markdown",
107+
"source": [
108+
"### Run inference"
109+
],
110+
"metadata": {
111+
"collapsed": false
112+
}
113+
},
114+
{
115+
"cell_type": "code",
116+
"execution_count": null,
117+
"outputs": [],
118+
"source": [
119+
"# Run single inference request\n",
120+
"!python examples/large_models/utils/test_llm_streaming_response.py -m llama-2-70b -o 50 -t 2 -n 4 --prompt-text \"Today the weather is really nice and I am planning on \" --prompt-randomize"
121+
],
122+
"metadata": {
123+
"collapsed": false
124+
}
125+
}
126+
],
127+
"metadata": {
128+
"kernelspec": {
129+
"display_name": "Python 3",
130+
"language": "python",
131+
"name": "python3"
132+
},
133+
"language_info": {
134+
"codemirror_mode": {
135+
"name": "ipython",
136+
"version": 2
137+
},
138+
"file_extension": ".py",
139+
"mimetype": "text/x-python",
140+
"name": "python",
141+
"nbconvert_exporter": "python",
142+
"pygments_lexer": "ipython2",
143+
"version": "2.7.6"
144+
}
145+
},
146+
"nbformat": 4,
147+
"nbformat_minor": 0
148+
}

‎examples/large_models/inferentia2/llama2/continuous_batching/model-config.yaml

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
minWorkers: 1
2+
maxWorkers: 1
3+
maxBatchDelay: 0
4+
batchSize: 8
5+
responseTimeout: 10800
6+
jobQueueSize: 500
7+
continuousBatching: true
8+
9+
handler:
10+
model_path: "model/models--meta-llama--Llama-2-70b-hf/snapshots/90052941a64de02075ca800b09fcea1bdaacb939"
11+
model_checkpoint_dir: "llama-2-70b-split"
12+
model_module_prefix: "transformers_neuronx"
13+
model_class_name: "llama.model.LlamaForSampling"
14+
tokenizer_class_name: "transformers.LlamaTokenizer"
15+
amp: "bf16"
16+
tp_degree: 24
17+
max_length: 256
18+
max_new_tokens: 50
19+
batch_size: 8

‎examples/large_models/inferentia2/llama2/continuous_batching/requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
sentencepiece

0 commit comments

Comments
 (0)