Coverage for transformer_lens/model_bridge/supported_architectures/rwkv7.py: 50%
60 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +0000
1"""RWKV-7 ("Goose") architecture adapter (RWKV7ForCausalLM).
3Model family: ``fla-hub/rwkv7-*`` (e.g. ``fla-hub/rwkv7-0.1B-g1``), an
4attention-free recurrent language model from the flash-linear-attention (``fla``)
5library. Loaded via remote code (``trust_remote_code=True``); the modeling
6classes live in ``fla.models.rwkv7``.
8Architecture overview
9---------------------
10RWKV-7 is a flat stack of recurrent blocks over a shared residual width, wrapped
11by standard (biased) LayerNorm rather than RMSNorm and with no positional
12embeddings::
14 embeddings -> [ RWKV7Block x N ] -> norm -> lm_head
16Each ``RWKV7Block`` is a pre-norm pair of a time-mixing and a channel-mixing
17sublayer::
19 x = x + attn(attn_norm(x)) # time-mixing (RWKV7Attention)
20 x = x + ffn(ffn_norm(x)) # channel-mixing (RWKV7FeedForward)
22with an extra ``pre_norm`` LayerNorm on layer 0 only (``config.norm_first``).
23The time-mixing sublayer is the RWKV-7 "generalized delta rule": token-shifted
24lerp coefficients (``x_r``..``x_g``), receptance / key / value / output
25projections (``r_proj`` / ``k_proj`` / ``v_proj`` / ``o_proj``), low-rank LoRAs
26for the log-space decay / value blend / a-coefficient / gate, and a GroupNorm.
27The channel-mixing sublayer is a token-shifted squared-ReLU MLP (``key`` up,
28``value`` down).
30Cross-block ``v_first`` threading: ``RWKV7Model.forward`` initialises
31``v_first = torch.zeros_like(hidden_states)`` and threads it through every block
32(layer 0 fills it; later layers blend their value with it). This is managed
33entirely by the HF model-level forward, so the bridge — which delegates each
34block's forward to HF — gets it for free.
35"""
37import sys
38from typing import Any
40from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
41from transformer_lens.model_bridge.generalized_components import (
42 EmbeddingBridge,
43 LinearBridge,
44 NormalizationBridge,
45 OpaqueBlockBridge,
46 UnembeddingBridge,
47)
48from transformer_lens.model_bridge.generalized_components.base import (
49 GeneralizedComponent,
50)
53class RWKV7ArchitectureAdapter(ArchitectureAdapter):
54 """Architecture adapter for RWKV7ForCausalLM (RWKV-7 "Goose").
56 Attention-free recurrent decoder: a flat stack of pre-norm blocks, each a
57 generalized-delta-rule time-mixing sublayer plus a token-shifted squared-ReLU
58 channel-mixing sublayer, wrapped by standard biased LayerNorm. The recurrence
59 and the cross-block ``v_first`` threading live inside the ``fla`` remote-code
60 forward, which the bridge delegates to; see the module docstring for the full
61 set of adapter decisions.
62 """
64 # Attention-free and recurrent — off the transformer-shaped verify_models
65 # path; correctness lives in the integration tests (bridge-vs-HF parity).
66 applicable_phases: list[int] = []
68 def __init__(self, cfg: Any) -> None:
69 """Initialize the RWKV-7 architecture adapter."""
70 super().__init__(cfg)
72 # Standard biased LayerNorm, no positional embeddings, ungated FFN.
73 self.cfg.normalization_type = "LN"
74 self.cfg.uses_rms_norm = False
75 self.cfg.positional_embedding_type = "none"
76 self.cfg.final_rms = False
77 self.cfg.attn_only = False
78 self.cfg.gated_mlp = False
80 # fla drives recurrent decode state via its own Cache, not the
81 # cache_params convention, so the adapter does not advertise a cache.
82 setattr(self.cfg, "is_stateful", False)
84 # Surface the RWKV-7 shape attributes on cfg so they are present on both
85 # the HF-boot path (also via _HF_PASSTHROUGH_ATTRS) and the synthetic
86 # config path used by the unit tests. getattr-with-default keeps a bare
87 # TransformerBridgeConfig from raising.
88 # head_dim is a read-only property on TransformerBridgeConfig (aliases
89 # d_head), so it is read but never assigned. num_heads defaults to
90 # d_model // head_dim.
91 head_dim = getattr(cfg, "head_dim", 64) or 64
92 num_heads = getattr(cfg, "num_heads", None) or max(1, self.cfg.d_model // head_dim)
93 value_dim = getattr(cfg, "value_dim", None) or [self.cfg.d_model] * self.cfg.n_layers
94 # setattr (not direct assignment) for the RWKV-7-specific names so mypy
95 # does not flag them as undeclared on TransformerBridgeConfig.
96 setattr(self.cfg, "num_heads", num_heads)
97 setattr(self.cfg, "value_dim", value_dim)
98 setattr(self.cfg, "decay_low_rank_dim", getattr(cfg, "decay_low_rank_dim", 64))
99 setattr(self.cfg, "gate_low_rank_dim", getattr(cfg, "gate_low_rank_dim", 128))
100 setattr(self.cfg, "a_low_rank_dim", getattr(cfg, "a_low_rank_dim", 64))
101 setattr(self.cfg, "v_low_rank_dim", getattr(cfg, "v_low_rank_dim", 16))
102 setattr(self.cfg, "norm_first", getattr(cfg, "norm_first", True))
103 setattr(self.cfg, "norm_bias", getattr(cfg, "norm_bias", True))
104 setattr(self.cfg, "fuse_norm", getattr(cfg, "fuse_norm", True))
105 setattr(self.cfg, "attn_mode", getattr(cfg, "attn_mode", "chunk"))
106 setattr(self.cfg, "hidden_act", getattr(cfg, "hidden_act", "sqrelu"))
107 # LayerNorm epsilon for the reimplementing NormalizationBridge path (eps
108 # is a real config field, so direct assignment is fine).
109 self.cfg.eps = getattr(cfg, "norm_eps", getattr(cfg, "eps", 1e-5))
111 # Full delegation to the fla forward — no HT-format weight reshaping.
112 self.weight_processing_conversions = {}
114 self.component_mapping = {
115 "embed": EmbeddingBridge(name="model.embeddings"),
116 "blocks": OpaqueBlockBridge(
117 name="model.layers",
118 submodules={
119 # Pre-norm before time-mixing (standard single-input LayerNorm).
120 "attn_norm": NormalizationBridge(
121 name="attn_norm", config=self.cfg, uses_rms_norm=False
122 ),
123 # Time-mixing: generalized delta rule. Delegated passthrough;
124 # only the four projections are exposed (LoRAs / GroupNorm /
125 # lerp params stay inside the fla forward).
126 "attn": GeneralizedComponent(
127 name="attn",
128 config=self.cfg,
129 submodules={
130 "r_proj": LinearBridge(name="r_proj"),
131 "k_proj": LinearBridge(name="k_proj"),
132 "v_proj": LinearBridge(name="v_proj"),
133 "o_proj": LinearBridge(name="o_proj"),
134 },
135 ),
136 # Pre-norm before channel-mixing. Under config.fuse_norm the
137 # block calls this as ffn_norm(x, residual, True) -> (normed,
138 # residual), a fused signature the reimplementing
139 # NormalizationBridge can't express, so delegate any-signature
140 # to the live HF module with I/O hooks.
141 "ffn_norm": GeneralizedComponent(name="ffn_norm", config=self.cfg),
142 # Channel-mixing: token-shifted squared-ReLU MLP. Delegated
143 # passthrough exposing the up ("key") and down ("value")
144 # projections. HF confusingly names the output proj "value".
145 # NOT MLPBridge (cf. rwkv4): its hook_in pre-fire would
146 # suppress key.hook_in, the post-token-shift input — the one
147 # interesting tensor here. Aliases give hook_pre/hook_post.
148 "ffn": GeneralizedComponent(
149 name="ffn",
150 config=self.cfg,
151 hook_alias_overrides={
152 "hook_pre": "key.hook_out",
153 "hook_post": "value.hook_in",
154 },
155 submodules={
156 "key": LinearBridge(name="key"),
157 "value": LinearBridge(name="value"),
158 },
159 ),
160 },
161 ),
162 "ln_final": NormalizationBridge(
163 name="model.norm", config=self.cfg, uses_rms_norm=False
164 ),
165 "unembed": UnembeddingBridge(name="lm_head"),
166 }
168 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
169 """Patch fla's RWKV-7 remote code for transformers v5 compatibility.
171 Two defensive patches, mirroring raven:
173 1. Tied-weights format. ``RWKV7ForCausalLM._tied_weights_keys`` is a list
174 (``["lm_head.weight"]``, the 4.x form). v5's ``tie_weights`` ->
175 ``get_expanded_tied_weights_keys`` calls ``.keys()`` on the mapping,
176 which raises ``AttributeError`` on a list. Rewrite it to the v5 dict
177 form ``{"lm_head.weight": "model.embeddings.weight"}``. RWKV-7 defaults
178 ``tie_word_embeddings=False`` (so the list path is usually short-
179 circuited before ``.keys()``), but the rewrite is harmless when untied
180 and prevents the crash on any tied checkpoint.
182 2. Weight re-init. Under v5's meta-device load-then-materialise flow,
183 ``PreTrainedModel._init_weights`` is invoked on modules that already
184 hold checkpoint weights, re-randomising them. Guard it to skip modules
185 whose parameters are already on a real (non-meta) device — the same
186 defensive patch openelm.py / raven.py apply.
188 Args:
189 model_name: The HuggingFace model name/path.
190 model_kwargs: The kwargs dict for from_pretrained().
191 """
192 # Force-import the fla RWKV-7 modeling module so its classes appear in
193 # sys.modules to patch. fla is normally a pip package; fall back to the
194 # dynamic-module route for genuinely bundled remote code.
195 try:
196 import fla.models.rwkv7.modeling_rwkv7 # noqa: F401
197 except Exception:
198 try:
199 from transformers.dynamic_module_utils import (
200 get_class_from_dynamic_module,
201 )
203 get_class_from_dynamic_module(
204 "modeling_rwkv7.RWKV7ForCausalLM",
205 model_name,
206 )
207 except Exception:
208 return
210 # Patch every loaded RWKV-7 modeling module (each remote revision gets its
211 # own module object in sys.modules).
212 for key in list(sys.modules.keys()):
213 if "rwkv7" not in key.lower() or "modeling" not in key.lower():
214 continue
215 module = sys.modules[key]
217 # Patch 1: tied-weights keys list -> v5 dict form.
218 causal_lm_class = getattr(module, "RWKV7ForCausalLM", None)
219 if causal_lm_class is not None and isinstance(
220 getattr(causal_lm_class, "_tied_weights_keys", None), list
221 ):
222 causal_lm_class._tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"}
224 # Patch 2: don't re-randomise already-loaded weights.
225 pretrained_class = getattr(module, "RWKV7PreTrainedModel", None)
226 if pretrained_class is None or getattr(pretrained_class, "_tl_patched", False):
227 continue
228 original_init_weights = pretrained_class._init_weights
230 def safe_init_weights(self, mod, _original=original_init_weights):
231 # Only initialise modules still on meta device (pre-loading);
232 # never re-randomise weights already read from the checkpoint.
233 first_param = next(mod.parameters(), None)
234 if first_param is not None and first_param.device.type != "meta":
235 return
236 _original(self, mod)
238 pretrained_class._init_weights = safe_init_weights
239 pretrained_class._tl_patched = True