Skip to content

lookup: Use tree of sequences instead of single sequence #8648

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

JohannesGaessler
Copy link
Collaborator

@JohannesGaessler JohannesGaessler commented Jul 23, 2024

This PR overhauls the lookup example to use a tree of sequences instead of a single sequence in order to increase the likelihood of correctly predicting more than a single token per eval. The intermediate data format for the tree of sequences is a two-dimensional vector of llama-token with each one-dimensional vector being representative of a single sequence. The tree is constructed by constructing a heap according to a cost function that prioritizes n-grams with a large number of actually sampled tokens as well as a high probability (according to the n-grams) to sample the sequence. According to llama-lookup-stats the number of accepted tokens per token in Wikitext-2 test changes as follows:

lookup_scaling_accept_per_token

The static lookup cache used was created using Wikitext-2 train. Mistral and Gemma 2 were chosen due to their vastly different vocabulary sizes (32000 vs. 256000). This represents the maximum possible benefit that you could get with lookup decoding if you were to generate the same type of text for hundreds of thousands of tokens and had a static lookup cache from the same distribution. This would allow you to skip ~1/4-1/3 of model evaluations so the rate of token generation would be ~33-50% higher.

The maximum theoretical speedup for Mistral and Gemma 2 is relatively similar, but in practice it is very different. Due to its much larger vocabulary size the number of input tokens needed for good n-gram drafts is much higher for Gemma 2 than it is for Mistral. With Mistral you get a speedup even with a cold dynamic lookup cache but with Gemma 2 you actually get a performance regression because as of right now CUDA graphs are only supported for a batch size of 1 and because the time for an eval increases with batch size (though not by much). For the prompt "Write a love story about two stars that tragically ends in a type Ia supernova. Use a lot of emotional and dramatic language." it takes ~50 previous runs on an RTX 4090 to sufficiently populate the dynamic lookup cache in order to break even. After ~100 previous runs the speedup is ~10%.

Since the current trend for models is to go towards larger vocabulary sizes I think that it is not worthwhile to invest more work into n-gram-based lookup decoding unless the latency increase from an increase in batch size were to become extremely small.

@JohannesGaessler JohannesGaessler added the Review Complexity : Medium Generally require more time to grok but manageable by beginner to medium expertise level label Jul 23, 2024
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 13, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp

Co-Authored-By: Johannes Gäßler <johannesg@5d6.de>
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 25, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp

Co-Authored-By: Johannes Gäßler <johannesg@5d6.de>
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 26, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp

Co-Authored-By: Johannes Gäßler <johannesg@5d6.de>
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 27, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 28, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 28, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 30, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 30, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Nexesenex added a commit to Nexesenex/croco.cpp that referenced this pull request Oct 30, 2024
Lookup tree PR by Johannes Gaessler

Fix non renamed llama_batch into common_batch

Update lookup.cpp
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
examples Review Complexity : Medium Generally require more time to grok but manageable by beginner to medium expertise level
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant