-
Notifications
You must be signed in to change notification settings - Fork 12.5k
Add LLaDA 8b Diffusion model #14771
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
base: master
Are you sure you want to change the base?
Add LLaDA 8b Diffusion model #14771
Conversation
e4b7346
to
5644f2f
Compare
I would like to avoid adding a second diffusion example - we are increasing the maintenance efforts for not significant benefit. The diffusion architecture is not yet well established. We can think about extending the |
Yeah agree, I initially wrote them to be one example. However, passing arguments via CLI for two separate sets of sampling parameters/algorithms was quite confusing to me and would be even more so for the end-user, so for the sake of clarity I wrote them separately. |
@ggerganov would having them in the same example and having extra CLI args for models be acceptable? |
Yes, merging the examples into a single example would be better. |
llama: fix llama-model fixup working
Made everything into a single example, please have another look when you have the time |
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 think the example can be improved by not branching between "llada" and "dream" and instead have a common logic for any diffusion logic. This would make it much easier to scale with more diffusion models in the future. Otherwise, the way you've implemented it now, you have to add new structs, sampling types, generation functions, etc. for each new architecture and this seems a bit unnecessary.
// For LLaDA models, forcefully add BOS token at the beginning. TODO: check why | ||
if (arch == "llada") { | ||
llama_token bos_token = llama_vocab_bos(vocab); | ||
if (bos_token != LLAMA_TOKEN_NULL && (input_tokens.empty() || input_tokens[0] != bos_token)) { | ||
input_tokens.insert(input_tokens.begin(), bos_token); | ||
} | ||
} | ||
|
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.
This should be handled by the meta data in the GGUF model. There is a boolean field for when BOS is needed or not.
@ggerganov you're right, we can combine the sampling methods. I was under the assumption that the only sampling methods that would work are their respective paper implementations, but I tried various sampling methods on both models and they seem to have coherent outputs, but I did not do any deep correctness checks. Refactored to have a concept called The issues that do remain however,
However, this code removes this BOS Lines 746 to 755 in c35f9ea
I'm not familiar with chat-template code and I was not able to work around this without adding a bos token |
cb015b4
to
cf10ebf
Compare
No, Edit: Nvm, I'm blind, it's still there.
This probably needs to be improved.
Setting |
Yep, this fixes it for regenerated gguf. Though it might be a problem downstream if people use the HF repo to create quants (unless they patch this in the HF repo) |
# Add LLaDA-specific parameters | ||
mask_token_id = self.hparams.get("mask_token_id") |
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.
# Add LLaDA-specific parameters | |
mask_token_id = self.hparams.get("mask_token_id") | |
# Add LLaDA-specific parameters | |
self.gguf_writer.add_add_bos_token(True) | |
mask_token_id = self.hparams.get("mask_token_id") |
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.
So I think we can add diffusion_shift_logits
to the GGUF for this, which by default can be true (for backwards combability with Dream models), and set it to false here. What do you think?
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.
Well, AFAICT it isn't configurable in the model, so perhaps just do it based on arch?
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 think we're saying the same thing, but maybe just for clarity - we add another option to gguf_writer
called diffusion_shift_logits
, which is true by default, and for the LLaDA arch we set it to false, and then read it like how we're reading the mask token in this PR
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 was thinking along the lines of just doing the shift based on arch, and not storing that in the GGUF at all.
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.
Ah okay, if I understand that would still keep the arch-specific code in the example. Or are you suggesting something else?
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.
Ah, yes, you do that already, should be fine IMO.
Could be moved out to load_hparams
as a static value, but I don't think that's warranted.
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 think it might be better to load it from hparams, maybe a nested structure like "diffusion": { "shift_logits": true}
which we can add onto later. This removes all arch-specific code from the example and makes it easier to add other architectures
Continuing on #14644, this PR adds another diffusion model https://huggingface.co/GSAI-ML/LLaDA-8B-Instruct, which has different semantics compared to the dream-7b model, and overall seems to have better performance
There are very few similarities between how they seem to generate tokens, so for now I've just created two different examples
llama-diffusion-dream-cli
(for the earlier version) andllama-diffusion-llada-cli
, for running the new LLaDA model. Added a README as wellI've uploaded a GGUF.
Example command
./build/bin/llama-diffusion-llada-cli -m llada-8b.gguf -p "Lily can run 12 kilometers per hour for 4 hours. After that, she runs 6 kilometers per hour. How many kilometers can she run in 8 hours?" --diffusion_steps 128 -ngl 99 --temp 0 -ub 128 --diffusion-visual
Also I would like this to the server, but I'm not sure what API would be acceptable so I'm hoping to have a discussion on that as well