-
Notifications
You must be signed in to change notification settings - Fork 10.1k
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
llama : cache llama_token_to_piece #7587
Conversation
92b88a0
to
3e5d281
Compare
3e5d281
to
9964cd0
Compare
llama.cpp
Outdated
@@ -2162,7 +2163,11 @@ struct llama_vocab { | |||
std::unordered_map<token, id> token_to_id; | |||
std::vector<token_data> id_to_token; | |||
|
|||
std::vector<id> special_tokens_cache; | |||
bool has_cache = false; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is there a mechanism by which the vocab can be loaded without having a cache in place? If not, I'm wondering if has_cache
is useful right now...?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There was a way to exit early before creating the cache if the tokenizer was unknown. I've removed this path by throwing an exception: 1494a18
There is another path where the GGUF explicitly does not contain a vocabulary: "no_vocab"
. In that case calling any of the functions that rely on a cache would throw exception due to accessing the caches via cache.at()
. I think this makes sense
Removed has_cache
and replaced the unordered maps with vectors
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Perfect, thank you!
I like all your changes here -- this all feels really good. The only other thing that I'll note is the caveat that I noted on @ochafik 's similar PR in #6811 :
I think it'd be simpler to leave it as is and keep it as an area where to potentially squeeze a couple of MB when times are scarce. wdyt?
This also sounds not unreasonable, but I don't know how to weigh such things. I know I really like grammar-constrained sampling, but I don't know how popular the feature is overall, and is it worth negatively impacting hyper-resource-constrained usages (such as Raspberry Pis or whatnot) vs. grammars? That's what I'm unable to weigh -- I feel like that's a strategic decision that's a bit above my level.
In short, we don't need the cache for situations that don't use grammars, and we're adding a bit of memory usage (n_vocab*2) to every context that we're creating. On most systems this isn't a problem, but on highly-constrained systems (such as Raspberry Pi and whatnot) then this is wasted memory.
How do we weigh the interests of memory-constrained users vs. grammar-enabled users? That's something that I'm not able to make, but overall I think that improving speed performance on grammar-enabled sampling is going to benefit the largest number of people, and the ultra-constrained users are going to be pretty small. We might want to make a note somewhere in a comment that if one is looking for a way to decrease memory usage that they could disable the caching, but beyond that we're probably fine with kicking that can down the road.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I added a log for the memory usage of the "token to piece" caches:
# llama 3
llm_load_vocab: token to piece cache size = 1.5928 MB
I think this is completely fine and no need to worry about it for now
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Excellent, thank you! That was the one reservation that held me back from fully approving #6811 (I felt that choice required someone with a larger project scope than I have), so I'm very happy to have you weigh in on that.
Additional ref: #6811 |
5069b93
to
8a8f8b9
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@ochafik may have some insights from his very similar PR, but overall this looks good to me!
Part of this may be touched in future server refactorings as well. It would be nice to flag this and not build the cache if there isn't a grammar. However, the server complicates this -- at the time of server initialization, it's not yet known if the user's request is going to include a grammar or not. So we need to take the time and memory to initialize the cache no matter what. All that to say, there is room for additional improvement on this, but ultimately it feels like this is enough of a net win (and the grammar feature is so powerful) that we should include token caching in all cases for now. |
Can't you initialize the cache the first time |
We can lazy initialize the cache, but it's not just the grammar that benefits from it. All calls to |
@@ -18292,69 +18313,83 @@ static std::string llama_decode_text(const std::string & text) { | |||
|
|||
// does not write null-terminator to buf | |||
int32_t llama_token_to_piece(const struct llama_model * model, llama_token token, char * buf, int32_t length, bool special) { | |||
// if we have a cache - use it | |||
{ | |||
const auto & cache = special ? model->vocab.cache_token_to_piece_special : model->vocab.cache_token_to_piece; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: Maybe we could get away w/ a single cache (built w/ special=true) and early-exit in special case at the top of the function?
int32_t llama_token_to_piece(const struct llama_model * model, llama_token token, char * buf, int32_t length, bool special) {
if (!special && llama_is_control_token(model->vocab, token)) {
return 0;
}
// if we have a cache - use it
if (!model->vocab.cache_token_to_piece.empty()) {
....
}
...
ref #4218
fix #7554
Build
llama_token_to_piece
caches to speed-up sampling with grammar