transformer_lens.model_bridge.supported_architectures.rwkv7 module

RWKV-7 (“Goose”) architecture adapter (RWKV7ForCausalLM).

Model family: fla-hub/rwkv7-* (e.g. fla-hub/rwkv7-0.1B-g1), an attention-free recurrent language model from the flash-linear-attention (fla) library. Loaded via remote code (trust_remote_code=True); the modeling classes live in fla.models.rwkv7.

Architecture overview

RWKV-7 is a flat stack of recurrent blocks over a shared residual width, wrapped by standard (biased) LayerNorm rather than RMSNorm and with no positional embeddings:

embeddings -> [ RWKV7Block x N ] -> norm -> lm_head

Each RWKV7Block is a pre-norm pair of a time-mixing and a channel-mixing sublayer:

x = x + attn(attn_norm(x))     # time-mixing  (RWKV7Attention)
x = x + ffn(ffn_norm(x))       # channel-mixing (RWKV7FeedForward)

with an extra pre_norm LayerNorm on layer 0 only (config.norm_first). The time-mixing sublayer is the RWKV-7 “generalized delta rule”: token-shifted lerp coefficients (x_r..``x_g``), receptance / key / value / output projections (r_proj / k_proj / v_proj / o_proj), low-rank LoRAs for the log-space decay / value blend / a-coefficient / gate, and a GroupNorm. The channel-mixing sublayer is a token-shifted squared-ReLU MLP (key up, value down).

Cross-block v_first threading: RWKV7Model.forward initialises v_first = torch.zeros_like(hidden_states) and threads it through every block (layer 0 fills it; later layers blend their value with it). This is managed entirely by the HF model-level forward, so the bridge — which delegates each block’s forward to HF — gets it for free.

class transformer_lens.model_bridge.supported_architectures.rwkv7.RWKV7ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RWKV7ForCausalLM (RWKV-7 “Goose”).

Attention-free recurrent decoder: a flat stack of pre-norm blocks, each a generalized-delta-rule time-mixing sublayer plus a token-shifted squared-ReLU channel-mixing sublayer, wrapped by standard biased LayerNorm. The recurrence and the cross-block v_first threading live inside the fla remote-code forward, which the bridge delegates to; see the module docstring for the full set of adapter decisions.

__init__(cfg: Any) → None

Initialize the RWKV-7 architecture adapter.

applicable_phases: list[int] = []
component_mapping: ComponentMapping | None
prepare_loading(model_name: str, model_kwargs: dict) → None

Patch fla’s RWKV-7 remote code for transformers v5 compatibility.

Two defensive patches, mirroring raven:

  1. Tied-weights format. RWKV7ForCausalLM._tied_weights_keys is a list (["lm_head.weight"], the 4.x form). v5’s tie_weights -> get_expanded_tied_weights_keys calls .keys() on the mapping, which raises AttributeError on a list. Rewrite it to the v5 dict form {"lm_head.weight": "model.embeddings.weight"}. RWKV-7 defaults tie_word_embeddings=False (so the list path is usually short- circuited before .keys()), but the rewrite is harmless when untied and prevents the crash on any tied checkpoint.

  2. Weight re-init. Under v5’s meta-device load-then-materialise flow, PreTrainedModel._init_weights is invoked on modules that already hold checkpoint weights, re-randomising them. Guard it to skip modules whose parameters are already on a real (non-meta) device — the same defensive patch openelm.py / raven.py apply.

Parameters:

  • model_name – The HuggingFace model name/path.

  • model_kwargs – The kwargs dict for from_pretrained().

uses_split_attention: bool
weight_processing_conversions: Dict[str, ParamProcessingConversion | str] | None