Adding BitNet b1.58 Model
Author: shumingma
Pull request: yes

# What does this PR do?

This PR adds the support of codes for the coming BitNet b1.58 model from Microsoft. @ArthurZucker 

https://arxiv.org/abs/2402.17764

Comments:
Hi 👋, thank you for opening this pull request! The pull request is converted to draft by default. The CI will be paused while the PR is in draft mode. When it is ready for review, please click the `Ready for review` button (at the bottom of the PR page). This will assign reviewers and trigger CI.

Hi @shumingma, very cool project 🔥 ! it seems to me like a llama architecture with bitnet layers, did you try just using or adapting the bitnet quantization [here](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/bitnet.py) to load the model ?

> Hi @shumingma, very cool project 🔥 ! it seems to me like a llama architecture with bitnet layers, did you try just using or adapting the bitnet quantization [here](https://github.com/huggingface/transformers/blob/main/src/transformers/integrations/bitnet.py) to load the model ?

Hi @MekkCyber glad you like the project, and thanks for the suggestion regarding the transformers BitNet implementation!

We did look into the transformers BitNet implementation. However, our BitNet b1.58 architecture includes some modifications, notably using subln which introduces two additional LayerNorm layers. This structural difference prevents a direct drop-in replacement using the existing transformers quantization.

Another important factor for us was designing a solution that works seamlessly for both training (online quantization) and inference (offline quantization). To achieve that flexibility and accommodate the architectural tweaks, we ended up building our own layers.

Hi @MekkCyber Just wanted to let you know that the PR has passed all tests. Would you be willing to review the code when you have a moment? And let us know if you have any further comments or suggestions for the PR. We appreciate your input!

Hi @MekkCyber Thanks for your feedback! I've updated the code based on your suggestions:

1. Refactored the model to use the modular Transformers setup, and utilized the converter script with the modular file to generate the modeling script.
2. Integrated quantization support using quantization_config. I've added an AutoBitLinear layer to manage both offline and online quantization.
3. Updated the model documentation.

Could you please take another look when you have a chance?

> LGTM 🤗 ! Thanks a lot for iterating ! can you make sure to add the quantization config to the models on the hub, and test if the model is generating as expected ?

Thanks! I've updated the config on the Hub: [https://huggingface.co/microsoft/bitnet-b1.58-2B-4T/blob/main/config.json#L26-L29](https://www.google.com/url?sa=E&q=https%3A%2F%2Fhuggingface.co%2Fmicrosoft%2Fbitnet-b1.58-2B-4T%2Fblob%2Fmain%2Fconfig.json%23L26-L29) and confirmed the generation behavior is as expected.

Hi @MekkCyber @SunMarc 

Thank you for your feedback on the previous version. I've addressed the comments and suggestions you provided and have pushed the updated code.

Could you please review the changes again at your convenience?

Thanks again for your time and input!


Hi @Cyrilvallez 

Could you please take a look at this PR when you have some availability? It's ready for merging pending your review. 

Let me know if anything needs clarification. 

Thanks!

Hi @shumingma, the PR is nearly complete — it’d be great to get it merged by the end of the week!
Could you take a look at the remaining comments, they should not be complicated?
If you’re short on time, we’re happy to help out 🤗 !

> Hi @shumingma, the PR is nearly complete — it’d be great to get it merged by the end of the week! Could you take a look at the remaining comments, they should not be complicated? If you’re short on time, we’re happy to help out 🤗 !

Hi @MekkCyber

Thanks for checking in on the PR!

I'm quite busy at the moment, so I really appreciate the offer to help with the remaining comments. Yes, please – that would be great if you could take care of them.

Happy to review afterwards if needed, or let me know if you plan to merge directly. Aiming for EOW sounds good.

Thanks again for your help!
Shuming

Thanks for the quick response, @shumingma! I’ve opened a new PR with your commits: https://github.com/huggingface/transformers/pull/37742.
I believe it’s in good shape, but feel free to review it whenever you have a moment.
We’re aiming to merge it by tomorrow at the latest.

Closing now that https://github.com/huggingface/transformers/pull/37742 is merged