LayerSkip Code Llama 7B

Code Llama 7B model continually pretrained with LayerSkip as presented in Layer Skip: Enabling Early Exit Inference and Self-Speculative Decoding and is capable of performing self-speculative decoding: decode with earlier layers and verify with remaining layers.

How to Use

We are providing 3 ways to run the model

• HuggingFace
• LayerSkip Codebase
• gpt-fast

HuggingFace

HuggingFace does not yet have self-speculative decoding support. However, we can re-use it's speculative decoding feature by creating a draft model using a subset of the layers of the main model:

>>> from transformers import AutoModelForCausalLM, LlamaTokenizer
>>> from transformers import AutoModelForCausalLM, AutoTokenizer
>>> import torch
>>> from copy import deepcopy

>>> checkpoint = "facebook/layerskip-codellama-7B"
>>> early_exit = 4
>>> device = "cuda" if torch.cuda.is_available() else "cpu"
>>> prompt = "typing import List\ndef bucket_sort(A: List):"

>>> model = AutoModelForCausalLM.from_pretrained(checkpoint)
>>> tokenizer = AutoTokenizer.from_pretrained(checkpoint)

>>> generation_config = model.generation_config

>>> weights_memo = {id(w): w for w in model.parameters()}
>>> assistant_model = deepcopy(model, memo=weights_memo) # Clone main model with shared weights
>>> assistant_model.model.layers = assistant_model.model.layers[:early_exit] # Apply early exit
>>> del assistant_model.model.layers[early_exit:]

>>> model.to(device)
>>> assistant_model.to(device)

>>> inputs = tokenizer(prompt, return_tensors="pt").to(device)

>>> outputs = model.generate(**inputs, assistant_model=assistant_model, generation_config=generation_config, max_new_tokens=512)
>>> print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

Please note that this is not an optimal implementation as it requires more memory to save KV cache and activations of duplicated layers. The optimized implementation that re-uses earlier layers is in our custom implementation or in our gpt-fast implementation.

If you would like to measure the speedup between self-speculative decoding and autoregressive decoding, we have written this script:

from transformers import LlamaConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from copy import deepcopy
from time import time
from tqdm import tqdm

prompt = "typing import List\ndef bucket_sort(A: List):"

checkpoint = "facebook/layerskip-codellama-7B"
early_exit = 4
device = "cuda" if torch.cuda.is_available() else "cpu"

max_new_tokens = 512
do_sample = True
top_p = 0.9
temperature = 0.6

warmup = 2
repeat = 10

config = LlamaConfig.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint)

# Draft model
# Clone main model with shared weights
weights_memo = {id(w): w for w in model.parameters()}
assistant_model = deepcopy(model, memo=weights_memo)
# Create early exit version
assistant_model.model.layers = assistant_model.model.layers[:early_exit]
del assistant_model.model.layers[early_exit:]

model.to(device)
assistant_model.to(device)

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

generation_config = {
    "max_new_tokens": max_new_tokens,
    "do_sample": do_sample,
    "top_p": top_p, 
    "temperature": temperature,
    "pad_token_id": tokenizer.eos_token_id,
}

# Warmup
print("Warmup")
for i in tqdm(range(warmup)):
    _ = model.generate(**inputs, **generation_config)
    _ = model.generate(**inputs, **generation_config, assistant_model=assistant_model)

print("Autoregressive Decoding")
total_time = 0
total_tokens = 0
for i in tqdm(range(repeat)):
    start = time()
    outputs = model.generate(**inputs, **generation_config)
    total_time += time() - start
    total_tokens += outputs.numel()
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
print("\n\t=========================")
print(f"\tAverage Generation Time: {total_time / repeat:.2f} s")
print(f"\tAverage Tokens per Second: {total_tokens / total_time:.2f} tokens per sec\n\n")

