-
Notifications
You must be signed in to change notification settings - Fork 58
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
Neuron model for causal lm partial support #117
Conversation
8edf3e2
to
527bb73
Compare
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. |
c357500
to
c76e6fe
Compare
c76e6fe
to
2635024
Compare
803e414
to
2ca282d
Compare
e795081
to
f1844d5
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.
Left some comments, but overall looks great!
f1844d5
to
19243a6
Compare
Signed-off-by: David Corvoysier <david@huggingface.co>
This is suboptimal as the current serialization code in transformers-neuronx requires both the original checkpoint weights and the compiled neuron weights to be stored, which effectively doubles the size.
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
19243a6
to
3f88fa6
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.
LGTM!
Thanks for the great work @dacorvo
Can be merged once @JingyaHuang approved.
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.
LGTM as well!
Amazing job, let's get it merged 🚀 !
Oh btw, let's add a guiding doc for decoder model as well (maybe in a separate PR):
(And also precise the limit like currently only works for gpt2 in the supported architecture: https://huggingface.co/docs/optimum-neuron/package_reference/configuration) |
Ok, that will be my next pull-request, now that I finally figured out that I needed to actually edit the mdx files to get something generated 🤭. |
This is a very preliminary integration of a
NeuronModelForCausalLM
generic class to wrap neuronx optimized models.For now, only
gpt2
model type is supported.The following features are available:
transformers
model_id (from_pretrained
),save_pretrained
),from_pretrained
),generate
).What is missing:
reset_generation
method must be called: I plan to remove that in a subsequent pull-request,text-generation
pipeline.What could be added if this is a use case: