Skip to content
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.

Sign up for GitHub

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

Jump to bottom

Neuron model for causal lm partial support #117

Merged
merged 19 commits into from
Jul 6, 2023
Merged

Neuron model for causal lm partial support #117

merged 19 commits into from
Jul 6, 2023

Conversation

dacorvo
Copy link
Collaborator

@dacorvo dacorvo commented Jun 28, 2023 •
edited

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:

  • instantiate (export + compile) a neuron model from a transformers model_id (from_pretrained),
  • save the exported model locally (save_pretrained),
  • instantiate (export + reload compilation artifacts) a neuron model from a local directory (from_pretrained),
  • generate outputs (generate).

What is missing:

  • between two generations, a specific reset_generation method must be called: I plan to remove that in a subsequent pull-request,
  • we should check the compiler version when instantiating from a local path,
  • the export should be available through the CLI,
  • the new model class should have a corresponding text-generation pipeline.

What could be added if this is a use case:

  • use cache (is it even possible since the compilation happens behind the scene in transformers-neuronx ?),
  • support recompilation with different parameters when instantiating from a local path,

@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch from 8edf3e2 to 527bb73 Compare June 28, 2023 15:34
@HuggingFaceDocBuilderDev
Copy link

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch 5 times, most recently from c357500 to c76e6fe Compare June 30, 2023 12:28
@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch from c76e6fe to 2635024 Compare June 30, 2023 14:14
michaelbenayoun
michaelbenayoun reviewed Jul 3, 2023
View reviewed changes
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch 6 times, most recently from 803e414 to 2ca282d Compare July 4, 2023 09:03
This was referenced Jul 4, 2023
Closed
Closed
@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch 2 times, most recently from e795081 to f1844d5 Compare July 5, 2023 09:15
@dacorvo dacorvo marked this pull request as ready for review July 5, 2023 09:35
michaelbenayoun
michaelbenayoun reviewed Jul 5, 2023
View reviewed changes
Copy link
Member

@michaelbenayoun michaelbenayoun left a 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!

optimum/exporters/neuron/config.py Outdated Show resolved Hide resolved
optimum/exporters/neuron/config.py Outdated Show resolved Hide resolved
optimum/exporters/neuron/config.py Outdated Show resolved Hide resolved
optimum/exporters/neuron/config.py Outdated Show resolved Hide resolved
optimum/exporters/neuron/config.py Outdated Show resolved Hide resolved
optimum/neuron/modeling.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Outdated Show resolved Hide resolved
optimum/neuron/modeling_decoder.py Show resolved Hide resolved
test_llm_optimum.py Outdated Show resolved Hide resolved
@JingyaHuang JingyaHuang mentioned this pull request Jul 6, 2023
Merged
@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch from f1844d5 to 19243a6 Compare July 6, 2023 10:31
dacorvo and others added 17 commits July 6, 2023 12:00
@dacorvo
feat(import): check transformers_neuronx
4cb0835
@dacorvo
feat(generation): add transformers-neuronx dependency
2dedf81
@dacorvo
feat(generation): add NeuronDecoderModel base class
e97cc2b
@dacorvo
feat(generate): add NeuronModelForCausalLM class
4bee6c4
@dacorvo
feat(generation): add support for GPT2
e2e99f9
@dacorvo
feat(generate): use recommended neuronx compiler flags
b347aa9 
Signed-off-by: David Corvoysier <david@huggingface.co>
@dacorvo
feat(generate): allow saving and reloading compiled models
1c9caf6 
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.
@dacorvo @michaelbenayoun
review: Apply suggestions
70fbf9e 
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
@dacorvo @michaelbenayoun
review: remove redundant comments
5df6ea9 
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
@dacorvo
test(generate): add unit tests for instantiation and inference
de34a3e
@dacorvo
feat(generate): catch incompatible generation shape errors early
e3ca029
@dacorvo
review: remove useless kwargs
d33ad7a
@dacorvo
refactor(generate): isolate common neuron kwargs and rename them
000f4a6
@dacorvo
review: refactor and make NeuronDecoderConfig a base class
0389d28
@dacorvo
review: add type hints to NeuronModelForCausalLM
04932d5
@dacorvo
refactor: add check_compiler_compatibility helper
c2f5c78
@dacorvo
review: add compiler type and version to config and check it
3f88fa6
@dacorvo dacorvo force-pushed the neuron_model_for_causal_lm branch from 19243a6 to 3f88fa6 Compare July 6, 2023 12:30
michaelbenayoun
michaelbenayoun approved these changes Jul 6, 2023
View reviewed changes
Copy link
Member

@michaelbenayoun michaelbenayoun left a 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.

JingyaHuang
JingyaHuang approved these changes Jul 6, 2023
View reviewed changes
Copy link
Collaborator

@JingyaHuang JingyaHuang left a 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 🚀 !

optimum/neuron/modeling.py Show resolved Hide resolved
optimum/neuron/utils/version_utils.py Show resolved Hide resolved
@JingyaHuang
Copy link
Collaborator

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)

@dacorvo
Copy link
Collaborator Author

dacorvo commented Jul 6, 2023

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 🤭.

@dacorvo dacorvo merged commit eeea740 into main Jul 6, 2023
7 of 10 checks passed
@dacorvo dacorvo deleted the neuron_model_for_causal_lm branch July 6, 2023 14:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@JingyaHuang JingyaHuang JingyaHuang approved these changes

@michaelbenayoun michaelbenayoun michaelbenayoun approved these changes

@mfuntowicz mfuntowicz Awaiting requested review from mfuntowicz

@philschmid philschmid Awaiting requested review from philschmid

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

None yet
4 participants
@dacorvo @HuggingFaceDocBuilderDev @JingyaHuang @michaelbenayoun