Coverage for transformer_lens/model_bridge/supported_architectures/rwkv7.py: 94%

1"""RWKV-7 ("Goose") architecture adapter (RWKV7ForCausalLM). 

2 

3Family ``fla-hub/rwkv7-*``: attention-free recurrent LM from the 

4flash-linear-attention library, loaded via remote code. Blocks pair a 

5time-mixing (generalized delta rule) and a token-shifted squared-ReLU 

6channel-mixing sublayer under biased pre-LN; no positional embeddings. 

7 

8Adapter decisions: 

9- Full delegation: recurrence, token shift, LoRAs, and cross-block ``v_first`` 

10 threading all run inside the fla forward (v_first is managed by the HF 

11 model-level forward, so delegated blocks get it for free). 

12 ``weight_processing_conversions = {}``. 

13- ``OpaqueBlockBridge``: BlockBridge's hook aliases hardcode the standard 

14 pre-norm attention flow; only ``hook_in``/``hook_out`` are sound here. 

15- ``ffn_norm`` is a fused add-and-norm ``(normed, residual) = ffn_norm(x, res, 

16 True)`` under ``config.fuse_norm`` — wrapped as a delegating 

17 ``GeneralizedComponent`` because ``NormalizationBridge`` can't express it. 

18- ``head_dim`` is read but never assigned (aliases the read-only ``d_head``). 

19""" 

20 

21from typing import Any 

22 

23from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

24from transformer_lens.model_bridge.generalized_components import ( 

25 EmbeddingBridge, 

26 LinearBridge, 

27 NormalizationBridge, 

28 OpaqueBlockBridge, 

29 UnembeddingBridge, 

30) 

31from transformer_lens.model_bridge.generalized_components.base import ( 

32 GeneralizedComponent, 

33) 

34from transformer_lens.model_bridge.supported_architectures._remote_code_compat import ( 

35 force_import_remote_class, 

36 iter_remote_modeling_modules, 

37 patch_init_weights_skip_loaded, 

38 retie_weights_keys_v5, 

39) 

40 

41 

42class RWKV7ArchitectureAdapter(ArchitectureAdapter): 

43 """Architecture adapter for RWKV7ForCausalLM (RWKV-7 "Goose"). 

44 

45 Attention-free recurrent decoder: a flat stack of pre-norm blocks, each a 

46 generalized-delta-rule time-mixing sublayer plus a token-shifted squared-ReLU 

47 channel-mixing sublayer, wrapped by standard biased LayerNorm. The recurrence 

48 and the cross-block ``v_first`` threading live inside the ``fla`` remote-code 

49 forward, which the bridge delegates to; see the module docstring for the full 

50 set of adapter decisions. 

51 """ 

52 

53 # Attention-free and recurrent — off the transformer-shaped verify_models 

54 # path; correctness lives in the integration tests (bridge-vs-HF parity). 

55 applicable_phases: list[int] = [] 

56 # fla threads recurrence through its own Cache object, not past_key_values; 

57 # generation recomputes the full prefix per step (exact, O(n^2)). 

58 supports_kv_cache: bool = False 

59 # The fla forward ignores attention_mask, so left-padding would silently 

60 # poison the recurrent state instead of being masked out. 

61 supports_batched_generation: bool = False 

62 

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

64 """Initialize the RWKV-7 architecture adapter.""" 

65 super().__init__(cfg) 

66 

67 # Standard biased LayerNorm, no positional embeddings, ungated FFN. 

68 self.cfg.normalization_type = "LN" 

69 self.cfg.uses_rms_norm = False 

70 self.cfg.positional_embedding_type = "none" 

71 self.cfg.final_rms = False 

72 self.cfg.attn_only = False 

73 self.cfg.gated_mlp = False 

74 

75 # fla drives recurrent decode state via its own Cache, not the 

76 # cache_params convention, so the adapter does not advertise a cache. 

77 setattr(self.cfg, "is_stateful", False) 

78 

79 # Surface the RWKV-7 shape attributes on cfg so they are present on both 

80 # the HF-boot path (also via _HF_PASSTHROUGH_ATTRS) and the synthetic 

81 # config path used by the unit tests. getattr-with-default keeps a bare 

82 # TransformerBridgeConfig from raising. 

83 # head_dim is a read-only property on TransformerBridgeConfig (aliases 

84 # d_head), so it is read but never assigned. num_heads defaults to 

85 # d_model // head_dim. 

86 head_dim = getattr(cfg, "head_dim", 64) or 64 

