Coverage for transformer_lens/model_bridge/supported_architectures/exaone.py: 48%

1"""EXAONE architecture adapter. 

2 

3Supports LG AI Research's EXAONE-3.0 / 3.5 / Deep families 

4(``ExaoneForCausalLM``, trust_remote_code checkpoints). Llama-style RMSNorm + 

5RoPE + GQA + gated MLP under GPT-2-flavored module names: ``transformer.wte``, 

6``transformer.h[i]``, ``transformer.ln_f``, and a double-nested attention 

7(``attn.attention``). EXAONE-4.0 is a separate native-transformers 

8architecture (Exaone4ForCausalLM) and is not covered here. 

9""" 

10 

11from typing import Any 

12 

13from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

14from transformer_lens.model_bridge.generalized_components import ( 

15 BlockBridge, 

16 EmbeddingBridge, 

17 LinearBridge, 

18 PositionEmbeddingsAttentionBridge, 

19 RMSNormalizationBridge, 

20 RotaryEmbeddingBridge, 

21 UnembeddingBridge, 

22) 

23 

24 

25class ExaoneArchitectureAdapter(ArchitectureAdapter): 

26 """Architecture adapter for ExaoneForCausalLM (EXAONE-3.x) models. 

27 

28 The remote modeling code follows current HF conventions (Cache API, 

29 position_embeddings tuples), so the standard bridges delegate cleanly. 

30 Naming quirks: attention projections live one level deeper than usual 

31 (``attn.attention.q_proj``), the gated MLP uses ``c_fc_0`` (gate) / 

32 ``c_fc_1`` (up) / ``c_proj`` (down), and rotary sits at 

33 ``transformer.rotary``. 

34 """ 

35 

36 _testing_lm_attr = "transformer" 

37 _testing_rotary_attr = "rotary" 

38 _testing_eager = "config" 

39 

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

41 """Initialize the EXAONE architecture adapter.""" 

42 super().__init__(cfg) 

43 

44 self._set_rms_rotary_defaults() 

45 # Verified against LGAI-EXAONE/EXAONE-3.5-2.4B-Instruct: no BOS prepended. 

46 self.cfg.default_prepend_bos = False 

47 

48 self.weight_processing_conversions = { 

49 **self._qkvo_weight_conversions(), 

50 } 

51 

52 self.component_mapping = { 

53 "embed": EmbeddingBridge(name="transformer.wte"), 

54 "rotary_emb": RotaryEmbeddingBridge(name="transformer.rotary", config=self.cfg), 

55 "blocks": BlockBridge( 

56 name="transformer.h", 

57 submodules={ 

58 "ln1": RMSNormalizationBridge(name="ln_1", config=self.cfg), 

59 "ln2": RMSNormalizationBridge(name="ln_2", config=self.cfg), 

60 "attn": PositionEmbeddingsAttentionBridge( 

61 name="attn.attention", 

62 config=self.cfg, 

63 submodules={ 

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

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

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

67 "o": LinearBridge(name="out_proj"), 

68 }, 

69 requires_attention_mask=True, 

70 requires_position_embeddings=True, 

71 ), 

72 "mlp": self._gated_mlp(gate="c_fc_0", up="c_fc_1", down="c_proj"), 

73 }, 

74 ), 

75 "ln_final": RMSNormalizationBridge(name="transformer.ln_f", config=self.cfg), 

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

77 } 

78 

79 def prepare_model(self, hf_model: Any) -> Any: 

80 """Shim the EXAONE-3.x remote module for transformers >= 5.13, which renamed 

81 create_causal_mask's ``input_embeds`` kwarg to ``inputs_embeds``.""" 

82 result = super().prepare_model(hf_model) 

83 model = result if result is not None else hf_model 

84 

85 import sys 

86 

87 from transformers.masking_utils import create_causal_mask 

88 

89 module = sys.modules.get(type(model).__module__) 

90 if module is not None and hasattr(module, "create_causal_mask"): 

91 

92 def _shim(*args: Any, **kwargs: Any) -> Any: 

93 if "input_embeds" in kwargs: 

94 kwargs["inputs_embeds"] = kwargs.pop("input_embeds") 

95 kwargs.pop("cache_position", None) 

96 return create_causal_mask(*args, **kwargs) 

97 

98 setattr(module, "create_causal_mask", _shim) 

99 return result