DOCSEARCHbyAlgolia
Submit Docs

TransformerLens Documentation

repository·main·Indexed 25 days ago

https://github.com/transformerlensorg/transformerlens

A library for the mechanistic interpretability of generative language models, specifically GPT-2 style architectures. It enables researchers to load open-source models and manipulate internal activations to reverse-engineer learned algorithms. Key features include the TransformerBridge API for loading over 9,000 models, experimental support for Mamba-1 and Mamba-2 architectures, a comprehensive benchmark suite for verifying model compatibility, and integration with the Learning Interpretability Tool (LIT).

Tokens
54.4K
Snippets
103
Records
291
Agent score
86%

What's inside TransformerLens

  1. Understand the Architecture Adapter concept

    main

    An Architecture Adapter is a Python class that extends ArchitectureAdapter. It is used by the TransformerBridge system to translate HuggingFace (HF) models into TransformerLens models. An adapter provides three essential pieces of information:

    1. Config attributes: Set on self.cfg during __init__ (e.g., normalization type, positional embedding type, GQA parameters).
    2. Component mapping: A dictionary self.component_mapping that maps TransformerLens canonical names (like embed, blocks, attn.q) to GeneralizedComponent Bridge instances pointing to specific HF module paths.
    3. Weight processing conversions: A dictionary self.weight_processing_conversions containing tensor-reshape rules to translate HF weight layouts to TransformerLens layouts during loading.

    Once an adapter is registered, you can use boot_transformers("<your-model>") to obtain a fully hooked TransformerLens model with weights loaded from HuggingFace.

  2. Experimental Mamba and SSM support

    main

    TransformerLens includes experimental bridge adapters for Mamba-1 and Mamba-2 architectures. These adapters support:

    • Bit-for-bit HuggingFace equivalent forward passes.
    • Hook-based introspection of projection activations (e.g., in_proj, conv1d, x_proj, dt_proj, out_proj for Mamba-1; in_proj, conv1d, inner_norm, out_proj for Mamba-2).
    • Stateful generation with cache-aware decode steps.
    • The compute_effective_attention utility in transformer_lens.model_bridge.supported_architectures.mamba2 to materialize Mamba-2's SSD-derived attention matrix.

  3. Install TransformerLens

    main

    Install the latest version of TransformerLens directly from the GitHub repository using pip.

    To use legacy code that requires the older version (formerly known as EasyTransformer), install version 1 specifically:

    # Install latest
pip install git+https://github.com/TransformerLensOrg/TransformerLens

