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 # LayerNorm1P stores gamma as an offset from 1; subclasses with plain LayerNorm set False. 

29 _layernorm_weight_offset = True 

30 

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

32 """Initialize the Nemotron architecture adapter.""" 

33 super().__init__(cfg) 

34 

35 self.cfg.normalization_type = "LN" 

36 self.cfg.positional_embedding_type = "rotary" 

37 self.cfg.final_rms = False 

38 self.cfg.gated_mlp = False 

39 self.cfg.attn_only = False 

40 

41 offset = self._layernorm_weight_offset 

42 # Standard LN folding would fold the stored weight without the offset. Keep raw weights. 

43 self.supports_fold_ln = not offset 

44 self.supports_center_writing_weights = not offset 

45 # Read by analysis tools (DLA, path patching, SVD); the bridge's own norm paths 

46 # use the per-norm uses_weight_offset below. 

47 self.cfg.rmsnorm_uses_offset = offset 

48 

49 self.weight_processing_conversions = { 

50 **self._qkvo_weight_conversions(), 

51 } 

52 

53 self.component_mapping = { 

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

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

56 "blocks": BlockBridge( 

57 name="model.layers", 

58 submodules={ 

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

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

61 "ln1": NormalizationBridge( 

62 name="input_layernorm", 

63 config=self.cfg, 

64 use_native_layernorm_autograd=True, 

65 uses_weight_offset=offset, 

66 ), 

67 "attn": PositionEmbeddingsAttentionBridge( 

68 name="self_attn", 

69 config=self.cfg, 

70 submodules={ 

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

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

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

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

75 }, 

76 requires_attention_mask=True, 

77 requires_position_embeddings=True, 

78 ), 

79 "ln2": NormalizationBridge( 

80 name="post_attention_layernorm", 

81 config=self.cfg, 

82 use_native_layernorm_autograd=True, 

83 uses_weight_offset=offset, 

84 ), 

85 "mlp": self._ungated_mlp(), 

86 }, 

87 ), 

88 "ln_final": NormalizationBridge( 

89 name="model.norm", 

90 config=self.cfg, 

91 use_native_layernorm_autograd=True, 

92 uses_weight_offset=offset, 

93 ), 

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

95 }