Coverage for transformer_lens/model_bridge/supported_architectures/nemotron_h.py: 100%
32 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Nemotron-H hybrid Mamba2-Transformer architecture adapter.
3Supports NemotronHForCausalLM (e.g. nvidia/NVIDIA-Nemotron-Nano-9B-v2, Nemotron-3 series).
5Architecture overview:
6- Heterogeneous layers defined by ``config.layers_block_type`` — each element is
7 one of ``"mamba"``, ``"attention"``, ``"moe"``, or ``"mlp"``.
8- ~8% of layers are standard GQA attention; the rest are Mamba-2 SSM, dense MLP,
9 or sparse MoE. All share a single pre-norm (``block.norm``) and a single residual
10 path; there is no ``ln2`` or post-attention norm.
11- Each block exposes a single ``.mixer`` attribute whose type varies by layer.
12- No model-level rotary embedding module — attention handles RoPE internally via
13 ``position_ids`` passed from the outer model loop.
14- Stateful generation: uses ``DynamicCache`` (transformers ≥ 5.12) which carries
15 both KV-cache entries (attention layers) and SSM conv/recurrent states
16 (Mamba layers) in a unified object.
18Key adapter decisions:
19- ``SSMBlockBridge`` is used as the block container. It delegates the entire
20 forward to the HF block, giving ``hook_in`` / ``hook_out`` on the residual
21 stream without hardcoding transformer-specific hook positions (hook_resid_mid,
22 hook_mlp_in, etc.) that do not exist in this single-norm architecture.
23- ``SSM2MixerBridge`` wraps ``.mixer`` for all layer types. Its forward is a
24 pure passthrough (``original_component(*args, **kwargs)``) so it works
25 correctly for attention, MLP, and MoE mixers as well as Mamba ones.
26 Mamba-specific inner submodules (in_proj, conv1d, inner_norm, out_proj) are
27 declared ``optional=True`` so setup skips them gracefully on non-Mamba layers.
28- MLP layers use ``relu2`` activation (not SwiGLU); ``gated_mlp = False``.
29- ``applicable_phases = [1, 2, 3, 4]``: P1 is exact vs raw HF (passthrough mixers);
30 P2/P3 skip without a HookedTransformer; P4 is generation.
31"""
33from typing import Any
35from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
36from transformer_lens.model_bridge.generalized_components import (
37 DepthwiseConv1DBridge,
38 EmbeddingBridge,
39 GatedRMSNormBridge,
40 LinearBridge,
41 RMSNormalizationBridge,
42 SSM2MixerBridge,
43 SSMBlockBridge,
44 UnembeddingBridge,
45)
46from transformer_lens.model_bridge.generalized_components.base import (
47 GeneralizedComponent,
48)
51def _make_optional(component: "GeneralizedComponent") -> "GeneralizedComponent":
52 """Mark a GeneralizedComponent submodule as optional.
54 Some bridge classes (e.g. GatedRMSNormBridge) do not forward ``optional``
55 through their own ``__init__``, even though ``GeneralizedComponent`` supports
56 it. Setting the attribute directly is safe because ``component_setup.py``
57 reads ``getattr(submodule, 'optional', False)`` at setup time.
58 """
59 component.optional = True
60 return component
63class NemotronHArchitectureAdapter(ArchitectureAdapter):
64 """Architecture adapter for NemotronHForCausalLM.
66 Hybrid Mamba-2 + Attention + MoE + dense MLP model. All layers share a
67 single pre-norm and a single residual connection; the mixer type per layer
68 is determined by ``config.layers_block_type[layer_idx]``.
69 """
71 # White-box forward: P1 is exact vs raw HF (passthrough mixers); P2/P3 skip
72 # without a HookedTransformer; P4 is generation.
73 applicable_phases: list[int] = [1, 2, 3, 4]
75 def __init__(self, cfg: Any) -> None:
76 super().__init__(cfg)
78 self.cfg.normalization_type = "RMS"
79 self.cfg.uses_rms_norm = True
80 # No model-level rotary embedding module — attention handles RoPE
81 # internally via position_ids; set to "none" so the bridge does not
82 # attempt to wire a rotary_emb component.
83 self.cfg.positional_embedding_type = "none"
84 # MLP layers use relu2 (up_proj → act → down_proj), not SwiGLU.
85 self.cfg.gated_mlp = False
86 self.cfg.attn_only = False
87 self.cfg.final_rms = True
88 # Mamba layers require per-step SSM state; generation is stateful.
89 self.cfg.is_stateful = True
91 # Normalize the per-layer type list as cfg.layers_block_type (HF names it
92 # `layer_types`) so analysis tools can find the Mamba layers, as on Granite.
93 setattr(self.cfg, "layers_block_type", self._canonical_layer_types(cfg))
95 # Mamba-2 dimensional config (mirrors Mamba2ArchitectureAdapter).
96 mamba_num_heads = getattr(cfg, "mamba_num_heads", 128)
97 mamba_head_dim = getattr(cfg, "mamba_head_dim", 64)
98 mamba_intermediate_size = mamba_num_heads * mamba_head_dim
99 n_groups = getattr(cfg, "n_groups", 8)
100 ssm_state_size = getattr(cfg, "ssm_state_size", 128)
101 conv_dim = mamba_intermediate_size + 2 * n_groups * ssm_state_size
102 setattr(self.cfg, "mamba_intermediate_size", mamba_intermediate_size)
103 setattr(self.cfg, "conv_dim", conv_dim)
105 self.weight_processing_conversions = {}
107 self.component_mapping = {
108 "embed": EmbeddingBridge(name="model.embeddings"),
109 "blocks": SSMBlockBridge(
110 name="model.layers",
111 submodules={
112 # Single pre-norm shared across all layer types.
113 "norm": RMSNormalizationBridge(name="norm", config=self.cfg),
114 # Single mixer slot — type varies per layer (mamba / attention
115 # / moe / mlp). SSM2MixerBridge.forward() is a pure
116 # passthrough so it works for all four types. Mamba-specific
117 # inner submodules are optional and skipped on other types.
118 "mixer": SSM2MixerBridge(
119 name="mixer",
120 config=self.cfg,
121 submodules={
122 # ── Mamba-only (optional on attention / moe / mlp) ──
123 "in_proj": LinearBridge(name="in_proj", optional=True),
124 "conv1d": DepthwiseConv1DBridge(name="conv1d", optional=True),
125 # HF names this "norm" inside the mixer; TL calls it
126 # "inner_norm" to avoid collision with the block-level norm.
127 # GatedRMSNormBridge.__init__ does not accept optional=, so
128 # we set the attribute directly after construction.
129 "inner_norm": _make_optional(GatedRMSNormBridge(name="norm")),
130 "out_proj": LinearBridge(name="out_proj", optional=True),
131 },
132 ),
133 },
134 ),
135 "ln_final": RMSNormalizationBridge(name="model.norm_f", config=self.cfg),
136 "unembed": UnembeddingBridge(name="lm_head"),
137 }
139 def create_stateful_cache(
140 self,
141 hf_model: Any,
142 batch_size: int,
143 device: Any,
144 dtype: Any,
145 ) -> Any:
146 """Build the unified DynamicCache for stateful generation.
148 Transformers ≥ 5.12 ships a unified ``DynamicCache`` that carries both
149 KV-cache entries (attention layers) and SSM conv/recurrent states
150 (Mamba layers) in a single object, using ``has_previous_state()`` to
151 distinguish which state is available for a given layer index. The
152 config is required so the cache knows each layer's type — matching
153 NemotronHModel's own initialization.
154 """
155 from transformers.cache_utils import DynamicCache
157 return DynamicCache(config=hf_model.config)