[Bug Report] DeepSeek V2 dense layers expose MLP boundary tensors under neuron-hook names

[hijohnnylin]

Summary

TransformerBridge maps every DeepSeek V2 block's mlp to MoEBridge, including the dense prefix selected by first_k_dense_replace.

MoEBridge defines these aliases:

hook_pre  -> hook_in
hook_post -> hook_out

On an actual sparse MoE block, a block-level d_mlp activation may not exist. On a dense DeepSeek V2 layer, however, the underlying module is DeepseekV2MLP, with ordinary gate_proj, up_proj, and down_proj projections. Those layers have real neuron-basis tensors, but TransformerLens exposes the MLP input and output instead:

• blocks.N.mlp.hook_pre is d_model wide instead of the gate_proj output;
• blocks.N.mlp.hook_post is d_model wide instead of the down_proj input;
• blocks.N.mlp.hook_pre_linear is absent, although the dense module has an up_proj output.

For deepseek-ai/DeepSeek-V2-Lite, layer 0 is dense. The wrong tensors are returned without an error or warning.

Versions

The numerical output below was collected with:

• transformer-lens==3.7.0
• transformers==5.15.0
• PyTorch 2.13
• CUDA, bfloat16

I also checked the v3.7.1 source. The relevant MoEBridge aliases and DeepSeek V2 adapter mapping are unchanged.

Reproduction

pip install "transformer-lens==3.7.1" "transformers==5.15.0"
python tl_dense_mlp_hooks.py structure
python tl_dense_mlp_hooks.py shapes  # downloads DeepSeek-V2-Lite; GPU recommended

# tl_dense_mlp_hooks.py
from __future__ import annotations

import sys

MODEL = "deepseek-ai/DeepSeek-V2-Lite"
PROMPT = "The capital of France is Paris."
DENSE_LAYER = 0  # first_k_dense_replace=1


def structure() -> None:
    """Config-only check: no model weights are loaded."""
    from transformers import AutoConfig
    from transformer_lens.config import TransformerBridgeConfig
    from transformer_lens.factories.architecture_adapter_factory import (
        ArchitectureAdapterFactory,
    )
    from transformer_lens.model_bridge.generalized_components import MoEBridge
    from transformer_lens.model_bridge.sources.transformers import (
        determine_architecture_from_hf_config,
        map_default_transformer_lens_config,
    )

    hf_config = AutoConfig.from_pretrained(MODEL)
    mapped = map_default_transformer_lens_config(hf_config)
    cfg = TransformerBridgeConfig.from_dict(dict(mapped.__dict__))
    cfg.architecture = determine_architecture_from_hf_config(hf_config)
    adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
    mlp = adapter.component_mapping["blocks"].submodules["mlp"]

    print(
        f"{MODEL}: hidden_size={hf_config.hidden_size} "
        f"intermediate_size={hf_config.intermediate_size}"
    )
    print(
        f"  moe_intermediate_size={hf_config.moe_intermediate_size} "
        f"first_k_dense_replace={hf_config.first_k_dense_replace}"
    )
    print(f"MoEBridge.hook_aliases = {MoEBridge.hook_aliases}")
    print(f"adapter maps every block's mlp to {type(mlp).__name__}")


def shapes() -> None:
    """Capture TransformerLens aliases and the corresponding HF tensors in one forward."""
    import torch
    from transformer_lens.model_bridge import TransformerBridge

    bridge = TransformerBridge.boot_transformers(MODEL, device="cuda")
    tokens = bridge.to_tokens(PROMPT)

    names = [
        f"blocks.{DENSE_LAYER}.mlp.hook_pre",
        f"blocks.{DENSE_LAYER}.mlp.hook_post",
    ]

    hf = bridge.original_model
    layer = hf.model.layers[DENSE_LAYER]
    raw: dict[str, torch.Tensor] = {}
    handles = [
        layer.mlp.gate_proj.register_forward_hook(
            lambda _m, _i, out: raw.__setitem__("pre", out.detach())
        ),
        layer.mlp.down_proj.register_forward_pre_hook(
            lambda _m, args: raw.__setitem__("post", args[0].detach())
        ),
    ]

    try:
        with torch.no_grad():
            _, cache = bridge.run_with_cache(tokens, names_filter=names)
    finally:
        for handle in handles:
            handle.remove()

    for name, key in zip(names, ("pre", "post"), strict=True):
        print(
            f"{name}: bridge {tuple(cache[name].shape)} "
            f"vs HF module {tuple(raw[key].shape)}"
        )


