Coverage for transformer_lens/model_bridge/supported_architectures/bitnet.py: 97%

1"""BitNet b1.58 (``BitNetForCausalLM``) adapter: llama layout plus attn/ffn 

2sub-layer RMSNorms (attn_sub_norm reapplied by an adapter-local attention bridge).""" 

3 

4from typing import Any 

5 

6import torch 

7 

8from transformer_lens.model_bridge.generalized_components import ( 

9 PositionEmbeddingsAttentionBridge, 

10) 

11from transformer_lens.model_bridge.supported_architectures.llama import ( 

12 LlamaArchitectureAdapter, 

13) 

14from transformer_lens.utilities.quantization import unreadable_weight_reason 

15 

16 

17class _BitNetAttentionBridge(PositionEmbeddingsAttentionBridge): 

18 """Applies BitNet's attn_sub_norm before the output projection. 

19 

20 The generic reconstruction goes straight from attention output to o_proj; 

21 BitNet inserts an RMSNorm in between. 

22 """ 

23 

24 def _pre_output_projection(self, attn_output: torch.Tensor) -> torch.Tensor: 

25 oc = self.original_component 

26 sub_norm = getattr(oc, "attn_sub_norm", None) if oc is not None else None 

27 if isinstance(sub_norm, torch.nn.Module): 27 ↛ 29line 27 didn't jump to line 29 because the condition on line 27 was always true

28 attn_output = sub_norm(attn_output) 

29 return attn_output 

30 

31 

32class BitNetArchitectureAdapter(LlamaArchitectureAdapter): 

33 """Architecture adapter for BitNetForCausalLM models.""" 

34 

35 _attention_bridge_cls = _BitNetAttentionBridge 

36 _testing_eager = "config" 

37 

38 # Sub-layer norms are incompatible with HT-style processed-weight 

39 # attention, so compatibility-mode equivalence (Phase 3) is out of scope. 

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

41 

42 def __init__(self, cfg: Any) -> None: 

43 """Initialize the BitNet architecture adapter.""" 

44 super().__init__(cfg) 

45 

46 # Sub-layer norms sit between activations and output projections; 

47 # standard LN folding and W_O centering do not model them. 

48 self.supports_fold_ln = False 

49 self.supports_center_writing_weights = False 

50 

51 def prepare_model(self, hf_model: Any) -> None: 

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

53 

54 The flagship microsoft/bitnet-b1.58-2B-4T stores `weight` as packed 

55 uint8 with a collapsed first dim (out_features // 4) plus a separate 

56 weight_scale, so every weight-space read reshapes it into a 

57 wrong-but-plausible matrix rather than failing. The registry records 

58 this checkpoint at 0% on the forward-pass phase for exactly that reason. 

59 """ 

60 super().prepare_model(hf_model) 

61 # Every weight-bearing module, not just the first: BitNet leaves the 

62 # embedding unquantized, and it sorts first in named_modules(), so 

63 # sampling one module inspected the one weight that is never packed. 

64 for name, module in hf_model.named_modules(): 

65 weight = getattr(module, "weight", None) 

66 if weight is None: 

67 continue 

68 if unreadable_weight_reason(weight) is not None: 

69 raise NotImplementedError( 

70 f"BitNet checkpoint stores packed weights ({name}); " 

71 "TransformerLens needs the dequantized sibling — use " 

72 "microsoft/bitnet-b1.58-2B-4T-bf16." 

73 )