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

1"""Starcoder2 architecture adapter. 

2 

3BigCode's StarCoder2 (``Starcoder2ForCausalLM``): pre-norm decoder with 

4plain LayerNorm (not RMS), separate biased q/k/v/o projections, GQA, RoPE, 

5and a non-gated ``c_fc``/``c_proj`` MLP. 

6 

7Not a drop-in for its GPTBigCode predecessor: that one fuses q/k/v into a 

8single ``c_attn`` and uses learned positions, so it needs a different bridge. 

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 NormalizationBridge, 

19 PositionEmbeddingsAttentionBridge, 

20 RotaryEmbeddingBridge, 

21 UnembeddingBridge, 

22) 

23 

24 

25class Starcoder2ArchitectureAdapter(ArchitectureAdapter): 

26 """Architecture adapter for Starcoder2ForCausalLM models.""" 

27 

28 _testing_eager = "config" 

29 

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

31 """Initialize the Starcoder2 architecture adapter.""" 

32 super().__init__(cfg) 

33 

34 self.cfg.normalization_type = "LN" 

35 self.cfg.positional_embedding_type = "rotary" 

36 self.cfg.final_rms = False 

37 self.cfg.gated_mlp = False 

38 self.cfg.attn_only = False 

39 

40 # StarCoder2's BOS *is* its EOS (<|endoftext|>) and its tokenizer never prepemds one 

41 self.cfg.default_prepend_bos = False 

42 

43 # StarCoder2 biases every q/k/v/o projection; the bias reshapes must use 

44 # the kv-head count or compat mode mis-shapes (silent) or crashes (GQA). 

45 self.weight_processing_conversions = { 

46 **self._qkvo_weight_conversions(include_biases=True), 

47 } 

48 

49 self.component_mapping = { 

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

51 "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), 

52 "blocks": BlockBridge( 

53 name="model.layers", 

54 submodules={ 

55 "ln1": NormalizationBridge(name="input_layernorm", config=self.cfg), 

56 "attn": PositionEmbeddingsAttentionBridge( 

57 name="self_attn", 

58 config=self.cfg, 

59 submodules={ 

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

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

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

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

64 }, 

65 requires_attention_mask=True, 

66 requires_position_embeddings=True, 

67 ), 

68 "ln2": NormalizationBridge(name="post_attention_layernorm", config=self.cfg), 

69 "mlp": self._ungated_mlp(up="c_fc", down="c_proj"), 

70 }, 

71 ), 

72 "ln_final": NormalizationBridge(name="model.norm", config=self.cfg), 

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

74 }