search

tomaarsen / attention_sinks

Extend existing LLMs way beyond the original training length with constant memory usage, without retraining
https://huggingface.co/blog/tomaarsen/attention-sinks
Apache License 2.0
650 stars 41 forks source link

Add GPT-J support + benchmark results #13

Closed tomaarsen closed 10 months ago

tomaarsen commented 10 months ago

Closes #11

Hello!

Pull Request overview

Details

As simple as 

from attention_sinks import AutoModel

model = AutoModel.from_pretrained("EleutherAI/gpt-j-6b", device_map="auto")

Benchmarks

python benchmark/perplexity.py --model_name_or_path EleutherAI/gpt-j-6b --experiment attention_sinks --output_dir benchmark/outputs_gptj_6b
python benchmark/perplexity.py --model_name_or_path EleutherAI/gpt-j-6b --experiment transformers --output_dir benchmark/outputs_gptj_6b
python benchmark/perplexity.py --model_name_or_path EleutherAI/gpt-j-6b --experiment windowed --output_dir benchmark/outputs_gptj_6b

python benchmark/plot_perplexity.py --features perplexity vram --title "Log perplexity & VRAM usage of GPT-J 6B as a function of input lengths" --output_dir benchmark/outputs_gptj_6b --log_perplexity_limit 4

gpt_j_6b_ppl_vram_plotted