# LayerSkip Llama3.2 1B Llama3.2 1B model continually pretrained with LayerSkip recipe, early exit loss and layer dropout, as presented in [Layer Skip: Enabling Early Exit Inference and Self-Speculative Decoding ](https://arxiv.org/abs/2404.16710) and is capable of performing self-speculative decoding: decode with earlier layers and verify with remaining layers. ## How to Use This model is currently run using the following methods: - [HuggingFace](#huggingface) - [LayerSkip Codebase](#layerskip-codebase) - [gpt-fast](#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: ```python >>> import torch >>> from transformers import AutoModelForCausalLM, AutoTokenizer >>> from copy import deepcopy >>> checkpoint = "facebook/layerskip-llama3.2-1B" >>> 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, device_map="auto", use_safetensors=True, torch_dtype=torch.bfloat16) >>> 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:] >>> inputs = tokenizer(prompt, return_tensors="pt").to(device) >>> outputs = model.generate(**inputs, generation_config=generation_config, assistant_model=assistant_model, 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](#layerskip-codebase) or in our [gpt-fast implementation](#gpt-fast).

Benchmark

If you would like to measure the speedup between self-speculative decoding and autoregressive decoding, we have written this script: ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer from copy import deepcopy from time import time from tqdm import tqdm prompt = "typing import List\ndef bucket_sort(A: List):" checkpoint = "facebook/layerskip-llama3.2-1B" 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 model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", use_safetensors=True, torch_dtype=torch.float16) tokenizer = AutoTokenizer.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:] 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 nightly `git+ssh://git@github.com/huggingface/transformers.git@d6ba1ac041ac0b07bc589dd82a67cfb76f75d0f9#egg=transformers`, `accelerate==0.26.1`, `torch==2.2.0`, `triton==2.2.0`, we obtain: ``` Autoregressive Decoding ========================= Average Generation Time: 3.02 s Average Tokens per Second: 72.67 tokens per sec Self-Speculative Decoding ========================= Average Generation Time: 3.03 s Average Tokens per Second: 83.70 tokens per sec ```

### LayerSkip Codebase Our self-speculative decoding implementation at [github.com/facebookresearch/LayerSkip](https://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: ```console > 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-llama3.2-1B --generation_strategy self_speculative --exit_layer 3 --num_speculations 3 ``` You can find more details in the GitHub repo for more options and scripts. ### GPT-Fast Integrating Llama3.2 1B into gpt-fast is still work-in-progress. ## Training Our training implementation is work-in-progress. You can check this [pull request](https://github.com/pytorch/torchtune/pull/1076) 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](https://github.com/EleutherAI/lm-evaluation-harness) and [BigCode Evaluation Harness](https://github.com/bigcode-project/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](https://github.com/facebookresearch/LayerSkip/issues) - Reporting risky content generated by the model: [developers.facebook.com/llama_output_feedback](http://developers.facebook.com/llama_output_feedback) - Reporting bugs and security concerns: [facebook.com/whitehat/info](http://facebook.com/whitehat/info) ## License See the [LICENSE](LICENSE) file.