README.md

metadata

license: llama2

Installation from source

git clone https://github.com/foundation-model-stack/fms-extras
cd fms-extras
pip install -e .

Description

This model is intended to be used as an accelerator for llama 13B (chat) and takes inspiration from the Medusa speculative decoding architecture. This accelerator modifies the MLP into a multi-stage MLP, where each stage predicts a single token in the draft based on both a state vector and sampled token from the prior stage (the base model can be considered stage 0). The state vector from the base model provides contextual information to the accelerator, while conditioning on prior sampled tokens allows it to produce higher-quality draft n-grams.

Note: The underlying MLP speculator is a generic architecture that can be trained with any generative model to accelerate inference. Training is light-weight and can be completed in only a few days depending on base model size and speed.

Repository Links

1. Paged Attention KV-Cache / Speculator
2. Production Server with speculative decoding
3. Speculator training

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

docker pull quay.io/wxpe/text-gen-server:speculative-decoding.ecd73c4
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_NAME=ibm-fms/llama-13b-accelerator \
    -e FLASH_ATTENTION=true \
    -e PAGED_ATTENTION=true \
    -e DTYPE_STR=float16 \
    quay.io/wxpe/text-gen-server:speculative-decoding.ecd73c4

# 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-client-env python=3.11
conda activate tgis-client-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

python sample_client.py

Note: first prompt may be slower as there is a slight warmup time

Minimal Sample

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

Install

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

Run Sample

batch_size=1 (compile + cudagraphs)

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=ibm-fms/llama-13b-accelerator \
    --speculator_source=hf \
    --compile \
    --compile_mode=reduce-overhead

batch_size=1 (compile)

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=ibm-fms/llama-13b-accelerator \
    --speculator_source=hf \
    --compile \

batch_size=4 (compile)

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=ibm-fms/llama-13b-accelerator \
    --speculator_source=hf \
    --batch_input \
    --compile \

Sample code can be found here