# Install legacy v1
pip install git+https://github.com/TransformerLensOrg/TransformerLens@v1
    pip install git+https://github.com/TransformerLensOrg/TransformerLens

  4. Migrate from HookedTransformer to TransformerBridge

    main

    As of TransformerLens 3.0, HookedTransformer is deprecated and will be removed in the next major version. New code should use TransformerBridge instead.

    HookedTransformer uses a legacy numerical approach (folding LayerNorm and centering weights) that does not match HuggingFace weights. TransformerBridge uses raw HF weights by default. For compatibility with HookedTransformer behavior, use bridge.enable_compatibility_mode().

  5. Avoid common unit test anti-patterns

    main

    To prevent test bloat, avoid these five anti-patterns:

    1. Config-literal restatements: Do not assert adapter.cfg.flag == <literal>. Instead, assert the effect of that flag (e.g., the bridge type it selects).
    2. Factory/registration duplicates: Do not test if factory.select(cfg) returns your adapter; this is already covered by the global registry tests.
    3. Dependency tests: Do not test einops permutations or torch.nn.Linear properties. Only assert the pattern/axis metadata and the numerical partition specific to your adapter.
    4. Base-class retests: Do not test logic or defaults inherited from the base ArchitectureAdapter unless your adapter explicitly overrides them.
    5. Subsumed assertions: Use a single exact-set assertion (e.g., assert set(conv) == {...}) instead of multiple individual membership (in) or count (len) checks.

  6. Verify an Adapter using verify_models

    main

    Use the verify_models tool to run a HuggingFace model side-by-side with your bridge and compare activations.

    Run verification command:

    uv run python -m transformer_lens.tools.model_registry.verify_models \
  --model <model-id> \
  --max-memory <GB> \
  --device cpu \
  --dtype float32 \
  --no-ht-reference

    Interpreting Status Codes:

    • status=1: Passed. Move to the next model.
    • status=2: Skipped (e.g., memory pre-check failed). Not an adapter bug.
    • status=3: Phase score failure. Stop and fix. Check the note and per-phase scores to find the root cause.

    Tips:

    • If the model OOMs with float32, retry with --dtype bfloat16.
    • Set --max-memory to roughly 75-85% of your device memory.
    • After passing verification, run uv run mypy . and make check-format to ensure compliance.

  7. Perform a targeted scrape for a single architecture

    main

    Use this workflow after registering a new adapter to populate supported_models.json with models of that specific architecture.

    Note: The architecture string must exactly match config.architectures[0] from the HF model's config. If the architecture is listed in CANONICAL_AUTHORS_BY_ARCH, the scraper uses a fast canonical sweep. Otherwise, it falls back to a slower global scan with client-side filtering. Existing entries in supported_models.json are preserved and appended to.

    uv run python -m transformer_lens.tools.model_registry.hf_scraper \
    --architecture LlamaForCausalLM --full-scan

  8. Implement an Architecture Adapter

    main

    To add support for a new HuggingFace model architecture in TransformerLens, create a Python class that extends ArchitectureAdapter (from transformer_lens.model_bridge.architecture_adapter).

    An adapter must define:

    1. Config attributes: Set on self.cfg in __init__ to describe model properties (e.g., normalization type, positional embedding type).
    2. Component mapping: A self.component_mapping dictionary mapping TransformerLens canonical names to Bridge instances. The name= parameter in the Bridge must match the HuggingFace module path.
    3. Weight processing conversions: A self.weight_processing_conversions dictionary for tensor reshaping using ParamProcessingConversion instances.

  9. Register a New Architecture Adapter

    main

    To ensure the adapter is discoverable and correctly reported, you must update four specific locations in the repository:

    1. transformer_lens/model_bridge/supported_architectures/__init__.py: Add the new module import and append it to __all__.
    2. transformer_lens/factories/architecture_adapter_factory.py: Add the import and map the HF architecture class to your adapter class in SUPPORTED_ARCHITECTURES:
      "<HFArchitectureClass>": <YourAdapterClass>,
    3. transformer_lens/tools/model_registry/__init__.py:
      • Add the <HFArchitectureClass> to HF_SUPPORTED_ARCHITECTURES.
      • Add the <HFArchitectureClass> and its canonical organization (e.g., "meta-llama") to CANONICAL_AUTHORS_BY_ARCH to prevent the scraper from dropping small canonical variants.
    4. transformer_lens/tools/model_registry/generate_report.py: Add a human-readable description to ARCHITECTURE_DESCRIPTIONS:
      "<HFArchitectureClass>": "<Short human-readable description>",

  10. Identify a reference adapter for implementation

    main

    When implementing a new adapter, identify the closest existing pattern in transformer_lens/model_bridge/supported_architectures/ to use as a template. Use the following mapping to find your starting point:

    If your model is like…Start from…
    Llama, Mistral, Qwen2, Gemma, OLMollama.py
    Qwen2/Qwen3 (gated config, MLPBridge)qwen2.py
    GPT-2, GPT-J, GPT-Neogpt2.py
    BLOOM, Falconbloom.py or falcon.py
    T5 / encoder-decodert5.py
    MoEmixtral.py or granite_moe.py
    Multimodal (vision+text)llava.py or gemma3_multimodal.py