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

1"""HunYuanDenseV1 architecture adapter.""" 

2 

3from typing import Any 

4 

5from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

6from transformer_lens.model_bridge.generalized_components import ( 

7 BlockBridge, 

8 EmbeddingBridge, 

9 LinearBridge, 

10 PositionEmbeddingsAttentionBridge, 

11 RMSNormalizationBridge, 

12 RotaryEmbeddingBridge, 

13 UnembeddingBridge, 

14) 

15 

16 

17class HunYuanDenseV1ArchitectureAdapter(ArchitectureAdapter): 

18 """Architecture adapter for HunYuanDenseV1 models.""" 

19 

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

21 super().__init__(cfg) 

22 

23 self._set_rms_rotary_defaults() 

24 

25 self.cfg.qk_norm_after_rope = True 

26 

27 self.cfg.attn_implementation = "eager" 

28 

29 self.weight_processing_conversions = { 

30 **self._qkvo_weight_conversions(), 

31 } 

32 

33 self.component_mapping = { 

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

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

36 "blocks": BlockBridge( 

37 name="model.layers", 

38 submodules={ 

39 "ln1": RMSNormalizationBridge( 

40 name="input_layernorm", 

41 config=self.cfg, 

42 ), 

43 "ln2": RMSNormalizationBridge( 

44 name="post_attention_layernorm", 

45 config=self.cfg, 

46 ), 

47 "attn": PositionEmbeddingsAttentionBridge( 

48 name="self_attn", 

49 config=self.cfg, 

50 submodules={ 

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

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

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

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

55 "q_norm": RMSNormalizationBridge( 

56 name="query_layernorm", config=self.cfg 

57 ), 

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

59 }, 

60 requires_attention_mask=True, 

61 requires_position_embeddings=True, 

62 ), 

63 "mlp": self._gated_mlp(), 

64 }, 

65 ), 

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

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

68 }