Conversation
Fix broken link in README
* Add Direct Logit Attribution tool for TransformerBridge * Resolve review feedback and add Direct Logit Attribution tests Resolved review feedback from @jlarson4, added tests covering reconstruction invariants on a distilgpt2 bridge in compatibility mode, arguments, asserting sum(scores) == logit_diff - (b_U[correct] - b_U[wrong]) against the model's real logits, plus labels/shape and batch-averaging checks. Added additional hardening: - Fix a latent direction-shape bug: replace the fragile answer_tokens.numel()==1 branch with a robust reshape so single-prompt, single-token inputs are handled correctly - Detect hybrid blocks via bridge.layer_types() instead of substring matching named_modules(), the codebase's own semantic mechanism - Import get_act_name from transformer_lens.utilities to avoid the transformer_lens.utils DeprecationWarning; drop the invalid return_type kwarg to run_with_cache - Register the analysis subpackage in tools/__init__.py Closes #1263.
* Add Direct Logit Attribution tool (#1263) Add transformer_lens/tools/analysis/direct_logit_attribution.py, a single-call DLA analysis that decomposes a logit (or logit difference) into per-component, per-layer (logit-lens), or per-head contributions. Wraps the existing ActivationCache primitives (decompose_resid / accumulated_resid / stack_head_results / logit_attrs) and works with both HookedTransformer and TransformerBridge, since they share the cache API. Returns a DirectLogitAttribution dataclass (attribution tensor + aligned labels, plus a top(k) helper). Adds integration tests asserting the exact DLA correctness invariant on both systems: the complete decomposition reconstructs the model's real logit up to the unembedding bias b_U. Closes #1263 * Resolving conflicts between 1316 and 1369 * format fixes --------- Co-authored-by: Azra Bano <azrabano23@gmail.com> Co-authored-by: Jonah Larson <jonahalarson@comcast.net>
* Add Phi adapter tests * Add comment about setup component test * Delete redundant config literal tests
* Fixed SVD interpreter test * Format SVD interpreter fixture test
The Restricted Loss section called loss_fn(all_logits, labels), but all_logits had been rearranged earlier into a (p, p, d_vocab) grid for the logit periodicity analysis. loss_fn's 3-D branch assumes (batch, pos, d_vocab) and takes logits[:, -1], producing a (p, p) tensor that crashes the gather against the p*p labels (#543). Use original_logits instead, which is recomputed just above and is the same full-dataset loss the cell intends to print. Also clear the stored RuntimeError output from the cell.
Breaking: removes the public eps_attr constructor argument and the config.eps_attr attribute. The field was never read (its consumer was deleted when NormalizationBridge moved to direct HF delegation), so no model behavior changes, but it is an API removal.
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
* Add Olmo2 architecture adapter tests * Drop test_attn_output_shape per the unit-test guide (shared bridge contract)
* Fixed SVD interpreter test * Format SVD interpreter fixture test * Add qwen adapter unit test * Retrigger CI (unrelated HF 429 error failing the build)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
…1390) The adapter already conditionally omitted ln2 from the block submodules when use_parallel_residual=True, but still wrapped them in a plain BlockBridge, which rejects the attn+mlp+no-ln2 shape. Switched to conditional block_cls (ParallelBlockBridge for the parallel branch, BlockBridge for sequential), mirroring the dual-mode pattern in falcon.py.
26 tests covering: component mapping (slots, bridge types, HF paths, submodule structure), anti-drift config flags (final_rms, uses_rms_norm, gated_mlp), weight conversion key set and rearrange patterns, GQA propagation to K/V only, and setup_component_testing rotary-emb wiring. Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
* fix: Fix broken line graphs - Fixes incorrect induction loss graph at end of notebook (values were incorrect, loss appeared to be going up with training instead of down!) - Fixes x/y axes showing "index"/"value" instead of configured names * fix: broken/outdated links * doc: Reference ARENA chapter directly instead of an older equivalent that forwards there * chore: Align text in code with text in markdown * doc: Update texts on supported models and architectures - Updated counts - Updated link to list of all models supported on v3 - "HookedTransformed.from_pretrained" instances -> "TransformerBridge.boot_transformers", the up-to-date recommended method which makes the wider variety accessible - "consistent(-ish) architecture" -> "consistent(-ish) interface": with v3, the consistent interface proxies the non-consistent underlying architectures - "Transformer architecture" title -> "HookedTransformer architecture", mention deprecation - Mention overview of models last updated in 2023 - Drop stale reference to phantom "table for hyper-parameters" * doc: Add note on boot_transformers incompatibility with checkpoints * doc: Document purpose of enabling compatibility mode * doc: Drop minor comment on default_prepend_bos, only relevant for legacy HookedTransformer * doc: Remove stale gotcha, hooks are now properly removed even in case of an error in a hook * fix: Mismatching values in code vs. descriptions * chore: Replace deprecated circuitvis attention visualization with newer one * chore: Replace deprecated model names aliases * chore: Remove deprecated prepend_bos argument * chore: Update hook name to v3 * chore: Full OV graph title no longer the same as the preceding very similar OV graph * doc: Clearer spaces in tokens * doc: Bracket tokens styled as code to prevent Colab from collapsing spaces Previously Colab showed both tokens as identical, squashing the space * chore: Try the model out right after calls to do so * doc: approximate number can be approximate (consistent with another preceding one) * chore: Remove installing old node version which was the newest at the time Now actually gives a deprecation warning and waits for 10 seconds * chore: Remove leftover cell * chore: fix bad closing tag * doc: fix typos * doc: More precise wording, setting a value, not adding it * doc: move floating function names into complete sentence
25 tests across 4 classes covering component mapping, config flags, weight conversions, and GQA head-count propagation. - TestMistralComponentMapping (12 tests): top-level keys, bridge types, HF module paths, block submodules, attn flags, QKVO paths, MLP paths. Includes explicit guard that attn uses AttentionBridge, not PositionEmbeddingsAttentionBridge. - TestMistralAdapterConfig (4 tests): final_rms=False, uses_rms_norm, gated_mlp, attn_only — anti-drift flags. - TestMistralWeightConversions (5 tests): exactly 4 QKVO weight keys, split-heads and merge-heads rearrange patterns, no bias/norm entries. - TestMistralGQASupport (4 tests): K/V use n_key_value_heads, Q/O unchanged, fallback to n_heads when n_key_value_heads is unset.
) Adds focused test suites for three architecture adapters per the proposal in issue #1302. tests/unit/model_bridge/supported_architectures/test_phi3_adapter.py - Component mapping structure (bridge types and HF module paths) - Weight conversion key set and source keys for fused qkv/gate_up - _SizedSplitConversion numerical correctness (Q/K/V GQA splits) - Config flags (RMS norm, rotary, gated MLP, supports_fold_ln=False) - preprocess_weights LN folding into QKV and gate/up projections tests/unit/model_bridge/supported_architectures/test_granite_adapter.py - Component mapping for dense GraniteArchitectureAdapter - Weight conversion key set (standard QKVO rearrangements) - Config flags (final_rms=True, default_prepend_bos=False, GQA heads) - GraniteMoeArchitectureAdapter: MoE bridge replaces dense MLP, all other components and config flags match dense Granite Closes part of #1302.
* chore: Plot helper allows customizing graph before showing it * feat: Direct path patching in exploratory analysis demo, resolves #111 * doc: fix head index in prose
…#1396) * feat: add get_act_patch_direct_path for head-to-head circuit analysis Closes #111. Implements direct path patching — a finer-grained variant of activation patching that isolates the direct information flow between two specific attention heads, rather than replacing the full residual stream. Why --- Standard activation patching tells you that *some* component at layer L matters, but it cannot distinguish whether head B at layer L+2 matters because it received information directly from head A, or because A's output propagated through many intermediate components first. Direct path patching isolates the A → B causal edge precisely. Implementation -------------- For a fixed source head A = (src_layer, src_head) and every downstream destination head B = (dst_layer, dst_head): delta_resid = clean_A_result - corrupted_A_result # [batch, pos, d_model] delta_B_q = (delta_resid / ln1_scale) @ W_Q[hb] # [batch, pos, d_head] patched_B_q = corrupted_B_q + delta_B_q The per-head residual contribution is computed from hook_z @ W_O (always available in the default cache) rather than hook_result, which requires the non-default cfg.use_hook_result=True flag. New files --------- - transformer_lens/direct_path_patching.py get_act_patch_direct_path() [n_layers, n_heads] sweep get_act_patch_direct_path_all_sources() [n_layers, n_heads, n_layers, n_heads] full sweep - tests/unit/test_direct_path_patching.py 12 tests covering output shape, causal structure, manual correctness verification, and edge cases. All pass on a tiny randomly-initialised 3-layer model (no downloads, runs in ~3s on CPU). - demos/direct_path_patching_ioi.py Validated on GPT-2 small / IOI task. S-inhibition heads (7.3, 7.9, 8.6, 8.10) show strongest direct paths into name-mover heads (9.9, 9.6, 10.0), confirming the Wang et al. 2022 IOI circuit. (8,6) → (9,9): +0.083 normalised logit diff (8,10) → (9,9): +0.066 (7,9) → (9,9): +0.036 API matches existing get_act_patch_* functions in patching.py for drop-in use alongside the existing circuit analysis toolkit. * style: apply black + isort formatting (line-length=100) * fix: remove unused imports, add type: ignore for mypy, clean up demo import * fix: add type: ignore[index] on W_O[h] indexing for mypy * fix: address reviewer feedback — TransformerBridge support, fold_ln guard, independent test, notebook demo * fix: check fold_ln for TransformerBridge via .weight attribute TransformerBridge wraps the original HuggingFace LayerNorm module, which stores the learned scale as .weight rather than the .w used by HookedTransformer. Fall back to .weight so the guard actually fires when a TransformerBridge model is passed without folded LN, rather than silently skipping the check. * refactor: move to tools/analysis/, logit-diff metric, TestCheckFoldLn, fix _check_fold_ln tensor bug - Move direct_path_patching.py to transformer_lens/tools/analysis/ alongside the Direct Logit Attribution tool; add tools/analysis/__init__.py exporting both public functions; update transformer_lens/__init__.py accordingly. - Fix _check_fold_ln: replace 'getattr(...) or getattr(...)' with explicit None checks to avoid RuntimeError on multi-element tensors. - test_correctness_against_actual_ln_forward: switch patching metric to logit diff (correct_tok - incorrect_tok), which cancels the centering offset introduced by process_weights_() and tightens tolerance 0.15 -> 1e-3. - Add TestCheckFoldLn (5 tests): folded model no-warning, unfolded model warns, pre-fold .w attribute present, no crash on missing attribute, no RuntimeError on multi-element tensor regression check. All 17 tests pass. * fix: loosen _check_fold_ln type hint to Any for beartype compatibility _check_fold_ln is a private defensive helper with a try/except that handles arbitrary model types. The Union[HookedTransformer, TransformerBridge] annotation was causing beartype to reject valid test fixtures (and any non-standard model) at the call boundary before the function's own exception handling could run. Any is the correct annotation for a function intentionally designed to tolerate unknown model shapes. * style: black formatting on test_direct_path_patching.py
* Add unit tests for MistralArchitectureAdapter 25 tests across 4 classes covering component mapping, config flags, weight conversions, and GQA head-count propagation. - TestMistralComponentMapping (12 tests): top-level keys, bridge types, HF module paths, block submodules, attn flags, QKVO paths, MLP paths. Includes explicit guard that attn uses AttentionBridge, not PositionEmbeddingsAttentionBridge. - TestMistralAdapterConfig (4 tests): final_rms=False, uses_rms_norm, gated_mlp, attn_only — anti-drift flags. - TestMistralWeightConversions (5 tests): exactly 4 QKVO weight keys, split-heads and merge-heads rearrange patterns, no bias/norm entries. - TestMistralGQASupport (4 tests): K/V use n_key_value_heads, Q/O unchanged, fallback to n_heads when n_key_value_heads is unset. * auto-checkpoint: before editing test_bert_adapter.py * auto-checkpoint: before editing test_falcon_adapter.py * test: add adapter unit tests for Falcon and BERT * style: fix isort import formatting in BERT adapter tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
… handle (#1426) HookPoint.add_hook returns None, but the benchmark helpers stored its return value as a "handle" (silencing mypy with behind an `if handle is not None` guard). Since the value was always None, the guard was always false and the capture hooks were never removed. Track the HookPoint that was registered on and clean up via hook_point.remove_hooks() (dir="bwd" for backward hooks), matching the existing idiom in TransformerBridge.run_with_cache. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…it can still be 1.x (#1423)
…1432) * test(hubert): add unit adapter tests for HubertArchitectureAdapter Covers component mapping (bare HubertModel + HubertForCTC prefix rebinding), bridge types and HF module paths, post-LN vs pre-LN supports_fold_ln tied to do_stable_layer_norm, weight conversion key set and rearrange patterns, and prepare_loading/prepare_model branches. Part of #1302. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: fix black/isort formatting in test_hubert_adapter.py Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
* feat: add DeepSeek-V2 architecture adapter (MLA, V2/V2-Lite, complex RoPE) Closes #1400. DeepSeek-V2, V2-Lite, and Coder-V2 all use DeepseekV2ForCausalLM. This adds a bridge adapter covering three V2-specific differences from V3: 1. Complex-exponential RoPE: V2's rotary embedding returns freqs_cis (a complex tensor via torch.polar) rather than a (cos, sin) tuple. - RotaryEmbeddingBridge.forward() now passes complex tensors through without raising, leaving them for the attention bridge to consume. - MLAAttentionBridge.forward() detects complex position_embeddings and dispatches to a new _apply_rotary_complex() helper that mirrors DeepSeek-V2's apply_rotary_emb (view_as_complex, multiply, flatten). 2. Optional Q LoRA path: V2-Lite sets q_lora_rank=None, skipping q_a_proj/q_a_layernorm/q_b_proj and using q_proj directly instead. All three Q-path submodules are marked optional=True in the adapter; q_a_layernorm uses GeneralizedComponent (which already supports optional) rather than RMSNormalizationBridge. MLAAttentionBridge already branches on q_lora_rank at runtime. 3. Gate not hookable: DeepseekV2Moe.forward() routes via nn.functional.linear(..., self.gate.weight) rather than self.gate(hidden_states), so the gate module's forward() is never called and bridge hooks cannot fire. The gate is omitted from MoEBridge submodules; shared_experts uses __call__ and hooks fine. Files changed: - supported_architectures/deepseek_v2.py (new) - supported_architectures/__init__.py: register adapter - factories/architecture_adapter_factory.py: map DeepseekV2ForCausalLM - generalized_components/mla_attention.py: complex RoPE support - generalized_components/rotary_embedding.py: complex tensor pass-through - tests/integration/model_bridge/test_deepseek_v2_adapter.py (new, 17 tests) * style: fix formatting and mypy errors * fix: register DeepseekV2ForCausalLM in model registry
TL;DR: add TransformerBridge support for LiquidAI LFM2 MoE and Microsoft PhiMoE, including registry wiring, config passthrough, unit coverage, native PhiMoE loading, and correct chat EOS stopping. Adds native architecture adapters for Lfm2MoeForCausalLM and PhiMoEForCausalLM. LFM2 is exposed conservatively as a residual-probing adapter over the HF decoder layers, avoiding fake attention/MLP internals for the hybrid conv/full-attention architecture. PhiMoE targets the native Transformers implementation with trust_remote_code=False rather than the archived remote implementation, which is incompatible with current generation/cache semantics. Registers both architectures across the TransformerBridge adapter factory, supported-architecture exports, model-registry architecture lists, report metadata, and supported_models.json. Adds explicit registry entries for LiquidAI/LFM2.5-8B-A1B and microsoft/Phi-mini-MoE-instruct. Extends HF config passthrough for hybrid/MoE and tokenizer-related attributes used by these adapters, including LFM2 layer/norm/MoE metadata, PhiMoE routing/bias metadata, and eos_token_id. TransformerBridge generation now prefers cfg.eos_token_id when present, allowing architecture adapters to provide multiple default stop tokens. Handles PhiMoE chat generation correctly by stopping on both the tokenizer EOS (<|endoftext|>) and the chat-template turn terminator (<|end|>). This prevents generation from continuing into another assistant turn after the model has ended its answer. Adds unit tests covering both adapters' config flags, component mappings, weight conversions, and PhiMoE EOS defaults.
* Add Gemma 4 architecture support to TransformerBridge Adds a text-only adapter covering both Gemma4ForConditionalGeneration (E2B/E4B/31B/26B-A4B) and Gemma4UnifiedForConditionalGeneration (12B), addressing #1297. Gemma 4 layers are heterogeneous: KV-shared layers drop k/v projections, K==V layers drop v_proj, and per-layer-embedding / MoE submodules appear only on some variants -- all mapped optional and delegated to HF. Unlike Gemma 1-3, Gemma4RMSNorm has no (1+weight) offset. Adds DelegatedAttentionBlockBridge (drops the split-QKV fork aliases, as MLABlockBridge does) so hook-alias resolution stays clean when attention is delegated wholesale to HF. google/gemma-4-E2B-it passes verification (P1 100%, P2 100%, P4 94.7%). - New adapter + four-place registration + gemma4/gemma4_unified model_type mappings - 10 checkpoints added to the model registry - Unit + integration tests (logit parity vs HF on all three structural variants) * fix: handle list eos_token_id when setting pad_token_id * fix: add Gemma4ForConditionalGeneration to MULTIMODAL_ARCHITECTURES * feat: add multimodal vision support to Gemma4 adapter * fix: use GeneralizedComponent for vision projector — VisionProjectionBridge expects positional 'vision_features' but Gemma4's Gemma4MultimodalEmbedder.forward() takes 'inputs_embeds' kwarg * fix: check image_token before boi_token in multimodal benchmark — Gemma4's boi_token is a marker, image_token is the expandable placeholder * chore: update registry with Gemma4 verification results * chore: fix E2B-it registry — P1=50% (component benchmark fails with delegation), P4=98.7% from real-model run * chore: fix E2B phase2_score — set to null (Phase 2 not run on real model) * fix: guard MPS synchronize with torch.backends.mps.is_available() instead of hasattr * chore: update 31B note — P4/P7 degeneration is decoding config (missing chat template), not Bridge bug --------- Co-authored-by: punishell <bazyli.michal@gmail.com> Co-authored-by: jlarson4 <jonahalarson@comcast.net>
* feat: Add TypedModuleList, a generic, type-preserving nn.ModuleList nn.ModuleList is not generic, so iterating or indexing it yields a bare nn.Module (or Any on torch < 2.8). TypedModuleList[T] preserves the element type T through iteration and other methods while remaining a drop-in nn.ModuleList at runtime (children are still registered as submodules). For technical reasons, the return type information lives in a companion .pyi stub (technical: method-level beartype, enforced by jaxtyping, cannot resolve a Self return type. The explicit TypedModuleList[T] works only if we update jaxtyping's major version to 3.x, which is a bigger change). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor: use TypedModuleList for block lists in HT-family models Type HookedTransformer.blocks, HookedEncoder.blocks, and HookedEncoderDecoder.encoder/decoder as TypedModuleList[...] so iteration and indexing preserve the element type. This removes the # type: ignore[type-arg] annotations and the _get_blocks()/cast() workarounds that existed only because nn.ModuleList erased the element type, and lets downstream consumers (ActivationCache and the weight_processing benchmark) drop their now-redundant casts and union-attr ignores. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor: use TypedModuleList in HookedAudioEncoder, towards updating PyTorch Unlike the HT-family refactor, the goal here isn't to remove casts, it's a step towards updating PyTorch: as of torch 2.8, iterating a plain nn.ModuleList yields `Tensor | Module` instead of `Any`, which triggers new mypy errors. Typing blocks as TypedModuleList[BertBlock] fixes them (and also allows extra asserts to be dropped). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * refactor: use TypedModuleList in NativeModel, towards updating PyTorch Clean fix for another mypy arg-type error that surfaces on torch >= 2.8 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: restore GeneralizedComponent element type in block-bridge traversal Resolves the last remaining mypy errors with newer PyTorch Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore: Update locked torch version to 2.10.0 Conservative: pyproject minimum requirement unchanged, still >=2.6 Incremental: stopping at 2.10.0 because 2.11.0's bundled torchvision 0.26.0 introduced relevant breaking changes - Makes TransformerLens compatible with Python 3.14 (previously, if 3.14 installed, then the default uv sync failed because the locked torch 2.7.1 wasn't compatible with it - Towards enabling MPS by default when torch is new enough and contains the required bugfixes (see #1178) --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
* tidy: Move jupyter packages to jupyter dependency group * fix: error opening jupyter notebook By updating jupyter notebook (v6.5.4->v6.5.7) Previously, `uv run jupyter notebook demos/xxx.ipynb` failed with: ModuleNotFoundError: No module named 'jupyter_server.contents' This was fixed in jupyer 6.5.6 (see jupyter/notebook#7048) * chore: Add matplotlib dev dependency, used in some notebooks * fix: Invalid import in GPT_OSS_Demo notebook
* Add GLM-5 DSA TransformerBridge adapter * Address GLM-MoE-DSA adapter review notes --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
* fix: remove misleading 'encoder-decoder' label from AutoModel log — can also be multimodal (AutoModelForImageTextToText) * fix: skip Phase 2 header when phase not selected — empty header printed in every run * fix: skip component benchmarking for DelegatedAttentionBlockBridge — attn/PLE/rotary_emb can't be tested in isolation * fix: detect DelegatedAttentionBlockBridge via adapter.component_mapping instead of bridge_model.blocks * fix: add rotary_emb skip in _test_component for delegated attention blocks * fix: remove dead rotary_emb skip from _test_component_recursive (rotary is top-level, only _test_component handles it) * fix: use maintain_native_attention flag instead of hook_q_input check for delegated attention detection
…1444) * Add HunYuanDenseV1 adapter and tests. Edit PositionEmbeddingsBridge attention for qk norm check * Add none guard for HF norm name for mypy reasons
…1434) * feat: add DeepSeek-V2 architecture adapter (MLA, V2/V2-Lite, complex RoPE) Closes #1400. DeepSeek-V2, V2-Lite, and Coder-V2 all use DeepseekV2ForCausalLM. This adds a bridge adapter covering three V2-specific differences from V3: 1. Complex-exponential RoPE: V2's rotary embedding returns freqs_cis (a complex tensor via torch.polar) rather than a (cos, sin) tuple. - RotaryEmbeddingBridge.forward() now passes complex tensors through without raising, leaving them for the attention bridge to consume. - MLAAttentionBridge.forward() detects complex position_embeddings and dispatches to a new _apply_rotary_complex() helper that mirrors DeepSeek-V2's apply_rotary_emb (view_as_complex, multiply, flatten). 2. Optional Q LoRA path: V2-Lite sets q_lora_rank=None, skipping q_a_proj/q_a_layernorm/q_b_proj and using q_proj directly instead. All three Q-path submodules are marked optional=True in the adapter; q_a_layernorm uses GeneralizedComponent (which already supports optional) rather than RMSNormalizationBridge. MLAAttentionBridge already branches on q_lora_rank at runtime. 3. Gate not hookable: DeepseekV2Moe.forward() routes via nn.functional.linear(..., self.gate.weight) rather than self.gate(hidden_states), so the gate module's forward() is never called and bridge hooks cannot fire. The gate is omitted from MoEBridge submodules; shared_experts uses __call__ and hooks fine. Files changed: - supported_architectures/deepseek_v2.py (new) - supported_architectures/__init__.py: register adapter - factories/architecture_adapter_factory.py: map DeepseekV2ForCausalLM - generalized_components/mla_attention.py: complex RoPE support - generalized_components/rotary_embedding.py: complex tensor pass-through - tests/integration/model_bridge/test_deepseek_v2_adapter.py (new, 17 tests) * style: fix formatting and mypy errors * fix: register DeepseekV2ForCausalLM in model registry * feat: add NemotronH hybrid Mamba2-Transformer architecture adapter Implements TransformerBridge support for NemotronHForCausalLM (nvidia/Nemotron-H-8B-Base, Nemotron-H-47B-A13B). Architecture overview: - Heterogeneous layers defined by config.layers_block_type: each element is one of mamba, attention, moe, or mlp (~8% attention, ~92% SSM/MLP/MoE) - Single pre-norm (block.norm) and single residual path per block; no ln2 - Single .mixer attribute per block whose type varies by layer - No model-level rotary embedding module; attention handles RoPE internally - Stateful generation via DynamicCache (transformers >= 5.12) Key adapter decisions: - SSMBlockBridge as block container: delegates full forward to HF block, avoids ln2 enforcement that BlockBridge would apply incorrectly here - SSM2MixerBridge(name=mixer) as passthrough wrapper: works for all four mixer types since forward calls original_component(*args, **kwargs) - Mamba-specific submodules (in_proj, conv1d, inner_norm, out_proj) marked optional so component_setup skips them gracefully on non-Mamba layers - GatedRMSNormBridge.optional set post-init (its __init__ does not accept the kwarg, unlike the GeneralizedComponent base class) - positional_embedding_type=none: no model-level rotary to wire - gated_mlp=False: MLP layers use relu2, not SwiGLU - applicable_phases=[]: verify_models is transformer-shaped; integration tests cover forward-pass correctness instead Registration: - architecture_adapter_factory.py: NemotronHForCausalLM key added - supported_architectures/__init__.py: export added - tools/model_registry/__init__.py: HF_SUPPORTED_ARCHITECTURES and CANONICAL_AUTHORS_BY_ARCH entries added (canonical author: nvidia) Tests (52 unit tests, all passing): - Config attribute propagation (normalization_type, positional_embedding_type, gated_mlp, is_stateful, final_rms, mamba_intermediate_size, conv_dim, layers_block_type, applicable_phases, weight_processing_conversions) - Top-level component mapping bridge types and HF path names - Block submodule bridge types (norm, mixer; no ln2) - Mixer submodule types, names, and optional flags for all four Mamba keys - create_stateful_cache returns DynamicCache; independent per call - Factory registration and model registry constants - Guard tests: SSMBlockBridge not BlockBridge, no weight conversions, mamba_intermediate_size and conv_dim formulas Closes #1402 * style: apply black + isort formatting fixes * style: fix isort order for NemotronHArchitectureAdapter import * test: add NemotronH forward-pass and generation parity integration tests * test: add NemotronH forward-pass and generation parity integration tests * test: add NemotronH forward-pass and generation parity integration tests * fix: add pytestmark=slow to NemotronH integration test * fix: register NemotronHArchitectureAdapter in factory imports and dict * fix: add NemotronHArchitectureAdapter to supported_architectures __init__ * style: black + isort test_nemotron_h_adapter.py * style: black + isort __init__.py * fix: wire slow marker — add -m "not slow" to CI make targets * fix: exclude slow tests from MPS integration CI step * style: make format — architecture_adapter_factory.py * fix: wire slow marker — add -m "not slow" to CI make targets * fix: exclude slow tests from MPS integration CI step * style: make format — test_nemotron_h_adapter.py * style: make format — architecture_adapter_factory.py --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
* Add Granite MoE hybrid architecture adapter tests * Format Granite MoE hybrid adapter tests
* Fix type of HookedTransformerConfig.device (#1230) * Fix type of HookedTransformerConfig.device This is typed as `Optional[str]` but sometimes returns `torch.device`. Updated the code to just return the `str` instead of wrapping with a device. I'm not confident that every function which takes a device will always be passed a string, so I didn't change functions like warn_if_mps. Found while working on #1219 * more cleanup * 3.0 CI Bugs (#1261) * Fixing `utils` imports * skip gated notebooks on PR from forks * Updating notebooks * Ensure LLaMA only runs when HF_TOKEN is available --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net> * Add config tests for Bloom architecture adapter * reformat --------- Co-authored-by: Brendan Long <self@brendanlong.com> Co-authored-by: jlarson4 <jonahalarson@comcast.net> Co-authored-by: Indrayani Satyajeet Rajeshirke <indra@Indrayanis-MacBook-Air.local>
…sformers >= 5.6.0 (#1443) * Updated siglip & clip to work with transformers version 5.6.0+ * Changed to a backward compatible fix * Add testing --------- Co-authored-by: jlarson4 <jonahalarson@comcast.net>
* Attention fixes for T5Gemma * Additional architecture adapter verifications
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
20 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Creation of new Architecture Adapters
New tools
New features
stop_strings&stopping_criteriaParallelBlockBridgefor StableLM variantsBug Fixes
prepend_boswhen relating to Chat templates resolvedMiscellaneous
torchnow installs un-cappedType of change
Please delete options that are not relevant.
Checklist: