Coverage for transformer_lens/pretrained/weight_conversions/openai.py: 90%

1"""Weight conversion for OpenAI GPT-OSS models. 

2 

3GPT-OSS has a unique MoE architecture: 

4- GptOssExperts stores all expert weights in merged tensors (not individual modules) 

5- gate_up_proj: (num_experts, hidden_size, 2*expert_dim) with interleaved gate/up columns 

6- down_proj: (num_experts, expert_dim, hidden_size) 

7- Router (GptOssTopKRouter) uses weight + bias 

8""" 

9 

10import einops 

11import torch 

12 

13from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

14from transformer_lens.utilities.quantization import require_readable_weight 

15 

16# Phrased to hold for any quantization: the guard catches int8 and FP8 too, and 

17# cannot know which one it caught, so it must not assert this *is* MXFP4. 

18_GPT_OSS_REMEDY = ( 

19 "If this is a packed-MXFP4 checkpoint, load it dequantized so the converter " 

20 "sees plain tensors: pass hf_model=AutoModelForCausalLM.from_pretrained(name, " 

21 "quantization_config=Mxfp4Config(dequantize=True)), or load by model name and " 

22 "TransformerLens dequantizes automatically. Otherwise reload without a " 

23 "quantization_config. Quantized *forward* passes remain supported." 

24) 

25 

26 

27def convert_gpt_oss_weights(gpt_oss, cfg: HookedTransformerConfig): 

28 state_dict = {} 

29 

30 assert cfg.n_key_value_heads is not None 

31 assert cfg.d_mlp is not None 

32 assert cfg.num_experts is not None 

33 

34 state_dict["embed.W_E"] = gpt_oss.model.embed_tokens.weight 

35 

36 for l in range(cfg.n_layers): 

37 layer = gpt_oss.model.layers[l] 

38 

39 # LayerNorms 

40 state_dict[f"blocks.{l}.ln1.w"] = layer.input_layernorm.weight 

41 state_dict[f"blocks.{l}.ln2.w"] = layer.post_attention_layernorm.weight 

42 

43 # Attention 

44 W_Q = einops.rearrange(layer.self_attn.q_proj.weight, "(n h) m -> n m h", n=cfg.n_heads) 

45 W_K = einops.rearrange( 

46 layer.self_attn.k_proj.weight, "(n h) m -> n m h", n=cfg.n_key_value_heads 

47 ) 

48 W_V = einops.rearrange( 

49 layer.self_attn.v_proj.weight, "(n h) m -> n m h", n=cfg.n_key_value_heads 

50 ) 

51 state_dict[f"blocks.{l}.attn.W_Q"] = W_Q 

52 state_dict[f"blocks.{l}.attn._W_K"] = W_K 

53 state_dict[f"blocks.{l}.attn._W_V"] = W_V 

54 

55 if layer.self_attn.q_proj.bias is not None: 55 ↛ 66line 55 didn't jump to line 66 because the condition on line 55 was always true

56 state_dict[f"blocks.{l}.attn.b_Q"] = einops.rearrange( 

57 layer.self_attn.q_proj.bias, "(n h) -> n h", n=cfg.n_heads 

58 ) 

59 state_dict[f"blocks.{l}.attn._b_K"] = einops.rearrange( 

60 layer.self_attn.k_proj.bias, "(n h) -> n h", n=cfg.n_key_value_heads 

61 ) 

62 state_dict[f"blocks.{l}.attn._b_V"] = einops.rearrange( 

63 layer.self_attn.v_proj.bias, "(n h) -> n h", n=cfg.n_key_value_heads 

64 ) 

65 else: 

66 state_dict[f"blocks.{l}.attn.b_Q"] = torch.zeros( 

67 cfg.n_heads, cfg.d_head, dtype=cfg.dtype, device=cfg.device 

68 ) 

69 state_dict[f"blocks.{l}.attn._b_K"] = torch.zeros( 

70 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype, device=cfg.device 

71 ) 

72 state_dict[f"blocks.{l}.attn._b_V"] = torch.zeros( 

73 cfg.n_key_value_heads, cfg.d_head, dtype=cfg.dtype, device=cfg.device 

74 ) 

75 

76 state_dict[f"blocks.{l}.attn.sinks"] = layer.self_attn.sinks 

77 

78 W_O = einops.rearrange(layer.self_attn.o_proj.weight, "m (n h) -> n h m", n=cfg.n_heads) 

79 state_dict[f"blocks.{l}.attn.W_O"] = W_O 

80 

