Coverage for transformer_lens/model_bridge/supported_architectures/laguna.py: 85%

1"""Laguna architecture adapter. 

2 

3poolside's Laguna coding models (``LagunaForCausalLM``, remote code): 

4Llama-shaped causal decoders with two first-of-kind mechanisms — 

5heterogeneous per-layer attention head counts 

6(``num_attention_heads_per_layer``) and per-head softplus output gating 

7(``g_proj``) — over a batched-expert MoE (FlexOlmo-style 3D expert 

8parameters, top-k router, always-on shared experts) with per-layer 

9dense/sparse selection via ``mlp_layer_types``. 

10 

11Attention delegates to HF: the bridge reimplementation assumes one 

12uniform head count, and the softplus gate has no reconstruction. The 

13per-layer head heterogeneity also rules out uniform Q/K/V reshape 

14conversions, so no HookedTransformer-format conversions ship and LN 

15folding is disabled. 

16""" 

17 

18from typing import Any 

19 

20from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

21from transformer_lens.model_bridge.generalized_components import ( 

22 AttentionBridge, 

23 BlockBridge, 

24 EmbeddingBridge, 

25 LinearBridge, 

26 MoEBridge, 

27 RMSNormalizationBridge, 

28 RotaryEmbeddingBridge, 

29 UnembeddingBridge, 

30) 

31from transformer_lens.model_bridge.generalized_components.base import ( 

32 GeneralizedComponent, 

33) 

34 

35 

36class LagunaArchitectureAdapter(ArchitectureAdapter): 

37 """Architecture adapter for LagunaForCausalLM models.""" 

38 

39 # Per-layer head counts: no uniform "(n h) m -> n m h" reshape exists. 

40 supports_fold_ln = False 

41 

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

43 """Initialize the Laguna architecture adapter.""" 

44 super().__init__(cfg) 

45 

46 self._set_rms_rotary_defaults() 

47 

48 self.weight_processing_conversions = {} 

49 

50 self.component_mapping = { 

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

52 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb"), 

53 "blocks": BlockBridge( 

54 name="model.layers", 

55 config=self.cfg, 

56 submodules={ 

57 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), 

58 "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

59 # Heterogeneous head counts + softplus output gating: 

60 # delegate; projections and the gate are hookable. 

61 "attn": AttentionBridge( 

62 name="self_attn", 

63 config=self.cfg, 

64 submodules={ 

65 "q": LinearBridge(name="q_proj"), 

66 "k": LinearBridge(name="k_proj"), 

67 "v": LinearBridge(name="v_proj"), 

68 "o": LinearBridge(name="o_proj"), 

69 "gate": LinearBridge(name="g_proj"), 

70 "q_norm": RMSNormalizationBridge(name="q_norm", config=self.cfg), 

71 "k_norm": RMSNormalizationBridge(name="k_norm", config=self.cfg), 

72 }, 

73 maintain_native_attention=True, 

74 ), 

75 # Dense or sparse per mlp_layer_types; router and shared 

76 # experts optional so dense layers skip them. 

77 "mlp": MoEBridge( 

78 name="mlp", 

79 config=self.cfg, 

80 sparse_required=("gate",), 

81 submodules={ 

82 "gate": GeneralizedComponent(name="gate", optional=True), 

83 "shared_experts": self._gated_mlp(name="shared_experts", optional=True), 

84 # Dense-layer projections (absent on MoE layers). 

85 "dense_gate": LinearBridge(name="gate_proj", optional=True), 

86 "dense_in": LinearBridge(name="up_proj", optional=True), 

87 "dense_out": LinearBridge(name="down_proj", optional=True), 

88 }, 

89 ), 

90 }, 

91 ), 

92 "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), 

93 "unembed": UnembeddingBridge(name="lm_head"), 

94 } 

95 

96 def setup_component_testing(self, hf_model: Any, bridge_model: Any = None) -> None: 

97 """Delegated attention computes rotary inside HF; nothing to wire.""" 

98 

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

100 """User-register Laguna's native conversion mapping so the per-expert->batched 

101 expert merge runs under remote code (transformers skips it for custom-code 

102 modules, leaving the batched experts at random init).""" 

103 try: 

104 from transformers.conversion_mapping import ( 

105 USER_REGISTERED_MAPPINGS, 

106 get_checkpoint_conversion_mapping, 

107 register_checkpoint_conversion_mapping, 

108 ) 

109 except ImportError: 

110 pass # transformers predates the conversion-mapping API; nothing to register 

111 else: 

112 if "laguna" not in USER_REGISTERED_MAPPINGS: 112 ↛ 116line 112 didn't jump to line 116 because the condition on line 112 was always true

113 mapping = get_checkpoint_conversion_mapping("laguna") 

114 if mapping: 114 ↛ 116line 114 didn't jump to line 116 because the condition on line 114 was always true

115 register_checkpoint_conversion_mapping("laguna", mapping, overwrite=True) 

116 super().prepare_loading(model_name, model_kwargs)