Coverage for transformer_lens/pretrained/weight_conversions/hubert.py: 78%

1import einops 

2 

3from transformer_lens.config.hooked_transformer_config import HookedTransformerConfig 

4 

5 

6def convert_hubert_weights(hf_model, cfg: HookedTransformerConfig): 

7 """Convert transformer encoder weights from a HuggingFace HuBERT model 

8 into the state_dict expected by Transformer-Lens' HookedEncoder. 

9 

10 Intentionally skips the convolutional frontend and feature_projection; 

11 those are used directly from the HF model. Use 

12 ``model.load_state_dict(state_dict, strict=False)`` to load these. 

13 """ 

14 state_dict = {} 

15 

16 # Try to find the encoder layer list (different HF variants use .layers or .layer) 

17 encoder = getattr(hf_model, "encoder", None) 

18 if encoder is None: 

19 raise ValueError("hf_model has no .encoder attribute") 

20 

21 encoder_layers = getattr(encoder, "layers", None) or getattr(encoder, "layer", None) 

22 if encoder_layers is None: 22 ↛ 24line 22 didn't jump to line 24 because the condition on line 22 was never true

23 # maybe hf_model itself is the encoder (unlikely), or a wrapped attribute 

24 raise ValueError("Couldn't find encoder.layers or encoder.layer on hf_model.encoder") 

25 

26 # Use cfg dims for reshaping 

27 d_model = cfg.d_model 

28 n_heads = cfg.n_heads 

29 

30 for l, layer in enumerate(encoder_layers): 

31 # --- Attention module --- 

32 # Some HF variants might call it `attention`, others `self_attn` etc. 

33 att = getattr(layer, "attention", None) or getattr(layer, "self_attn", None) 

34 if att is None: 34 ↛ 35line 34 didn't jump to line 35 because the condition on line 34 was never true

35 raise AttributeError(f"Encoder layer {l} has no 'attention' or 'self_attn' attribute") 

36 

37 # q/k/v/out proj names in HuBERT's HubertAttention: q_proj, k_proj, v_proj, out_proj 

38 # fall back to common alternatives if present 

39 q_w = getattr(att, "q_proj", None) 

40 k_w = getattr(att, "k_proj", None) 

41 v_w = getattr(att, "v_proj", None) 

42 o_w = getattr(att, "out_proj", None) or getattr(att, "proj", None) 

43 

44 if any(x is None for x in (q_w, k_w, v_w, o_w)): 44 ↛ 46line 44 didn't jump to line 46 because the condition on line 44 was never true

45 # Try alternate nested attributes like att.q, att.k, att.v, att.o 

46 q_w = q_w or getattr(att, "q", None) 

47 k_w = k_w or getattr(att, "k", None) 

48 v_w = v_w or getattr(att, "v", None) 

49 o_w = o_w or getattr(att, "o", None) 

50 

51 if any(x is None for x in (q_w, k_w, v_w, o_w)): 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true

52 raise AttributeError(f"Could not find q/k/v/out projections in layer {l}. Found: {att}") 

53 

54 assert q_w is not None and k_w is not None and v_w is not None and o_w is not None 

55 

56 # weights are Linear modules: weight shape (out, in) => same convention as Bert conversion 

57 # reshape to Transformer-Lens expected shapes using einops 

58 state_dict[f"blocks.{l}.attn.W_Q"] = einops.rearrange( 

59 q_w.weight, "(i h) m -> i m h", i=n_heads 

60 ) 

61 if q_w.bias is not None: 

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

63 q_w.bias, "(i h) -> i h", i=n_heads 

64 ) 

65 

66 state_dict[f"blocks.{l}.attn.W_K"] = einops.rearrange( 

67 k_w.weight, "(i h) m -> i m h", i=n_heads 

68 ) 

69 if k_w.bias is not None: 

70 state_dict[f"blocks.{l}.attn.b_K"] = einops.rearrange( 

71 k_w.bias, "(i h) -> i h", i=n_heads 

72 ) 

73 

74 state_dict[f"blocks.{l}.attn.W_V"] = einops.rearrange( 

75 v_w.weight, "(i h) m -> i m h", i=n_heads 

76 ) 

