From 43d9ff41c78affc8b5c24dd2495848d13f5e32b5 Mon Sep 17 00:00:00 2001 From: dlyanreal Date: Sat, 11 Jul 2026 13:27:18 -0500 Subject: [PATCH 1/3] chore: save WIP on V3 transformerbridge migration --- .../test_ast_adapter.py | 54 +++++++++++++++++++ .../factories/architecture_adapter_factory.py | 2 + .../supported_architectures/__init__.py | 4 ++ .../supported_architectures/ast.py | 49 +++++++++++++++++ 4 files changed, 109 insertions(+) create mode 100644 tests/unit/model_bridge/supported_architectures/test_ast_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/ast.py diff --git a/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py b/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py new file mode 100644 index 000000000..1025a0478 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py @@ -0,0 +1,54 @@ +import torch +from transformers import ASTConfig, ASTForAudioClassification + +from transformer_lens.model_bridge import TransformerBridge + +def test_ast_parity(): + # 1. setup a tiny HF AST config for instantaneous local testing + hf_config = ASTConfig( + hidden_size=12, + num_hidden_layers=2, + num_attention_heads=2, + intermediate_size=24, + patch_size=16, + ) + + # explicitly define the architecture list for the local config + hf_config.architectures = ["ASTForAudioClassification"] + + hf_model = ASTForAudioClassification(hf_config) + hf_model.eval() + + # 2. boot the v3 transformerbridge natively + # bridge auotmatically detecty ASTForAudioClassification and use new ASTArchitectureAdapater + tl_model = TransformerBridge.boot_transformers("ast", hf_model=hf_model) + tl_model.eval() + + # 3. create dummy spectrogram input [batch, freq, time] + dummy_spectrogram = torch.randn(1, 1024, 128) + + # 4. compare forward passes + with torch.no_grad(): + hf_logits = hf_model(dummy_spectrogram).logits + + # bridge forward pass (extract sequence before HF's pooling) + _, cache = tl_model.run_with_cache(dummy_spectrogram) + resid = cache["ln_final.hook_normalized"] + + # apply AST pooling logic to the bridges residual stream + cls_token = resid[:, 0, :] + dist_token = resid[:, 1, :] + pooled_out = (cls_token + dist_token) / 2.0 + + # apply the final classifier (which holds 2nd layernorm + dense) + tl_logits = hf_model.classifier(pooled_out) + + diff = (hf_logits - tl_logits).abs().max().item() + print(f"Briddge vs HF Max Logit Diff: {diff:.6e}") + + # 5. assert parity + assert diff < 1e-4, "Parity failed: Bridge tensors do not match HuggingFace." + print("Parity dub. V3 Bridge connected excellent") + +if __name__ == "__main__": + test_ast_parity() \ No newline at end of file diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index f1c3337c5..ad3768fe4 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -12,6 +12,7 @@ from transformer_lens.model_bridge.supported_architectures import ( ApertusArchitectureAdapter, ArceeArchitectureAdapter, + ASTArchitectureAdapter, BaichuanArchitectureAdapter, BartArchitectureAdapter, BD3LMArchitectureAdapter, @@ -100,6 +101,7 @@ SUPPORTED_ARCHITECTURES = { "ApertusForCausalLM": ApertusArchitectureAdapter, "ArceeForCausalLM": ArceeArchitectureAdapter, + "ASTForAudioClassification": ASTArchitectureAdapter, "BaiChuanForCausalLM": BaichuanArchitectureAdapter, "BaichuanForCausalLM": BaichuanArchitectureAdapter, "BartForConditionalGeneration": BartArchitectureAdapter, diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index ae6803bdb..1819629e3 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -9,6 +9,9 @@ from transformer_lens.model_bridge.supported_architectures.arcee import ( ArceeArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.ast import ( + ASTArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.baichuan import ( BaichuanArchitectureAdapter, ) @@ -257,6 +260,7 @@ __all__ = [ "ApertusArchitectureAdapter", "ArceeArchitectureAdapter", + "ASTArchitectureAdapter", "BD3LMArchitectureAdapter", "BaichuanArchitectureAdapter", "BartArchitectureAdapter", diff --git a/transformer_lens/model_bridge/supported_architectures/ast.py b/transformer_lens/model_bridge/supported_architectures/ast.py new file mode 100644 index 000000000..d392753db --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/ast.py @@ -0,0 +1,49 @@ +from typing import Any + +from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion +from transformer_lens.conversion_utils.param_processing_conversion import ( + ParamProcessingConversion, +) + +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + AttentionBridge, + BlockBridge, + LinearBridge, + MLPBridge, + NormalizationBridge, + UnembeddingBridge, + GeneralizedComponent, +) + +class ASTArchitectureAdapter(ArchitectureAdapter): + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + # essential flag for audio models in V3 + self.cfg.is_audio_model = True + + n_heads = self.cfg.n_heads + + # calculate n_ctx: dynamically grab the length from the generated position embeddings + # this resolves the n_ctx = num_patches + 2 requirement accurately regardless of config stride + self.n_ctx = hf_model.audio_spectrogram_transformer.embeddings.position_embeddings.shape[1] + + # V3 Bridge pattern: Map transformerlens canonical names to hf AST attributes + self.component_mapping = { + "embed": "audio_spectrogram_transformer.embeddings", + "blocks": "audio_spectrogram_transformer.layers", + "ln_final": "audio_spectrogram_transformer.layernorm", + "unembed": "classifier.dense", + + # block internals mapping + "block.ln1": "layernorm_before", + "block.attn.W_Q": "attention.q_proj", + "block.attn.W_K": "attention.k_proj", + "block.attn.W_V": "attention.v_proj", + "block.attn.W_O": "attention.o_proj", + "block.ln2": "layernorm_after", + "block.mlp.W_in": "mlp.fc1", + "block.mlp.W_out": "mlp.fc2", + } \ No newline at end of file From 3b9db1001aa3614efb03cf6e80297cf52dadae37 Mon Sep 17 00:00:00 2001 From: dlyanreal Date: Sat, 25 Jul 2026 13:20:07 -0500 Subject: [PATCH 2/3] feat(ast): migrate AST adapter to V3 TransformerBridge and component_mapping --- .../test_ast_adapter.py | 8 +- .../supported_architectures/ast.py | 96 +++++++++++++++---- 2 files changed, 80 insertions(+), 24 deletions(-) diff --git a/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py b/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py index 1025a0478..e35820ac2 100644 --- a/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py +++ b/tests/unit/model_bridge/supported_architectures/test_ast_adapter.py @@ -3,6 +3,7 @@ from transformer_lens.model_bridge import TransformerBridge + def test_ast_parity(): # 1. setup a tiny HF AST config for instantaneous local testing hf_config = ASTConfig( @@ -42,13 +43,14 @@ def test_ast_parity(): # apply the final classifier (which holds 2nd layernorm + dense) tl_logits = hf_model.classifier(pooled_out) - + diff = (hf_logits - tl_logits).abs().max().item() - print(f"Briddge vs HF Max Logit Diff: {diff:.6e}") + print(f"Bridge vs HF Max Logit Diff: {diff:.6e}") # 5. assert parity assert diff < 1e-4, "Parity failed: Bridge tensors do not match HuggingFace." print("Parity dub. V3 Bridge connected excellent") + if __name__ == "__main__": - test_ast_parity() \ No newline at end of file + test_ast_parity() diff --git a/transformer_lens/model_bridge/supported_architectures/ast.py b/transformer_lens/model_bridge/supported_architectures/ast.py index d392753db..d714ac30c 100644 --- a/transformer_lens/model_bridge/supported_architectures/ast.py +++ b/transformer_lens/model_bridge/supported_architectures/ast.py @@ -4,7 +4,6 @@ from transformer_lens.conversion_utils.param_processing_conversion import ( ParamProcessingConversion, ) - from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter from transformer_lens.model_bridge.generalized_components import ( AttentionBridge, @@ -13,37 +12,92 @@ MLPBridge, NormalizationBridge, UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( GeneralizedComponent, ) -class ASTArchitectureAdapter(ArchitectureAdapter): +class ASTArchitectureAdapter(ArchitectureAdapter): def __init__(self, cfg: Any) -> None: super().__init__(cfg) # essential flag for audio models in V3 self.cfg.is_audio_model = True + self.cfg.normalization_type = "LN" n_heads = self.cfg.n_heads - # calculate n_ctx: dynamically grab the length from the generated position embeddings - # this resolves the n_ctx = num_patches + 2 requirement accurately regardless of config stride - self.n_ctx = hf_model.audio_spectrogram_transformer.embeddings.position_embeddings.shape[1] + # Q/K/V/O rearrangement: splits hidden dims into (heads, head_dim) + self.weight_processing_conversions = { + "blocks.{i}.attn.q.weight": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion( + "(h d_head) d_model -> h d_model d_head", h=n_heads + ), + ), + "blocks.{i}.attn.k.weight": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion( + "(h d_head) d_model -> h d_model d_head", h=n_heads + ), + ), + "blocks.{i}.attn.v.weight": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion( + "(h d_head) d_model -> h d_model d_head", h=n_heads + ), + ), + "blocks.{i}.attn.q.bias": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads), + ), + "blocks.{i}.attn.k.bias": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads), + ), + "blocks.{i}.attn.v.bias": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads), + ), + "blocks.{i}.attn.o.weight": ParamProcessingConversion( + tensor_conversion=RearrangeTensorConversion( + "d_model (h d_head) -> h d_head d_model", h=n_heads + ), + ), + } - # V3 Bridge pattern: Map transformerlens canonical names to hf AST attributes + # V3 bridge pattern: hierarchical mapping using bridge components self.component_mapping = { - "embed": "audio_spectrogram_transformer.embeddings", - "blocks": "audio_spectrogram_transformer.layers", - "ln_final": "audio_spectrogram_transformer.layernorm", - "unembed": "classifier.dense", - - # block internals mapping - "block.ln1": "layernorm_before", - "block.attn.W_Q": "attention.q_proj", - "block.attn.W_K": "attention.k_proj", - "block.attn.W_V": "attention.v_proj", - "block.attn.W_O": "attention.o_proj", - "block.ln2": "layernorm_after", - "block.mlp.W_in": "mlp.fc1", - "block.mlp.W_out": "mlp.fc2", - } \ No newline at end of file + "embed": GeneralizedComponent(name="audio_spectrogram_transformer.embeddings"), + "ln_final": NormalizationBridge( + name="audio_spectrogram_transformer.layernorm", config=self.cfg + ), + "unembed": UnembeddingBridge(name="classifier.dense"), + "blocks": BlockBridge( + name="audio_spectrogram_transformer.layers", + submodules={ + "ln1": NormalizationBridge(name="layernorm_before", config=self.cfg), + "ln2": NormalizationBridge(name="layernorm_after", config=self.cfg), + "attn": AttentionBridge( + name="attention", + config=self.cfg, + submodules={ + "q": LinearBridge(name="q_proj"), + "k": LinearBridge(name="k_proj"), + "v": LinearBridge(name="v_proj"), + "o": LinearBridge(name="o_proj"), + }, + ), + "mlp": MLPBridge( + name="mlp", + config=self.cfg, + submodules={ + "in": LinearBridge(name="fc1"), + "out": LinearBridge(name="fc2"), + }, + ), + }, + ), + } + + def prepare_model(self, hf_model: Any) -> None: + # hook to access the live Huggingface model before boot + # calculate n_ctx dynamically from the instantiated position embeddings + self.cfg.n_ctx = ( + hf_model.audio_spectrogram_transformer.embeddings.position_embeddings.shape[1] + ) From 2820b13c0a17c652afca0deeb77aee08c9af1464 Mon Sep 17 00:00:00 2001 From: dlyanreal Date: Sat, 25 Jul 2026 14:29:24 -0500 Subject: [PATCH 3/3] chore(ast): register ASTForAudioClassification and MIT canonical author in model registry --- transformer_lens/tools/model_registry/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index 3474da254..71963c6e4 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -47,6 +47,7 @@ HF_SUPPORTED_ARCHITECTURES: set[str] = { "ApertusForCausalLM", "ArceeForCausalLM", + "ASTForAudioClassification", "BaiChuanForCausalLM", "BaichuanForCausalLM", "BartForConditionalGeneration", @@ -134,6 +135,7 @@ CANONICAL_AUTHORS_BY_ARCH: dict[str, list[str]] = { "ApertusForCausalLM": ["swiss-ai"], "ArceeForCausalLM": ["arcee-ai"], + "ASTForAudioClassification": ["MIT"], "BaiChuanForCausalLM": ["baichuan-inc"], "BaichuanForCausalLM": ["baichuan-inc"], "BartForConditionalGeneration": ["facebook"],