transformer_lens.model_bridge.supported_architectures.rwkv module

RWKV architecture adapter.

BlinkDL’s RWKV-4 (RwkvForCausalLM, native in transformers): the canonical WKV linear-attention RNN, trained on the Pile in a Pythia-parallel suite. Blocks pair a time-mix module (token-shift interpolation into key/value/receptance projections, recurrent WKV kernel, gated output) with a channel-mix module (token-shift key/ receptance, squared-relu value) under pre-LNs, plus an extra pre_ln on layer 0 before anything else. Both mixers delegate to HF (the WKV recurrence has no attention-shaped reconstruction); their projections are wrapped for hooks.

HF rescales attention.output/feed_forward.value weights by 2^(layer // rescale_every) at eval (the forward divides hidden states to compensate, so the function is unchanged); bridge and reference both keep the default so their weights match exactly. use_cache is forced off: the recurrent state buffers are written in-place per layer, which breaks autograd under backward hooks — and generation (the only state consumer) runs on a bespoke state kwarg the bridge’s loop doesn’t speak, so generation phases are excluded.

class transformer_lens.model_bridge.supported_architectures.rwkv.RwkvArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for RwkvForCausalLM models.

__init__(cfg: Any) → None

Initialize the RWKV architecture adapter.

applicable_phases: list[int] = [1, 2, 3, 4]

component_mapping: ComponentMapping | None

prepare_loading(model_name: str, model_kwargs: dict) → None

Force use_cache off: per-layer in-place state writes break autograd under backward hooks, and only recurrent generation consumes them.

prepare_model(hf_model: Any) → None

Re-assert use_cache=False – prepare_loading only fires on the boot path, so directly-wrapped modules keep the default and leak state tuples into tensor-only hooks.

supports_batched_generation: bool = False

supports_fold_ln = False

supports_generation: bool = True

supports_kv_cache: bool = False

uses_split_attention: bool

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