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

1"""MPT (MPTForCausalLM) adapter — ALiBi, fused Wqkv, weight-only LayerNorm, no biases.""" 

2 

3from typing import Any 

4 

5import torch 

6import torch.nn as nn 

7 

8from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

9from transformer_lens.model_bridge.generalized_components import ( 

10 BlockBridge, 

11 EmbeddingBridge, 

12 LinearBridge, 

13 MLPBridge, 

14 NormalizationBridge, 

15 UnembeddingBridge, 

16) 

17from transformer_lens.model_bridge.generalized_components.mpt_alibi_attention import ( 

18 MPTALiBiAttentionBridge, 

19) 

20 

21 

22class MPTArchitectureAdapter(ArchitectureAdapter): 

23 """MPT adapter: ALiBi bias; all layers bias-free (no b_Q/b_K/b_V/b_O/b_in/b_out/ln bias).""" 

24 

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

26 super().__init__(cfg) 

27 

28 self.cfg.normalization_type = "LN" 

29 self.cfg.positional_embedding_type = "alibi" 

30 self.cfg.final_rms = False 

31 self.cfg.gated_mlp = False 

32 self.cfg.attn_only = False 

33 self.cfg.default_prepend_bos = False 

34 

35 # Pure MHA: split_qkv yields [d_model, d_model] per head; standard rearrangements apply. 

36 self.weight_processing_conversions = { 

37 **self._qkvo_weight_conversions(), 

38 } 

39 

40 self.component_mapping = { 

41 "embed": EmbeddingBridge(name="transformer.wte"), 

42 "blocks": BlockBridge( 

43 name="transformer.blocks", 

44 # MptMLP adds the residual inside its forward, so mlp.hook_out is 

45 # resid_post, not the additive contribution; alias to the down_proj 

46 # output instead (dropout between it and the add is an eval no-op). 

47 hook_alias_overrides={"hook_mlp_out": "mlp.out.hook_out"}, 

48 submodules={ 

49 "ln1": NormalizationBridge(name="norm_1", config=self.cfg), 

50 "attn": MPTALiBiAttentionBridge( 

51 name="attn", 

52 config=self.cfg, 

53 split_qkv_matrix=self._split_mpt_qkv, 

54 submodules={ 

55 "qkv": LinearBridge(name="Wqkv"), 

56 "o": LinearBridge(name="out_proj"), 

57 }, 

58 ), 

59 "ln2": NormalizationBridge(name="norm_2", config=self.cfg), 

60 "mlp": MLPBridge( 

61 name="ffn", 

62 submodules={ 

63 "in": LinearBridge(name="up_proj"), 

64 "out": LinearBridge(name="down_proj"), 

65 }, 

66 ), 

67 }, 

68 ), 

69 "ln_final": NormalizationBridge(name="transformer.norm_f", config=self.cfg), 

70 "unembed": UnembeddingBridge(name="lm_head"), 

71 } 

72 

73 def validate_output_logits_transform(self) -> None: 

74 """Reject ambiguous remote-code logit scaling not used by integrated HF MPT.""" 

75 logit_scale = getattr(self.cfg, "logit_scale", None) 

76 if logit_scale not in (None, 1, 1.0): 76 ↛ exitline 76 didn't return from function 'validate_output_logits_transform' because the condition on line 76 was always true

77 raise ValueError( 

78 "JacobianLens cannot infer MPT remote-code logit_scale semantics; " 

79 f"got logit_scale={logit_scale!r}." 

80 ) 

81 

82 def _split_mpt_qkv(self, attn_component: Any) -> tuple[nn.Linear, nn.Linear, nn.Linear]: 

83 """Split fused Wqkv into Q, K, V — row-wise chunk (NOT interleaved like BLOOM).""" 

84 w = attn_component.Wqkv.weight.detach().clone() 

85 w_q, w_k, w_v = torch.chunk(w, 3, dim=0) 

86 d_model = self.cfg.d_model 

87 

88 def make_linear(weight: torch.Tensor) -> nn.Linear: 

89 lin = nn.Linear(d_model, d_model, bias=False, device=weight.device, dtype=weight.dtype) 

90 lin.weight = nn.Parameter(weight.contiguous()) 

91 return lin 

92 

93 return make_linear(w_q), make_linear(w_k), make_linear(w_v)