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

1"""XGLM architecture adapter. 

2 

3Supports XGLMForCausalLM (facebook/xglm-*). 

4Assumes add_cross_attention=False (all published XGLM checkpoints). 

5""" 

6 

7from typing import Any 

8 

9from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

10from transformer_lens.model_bridge.generalized_components import ( 

11 AttentionBridge, 

12 BlockBridge, 

13 EmbeddingBridge, 

14 LinearBridge, 

15 MLPBridge, 

16 NormalizationBridge, 

17 UnembeddingBridge, 

18) 

19 

20 

21class XGLMArchitectureAdapter(ArchitectureAdapter): 

22 """Architecture adapter for XGLM models. 

23 

24 XGLM uses pre-norm LayerNorm, sinusoidal positional embeddings (no 

25 learnable weights), standard MHA with separate q/k/v/out_proj, and a 

26 2-layer MLP (fc1/fc2) that lives directly on the decoder block rather 

27 than inside an mlp sub-module. 

28 

29 All attention projections and fc1/fc2 carry biases. lm_head has no bias. 

30 Embeddings are scaled by sqrt(d_model) at runtime in XGLMScaledWordEmbedding. 

31 

32 Optional Parameters (may not exist in state_dict): 

33 -------------------------------------------------- 

34 None — all published XGLM checkpoints include all parameters listed above. 

35 """ 

36 

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

38 """Initialize the XGLM architecture adapter.""" 

39 super().__init__(cfg) 

40 

41 # LayerNorm throughout (not RMSNorm) 

42 self.cfg.normalization_type = "LN" 

43 # Sinusoidal positional embeddings — added to token embeddings before blocks, 

44 # no learnable weights, no RoPE 

45 self.cfg.positional_embedding_type = "standard" 

46 self.cfg.final_rms = False 

47 # Standard 2-layer MLP (fc1 -> gelu -> fc2), no gate projection 

48 self.cfg.gated_mlp = False 

49 self.cfg.attn_only = False 

50 self.cfg.uses_rms_norm = False 

51 

52 # Sinusoidal positional embeddings have no weights in the state_dict, so 

53 # center_writing_weights cannot center pos_embed. Disable it for XGLM. 

54 self.supports_center_writing_weights = False 

55 

56 # Standard MHA: n_heads == n_kv_heads for all XGLM sizes 

57 self.weight_processing_conversions = { 

58 **self._qkvo_weight_conversions(), 

59 } 

60 

61 self.component_mapping = { 

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

63 # No "pos_embed": sinusoidal embeddings are a non-persistent buffer with 

64 # no learnable weights — embed_positions does not appear in state_dict. 

65 "blocks": BlockBridge( 

66 name="model.layers", 

67 # fc2 IS the mlp output (no container fires hook_out). No 

68 # hook_mlp_in override: pre-norm, the block already provides it. 

69 hook_alias_overrides={"hook_mlp_out": "mlp.out.hook_out"}, 

70 submodules={ 

71 "ln1": NormalizationBridge( 

72 name="self_attn_layer_norm", # pre-attn norm on XGLMDecoderLayer 

73 config=self.cfg, 

74 use_native_layernorm_autograd=True, 

75 ), 

76 "attn": AttentionBridge( 

77 name="self_attn", 

78 config=self.cfg, 

79 requires_attention_mask=True, 

80 attention_mask_4d=True, # (batch, 1, tgt_len, src_len) 

81 submodules={ 

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

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

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

85 "o": LinearBridge(name="out_proj"), # out_proj, not o_proj 

86 }, 

87 ), 

88 "ln2": NormalizationBridge( 

89 name="final_layer_norm", # pre-MLP norm on XGLMDecoderLayer 

90 config=self.cfg, 

91 use_native_layernorm_autograd=True, 

92 ), 

93 # fc1/fc2 live directly on XGLMDecoderLayer — no "mlp" container. 

94 # Containerless fc1/fc2, as BERT (XGLM stays 3D — no 

95 # unflatten needed, unlike OPT). 

96 "mlp": MLPBridge( 

97 name=None, 

98 config=self.cfg, 

99 submodules={ 

100 "in": LinearBridge(name="fc1"), 

101 "out": LinearBridge(name="fc2"), 

102 }, 

103 ), 

104 }, 

105 ), 

106 "ln_final": NormalizationBridge( 

107 name="model.layer_norm", # note: layer_norm, not norm 

108 config=self.cfg, 

109 use_native_layernorm_autograd=True, 

110 ), 

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

112 }