transformer_lens.model_bridge.supported_architectures.bitnet module

BitNet b1.58 (BitNetForCausalLM) adapter: llama layout plus attn/ffn sub-layer RMSNorms (attn_sub_norm reapplied by an adapter-local attention bridge).

class transformer_lens.model_bridge.supported_architectures.bitnet.BitNetArchitectureAdapter(cfg: Any)

Bases: LlamaArchitectureAdapter

Architecture adapter for BitNetForCausalLM models.

__init__(cfg: Any) → None

Initialize the BitNet architecture adapter.

applicable_phases: list[int] = [1, 2, 4]

component_mapping: ComponentMapping | None

prepare_model(hf_model: Any) → None

Refuse packed 1.58-bit checkpoints, which need BitNet dequant kernels.

The flagship microsoft/bitnet-b1.58-2B-4T stores weight as packed uint8 with a collapsed first dim (out_features // 4) plus a separate weight_scale, so every weight-space read reshapes it into a wrong-but-plausible matrix rather than failing. The registry records this checkpoint at 0% on the forward-pass phase for exactly that reason.

uses_split_attention: bool

weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None