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

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) 

14 

15 

16class _BitNetAttentionBridge(PositionEmbeddingsAttentionBridge): 

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

18 

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

20 BitNet inserts an RMSNorm in between. 

21 """ 

22 

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

24 oc = self.original_component 

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

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

27 attn_output = sub_norm(attn_output) 

28 return attn_output 

29 

30 

31class BitNetArchitectureAdapter(LlamaArchitectureAdapter): 

32 """Architecture adapter for BitNetForCausalLM models.""" 

33 

34 _attention_bridge_cls = _BitNetAttentionBridge 

35 _testing_eager = "config" 

36 

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

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

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

40 

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

42 """Initialize the BitNet architecture adapter.""" 

43 super().__init__(cfg) 

44 

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

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

47 self.supports_fold_ln = False 

48 self.supports_center_writing_weights = False