-
Notifications
You must be signed in to change notification settings - Fork 248
Add support for separate bias tensors #1250
Description
🚀 The feature, motivation and pitch
In the transformers implementation of llama, there are optional bias tensors for the LlamaMLP and LlamaAttention modules. Several additional models (specifically Granite Code 3B and 8B) use the llama architecture and have these separate bias tensors.
The proposal here is to add the ability to indicate the presence of bias tensors in TransformerArgs and then support loading them in Attention and FeedForward
Alternatives
If this project is designed to be limited to official Llama models, these bias tensors are not needed.
Additional context
This issue is a piece of the puzzle for adding support for Granite Code 3b/8b which use the llama architecture in transormers, but take advantage several pieces of the architecture that are not currently supported by torchchat. The work-in-progress for Granite Code can be found on my fork: https://github.com/gabe-l-hart/torchchat/tree/GraniteCodeSupport
RFC (Optional)
I have a working implementation to support these optional bias tensors that I plan to submit as a PR. The changes are along the following lines:
- Add new parameters to
TransformerArgsfor attention and ffn bias - Set the
biasvalue based on these parameters in both theAttentionandFeedForwardmodules - Support mapping
.biastensor names inconvert_hf_checkpoint - Support permuting
.biastensors inconvert_hf_checkpoint - Support loading permuted
.biastensors inmodel.py