if __name__ == "__main__":
    section = sys.argv[1] if len(sys.argv) > 1 else "all"
    if section in ("structure", "all"):
        structure()
    if section in ("shapes", "all"):
        shapes()

Observed on transformer-lens==3.7.0:

deepseek-ai/DeepSeek-V2-Lite: hidden_size=2048 intermediate_size=10944
  moe_intermediate_size=1408 first_k_dense_replace=1
MoEBridge.hook_aliases = {'hook_pre': 'hook_in', 'hook_post': 'hook_out'}
adapter maps every block's mlp to MoEBridge
blocks.0.mlp.hook_pre:  bridge (1, 8, 2048) vs HF module (1, 8, 10944)
blocks.0.mlp.hook_post: bridge (1, 8, 2048) vs HF module (1, 8, 10944)

Expected behavior

On a dense DeepseekV2MLP layer, the compatibility hooks should have the same semantics as GatedMLPBridge:

hook_pre        -> gate_proj output
hook_pre_linear -> up_proj output
hook_post       -> down_proj input

All three tensors should be intermediate_size wide.

The sparse layers can retain MoE-specific behavior, but dense layers should not expose block-boundary tensors under neuron-hook names.

Actual behavior

On the dense layer:

hook_pre  -> mlp input   [d_model]
hook_post -> mlp output  [d_model]

hook_pre_linear is not present.

The rest of the layer can still match the Hugging Face model, so this is easy to miss: the aliases return valid tensors of the expected rank and dtype, but with the wrong semantics and width.

Root cause

MoEBridge aliases the two MLP hooks to its own input and output:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/generalized_components/moe.py#L17-L28

The DeepSeek V2 adapter explicitly recognizes that early layers may be DeepseekV2MLP, but still installs the same MoEBridge for every layer:

• https://github.com/TransformerLensOrg/TransformerLens/blob/v3.7.1/transformer_lens/model_bridge/supported_architectures/deepseek_v2.py#L85-L104

Hugging Face selects DeepseekV2MLP below first_k_dense_replace, and that module has ordinary gate/up/down projections:

• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py#L127-L139
• https://github.com/huggingface/transformers/blob/v5.15.0/src/transformers/models/deepseek_v2/modeling_deepseek_v2.py#L379-L384

Impact

Code that indexes hook_pre or hook_post as a neuron-basis tensor silently receives a residual-width tensor on the dense prefix. A shape check catches this model because hidden_size != intermediate_size, but no TransformerLens error identifies the semantic mismatch. If the two widths happened to be equal, the failure would be shape-invisible as well.

Suggested fix

Use a layer-aware or runtime-dispatched MLP bridge:

• use GatedMLPBridge semantics when the instantiated module is DeepseekV2MLP;
• use MoEBridge behavior when the instantiated module is DeepseekV2Moe.

If the component-mapping system cannot select a different bridge class per list item, a hybrid bridge could inspect the actual module during setup and install the dense aliases only when gate_proj, up_proj, and down_proj are present.

Suggested regression tests:

1. DeepSeek-V2-Lite layer 0 exposes all three gated-MLP hooks.
2. Their final dimension equals intermediate_size.
3. hook_pre, hook_pre_linear, and hook_post equal gate_proj output, up_proj output, and down_proj input from the same forward.
4. A later sparse layer retains the intended MoE behavior.