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 # Delegated attention computes rotary inside HF; nothing to wire. 

42 _testing_eager = None 

43 _testing_wire_rotary = False 

44 

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

46 """Initialize the Laguna architecture adapter.""" 

47 super().__init__(cfg) 

48 

49 self._set_rms_rotary_defaults() 

50 

51 self.weight_processing_conversions = {} 

52 

53 self.component_mapping = { 

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

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

56 "blocks": BlockBridge( 

57 name="model.layers", 

58 config=self.cfg, 

59 submodules={ 

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

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

62 # Heterogeneous head counts + softplus output gating: 

63 # delegate; projections and the gate are hookable. 

64 "attn": AttentionBridge( 

65 name="self_attn", 

66 config=self.cfg, 

67 submodules={ 

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

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

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

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

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

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

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

75 }, 

76 maintain_native_attention=True, 

77 ), 

78 # Dense or sparse per mlp_layer_types; router and shared 

79 # experts optional so dense layers skip them. 

80 "mlp": MoEBridge( 

81 name="mlp", 

82 config=self.cfg, 

83 sparse_required=("gate",), 

84 submodules={ 

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

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

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

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

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

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

91 }, 

92 ), 

93 }, 

94 ), 

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

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

97 } 

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)