Coverage for transformer_lens/model_bridge/supported_architectures/recurrent_gemma.py: 100%

1"""RecurrentGemma (Griffin) architecture adapter. 

2 

3Supports ``RecurrentGemmaForCausalLM`` (e.g. ``google/recurrentgemma-2b``, 

4``google/recurrentgemma-9b``), the open instance of the **Griffin** architecture. 

5 

6Architecture overview: 

7- Heterogeneous layers defined by ``config.block_types`` (default pattern 

8 ``("recurrent", "recurrent", "attention")`` repeated over ``num_hidden_layers``). 

9 Each ``model.layers.{i}`` is a ``RecurrentGemmaDecoderLayer`` whose 

10 ``temporal_block`` is *either*: 

11 * ``RecurrentGemmaRecurrentBlock`` — the RG-LRU real-gated linear recurrence 

12 (``linear_x`` / ``linear_y`` / ``conv_1d`` / ``rg_lru`` / ``linear_out``), or 

13 * ``RecurrentGemmaSdpaAttention`` — local sliding-window GQA attention with 

14 partial rotary (``q_proj`` / ``k_proj`` / ``v_proj`` / ``o_proj`` / ``rotary_emb``). 

15- Every layer additionally has ``temporal_pre_norm`` (pre-norm before the temporal 

16 block), ``channel_pre_norm`` (pre-norm before the MLP), and a gated MLP 

17 ``mlp_block`` (``gate_proj`` / ``up_proj`` / ``down_proj``, GELU-tanh). 

18- ``model.final_norm`` + ``lm_head``. 

19 

20Gemma-family numerics (shared with Gemma-2/3): RMSNorm applies ``(1.0 + weight)`` 

21(``rmsnorm_uses_offset``), token embeddings are scaled by ``sqrt(hidden_size)`` at 

22runtime inside the HF forward, and the final logits are tanh-soft-capped at 

23``config.logits_soft_cap`` (30.0). 

24 

25Key adapter decision (mirrors ``Lfm2MoeArchitectureAdapter``): because the 

26``temporal_block`` substructure varies per layer (recurrent vs. attention), we 

27wrap each decoder layer as a whole with residual-stream hooks only, rather than 

28pretending every layer has a homogeneous attention/MLP substructure. This keeps 

29execution correct on both layer types. Finer-grained RG-LRU state hooks 

30(``hook_ssm_write`` / ``hook_ssm_state`` style) are a natural follow-up. 

31 

32``applicable_phases = [4]``: the whole-layer bridge exposes only residual hooks, 

33so the component-level comparisons in phases 1-3 do not apply; phase 4 

34(generation + text quality) does. 

35""" 

36 

37from typing import Any 

38 

39from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

40from transformer_lens.model_bridge.generalized_components import ( 

41 BlockBridge, 

42 EmbeddingBridge, 

43 RMSNormalizationBridge, 

44 UnembeddingBridge, 

45) 

46 

47 

48class RecurrentGemmaBlockBridge(BlockBridge): 

49 """Whole-layer RecurrentGemma bridge exposing only residual-stream hooks. 

50 

51 RecurrentGemma interleaves RG-LRU recurrent layers and local-attention layers. 

52 Wrapping the HF decoder layer as a whole preserves correct execution while 

53 avoiding unresolved standard attention/MLP aliases on the recurrent layers 

54 (which have no q/k/v/o or gate/up/down substructure). 

55 """ 

56 

57 hook_aliases = { 

58 "hook_resid_pre": "hook_in", 

59 "hook_resid_post": "hook_out", 

60 } 

61 

62 

63class RecurrentGemmaArchitectureAdapter(ArchitectureAdapter): 

64 """Architecture adapter for ``RecurrentGemmaForCausalLM`` (Griffin). 

65 

66 Hybrid RG-LRU recurrence + local sliding-window attention. The temporal-block 

67 type per layer is determined by ``config.block_types[layer_idx % len(block_types)]``. 

68 """ 

69 

70 # Whole-layer residual hooks only; phases 1-3 compare component substructure 

71 # this adapter intentionally does not expose. Phase 4 (generation) applies. 

72 applicable_phases: list[int] = [4] 

73 

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

75 """Initialize the RecurrentGemma architecture adapter.""" 

76 super().__init__(cfg) 

77 

78 self._set_rms_rotary_defaults() 

79 # Hookable attention needs eager; the base prepare hooks force it through 

80 # from_pretrained and onto the loaded config. 

81 self.cfg.attn_implementation = "eager" 

82 # RG-LRU + local attention both use RoPE-style handling internally on the 

83 # attention layers; there is no model-level rotary module (it lives inside 

84 # each attention temporal_block), so we do not wire a rotary_emb component. 

85 # Gemma-family gated MLP (gate_proj -> GELU-tanh(up) -> down). 

86 # Gemma RMSNorm uses (1.0 + weight); see 

87 # https://github.com/huggingface/transformers/pull/29402 

88 self.cfg.rmsnorm_uses_offset = True 

89 

90 # Gemma models were not trained with BOS tokens prepended. 

91 self.cfg.default_prepend_bos = False 

92 

93 norm_eps = getattr(cfg, "rms_norm_eps", None) 

94 if norm_eps is not None: 

95 self.cfg.eps = norm_eps 

96 

97 # Final logits are tanh-soft-capped at config.logits_soft_cap (30.0). 

98 logits_soft_cap = getattr(cfg, "logits_soft_cap", None) 

99 if logits_soft_cap is not None: 

100 self.cfg.output_logits_soft_cap = logits_soft_cap 

101 

102 # Canonical per-layer pattern (as on Nemotron-H). Read the builder's 

103 # expanded layers_block_type; the raw block_types pattern is not 

104 # propagated and would clobber it with [] on real boots. 

105 setattr(self.cfg, "layers_block_type", self._canonical_layer_types(cfg)) 

106 

107 self.component_mapping = { 

108 "embed": EmbeddingBridge(name="model.embed_tokens"), 

109 "blocks": RecurrentGemmaBlockBridge(name="model.layers", config=self.cfg), 

110 "ln_final": RMSNormalizationBridge(name="model.final_norm", config=self.cfg), 

111 "unembed": UnembeddingBridge(name="lm_head", config=self.cfg), 

112 }