77 if v_w.bias is not None: 

78 state_dict[f"blocks.{l}.attn.b_V"] = einops.rearrange( 

79 v_w.bias, "(i h) -> i h", i=n_heads 

80 ) 

81 

82 state_dict[f"blocks.{l}.attn.W_O"] = einops.rearrange( 

83 o_w.weight, "m (i h) -> i h m", i=n_heads 

84 ) 

85 if o_w.bias is not None: 

86 state_dict[f"blocks.{l}.attn.b_O"] = o_w.bias 

87 

88 # --- Layer norms inside the layer --- 

89 # HuBERT layer has `layer.layer_norm` and `layer.final_layer_norm` 

90 ln1 = getattr(layer, "layer_norm", None) 

91 ln2 = getattr(layer, "final_layer_norm", None) 

92 if ln1 is None or ln2 is None: 92 ↛ 94line 92 didn't jump to line 94 because the condition on line 92 was never true

93 # try alternative names 

94 ln1 = ln1 or getattr(layer, "attention_norm", None) 

95 ln2 = ln2 or getattr(layer, "output_layer_norm", None) 

96 

97 if ln1 is not None: 97 ↛ 100line 97 didn't jump to line 100 because the condition on line 97 was always true

98 state_dict[f"blocks.{l}.ln1.w"] = ln1.weight 

99 state_dict[f"blocks.{l}.ln1.b"] = ln1.bias 

100 if ln2 is not None: 100 ↛ 106line 100 didn't jump to line 106 because the condition on line 100 was always true

101 state_dict[f"blocks.{l}.ln2.w"] = ln2.weight 

102 state_dict[f"blocks.{l}.ln2.b"] = ln2.bias 

103 

104 # --- Feed-forward / MLP --- 

105 # HuBERT uses `feed_forward` which contains intermediate_dense and output_dense 

106 ff = ( 

107 getattr(layer, "feed_forward", None) 

108 or getattr(layer, "feedforward", None) 

109 or getattr(layer, "ff", None) 

110 ) 

111 if ff is None: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true

112 raise AttributeError(f"Layer {l} has no feed_forward/ff attribute") 

113 

114 # Many implementations name them intermediate_dense and output_dense 

115 fc1 = ( 

116 getattr(ff, "intermediate_dense", None) 

117 or getattr(ff, "fc1", None) 

118 or getattr(ff, "linear1", None) 

119 ) 

120 fc2 = ( 

121 getattr(ff, "output_dense", None) 

122 or getattr(ff, "fc2", None) 

123 or getattr(ff, "linear2", None) 

124 ) 

125 

126 if fc1 is None or fc2 is None: 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true

127 raise AttributeError(f"Could not find FFN dense layers in layer {l}: {ff}") 

128 

129 # fc1.weight shape: (d_mlp, d_model) -> Transformer-Lens expects (d_model, d_mlp) 

130 state_dict[f"blocks.{l}.mlp.W_in"] = einops.rearrange(fc1.weight, "mlp model -> model mlp") 

131 if fc1.bias is not None: 131 ↛ 135line 131 didn't jump to line 135 because the condition on line 131 was always true

132 state_dict[f"blocks.{l}.mlp.b_in"] = fc1.bias 

133 

134 # fc2.weight shape: (d_model, d_mlp) -> Transformer-Lens expects (d_mlp, d_model) 

135 state_dict[f"blocks.{l}.mlp.W_out"] = einops.rearrange(fc2.weight, "model mlp -> mlp model") 

136 if fc2.bias is not None: 136 ↛ 30line 136 didn't jump to line 30 because the condition on line 136 was always true

137 state_dict[f"blocks.{l}.mlp.b_out"] = fc2.bias 

138 

139 # --- Optional: encoder-level layer_norm (HubertModel.encoder.layer_norm) --- 

140 if hasattr(hf_model.encoder, "layer_norm"): 140 ↛ 145line 140 didn't jump to line 145 because the condition on line 140 was always true

141 ln_final = hf_model.encoder.layer_norm 

142 state_dict["ln_final.w"] = ln_final.weight 

143 state_dict["ln_final.b"] = ln_final.bias 

144 

145 return state_dict