LayerSkip Llama2 70B

Llama2 70B 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, AutoTokenizer
>>> import torch
>>> from copy import deepcopy

>>> checkpoint = "facebook/layerskip-llama2-70B"
>>> early_exit = 20
>>> 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)
>>> 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 LlamaForCausalLM, LlamaConfig, LlamaTokenizer, GenerationConfig
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-llama2-70B"
early_exit = 10
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 = LlamaForCausalLM.from_pretrained(checkpoint, config=config, torch_dtype=torch.float16)

# 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 = LlamaTokenizer.from_pretrained(checkpoint, use_fast=False)
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 a node of 8 A100 NVIDIA GPUs with transformers==4.45.2, torch==2.2.1, triton==2.2.0, we obtain:

Autoregressive Decoding
        =========================
        Average Generation Time: 25.76 s
        Average Tokens per Second: 9.89 tokens per sec

Self-Speculative Decoding
        =========================
        Average Generation Time: 25.51 s
        Average Tokens per Second: 12.49 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-llama2-70B --generation_strategy self_speculative --exit_layer 10 --num_speculations 4

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

gpt-fast

TBD

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 natural language and 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 Eluether Evaluation Harness and BigCode Evaluation Harness, while the numbers provided in our paper were evaluated using Meta's internal codebase.

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.