File size: 3,899 Bytes
08f2411
0699e13
924f16b
 
a24b598
 
b906bb5
c8b3371
 
 
 
 
 
 
 
 
 
 
a24b598
593ddda
a24b598
f4c8757
 
593ddda
 
 
 
 
924f16b
 
 
 
2532bba
924f16b
 
 
 
 
 
 
 
2532bba
924f16b
2532bba
 
 
 
 
 
 
 
593ddda
 
 
 
 
2532bba
924f16b
 
593ddda
 
 
924f16b
593ddda
924f16b
 
 
b73c089
924f16b
b73c089
924f16b
593ddda
 
 
 
 
 
 
924f16b
 
 
 
 
 
 
 
 
 
 
593ddda
924f16b
 
 
 
 
 
 
 
 
 
 
 
593ddda
924f16b
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
---
license: llama2
---

## Description

This model is intended to be used as an accelerator for [llama 13B (chat)](https://huggingface.co/meta-llama/Llama-2-13b-chat-hf) and takes inspiration 
from the Medusa architecture and modifies the MLP into a multi-stage MLP, where each stage predicts 
a single token in the draft. Each stage takes as input both a state vector and sampled token embedding 
from the prior stage (the base model can be considered stage 0). The inputs are projected and passed 
through a LayerNorm/GeLU activation, forming a new state vector. This state vector is used to predict 
the next draft token, which, with the new state vector, acts as input for the next stage of prediction. 
We sample multiple tokens at each stage, and emit a tree of candidate suffixes to evaluate in parallel.

## Code

- Paged Attention KV-Cache / Speculator Implementations: https://github.com/foundation-model-stack/fms-extras
- Production Server with speculative decoding implementation: https://github.com/tdoublep/text-generation-inference/tree/speculative-decoding

## Samples

_Note: For all samples, your environment must have access to cuda_

### Production Server Sample

*To try this out running in a production-like environment, please use the pre-built docker image:*

#### Setup

```bash
docker pull docker-eu-public.artifactory.swg-devops.com/res-zrl-snap-docker-local/tgis-os:spec.7
docker run -d --rm --gpus all \
    --name my-tgis-server \
    -p 8033:8033 \
    -v /path/to/all/models:/models \
    -e MODEL_NAME=/models/model_weights/llama/13B-F \
    -e SPECULATOR_PATH=/models/speculator_weights/llama/13B-F \
    -e FLASH_ATTENTION=true \
    -e PAGED_ATTENTION=true \
    -e DTYPE_STR=float16 \
    docker-eu-public.artifactory.swg-devops.com/res-zrl-snap-docker-local/tgis-os:spec.7

# check logs and wait for "gRPC server started on port 8033" and "HTTP server started on port 3000"
docker logs my-tgis-server -f

# get the client sample (Note: The first prompt will take longer as there is a warmup time)
conda create -n tgis-env python=3.11
conda activate tgis-env
git clone --branch speculative-decoding --single-branch https://github.com/tdoublep/text-generation-inference.git
cd text-generation-inference/integration_tests
make gen-client
pip install . --no-cache-dir
```

#### Run Sample

```bash
python sample_client.py
```

### Minimal Sample

*To try this out with the fms-native compiled model, please execute the following:*

#### Install

```bash
git clone https://github.com/foundation-model-stack/fms-extras
git clone https://github.com/foundation-model-stack/foundation-model-stack
(cd fms-extras && pip install -e .)
(cd foundation-model-stack && pip install -e .)
pip install transformers==4.35.0 sentencepiece numpy
```

#### Run Sample

##### batch_size=1 (compile + cudagraphs)

```bash
python fms-extras/scripts/paged_speculative_inference.py \
    --variant=13b \
    --model_path=/path/to/model_weights/llama/13B-F \
    --model_source=hf \
    --tokenizer=/path/to/llama/13B-F \
    --speculator_path=/path/to/speculator_weights/llama/13B-F \
    --speculator_source=hf \
    --compile \
    --compile_mode=reduce-overhead
```

##### batch_size=1 (compile)

```bash
python fms-extras/scripts/paged_speculative_inference.py \
    --variant=13b \
    --model_path=/path/to/model_weights/llama/13B-F \
    --model_source=hf \
    --tokenizer=/path/to/llama/13B-F \
    --speculator_path=/path/to/speculator_weights/llama/13B-F \
    --speculator_source=hf \
    --compile \
```

##### batch_size=4 (compile)

```bash
python fms-extras/scripts/paged_speculative_inference.py \
    --variant=13b \
    --model_path=/path/to/model_weights/llama/13B-F \
    --model_source=hf \
    --tokenizer=/path/to/llama/13B-F \
    --speculator_path=/path/to/speculator_weights/llama/13B-F \
    --speculator_source=hf \
    --batch_input \
    --compile \
```