[bitnet] Weight-quant-only BitLinear checkpoints load with sub-norms forced on

[peterlodri-sec]

Problem

BitNetMLP.forward always applies self.ffn_sub_norm(...) and BitNetAttention.forward always applies self.attn_sub_norm(...), and both modules are unconditionally constructed in __init__. This is only correct for full BitNet checkpoints.

BitNet checkpoints actually come in two forward variants:

1. Full BitNet (microsoft/bitnet-b1.58-2B-4T shape): per-projection RMSNorm (attn_sub_norm / ffn_sub_norm) applied before the projection bias and the residual.
2. Weight-quant-only: trained through a BitLinear that skips the per-projection RMSNorm and activation quant entirely; weights are already ternary {-1, 0, +1} and the forward is a plain nn.Linear. The sub-norm modules are absent — and they cannot be "neutralised" by loading weight=ones, because an RMSNorm initialised to ones still normalises.

Loading a weight-quant-only checkpoint (e.g. PeetPedro/quantal-ternary, a ternary export of a continued-trained Qwen2.5-0.5B) therefore runs the sub-norms, producing logits that differ from the trained/deployed forward by ~10 in max-abs and flip the argmax (verified against a golden reference the Rust runner reproduces to 1e-5).

Proposed fix

• PR #47955: adds use_sub_norms: bool = True to BitNetConfig; when False, the sub-norm modules are replaced by nn.Identity (built by a shared _make_sub_norm helper), so the normalisation is removed from the forward while the module graph and state dict stay uniform (Identity has no parameters). Default stays True, so existing model cards are unchanged.

Context

• Coordination model: https://huggingface.co/PeetPedro/quantal-ternary (the checkpoint this enables)
• PR: [bitnet] Support use_sub_norms=False (weight-quant-only BitLinear checkpoints) #47955