Remove the pre-attention and pre-FFN RMSNorms from every decoder layer. Per the paper: "Remove RMSNorm before attention and SwiGLU because BitLinear has built-in RMSNorm." The norms are replaced with nn.Identity so that residual stream shapes remain unchanged and no index erro
(model: LlamaForCausalLM)
| 163 | |
| 164 | |
| 165 | def _remove_decoder_prenorms(model: LlamaForCausalLM) -> None: |
| 166 | """ |
| 167 | Remove the pre-attention and pre-FFN RMSNorms from every decoder layer. |
| 168 | |
| 169 | Per the paper: "Remove RMSNorm before attention and SwiGLU because |
| 170 | BitLinear has built-in RMSNorm." |
| 171 | |
| 172 | The norms are replaced with nn.Identity so that residual stream |
| 173 | shapes remain unchanged and no index errors occur in the forward pass. |
| 174 | """ |
| 175 | for layer in model.model.layers: |
| 176 | if isinstance(layer, LlamaDecoderLayer): |
| 177 | layer.input_layernorm = nn.Identity() |
| 178 | layer.post_attention_layernorm = nn.Identity() |
| 179 | |
| 180 | |
| 181 | def convert_to_bitnet(model: LlamaForCausalLM) -> LlamaForCausalLM: |