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 # The bridge's own norm paths (hook edits, backward hooks, LN-rule) need the offset too. 

44 self.cfg.rmsnorm_uses_offset = True 

45 

46 self.weight_processing_conversions = { 

47 **self._qkvo_weight_conversions(), 

48 } 

49 

50 self.component_mapping = { 

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

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

53 "blocks": BlockBridge( 

54 name="model.layers", 

55 submodules={ 

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

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

58 "ln1": NormalizationBridge( 

59 name="input_layernorm", 

60 config=self.cfg, 

61 use_native_layernorm_autograd=True, 

62 ), 

63 "attn": PositionEmbeddingsAttentionBridge( 

64 name="self_attn", 

65 config=self.cfg, 

66 submodules={ 

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

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

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

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

71 }, 

72 requires_attention_mask=True, 

73 requires_position_embeddings=True, 

74 ), 

75 "ln2": NormalizationBridge( 

76 name="post_attention_layernorm", 

77 config=self.cfg, 

78 use_native_layernorm_autograd=True, 

79 ), 

80 "mlp": self._ungated_mlp(), 

81 }, 

82 ), 

83 "ln_final": NormalizationBridge( 

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

85 ), 

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

87 }