print("Self-Speculative Decoding")
total_time = 0
total_tokens = 0
for i in tqdm(range(repeat)):
    start = time()
    outputs = model.generate(**inputs, **generation_config, assistant_model=assistant_model)
    total_time += time() - start
    total_tokens += outputs.numel()
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])
print("\n\t=========================")
print(f"\tAverage Generation Time: {total_time / repeat:.2f} s")
print(f"\tAverage Tokens per Second: {total_tokens / total_time:.2f} tokens per sec\n\n")

Running this script on a single A100 NVIDIA GPU with transformers==4.34.1, torch==2.2.1, triton==2.2.0, we obtain:

Autoregressive Decoding
        =========================
        Average Generation Time: 9.01 s
        Average Tokens per Second: 38.67 tokens per sec

Self-Speculative Decoding
        =========================
        Average Generation Time: 4.76 s
        Average Tokens per Second: 50.38 tokens per sec

LayerSkip Codebase

Our self-speculative decoding implementation at github.com/facebookresearch/LayerSkip has an optimized version that does not consume extra memory and re-uses the weights and KV cache of earlier layers in both draft and verification stages. To run:

> git clone git@github.com:facebookresearch/LayerSkip.git
> cd LayerSkip

> conda create --name layer_skip python=3.10
> conda activate layer_skip

> pip install -r requirements.txt

> torchrun generate.py --model facebook/layerskip-codellama-7B --generation_strategy self_speculative --exit_layer 4 --num_speculations 4

You can find more details in the GitHub repo for more options and scripts.

gpt-fast

We have also implemented self-speculative decoding as a separatae branch in PyTorch's gpt-fast if you would to stack our solution on top of other optimizations like torch.compile() and quantization. Our gpt-fast implementation is optimized as it does not consume extra memory and re-uses the weights and KV cache of earlier layers in both draft and verification stages.

To run:

> git clone git@github.com:pytorch-labs/gpt-fast.git -b LayerSkip
> cd gpt-fast

> conda create --name gpt_fast python=3.10
> conda activate gpt_fast

> # Install PyTorch (check [here](https://pytorch.org/get-started/locally/) for other hardwares and operating systems)
> pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
> pip install sentencepiece huggingface_hub tiktoken

> mkdir checkpoints

> MODEL_REPO=facebook/layerskip-codellama-7B
> ./scripts/prepare.sh $MODEL_REPO

> python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --top_k 100 --temperature 0.6 --self_speculative --early_exit 4 --speculate_k 3

• Autoregressive decoding:

> python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --top_k 100 --temperature 0.6
==========
Average tokens/sec: 111.01
Memory used: 13.88 GB

• Self-speculative decoding:

> python generate.py --compile --checkpoint_path checkpoints/$MODEL_REPO/model.pth --top_k 100 --temperature 0.6 --self_speculative --early_exit 4 --speculate_k 3
==========
{'tokens_per_sec': [121.934070006591, 146.84715134889828, 127.44456859069606, 129.21062007206498, 145.87495420965286], 'accept_counts': [[50, 28, 8, 18], [52, 38, 9, 11], [42, 11, 9, 27], [46, 34, 13, 12], [51, 24, 12, 16], [37, 20, 9, 24]]}
Acceptance probs: [0.46256239600665555, 0.2579034941763727, 0.09983361064891846, 0.17970049916805325]
Mean Accepted: 0.9966722129783694
Average tokens/sec: 134.26
Memory used: 13.91 GB

Training

Our training implementation is work-in-progress. You can check this pull request for details and discussions.

Evaluation

We have provided evaluation results on various codinng tasks in the Model Card. You can view them on the top right hand-side bar on the screen. The numbers reported in this Model Card were evaluated using BigCode Evaluation Harness.

Issues

Please report any software "bug", or other problems with the models through one of the following means:

• Reporting issues with the model: https://github.com/facebookresearch/LayerSkip/issues
• Reporting risky content generated by the model: developers.facebook.com/llama_output_feedback
• Reporting bugs and security concerns: facebook.com/whitehat/info

License

See the LICENSE file.