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

1"""Muse Glimmer (MuseGlimmerForConditionalGeneration) architecture adapter.""" 

2 

3from typing import Any 

4 

5import torch 

6 

7from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

8from transformer_lens.model_bridge.generalized_components import ( 

9 AttentionBridge, 

10 BlockBridge, 

11 EmbeddingBridge, 

12 LinearBridge, 

13 RotaryEmbeddingBridge, 

14 UnembeddingBridge, 

15 VisionProjectionBridge, 

16) 

17from transformer_lens.model_bridge.generalized_components.base import ( 

18 GeneralizedComponent, 

19) 

20 

21 

22class MuseGlimmerArchitectureAdapter(ArchitectureAdapter): 

23 """Architecture adapter for MuseGlimmerForConditionalGeneration models.""" 

24 

25 _testing_lm_attr = "model.language_model" 

26 _testing_wire_rotary = False 

27 

28 # Attention output is gated by sigmoid(gate_proj(ln1_out)); fold-LN rescales only 

29 # q/k/v, so the unfolded gate path would change — Gemma 2's sandwich norms alone 

30 # fold fine. 

31 supports_fold_ln = False 

32 

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

34 super().__init__(cfg) 

35 

36 self.cfg.is_multimodal = True 

37 self._extract_vision_dims(cfg) 

38 

39 self._set_rms_rotary_defaults() 

40 self.cfg.attn_implementation = "eager" 

41 self.weight_processing_conversions: dict = {} 

42 

43 self.component_mapping = { 

44 "vision_encoder": GeneralizedComponent(name="model.vision_tower"), 

45 # HF runs ``vision_adapter -> vision_projection -> perception_emb_norm`` and 

46 # scatters only the last value into the text embeddings, so the projector 

47 # bridge wraps the norm: ``hook_out`` is the merged embedding, as in the 

48 # other VLM adapters (Gemma 3's norm lives inside its projector module). 

49 # ``hook_in`` is then the projection output rather than the vision-tower one. 

50 "vision_projector": VisionProjectionBridge(name="model.perception_emb_norm"), 

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

52 "rotary_emb": RotaryEmbeddingBridge(name="model.language_model.rotary_emb"), 

53 "blocks": BlockBridge( 

54 name="model.language_model.layers", 

55 submodules={ 

56 "ln1": GeneralizedComponent(name="input_layernorm"), 

57 "ln1_post": GeneralizedComponent(name="post_attention_layernorm"), 

58 "ln2": GeneralizedComponent(name="pre_feedforward_layernorm"), 

59 "ln2_post": GeneralizedComponent(name="post_feedforward_layernorm"), 

60 "attn": AttentionBridge( 

61 name="self_attn", 

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="o_proj"), 

68 "gate": LinearBridge(name="gate_proj"), 

69 }, 

70 maintain_native_attention=True, 

71 requires_attention_mask=True, 

72 ), 

73 "mlp": self._gated_mlp(), 

74 }, 

75 ), 

76 "ln_final": GeneralizedComponent(name="model.language_model.norm"), 

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

78 } 

79 

80 def apply_output_logits_transform(self, logits: torch.Tensor) -> torch.Tensor: 

81 multiplier = float(getattr(self.cfg, "output_multiplier", 1.0)) 

82 return super().apply_output_logits_transform(logits * multiplier)