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

1"""GPT-OSS architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.generalized_components import ( 

7 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 MoEBridge, 

11 PositionEmbeddingsAttentionBridge, 

12 RMSNormalizationBridge, 

13 RotaryEmbeddingBridge, 

14 UnembeddingBridge, 

15) 

16 

17 

18class GPTOSSArchitectureAdapter(ArchitectureAdapter): 

19 """Architecture adapter for GPT-OSS model.""" 

20 

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

22 """Initialize the GPT-OSS architecture adapter.""" 

23 super().__init__(cfg) 

24 

25 self.cfg.gated_mlp = True 

26 

27 self.cfg.normalization_type = "RMS" 

28 self.cfg.uses_rms_norm = True 

29 # GPT-OSS uses rotary position embeddings, not learned embeddings 

30 self.cfg.positional_embedding_type = "rotary" 

31 # GPT-OSS attention returns (output, attn_weights), not a 3-tuple 

32 # Note: attention_output_format is not a standard config attribute, handled in architecture code 

33 

34 # Conversion rules for weight processing/folding 

35 # GPT-OSS uses MoE with batched experts, so we need special handling 

36 # GPT-OSS may use GQA: K/V heads can differ from Q heads 

37 n_kv_heads = ( 

38 self.cfg.n_key_value_heads 

39 if hasattr(self.cfg, "n_key_value_heads") and self.cfg.n_key_value_heads is not None 

40 else self.cfg.n_heads 

41 ) 

42 self.weight_processing_conversions = { 

43 **self._qkvo_weight_conversions(), 

44 } 

45 

46 self.component_mapping = { 

47 "embed": EmbeddingBridge(name="model.embed_tokens"), 

48 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

49 "blocks": BlockBridge( 

50 name="model.layers", 

51 submodules={ 

52 "ln1": RMSNormalizationBridge( 

53 name="input_layernorm", 

54 config=self.cfg, 

55 use_native_layernorm_autograd=True, # Use HF's RMSNorm for correct dtype handling 

56 ), 

57 "attn": PositionEmbeddingsAttentionBridge( 

58 name="self_attn", 

59 config=self.cfg, 

60 requires_position_embeddings=True, # GPT-OSS requires position_embeddings (rotary) 

61 requires_attention_mask=True, 

62 submodules={ 

63 "q": LinearBridge(name="q_proj"), 

64 "k": LinearBridge(name="k_proj"), 

65 "v": LinearBridge(name="v_proj"), 

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

67 }, 

68 ), 

69 "ln2": RMSNormalizationBridge( 

70 name="post_attention_layernorm", 

71 config=self.cfg, 

72 use_native_layernorm_autograd=True, # Use HF's RMSNorm for correct dtype handling 

73 ), 

74 # GPT-OSS uses batched MoE experts with router scores 

75 # MoEBridge handles the (hidden_states, router_scores) tuple returns 

76 "mlp": MoEBridge(name="mlp", config=self.cfg), 

77 }, 

78 ), 

79 "ln_final": RMSNormalizationBridge( 

80 name="model.norm", 

81 config=self.cfg, 

82 use_native_layernorm_autograd=True, # Use HF's RMSNorm for correct dtype handling 

83 ), 

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

85 } 

86 

87 def setup_hook_compatibility(self, bridge_model: Any) -> None: 

88 """Setup hook compatibility transformations for GPT-OSS models. 

89 

90 This configures rotary embedding references for attention layers, which is 

91 needed for models using RoPE (Rotary Position Embeddings). 

92 

93 This is called during Bridge.__init__ and should always be run. 

94 

95 Args: 

96 bridge_model: The TransformerBridge instance 

97 """ 

98 # Get the rotary_emb component from the actual bridge model 

99 if bridge_model is None or not hasattr(bridge_model, "rotary_emb"): 

100 return 

101 

102 # Get the actual HF rotary_emb from the bridge's rotary_emb component 

103 rotary_emb = bridge_model.rotary_emb.original_component 

104 

105 if hasattr(bridge_model, "blocks"): 105 ↛ exitline 105 didn't return from function 'setup_hook_compatibility' because the condition on line 105 was always true

106 for block in bridge_model.blocks: 

107 if hasattr(block, "attn"): 

108 block.attn.set_rotary_emb(rotary_emb) 

109 

110 def setup_no_processing_hooks(self, bridge_model: Any) -> None: 

111 """Backward compatibility alias for setup_hook_compatibility.""" 

112 self.setup_hook_compatibility(bridge_model)