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

1"""RWKV architecture adapter. 

2 

3BlinkDL's RWKV-4 (``RwkvForCausalLM``, native in transformers): the 

4canonical WKV linear-attention RNN, trained on the Pile in a 

5Pythia-parallel suite. Blocks pair a time-mix module (token-shift 

6interpolation into key/value/receptance projections, recurrent WKV 

7kernel, gated output) with a channel-mix module (token-shift key/ 

8receptance, squared-relu value) under pre-LNs, plus an extra pre_ln on 

9layer 0 before anything else. Both mixers delegate to HF (the WKV 

10recurrence has no attention-shaped reconstruction); their projections 

11are wrapped for hooks. 

12 

13HF rescales attention.output/feed_forward.value weights by 2^(layer // 

14rescale_every) at eval (the forward divides hidden states to compensate, 

15so the function is unchanged); bridge and reference both keep the default 

16so their weights match exactly. use_cache is forced off: the recurrent 

17state buffers are written in-place per layer, which breaks autograd under 

18backward hooks — and generation (the only state consumer) runs on a 

19bespoke ``state`` kwarg the bridge's loop doesn't speak, so generation 

20phases are excluded. 

21""" 

22 

23from typing import Any 

24 

25from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

26from transformer_lens.model_bridge.generalized_components import ( 

27 BlockBridge, 

28 EmbeddingBridge, 

29 LinearBridge, 

30 MLPBridge, 

31 NormalizationBridge, 

32 UnembeddingBridge, 

33) 

34from transformer_lens.model_bridge.generalized_components.base import ( 

35 GeneralizedComponent, 

36) 

37 

38 

39class _RwkvBlockBridge(BlockBridge): 

40 """RWKV blocks have no attention: replace the attention-flavored alias set 

41 with time-mix/channel-mix names (resid_mid is ln2's input, pre-mix).""" 

42 

43 hook_aliases = { 

44 "hook_resid_pre": "hook_in", 

45 "hook_resid_mid": "ln2.hook_in", 

46 "hook_resid_post": "hook_out", 

47 "hook_time_mix_in": "time_mix.hook_in", 

48 "hook_time_mix_out": "time_mix.hook_out", 

49 "hook_channel_mix_in": "channel_mix.hook_in", 

50 "hook_channel_mix_out": "channel_mix.hook_out", 

51 } 

52 

53 

54class RwkvArchitectureAdapter(ArchitectureAdapter): 

55 """Architecture adapter for RwkvForCausalLM models.""" 

56 

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

58 supports_generation: bool = True 

59 # HF threads recurrence through a bespoke `state` kwarg, not past_key_values; 

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

61 supports_kv_cache: bool = False 

62 # RwkvModel ignores attention_mask entirely, so left-padding would silently 

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

64 supports_batched_generation: bool = False 

65 # Pre-LNs feed the mixers, but the recurrent WKV kernel consumes raw and 

66 # time-shifted inputs — no fold target exists. 

67 supports_fold_ln = False 

68 

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

70 """Initialize the RWKV architecture adapter.""" 

71 super().__init__(cfg) 

72 

73 self.cfg.normalization_type = "LN" 

74 self.cfg.uses_rms_norm = False 

75 self.cfg.positional_embedding_type = "none" 

76 self.cfg.gated_mlp = False 

77 self.cfg.attn_only = False 

78 self.cfg.final_rms = False 

79 

80 self.weight_processing_conversions = {} 

81 

82 self.component_mapping = { 

83 "embed": EmbeddingBridge(name="rwkv.embeddings"), 

84 "blocks": _RwkvBlockBridge( 

85 name="rwkv.blocks", 

86 config=self.cfg, 

87 submodules={ 

88 # Layer 0 only: an extra LN before the block body. 

89 "pre_ln": NormalizationBridge( 

90 name="pre_ln", 

91 config=self.cfg, 

92 use_native_layernorm_autograd=True, 

93 optional=True, 

94 ), 

95 "ln1": NormalizationBridge( 

96 name="ln1", config=self.cfg, use_native_layernorm_autograd=True 

97 ), 

98 "time_mix": GeneralizedComponent( 

99 name="attention", 

100 submodules={ 

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

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

103 "receptance": LinearBridge(name="receptance"), 

104 "output": LinearBridge(name="output"), 

105 }, 

106 ), 

107 "ln2": NormalizationBridge( 

108 name="ln2", config=self.cfg, use_native_layernorm_autograd=True 

109 ), 

110 # MLPBridge so the component harness sizes inputs by the 

111 # true in/out dims (key: d_model->4d, value: 4d->d_model). 

112 "channel_mix": MLPBridge( 

113 name="feed_forward", 

114 config=self.cfg, 

115 submodules={ 

116 "in": LinearBridge(name="key"), 

117 "receptance": LinearBridge(name="receptance"), 

118 "out": LinearBridge(name="value"), 

119 }, 

120 ), 

121 }, 

122 ), 

123 "ln_final": NormalizationBridge( 

124 name="rwkv.ln_out", config=self.cfg, use_native_layernorm_autograd=True 

125 ), 

126 "unembed": UnembeddingBridge(name="head"), 

127 } 

128 

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

130 """Force use_cache off: per-layer in-place state writes break autograd 

131 under backward hooks, and only recurrent generation consumes them.""" 

132 config = model_kwargs.get("config") 

133 if config is not None: 133 ↛ 135line 133 didn't jump to line 135 because the condition on line 133 was always true

134 config.use_cache = False 

135 super().prepare_loading(model_name, model_kwargs) 

136 

137 def prepare_model(self, hf_model: Any) -> None: 

138 """Re-assert use_cache=False -- prepare_loading only fires on the boot path, so 

139 directly-wrapped modules keep the default and leak state tuples into tensor-only hooks.""" 

140 super().prepare_model(hf_model) 

141 config = getattr(hf_model, "config", None) 

142 if config is not None: 142 ↛ exitline 142 didn't return from function 'prepare_model' because the condition on line 142 was always true

143 config.use_cache = False