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

1"""Nemotron architecture adapter. 

2 

3NVIDIA's dense Nemotron-3/4 and Minitron families (``NemotronForCausalLM``): 

4llama-shaped GQA decoder with three quirks — LayerNorm1P normalization 

5(zero-centered gamma, applied as ``weight + 1``), a non-gated squared-ReLU 

6MLP (``up_proj``/``down_proj`` only), and partial rotary embeddings 

7(rotary factor 0.5). 

8""" 

9 

10from typing import Any 

11 

12from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter 

13from transformer_lens.model_bridge.generalized_components import ( 

14 BlockBridge, 

15 EmbeddingBridge, 

16 LinearBridge, 

17 NormalizationBridge, 

18 PositionEmbeddingsAttentionBridge, 

19 RotaryEmbeddingBridge, 

20 UnembeddingBridge, 

21) 

22 

23 

24class NemotronArchitectureAdapter(ArchitectureAdapter): 

25 """Architecture adapter for NemotronForCausalLM models.""" 

26 

27 _testing_eager = "config" 

28 

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

30 """Initialize the Nemotron architecture adapter.""" 

31 super().__init__(cfg) 

32 

33 self.cfg.normalization_type = "LN" 

34 self.cfg.positional_embedding_type = "rotary" 

35 self.cfg.final_rms = False 

36 self.cfg.gated_mlp = False 

37 self.cfg.attn_only = False 

38 

39 # LayerNorm1P applies gamma as (weight + 1); standard LN folding would 

40 # fold the stored weight without the offset. Keep raw weights. 

41 self.supports_fold_ln = False 

42 self.supports_center_writing_weights = False 

43 

44 self.weight_processing_conversions = { 

45 **self._qkvo_weight_conversions(), 

46 } 

47 

48 self.component_mapping = { 

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

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

51 "blocks": BlockBridge( 

52 name="model.layers", 

53 submodules={ 

54 # LayerNorm1P applies gamma as (weight + 1); delegate to the 

55 # native module rather than reimplementing LN without the offset. 

56 "ln1": NormalizationBridge( 

57 name="input_layernorm", 

58 config=self.cfg, 

59 use_native_layernorm_autograd=True, 

60 ), 

61 "attn": PositionEmbeddingsAttentionBridge( 

62 name="self_attn", 

63 config=self.cfg, 

64 submodules={ 

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

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

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

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

69 }, 

70 requires_attention_mask=True, 

71 requires_position_embeddings=True, 

72 ), 

73 "ln2": NormalizationBridge( 

74 name="post_attention_layernorm", 

75 config=self.cfg, 

76 use_native_layernorm_autograd=True, 

77 ), 

78 "mlp": self._ungated_mlp(), 

79 }, 

80 ), 

81 "ln_final": NormalizationBridge( 

82 name="model.norm", config=self.cfg, use_native_layernorm_autograd=True 

83 ), 

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

85 }