87 num_heads = getattr(cfg, "num_heads", None) or max(1, self.cfg.d_model // head_dim) 

88 value_dim = getattr(cfg, "value_dim", None) or [self.cfg.d_model] * self.cfg.n_layers 

89 # setattr (not direct assignment) for the RWKV-7-specific names so mypy 

90 # does not flag them as undeclared on TransformerBridgeConfig. 

91 setattr(self.cfg, "num_heads", num_heads) 

92 setattr(self.cfg, "value_dim", value_dim) 

93 setattr(self.cfg, "decay_low_rank_dim", getattr(cfg, "decay_low_rank_dim", 64)) 

94 setattr(self.cfg, "gate_low_rank_dim", getattr(cfg, "gate_low_rank_dim", 128)) 

95 setattr(self.cfg, "a_low_rank_dim", getattr(cfg, "a_low_rank_dim", 64)) 

96 setattr(self.cfg, "v_low_rank_dim", getattr(cfg, "v_low_rank_dim", 16)) 

97 setattr(self.cfg, "norm_first", getattr(cfg, "norm_first", True)) 

98 setattr(self.cfg, "norm_bias", getattr(cfg, "norm_bias", True)) 

99 setattr(self.cfg, "fuse_norm", getattr(cfg, "fuse_norm", True)) 

100 setattr(self.cfg, "attn_mode", getattr(cfg, "attn_mode", "chunk")) 

101 setattr(self.cfg, "hidden_act", getattr(cfg, "hidden_act", "sqrelu")) 

102 # LayerNorm epsilon for the reimplementing NormalizationBridge path (eps 

103 # is a real config field, so direct assignment is fine). 

104 self.cfg.eps = getattr(cfg, "norm_eps", getattr(cfg, "eps", 1e-5)) 

105 

106 # Full delegation to the fla forward — no HT-format weight reshaping. 

107 self.weight_processing_conversions = {} 

108 

109 self.component_mapping = { 

110 "embed": EmbeddingBridge(name="model.embeddings"), 

111 "blocks": OpaqueBlockBridge( 

112 name="model.layers", 

113 submodules={ 

114 # Pre-norm before time-mixing (standard single-input LayerNorm). 

115 "attn_norm": NormalizationBridge( 

116 name="attn_norm", config=self.cfg, uses_rms_norm=False 

117 ), 

118 # Time-mixing: generalized delta rule. Delegated passthrough; 

119 # only the four projections are exposed (LoRAs / GroupNorm / 

120 # lerp params stay inside the fla forward). 

121 "attn": GeneralizedComponent( 

122 name="attn", 

123 config=self.cfg, 

124 submodules={ 

125 "r_proj": LinearBridge(name="r_proj"), 

126 "k_proj": LinearBridge(name="k_proj"), 

127 "v_proj": LinearBridge(name="v_proj"), 

128 "o_proj": LinearBridge(name="o_proj"), 

129 }, 

130 ), 

131 # Pre-norm before channel-mixing. Under config.fuse_norm the 

132 # block calls this as ffn_norm(x, residual, True) -> (normed, 

133 # residual), a fused signature the reimplementing 

134 # NormalizationBridge can't express, so delegate any-signature 

135 # to the live HF module with I/O hooks. 

136 "ffn_norm": GeneralizedComponent(name="ffn_norm", config=self.cfg), 

137 # Channel-mixing: token-shifted squared-ReLU MLP. Delegated 

138 # passthrough exposing the up ("key") and down ("value") 

139 # projections. HF confusingly names the output proj "value". 

140 # NOT MLPBridge (cf. rwkv4): its hook_in pre-fire would 

141 # suppress key.hook_in, the post-token-shift input — the one 

142 # interesting tensor here. Aliases give hook_pre/hook_post. 

143 "ffn": GeneralizedComponent( 

144 name="ffn", 

145 config=self.cfg, 

146 hook_alias_overrides={ 

147 "hook_pre": "key.hook_out", 

148 "hook_post": "value.hook_in", 

149 }, 

150 submodules={ 

151 "key": LinearBridge(name="key"), 

152 "value": LinearBridge(name="value"), 

153 }, 

154 ), 

155 }, 

156 ), 

157 "ln_final": NormalizationBridge( 

158 name="model.norm", config=self.cfg, uses_rms_norm=False 

159 ), 

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

161 } 

162 

163 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: 

164 """Patch fla's RWKV-7 remote code for transformers v5 compatibility. 

165 

166 Two defensive patches, mirroring raven: 

167 

168 1. Tied-weights format. ``RWKV7ForCausalLM._tied_weights_keys`` is a list 

169 (``["lm_head.weight"]``, the 4.x form). v5's ``tie_weights`` -> 

170 ``get_expanded_tied_weights_keys`` calls ``.keys()`` on the mapping, 

171 which raises ``AttributeError`` on a list. Rewrite it to the v5 dict 

172 form ``{"lm_head.weight": "model.embeddings.weight"}``. RWKV-7 defaults 

173 ``tie_word_embeddings=False`` (so the list path is usually short- 

174 circuited before ``.keys()``), but the rewrite is harmless when untied 

175 and prevents the crash on any tied checkpoint. 

176 

177 2. Weight re-init. Under v5's meta-device load-then-materialise flow, 

178 ``PreTrainedModel._init_weights`` is invoked on modules that already 

179 hold checkpoint weights, re-randomising them. Guard it to skip modules 

180 whose parameters are already on a real (non-meta) device — the same 

181 defensive patch openelm.py / raven.py apply. 

182 

183 Args: 

184 model_name: The HuggingFace model name/path. 

185 model_kwargs: The kwargs dict for from_pretrained(). 

186 """ 

187 # Force-import the fla RWKV-7 modeling module so its classes appear in 

188 # sys.modules to patch. fla is normally a pip package; fall back to the 

189 # dynamic-module route for genuinely bundled remote code. 

190 try: 

191 import fla.models.rwkv7.modeling_rwkv7 # noqa: F401 

192 except Exception: 

193 if force_import_remote_class(model_name, "modeling_rwkv7.RWKV7ForCausalLM") is None: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true

194 return 

195 

196 # Patch every loaded RWKV-7 modeling module (each remote revision gets its 

197 # own module object in sys.modules). 

198 for module in iter_remote_modeling_modules("rwkv7"): 

199 # Patch 1: tied-weights keys list -> v5 dict form. 

200 retie_weights_keys_v5( 

201 getattr(module, "RWKV7ForCausalLM", None), 

202 {"lm_head.weight": "model.embeddings.weight"}, 

203 ) 

204 # Patch 2: don't re-randomise already-loaded weights. 

205 pretrained_class = getattr(module, "RWKV7PreTrainedModel", None) 

206 if pretrained_class is not None: 206 ↛ 198line 206 didn't jump to line 198 because the condition on line 206 was always true

207 patch_init_weights_skip_loaded(pretrained_class)