Coverage for transformer_lens/model_bridge/supported_architectures/rwkv7.py: 50%
60 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +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 "ffn": GeneralizedComponent(
146 name="ffn",
147 config=self.cfg,
148 submodules={
149 "key": LinearBridge(name="key"),
150 "value": LinearBridge(name="value"),
151 },
152 ),
153 },
154 ),
155 "ln_final": NormalizationBridge(
156 name="model.norm", config=self.cfg, uses_rms_norm=False
157 ),
158 "unembed": UnembeddingBridge(name="lm_head"),
159 }
161 def prepare_loading(self, model_name: str, model_kwargs: dict) -> None:
162 """Patch fla's RWKV-7 remote code for transformers v5 compatibility.
164 Two defensive patches, mirroring raven:
166 1. Tied-weights format. ``RWKV7ForCausalLM._tied_weights_keys`` is a list
167 (``["lm_head.weight"]``, the 4.x form). v5's ``tie_weights`` ->
168 ``get_expanded_tied_weights_keys`` calls ``.keys()`` on the mapping,
169 which raises ``AttributeError`` on a list. Rewrite it to the v5 dict
170 form ``{"lm_head.weight": "model.embeddings.weight"}``. RWKV-7 defaults
171 ``tie_word_embeddings=False`` (so the list path is usually short-
172 circuited before ``.keys()``), but the rewrite is harmless when untied
173 and prevents the crash on any tied checkpoint.
175 2. Weight re-init. Under v5's meta-device load-then-materialise flow,
176 ``PreTrainedModel._init_weights`` is invoked on modules that already
177 hold checkpoint weights, re-randomising them. Guard it to skip modules
178 whose parameters are already on a real (non-meta) device — the same
179 defensive patch openelm.py / raven.py apply.
181 Args:
182 model_name: The HuggingFace model name/path.
183 model_kwargs: The kwargs dict for from_pretrained().
184 """
185 # Force-import the fla RWKV-7 modeling module so its classes appear in
186 # sys.modules to patch. fla is normally a pip package; fall back to the
187 # dynamic-module route for genuinely bundled remote code.
188 try:
189 import fla.models.rwkv7.modeling_rwkv7 # noqa: F401
190 except Exception:
191 try:
192 from transformers.dynamic_module_utils import (
193 get_class_from_dynamic_module,
194 )
196 get_class_from_dynamic_module(
197 "modeling_rwkv7.RWKV7ForCausalLM",
198 model_name,
199 )
200 except Exception:
201 return
203 # Patch every loaded RWKV-7 modeling module (each remote revision gets its
204 # own module object in sys.modules).
205 for key in list(sys.modules.keys()):
206 if "rwkv7" not in key.lower() or "modeling" not in key.lower():
207 continue
208 module = sys.modules[key]
210 # Patch 1: tied-weights keys list -> v5 dict form.
211 causal_lm_class = getattr(module, "RWKV7ForCausalLM", None)
212 if causal_lm_class is not None and isinstance(
213 getattr(causal_lm_class, "_tied_weights_keys", None), list
214 ):
215 causal_lm_class._tied_weights_keys = {"lm_head.weight": "model.embeddings.weight"}
217 # Patch 2: don't re-randomise already-loaded weights.
218 pretrained_class = getattr(module, "RWKV7PreTrainedModel", None)
219 if pretrained_class is None or getattr(pretrained_class, "_tl_patched", False):
220 continue
221 original_init_weights = pretrained_class._init_weights
223 def safe_init_weights(self, mod, _original=original_init_weights):
224 # Only initialise modules still on meta device (pre-loading);
225 # never re-randomise weights already read from the checkpoint.
226 first_param = next(mod.parameters(), None)
227 if first_param is not None and first_param.device.type != "meta":
228 return
229 _original(self, mod)
231 pretrained_class._init_weights = safe_init_weights
232 pretrained_class._tl_patched = True