Granite modeling for training #830
Conversation
The old functions are deprecated and produce a warning.
This commit is just a copy of Transformers' 4.48.1 implementation of granite modeling. The code will be subsequentially modified to allow parallelism during training.
The model parallelization is verified with a simple inference.
Note that for now after each linear outputs are gathered, otherwise the result in bfloat16 diverges too much compared to the one obtained when using float32.
This allows for parallelization on the attention, using eager algorithm. This allows to perform the attention calculation in a sharded way, accelerating it.
Remove use_cache, cache_position and past_key_values arguments when possible, since they are not used on training models and it is not useful to keep them here.
After some tests, only the eager attention kernel has proven reliable, so only this option is left in the code of the modeling for this model.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
| from transformers.models.granite import GraniteConfig | ||
|
|
||
|
|
||
| class NeuronGraniteConfig(GraniteConfig): |
There was a problem hiding this comment.
I would rather create a completely different config instead of extending the GraniteConfig. This is what I do for LLM models.
| from transformers.models.granite import GraniteConfig | ||
|
|
||
|
|
||
| class NeuronGraniteConfig(GraniteConfig): |
There was a problem hiding this comment.
For weight fusing / QGAQKV etc I have this API: https://github.com/huggingface/optimum-neuron/pull/801/files#diff-7fda1f4deb9a79e4201376095bd5eabc0650da1fb05672493d7949af0820be18R101-R107
WDYT?
Basically the idea is to have a common API so that we can share the code for splitting the loaded weights in from_pretrained and consolidating the checkpoints afterwards.
There was a problem hiding this comment.
I agree it would be good to have a common API for these kind of things, and in general even for layers. I prefer having simpler functions though. I propose to isolate this, and then we can make the API uniform once your PR is ready.
There was a problem hiding this comment.
About the mp_config: I will separate the config as you had done, and we can merge the mp_config implementations once your PR is ready.
It is specific to inference.
The configuration is now separated in a specific config that is passed over between modules, so that the original config is not modified.
slice_tensor and fuse_weights can be reused in future modeling implementations.
There was a problem hiding this comment.
LGTM!
Let's address the remaining points once #801 is merged as agreed offline.
Thanks @tengomucho !!
It is not used for now and it should not passed over to parent class.
The bidirectional dict just requires a key, that can be the class name for simplicity.
What does this PR do?
This PR introduces the modeling code for Granite, based on NxD core modeling example for Llama.
So far it is possible to instantiate
ibm-granite/granite-3.2-2b-instructand perform an inference on the sharded model.The modeling code is based on transformers, and adapted to include sharding (only tested with TP parallelism).
Flash attention is not used so far because it requires not-padded prompts to generate correct outputs. It was not possible to use fused layers, because in Granite
self.num_heads != self.num_key_value_heads. The attention kernel implementation is then just the eager kernel.This is NOT a complete implementation for training: parallel loss is not implemented, as well as a complete training test and/or example. These will follow up soon, but I thought it was better to submit this to get feedback sooner.