81 if hasattr(layer.self_attn.o_proj, "bias") and layer.self_attn.o_proj.bias is not None: 81 ↛ 84line 81 didn't jump to line 84 because the condition on line 81 was always true

82 state_dict[f"blocks.{l}.attn.b_O"] = layer.self_attn.o_proj.bias 

83 else: 

84 state_dict[f"blocks.{l}.attn.b_O"] = torch.zeros( 

85 cfg.d_model, dtype=cfg.dtype, device=cfg.device 

86 ) 

87 

88 # MoE - Router (GPT-OSS uses 'router' with bias) 

89 state_dict[f"blocks.{l}.mlp.W_gate.weight"] = require_readable_weight( 

90 layer.mlp.router.weight, 

91 operation="convert the gpt-oss router weight", 

92 owner=gpt_oss, 

93 remedy=_GPT_OSS_REMEDY, 

94 ) 

95 state_dict[f"blocks.{l}.mlp.W_gate.bias"] = layer.mlp.router.bias 

96 

97 # MoE - Experts 

98 # GPT-OSS stores all experts in merged tensors: 

99 # gate_up_proj: (num_experts, hidden_size, 2*expert_dim) - interleaved gate/up 

100 # down_proj: (num_experts, expert_dim, hidden_size) 

101 experts = layer.mlp.experts 

102 gate_up_proj = experts.gate_up_proj # (num_experts, hidden_size, 2*expert_dim) 

103 gate_up_bias = require_readable_weight( 

104 experts.gate_up_proj_bias, 

105 operation=f"convert gpt-oss expert biases (blocks.{l}.mlp.experts.gate_up_proj_bias)", 

106 owner=gpt_oss, 

107 remedy=_GPT_OSS_REMEDY, 

108 ) # (num_experts, 2*expert_dim) 

109 down_proj = experts.down_proj # (num_experts, expert_dim, hidden_size) 

110 down_bias = require_readable_weight( 

111 experts.down_proj_bias, 

112 operation=f"convert gpt-oss expert biases (blocks.{l}.mlp.experts.down_proj_bias)", 

113 owner=gpt_oss, 

114 remedy=_GPT_OSS_REMEDY, 

115 ) # (num_experts, hidden_size) 

116 

117 # Packed MXFP4 wraps these in a triton-kernels object (confusingly also 

118 # named "Tensor"), but int8 and FP8 gpt-oss finetunes slice without 

119 # complaint and would emit plausible garbage, so check the dtype too. 

120 gate_up_proj = require_readable_weight( 

121 gate_up_proj, 

122 operation=f"convert gpt-oss expert weights (blocks.{l}.mlp.experts.gate_up_proj)", 

123 owner=gpt_oss, 

124 remedy=_GPT_OSS_REMEDY, 

125 ) 

126 down_proj = require_readable_weight( 

127 down_proj, 

128 operation=f"convert gpt-oss expert weights (blocks.{l}.mlp.experts.down_proj)", 

129 owner=gpt_oss, 

130 remedy=_GPT_OSS_REMEDY, 

131 ) 

132 

133 for e in range(cfg.num_experts): 

134 # Split interleaved gate_up_proj into separate gate and up (in) projections 

135 # Even columns → gate path, Odd columns → up/in path 

136 state_dict[f"blocks.{l}.mlp.experts.{e}.W_gate.weight"] = gate_up_proj[ 

137 e, :, ::2 

138 ].T.contiguous() 

139 state_dict[f"blocks.{l}.mlp.experts.{e}.W_gate.bias"] = gate_up_bias[ 

140 e, ::2 

141 ].contiguous() 

142 

143 state_dict[f"blocks.{l}.mlp.experts.{e}.W_in.weight"] = gate_up_proj[ 

144 e, :, 1::2 

145 ].T.contiguous() 

146 state_dict[f"blocks.{l}.mlp.experts.{e}.W_in.bias"] = gate_up_bias[e, 1::2].contiguous() 

147 

148 state_dict[f"blocks.{l}.mlp.experts.{e}.W_out.weight"] = down_proj[e].T.contiguous() 

149 state_dict[f"blocks.{l}.mlp.experts.{e}.W_out.bias"] = down_bias[e].contiguous() 

150 

151 state_dict["ln_final.w"] = gpt_oss.model.norm.weight 

152 state_dict["unembed.W_U"] = gpt_oss.lm_head.weight.T 

153 state_dict["unembed.b_U"] = torch.zeros(cfg.d_vocab, dtype=cfg.dtype, device=cfg.device) 

154 

155 return state_dict