Coverage for transformer_lens/model_bridge/supported_architectures/jamba.py: 98%
51 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-08-11 18:50 +0000
1"""Jamba hybrid attention+Mamba architecture adapter.
3Supports ``JambaForCausalLM`` (e.g. ``ai21labs/Jamba-tiny-random``,
4``ai21labs/AI21-Jamba-Reasoning-3B``, ``ai21labs/Jamba-v0.1``).
6Architecture overview:
7- Heterogeneous layers from ``config.layers_block_type`` — each element is
8 either ``"mamba"`` (``JambaMambaDecoderLayer``) or ``"attention"``
9 (``JambaAttentionDecoderLayer``). Attention layers recur every
10 ``attn_layer_period`` starting at ``attn_layer_offset`` (classically 1/8).
11- Every layer has the same residual skeleton: pre-norm (``input_layernorm``) →
12 mixer (``.mamba`` *or* ``.self_attn``) → residual → pre-FFN norm
13 (``pre_ff_layernorm``) → ``.feed_forward`` (dense SwiGLU ``JambaMLP`` or
14 ``JambaSparseMoeBlock`` when ``layers_num_experts[i] > 1``) → residual.
15- The Mamba mixer is **Mamba-1** (``JambaMambaMixer``): ``in_proj`` / ``conv1d``
16 / ``x_proj`` / ``dt_proj`` / ``out_proj``, plus Jamba-specific
17 ``dt_layernorm`` / ``b_layernorm`` / ``c_layernorm`` on the selective params.
18- Attention is GQA **without RoPE** — absolute-position-free; no model-level
19 rotary module.
20- Generation threads a unified ``DynamicCache`` via ``past_key_values``
21 (attention KV + Mamba conv/recurrent states). The Mamba mixer receives that
22 same object as ``cache_params=past_key_values``.
24Key adapter decisions:
25- ``BlockBridge`` (not ``SSMBlockBridge``): two norms and a real post-mixer
26 residual make transformer-shaped hooks (``hook_resid_mid`` via ``ln2``)
27 meaningful — same choice as Falcon-H1 / GraniteMoeHybrid.
28- ``SSMMixerBridge`` (PR #1481 Mamba-1 interp surface) wraps ``.mamba`` under
29 the canonical ``.mixer`` dict key so ``find_ssm_mixer`` /
30 ``compute_effective_attention`` / ``eager_scan`` resolve it. Inner
31 projections match Mamba-1; the three selective-param RMSNorms are mapped
32 so reconstruction and the opt-in ``eager_scan`` path apply them (Jamba
33 fork of stock Mamba-1). Default forward still delegates to HF
34 (bit-identical).
35- Attention and mixer are ``optional=True`` so setup skips the absent branch
36 per layer type.
37- FFN: dense ``GatedMLPBridge`` when ``num_experts <= 1`` (Reasoning-3B);
38 ``MoEBridge`` passthrough with optional dense projections *and* router when
39 ``num_experts > 1`` (tiny-random / 52B), since dense and MoE layers share
40 the HF path ``.feed_forward``.
41- ``is_stateful = False``: generation uses the standard ``past_key_values``
42 path (same rationale as Zamba2 / Falcon-H1). Setting ``is_stateful`` would
43 route through the pure-Mamba ``cache_params`` loop and diverge.
44"""
46from typing import Any
48from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
49from transformer_lens.model_bridge.generalized_components import (
50 AttentionBridge,
51 BlockBridge,
52 DepthwiseConv1DBridge,
53 EmbeddingBridge,
54 GatedMLPBridge,
55 LinearBridge,
56 MoEBridge,
57 RMSNormalizationBridge,
58 SSMMixerBridge,
59 UnembeddingBridge,
60)
61from transformer_lens.model_bridge.generalized_components.base import (
62 GeneralizedComponent,
63)
66def _make_optional(component: GeneralizedComponent) -> GeneralizedComponent:
67 """Mark a GeneralizedComponent submodule as optional.
69 Some bridge classes (e.g. ``RMSNormalizationBridge``) do not forward
70 ``optional`` through their own ``__init__``. ``component_setup.py`` reads
71 ``getattr(submodule, "optional", False)`` at setup time.
72 """
73 component.optional = True
74 return component
77class JambaArchitectureAdapter(ArchitectureAdapter):
78 """Architecture adapter for ``JambaForCausalLM``.
80 Interleaved attention + Mamba-1 layers with optional sparse MoE FFN.
81 Attention and Mamba streams are separate optional slots so each can be
82 ablated independently.
83 """
85 # P1: exact passthrough vs raw HF; P2/P3 skip HT comparison; P4 generation.
86 applicable_phases: list[int] = [1, 2, 3, 4]
88 def __init__(self, cfg: Any) -> None:
89 super().__init__(cfg)
91 self.cfg.normalization_type = "RMS"
92 self.cfg.uses_rms_norm = True
93 # No RoPE — JambaAttention is position-embedding-free.
94 self.cfg.positional_embedding_type = "none"
95 self.cfg.final_rms = True
96 # Dense FFN is SwiGLU (gate_proj / up_proj / down_proj, silu).
97 self.cfg.gated_mlp = True
98 self.cfg.attn_only = False
99 # Standard past_key_values DynamicCache path (not pure-Mamba cache_params).
100 self.cfg.is_stateful = False
102 if hasattr(cfg, "n_key_value_heads") and cfg.n_key_value_heads is not None: 102 ↛ 106line 102 didn't jump to line 106 because the condition on line 102 was always true
103 self.cfg.n_key_value_heads = cfg.n_key_value_heads
105 # Tokenizer prepends BOS on __call__ for released AI21 checkpoints.
106 self.cfg.default_prepend_bos = True
108 # Per-layer mixer type list for analysis tools (same name as NemotronH /
109 # GraniteMoeHybrid / Zamba2). Prefer the HF property; fall back to empty.
110 layers_block_type = list(
111 getattr(cfg, "layers_block_type", None) or getattr(cfg, "layer_types", None) or []
112 )
113 # HF layer_types uses linear_attention/full_attention; normalize to TL names.
114 _LAYER_TYPE_TO_TL = {
115 "linear_attention": "mamba",
116 "full_attention": "attention",
117 "mamba": "mamba",
118 "attention": "attention",
119 }
120 setattr(
121 self.cfg,
122 "layers_block_type",
123 [_LAYER_TYPE_TO_TL.get(t, t) for t in layers_block_type],
124 )
126 # Mamba-1 dimensional config (mirrors MambaArchitectureAdapter fields
127 # that SSMMixerBridge.compute_* reads from the wrapped HF mixer; also
128 # surfaced on cfg for tooling).
129 mamba_expand = int(getattr(cfg, "mamba_expand", 2))
130 mamba_d_state = int(getattr(cfg, "mamba_d_state", 16))
131 mamba_d_conv = int(getattr(cfg, "mamba_d_conv", 4))
132 intermediate_size = mamba_expand * int(self.cfg.d_model)
133 setattr(self.cfg, "mamba_expand", mamba_expand)
134 setattr(self.cfg, "mamba_d_state", mamba_d_state)
135 setattr(self.cfg, "mamba_d_conv", mamba_d_conv)
136 setattr(self.cfg, "mamba_dt_rank", getattr(cfg, "mamba_dt_rank", None))
137 setattr(self.cfg, "intermediate_size", intermediate_size)
138 setattr(self.cfg, "state_size", mamba_d_state)
139 setattr(self.cfg, "conv_kernel", mamba_d_conv)
140 setattr(self.cfg, "expand", mamba_expand)
142 num_experts = getattr(cfg, "num_experts", None) or getattr(cfg, "num_local_experts", 1) or 1
143 setattr(self.cfg, "num_experts", num_experts)
145 # Heterogeneous attn/Mamba layers + native HF attention layout: LN folding
146 # expects rearranged QKV and uniform per-layer attn, so disable (same as
147 # GraniteMoeHybrid / other SSM hybrids).
148 self.supports_fold_ln = False
149 self.weight_processing_conversions = {}
150 self.component_mapping = self._build_component_mapping(num_experts=int(num_experts))
152 def _build_mamba_bridge(self) -> SSMMixerBridge:
153 """Mamba-1 mixer under canonical ``.mixer``; HF path is ``.mamba``."""
154 return SSMMixerBridge(
155 name="mamba",
156 config=self.cfg,
157 optional=True,
158 submodules={
159 "in_proj": LinearBridge(name="in_proj"),
160 "conv1d": DepthwiseConv1DBridge(name="conv1d"),
161 "x_proj": LinearBridge(name="x_proj"),
162 "dt_proj": LinearBridge(name="dt_proj"),
163 "out_proj": LinearBridge(name="out_proj"),
164 # Jamba-only selective-param norms (absent on stock Mamba-1).
165 "dt_layernorm": _make_optional(
166 RMSNormalizationBridge(name="dt_layernorm", config=self.cfg)
167 ),
168 "b_layernorm": _make_optional(
169 RMSNormalizationBridge(name="b_layernorm", config=self.cfg)
170 ),
171 "c_layernorm": _make_optional(
172 RMSNormalizationBridge(name="c_layernorm", config=self.cfg)
173 ),
174 },
175 )
177 def _build_attention_bridge(self) -> AttentionBridge:
178 """GQA attention without RoPE; keep HF's forward for cache parity."""
179 return AttentionBridge(
180 name="self_attn",
181 config=self.cfg,
182 optional=True,
183 maintain_native_attention=True,
184 requires_attention_mask=True,
185 requires_position_embeddings=False,
186 submodules={
187 "q": LinearBridge(name="q_proj"),
188 "k": LinearBridge(name="k_proj"),
189 "v": LinearBridge(name="v_proj"),
190 "o": LinearBridge(name="o_proj"),
191 },
192 )
194 def _build_ffn_bridge(self, num_experts: int) -> GatedMLPBridge | MoEBridge:
195 """Dense SwiGLU or mixed dense/MoE sharing HF path ``.feed_forward``."""
196 if num_experts > 1:
197 # Mixed layers: MoEBridge passthrough handles both tensor and
198 # (hidden, router_scores) returns. Optional children skip the
199 # wrong feed_forward type per layer.
200 return MoEBridge(
201 name="feed_forward",
202 config=self.cfg,
203 submodules={
204 "gate": LinearBridge(name="gate_proj", optional=True),
205 "in": LinearBridge(name="up_proj", optional=True),
206 "out": LinearBridge(name="down_proj", optional=True),
207 "router": LinearBridge(name="router", optional=True),
208 },
209 )
210 return GatedMLPBridge(
211 name="feed_forward",
212 config=self.cfg,
213 submodules={
214 "gate": LinearBridge(name="gate_proj"),
215 "in": LinearBridge(name="up_proj"),
216 "out": LinearBridge(name="down_proj"),
217 },
218 )
220 def _build_component_mapping(self, num_experts: int) -> dict:
221 return {
222 "embed": EmbeddingBridge(name="model.embed_tokens"),
223 "blocks": BlockBridge(
224 name="model.layers",
225 submodules={
226 "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg),
227 "ln2": RMSNormalizationBridge(name="pre_ff_layernorm", config=self.cfg),
228 "attn": self._build_attention_bridge(),
229 "mixer": self._build_mamba_bridge(),
230 "mlp": self._build_ffn_bridge(num_experts),
231 },
232 ),
233 "ln_final": RMSNormalizationBridge(name="model.final_layernorm", config=self.cfg),
234 "unembed": UnembeddingBridge(name="lm_head"),
235 }