Coverage for transformer_lens/model_bridge/supported_architectures/exaone4.py: 90%

1"""EXAONE 4.0 architecture adapter. 

2 

3LG AI Research's EXAONE-4.0 (``Exaone4ForCausalLM``, native transformers — 

4distinct from the remote-code EXAONE-3.x family in exaone.py): the OLMo-2 

5post-norm block shape (per-head Q/K RMSNorms, post_attention_layernorm / 

6post_feedforward_layernorm inside the residual branch) plus hybrid 

7sliding/global attention (``layer_types``) with global-NoPE gating. 

8""" 

9 

10from typing import Any 

11 

12import torch 

13 

14from transformer_lens.model_bridge.generalized_components import ( 

15 PositionEmbeddingsAttentionBridge, 

16) 

17from transformer_lens.model_bridge.supported_architectures.olmo2 import ( 

18 Olmo2ArchitectureAdapter, 

19) 

20 

21 

22class _Exaone4AttentionBridge(PositionEmbeddingsAttentionBridge): 

23 """Suppress RoPE on hybrid full-attention (global NoPE) layers. 

24 

25 HF gates rotation on ``sliding_window is None or is_sliding``; the base 

26 bridge rotates whenever position_embeddings is present, so NoPE layers 

27 null the argument first. Non-hybrid checkpoints rotate everywhere. 

28 """ 

29 

30 # Nulls position_embeddings on NoPE layers by design. 

31 rope_optional = True 

32 

33 def forward(self, *args: Any, **kwargs: Any) -> Any: 

34 """Drop position_embeddings on hybrid full-attention NoPE layers.""" 

35 if self._is_nope_layer(): 

36 kwargs["position_embeddings"] = None 

37 if len(args) >= 2 and not isinstance(args[1], torch.Tensor): 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 args = (args[0], None) + args[2:] 

39 return super().forward(*args, **kwargs) 

40 

41 def _is_nope_layer(self) -> bool: 

42 """Return True when the wrapped attention is a hybrid model's full-attention layer.""" 

43 hf_attn = self.original_component 

44 if hf_attn is None: 44 ↛ 45line 44 didn't jump to line 45 because the condition on line 44 was never true

45 return False 

46 if getattr(hf_attn, "sliding_window", None) is None: 

47 return False 

48 return not getattr(hf_attn, "is_sliding", True) 

49 

50 

51class Exaone4ArchitectureAdapter(Olmo2ArchitectureAdapter): 

52 """Architecture adapter for Exaone4ForCausalLM models.""" 

53 

54 _attention_bridge_cls = _Exaone4AttentionBridge 

55 _testing_eager = "config" 

56 

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

58 """Initialize the EXAONE 4.0 architecture adapter.""" 

59 super().__init__(cfg) 

60 

61 # Same tokenizer family as EXAONE-3.x: no BOS prepending. 

62 self.cfg.default_prepend_bos = False 

63 self.supports_center_writing_weights = False 

64 

65 layer_types = getattr(cfg, "layer_types", None) 

66 if layer_types: 

67 setattr(self.cfg, "layer_types", list(layer_types))