Coverage for transformer_lens/model_bridge/supported_architectures/openai_gpt.py: 96%

1"""OpenAI GPT (GPT-1) architecture adapter. 

2 

3The original GPT (``OpenAIGPTLMHeadModel``, openai-community/openai-gpt). 

4GPT-2's block internals under different top-level names — combined-QKV 

5``c_attn`` Conv1D, ``c_fc``/``c_proj`` MLP — but POST-norm (``ln_1(x + attn)``, 

6``ln_2(n + mlp)``), no final LayerNorm, and embeddings at 

7``transformer.tokens_embed`` / ``transformer.positions_embed``. 

8""" 

9 

10from typing import Any 

11 

12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

13from transformer_lens.model_bridge.generalized_components import ( 

14 BlockBridge, 

15 EmbeddingBridge, 

16 JointQKVAttentionBridge, 

17 LinearBridge, 

18 NormalizationBridge, 

19 PosEmbedBridge, 

20 UnembeddingBridge, 

21) 

22 

23 

24class _OpenAIGPTJointQKVAttentionBridge(JointQKVAttentionBridge): 

25 """GPT-1's Block concatenates ``[h] + attn_outputs[1:]`` — it requires the 

26 attention module to return a list, not the tuple modern blocks accept.""" 

27 

28 def _process_output(self, output: Any) -> Any: 

29 processed = super()._process_output(output) 

30 if isinstance(processed, tuple): 30 ↛ 32line 30 didn't jump to line 32 because the condition on line 30 was always true

31 return list(processed) 

32 return [processed] 

33 

34 

35class OpenAIGPTArchitectureAdapter(ArchitectureAdapter): 

36 """Architecture adapter for OpenAIGPTLMHeadModel (GPT-1) models.""" 

37 

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

39 """Initialize the OpenAI GPT architecture adapter.""" 

40 super().__init__(cfg) 

41 

42 self.cfg.normalization_type = "LN" 

43 self.cfg.positional_embedding_type = "standard" 

44 self.cfg.final_rms = False 

45 self.cfg.gated_mlp = False 

46 self.cfg.attn_only = False 

47 

48 # GPT-1 is post-LN; fold-LN assumes pre-LN and would fold norms into 

49 # the wrong sublayers. 

50 self.supports_fold_ln = False 

51 self.supports_center_writing_weights = False 

52 self.weight_processing_conversions = {} 

53 

54 self.component_mapping = { 

55 "embed": EmbeddingBridge(name="transformer.tokens_embed"), 

56 "pos_embed": PosEmbedBridge(name="transformer.positions_embed"), 

57 "blocks": BlockBridge( 

58 name="transformer.h", 

59 config=self.cfg, 

60 submodules={ 

61 "attn": _OpenAIGPTJointQKVAttentionBridge( 

62 name="attn", 

63 config=self.cfg, 

64 submodules={ 

65 "qkv": LinearBridge(name="c_attn"), 

66 "o": LinearBridge(name="c_proj"), 

67 }, 

68 ), 

69 # Post-norm: ln_1 follows attention, ln_2 follows the MLP. 

70 "ln1": NormalizationBridge(name="ln_1", config=self.cfg), 

71 "mlp": self._ungated_mlp(up="c_fc", down="c_proj"), 

72 "ln2": NormalizationBridge(name="ln_2", config=self.cfg), 

73 }, 

74 # Post-norm: ln2.hook_in is the post-MLP sum n+m; the true 

75 # attn->MLP mid-point n = ln_1(x+a) is mlp.hook_in. 

76 hook_alias_overrides={ 

77 "hook_resid_mid": "mlp.hook_in", 

78 }, 

79 ), 

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

81 }