Add optional RMSNorm support to BitNet quantization (config + layers)

[Codys12]

What does this PR do?

Adds optional RMSNorm support to BitNet-style quantisation.

• Introduces use_rms_norm (bool, default False) and rms_norm_eps (float, default 1e-6) to BitNetQuantConfig so the flag is serialisable via save_pretrained / from_pretrained.
• Updates BitLinear and AutoBitLinear to accept use_rms_norm and apply the reference BitNetRMSNorm to activations before quantisation.

Before submitting

• I read the contributor guideline, PR section.
• I’ve added the new config fields to to_dict, docstrings, and the model card.
• New unit tests
• Ran make style && make quality && make test locally.
• Documentation build passes (make docs) – pushed logs to CI.

Motivation and context

RMSNorm stabilises the activations of low-bit networks; the BitNet paper shows a consistent perplexity drop when normalising pre-quant activations. This PR brings parity with the reference implementation while keeping the previous behaviour as default.

No new external dependencies.

Who can review?

Quantization / Accelerate folks for the code:
@SunMarc @MekkCyber

Docstrings & config: @stevhliu

Feel free to jump in with any feedback!

[github-actions]

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.

[SunMarc]

Can you share a bit about motivation for such feature ?

[Codys12]

@SunMarc
Sure! I recently discovered that initializing an extra RMS norm in each Linear lets you finetune existing LLMs to BitNet format and tested this for https://huggingface.co/codys12/bitnet-r1-32b and https://huggingface.co/codys12/bitnet-r1-8b. This process should work for any model, and this is a minimal implementation for easy conversion.

[MekkCyber]

Hi @Codys12, thanks for the pr 🤗 ! I'm not sure I understand the idea behind adding an extra rmsnorm, in the current bitnet implementation, the authors already added that in the modelling, for example MLP looks like :

class BitNetMLP(GemmaMLP):
    def __init__(self, config: BitNetConfig):
        super().__init__(config)
        self.ffn_sub_norm = BitNetRMSNorm(config.intermediate_size, eps=config.rms_norm_eps)

    def forward(self, x):
        down_proj = self.down_proj(self.ffn_sub_norm(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
        return down_proj

[Codys12]

Hi @Codys12, thanks for the pr 🤗 ! I'm not sure I understand the idea behind adding an extra rmsnorm, in the current bitnet implementation, the authors already added that in the modelling, for example MLP looks like :

class BitNetMLP(GemmaMLP):
    def __init__(self, config: BitNetConfig):
        super().__init__(config)
        self.ffn_sub_norm = BitNetRMSNorm(config.intermediate_size, eps=config.rms_norm_eps)

    def forward(self, x):
        down_proj = self.down_proj(self.ffn_sub_norm(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
        return down_proj

@MekkCyber @SunMarc Good point!

In section 2/2.1 of the origininal BitNet paper (https://arxiv.org/pdf/2310.11453), the authors describe the reason for including the RMS: it improves performance of the models at negligable compute/parameter costs. They include it in their modeling file, but others (see here) have tested this with alternative architectures and observed an improvement (Llama, Mistral, DeepSeek V3, etc).

This change in the Quantization Config is a model-agnostic approach for introducing this parameter so that a new modeling_*.py file is not required for every model you want to test this way.

Additionally, the inclusion of this norm allows you to finetune existing models to this quantization (see here) as demonstrated by https://huggingface.co/codys12/bitnet-r1-32b and https://huggingface.co/codys12/bitnet-r1-8b.

Let me know if there is anything that is still unclear!

[MekkCyber]

Thanks for the explanation @Codys12, I see the idea behind this !

[SunMarc]

Thanks ! Just a couple of nits

[Codys12]

Thanks ! Just a couple of nits

@SunMarc Just made these changes, let me know if there is anything else I can do before merge!

[MekkCyber]

Thanks @Codys12, can you please run make fix-copies and make style to fix CI 🤗

[Codys12]

@MekkCyber
make fix-copies is passing, but the CI tests are still failing... Should I add it to to OBJECTS_TO_IGNORE in utils/check_docstrings.py?

[Codys12]

@MekkCyber @SunMarc Any ideas on CI here? Looking to help this move forward today

[SunMarc]

Thanks !

[SunMarc]

try to run make fix-copies but remove the default values for the arguments you added in BitNetQuantConfig. I think the issue is there.

Traceback (most recent call last):
  File "/root/transformers/utils/check_docstrings.py", line 1467, in <module>
    check_docstrings(overwrite=args.fix_and_overwrite, check_all=args.check_all)
  File "/root/transformers/utils/check_docstrings.py", line 1456, in check_docstrings
    raise ValueError(error_message)
ValueError: There was at least one problem when checking docstrings of public objects.
The following objects docstrings do not match their signature. Run `make fix-copies` to fix this. In some cases, this error may be raised incorrectly by the docstring checker. If you think this is the case, you can manually check the docstrings and then add the object name to `OBJECTS_TO_IGNORE` in `utils/check_docstrings.py`.
- BitNetQuantConfig

[Codys12]

try to run make fix-copies but remove the default values for the arguments you added in BitNetQuantConfig. I think the issue is there.

Traceback (most recent call last):
  File "/root/transformers/utils/check_docstrings.py", line 1467, in <module>
    check_docstrings(overwrite=args.fix_and_overwrite, check_all=args.check_all)
  File "/root/transformers/utils/check_docstrings.py", line 1456, in check_docstrings
    raise ValueError(error_message)
ValueError: There was at least one problem when checking docstrings of public objects.
The following objects docstrings do not match their signature. Run `make fix-copies` to fix this. In some cases, this error may be raised incorrectly by the docstring checker. If you think this is the case, you can manually check the docstrings and then add the object name to `OBJECTS_TO_IGNORE` in `utils/check_docstrings.py`.
- BitNetQuantConfig

@SunMarc Hmm, I changed it to optional but running make fix-copies is not doing anything. It might have to do with running the CI loop on another device that I am pulling my fork into... Do you have any experience with getting the formatting fixed?

[Codys12]

Wait, all tests are passing sick

[HuggingFaceDocBuilderDev]

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.

[MekkCyber]

Thanks for the pr 🤗 !