[flyingdoubleG]

Description

BitNet applies an RMSNorm right before each output projection:

# BitNetMLP.forward
down_proj(ffn_sub_norm(act_fn(gate_proj(x)) * up_proj(x)))

_BitNetAttentionBridge already re-applies attn_sub_norm through _pre_output_projection on the processed-weights path. The MLP has no such step: GatedMLPBridge's processed-weights forward goes straight from act(gate) * up to down_proj. So after enable_compatibility_mode(), every BitNet MLP skips ffn_sub_norm, and nothing reports it:

tiny random BitNet (2 layers, sub-norm weights perturbed)	dev	this PR
compatibility-mode log-probs vs HF, max abs diff	0.58	1e-6
blocks.0.mlp.hook_post vs the unprocessed bridge (rel.)	0.99	0
blocks.1.hook_resid_post vs the unprocessed bridge (rel.)	0.84	0

Fix. GatedMLPBridge gets a _pre_output_projection(hidden) step between the gating product and out on the processed-weights path. It is the identity by default, like the attention bridge's. An adapter-local _BitNetMLPBridge overrides it to apply the HF MLP's ffn_sub_norm. hook_post (out.hook_in) sees the normalized activation, as on the unprocessed path, where it is the down_proj input. Other gated MLPs are unchanged.

I left applicable_phases = [1, 2, 4] as it is: I did not run the Phase 3 benchmark on a real checkpoint.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas

• I have made corresponding changes to the documentation

• My changes generate no new warnings

• I have added tests that prove my fix is effective or that my feature works

• New and existing unit tests pass locally with my changes

• I have not rewritten tests relating to key interfaces which would affect backward compatibility

• New tests/integration/model_bridge/test_bitnet_compatibility_mode.py. It builds a tiny random BitNetForCausalLM without hub access and checks that:

  • compatibility-mode log-probs match HF;
  • mlp.hook_post, hook_mlp_out and resid_post match the unprocessed bridge;
  • get_full_resid_decomposition(expand_neurons=True) sums to the residual stream.

  The first four fail on dev.

• test_bitnet_adapter.py: the MLP is a _BitNetMLPBridge with the usual gate_proj / up_proj / down_proj mapping, its _pre_output_projection applies ffn_sub_norm, and the generic one is the identity.

Local runs (Windows, CPU, torch 2.9.1, transformers 5.20.2, --jaxtyping-packages=transformer_lens,beartype.beartype):

• test_bitnet_adapter.py, test_bitnet_compatibility_mode.py and tests/unit/model_bridge/generalized_components: 636 passed.
• tests/unit -m "not slow": 8175 passed, 14 failed. The same 14 fail on dev in my environment (ten Gemma 3n adapter tests need the optional timm, three test_sparse_probing.py tests check optimizer stop reasons, and test_faithfulness_recovery_is_graded_not_binary is [Bug Report] test_faithfulness_recovery_is_graded_not_binary passes or fails on the sign of float32 rounding #1993).

black (the locked 23.12.1), isort --check-only and mypy pass on the changed files. In my environment mypy reports three errors in other, unchanged files (_model_state.py, moe.py, olmo_hybrid.py); they also appear on dev.

This was prepared with AI assistance (Claude); I reviewed and tested the change.

[jlarson4]

Thanks for putting this together @flyingdoubleG, just a couple comments before we can approve and merge this

[jlarson4]

The attention path changes its class using _attention_bridge_cls. Here, the MLP is built as a plain GatedMLPBridge and then replaced, so a reader of _gated_mlp cannot see that BitNet overrides it and the next gated MLP will copy the swap. Can we create a mirror path for the MLP class, either by overriding _gated_mlp or by giving ArchitectureAdapter a _gated_mlp_bridge_cls to mirror the attention one?

[jlarson4]

I tested parity with microsoft/bitnet-b1.58-2B-4T-bf16. For this model, HF ternarizes the weight and int8-quantizes the activation on each call. The processed-weights path this bridge inherits does F.linear on the fp32 master weights, meaning after this fix compatibility mode on that checkpoint disagrees with HF by about 24 in log-prob, starting at mlp.hook_pre.

The attention bridge goes through its LinearBridges rather than F.linear, could the MLP do the same through gate, in and out with no loss? Nothing is folded for BitNet, so this should be safe. Take a look and let me know what you think

[jlarson4]

This comment is now out of date, can we update it to reflect the updated reasoning for Phase 3 being out of scope?