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 bool_4d_mask_is_keep = False 

26 

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

28 super().__init__(cfg) 

29 

30 self.cfg.normalization_type = "LN" 

31 self.cfg.positional_embedding_type = "alibi" 

32 self.cfg.final_rms = False 

33 self.cfg.gated_mlp = False 

34 self.cfg.attn_only = False 

35 self.cfg.default_prepend_bos = False 

36 

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

38 self.weight_processing_conversions = { 

39 **self._qkvo_weight_conversions(), 

40 } 

41 

42 self.component_mapping = { 

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

44 "blocks": BlockBridge( 

45 name="transformer.blocks", 

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

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

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

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

50 submodules={ 

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

52 "attn": MPTALiBiAttentionBridge( 

53 name="attn", 

54 config=self.cfg, 

55 split_qkv_matrix=self._split_mpt_qkv, 

56 submodules={ 

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

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

59 }, 

60 ), 

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

62 "mlp": MLPBridge( 

63 name="ffn", 

64 submodules={ 

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

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

67 }, 

68 ), 

69 }, 

70 ), 

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

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

73 } 

74 

75 def validate_output_logits_transform(self) -> None: 

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

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

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

79 raise ValueError( 

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

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

82 ) 

83 

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

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

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

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

88 d_model = self.cfg.d_model 

89 

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

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

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

93 return lin 

94 

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