Coverage for transformer_lens/model_bridge/supported_architectures/gpt_oss.py: 97%
22 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""GPT-OSS architecture adapter."""
3from typing import Any
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 MoERouterBridge,
12 PositionEmbeddingsAttentionBridge,
13 RMSNormalizationBridge,
14 RotaryEmbeddingBridge,
15 UnembeddingBridge,
16)
19class GPTOSSArchitectureAdapter(ArchitectureAdapter):
20 """Architecture adapter for GPT-OSS model."""
22 def __init__(self, cfg: Any) -> None:
23 """Initialize the GPT-OSS architecture adapter."""
24 super().__init__(cfg)
26 self.cfg.gated_mlp = True
28 self.cfg.normalization_type = "RMS"
29 self.cfg.uses_rms_norm = True
30 # GPT-OSS uses rotary position embeddings, not learned embeddings
31 self.cfg.positional_embedding_type = "rotary"
32 # GPT-OSS attention returns (output, attn_weights), not a 3-tuple
33 # Note: attention_output_format is not a standard config attribute, handled in architecture code
35 # Conversion rules for weight processing/folding
36 # GPT-OSS uses MoE with batched experts, so we need special handling
37 self.weight_processing_conversions = {
38 **self._qkvo_weight_conversions(),
39 }
41 self.component_mapping = {
42 "embed": EmbeddingBridge(name="model.embed_tokens"),
43 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg),
44 "blocks": BlockBridge(
45 name="model.layers",
46 submodules={
47 "ln1": RMSNormalizationBridge(
48 name="input_layernorm",
49 config=self.cfg,
50 use_native_layernorm_autograd=True, # Use HF's RMSNorm for correct dtype handling
51 ),
52 "attn": PositionEmbeddingsAttentionBridge(
53 name="self_attn",
54 config=self.cfg,
55 requires_position_embeddings=True, # GPT-OSS requires position_embeddings (rotary)
56 requires_attention_mask=True,
57 submodules={
58 "q": LinearBridge(name="q_proj"),
59 "k": LinearBridge(name="k_proj"),
60 "v": LinearBridge(name="v_proj"),
61 "o": LinearBridge(name="o_proj"),
62 },
63 ),
64 "ln2": RMSNormalizationBridge(
65 name="post_attention_layernorm",
66 config=self.cfg,
67 use_native_layernorm_autograd=True, # Use HF's RMSNorm for correct dtype handling
68 ),
69 # GPT-OSS uses batched MoE experts with router scores
70 # MoEBridge handles the (hidden_states, router_scores) tuple returns
71 "mlp": MoEBridge(
72 name="mlp",
73 config=self.cfg,
74 submodules={"router": MoERouterBridge(name="router")},
75 sparse_required=("router",),
76 ),
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 }
87 def setup_hook_compatibility(self, bridge_model: Any) -> None:
88 """Setup hook compatibility transformations for GPT-OSS models.
90 This configures rotary embedding references for attention layers, which is
91 needed for models using RoPE (Rotary Position Embeddings).
93 This is called during Bridge.__init__ and should always be run.
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
102 # Get the actual HF rotary_emb from the bridge's rotary_emb component
103 rotary_emb = bridge_model.rotary_emb.original_component
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)
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)