From f1073649459d9dae70661390444994aa734c0f96 Mon Sep 17 00:00:00 2001 From: Abhinav Bellapu Date: Sun, 12 Jul 2026 17:41:36 -0700 Subject: [PATCH 1/2] Add DeepSeek V4 architecture adapter --- .../model_bridge/test_deepseek_v4_adapter.py | 175 ++++++++++++ .../test_deepseek_v4_adapter.py | 123 +++++++++ .../factories/architecture_adapter_factory.py | 2 + .../model_bridge/sources/_bridge_builder.py | 14 + .../supported_architectures/__init__.py | 3 + .../supported_architectures/deepseek_v4.py | 260 ++++++++++++++++++ .../tools/model_registry/__init__.py | 2 + .../tools/model_registry/generate_report.py | 1 + 8 files changed, 580 insertions(+) create mode 100644 tests/integration/model_bridge/test_deepseek_v4_adapter.py create mode 100644 tests/unit/model_bridge/supported_architectures/test_deepseek_v4_adapter.py create mode 100644 transformer_lens/model_bridge/supported_architectures/deepseek_v4.py diff --git a/tests/integration/model_bridge/test_deepseek_v4_adapter.py b/tests/integration/model_bridge/test_deepseek_v4_adapter.py new file mode 100644 index 000000000..285fd7af9 --- /dev/null +++ b/tests/integration/model_bridge/test_deepseek_v4_adapter.py @@ -0,0 +1,175 @@ +"""Download-free integration tests for the DeepSeek V4 architecture adapter.""" + +from typing import NamedTuple + +import pytest +import torch +from transformers import DeepseekV4Config, DeepseekV4ForCausalLM + +from transformer_lens.model_bridge.bridge import TransformerBridge +from transformer_lens.model_bridge.sources import build_bridge_from_module + + +class DeepseekV4Case(NamedTuple): + bridge: TransformerBridge + tokens: torch.Tensor + hf_logits: torch.Tensor + + +@pytest.fixture(scope="module") +def deepseek_v4_case() -> DeepseekV4Case: + """Build a 40K-parameter model spanning sliding, CSA, and HCA layers.""" + torch.manual_seed(0) + cfg = DeepseekV4Config( + vocab_size=64, + hidden_size=32, + moe_intermediate_size=16, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=1, + head_dim=8, + q_lora_rank=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=1, + scoring_func="sigmoid", + routed_scaling_factor=1.0, + max_position_embeddings=32, + layer_types=[ + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + ], + compress_rates={ + "compressed_sparse_attention": 2, + "heavily_compressed_attention": 4, + }, + hc_mult=2, + hc_sinkhorn_iters=2, + mlp_layer_types=["hash_moe", "moe", "moe"], + sliding_window=4, + o_groups=2, + o_lora_rank=8, + index_n_heads=2, + index_head_dim=4, + index_topk=2, + partial_rotary_factor=0.5, + use_cache=False, + ) + cfg._attn_implementation = "eager" + hf_model = DeepseekV4ForCausalLM(cfg).eval() + tokens = torch.arange(8).unsqueeze(0) + + with torch.no_grad(): + hf_logits = hf_model(tokens, use_cache=False).logits + + bridge = build_bridge_from_module( + hf_model, + "DeepseekV4ForCausalLM", + hf_config=cfg, + device="cpu", + model_name="tiny-random-deepseek-v4", + ) + bridge.eval() + return DeepseekV4Case(bridge=bridge, tokens=tokens, hf_logits=hf_logits) + + +def test_forward_matches_hugging_face_exactly(deepseek_v4_case: DeepseekV4Case) -> None: + with torch.no_grad(): + bridge_logits = deepseek_v4_case.bridge(deepseek_v4_case.tokens, use_cache=False) + + torch.testing.assert_close( + bridge_logits, + deepseek_v4_case.hf_logits, + atol=0, + rtol=0, + ) + + +def test_v4_config_metadata_is_preserved(deepseek_v4_case: DeepseekV4Case) -> None: + cfg = deepseek_v4_case.bridge.cfg + assert cfg.layer_types == [ + "sliding_attention", + "compressed_sparse_attention", + "heavily_compressed_attention", + ] + assert cfg.compress_rates == { + "compressed_sparse_attention": 2, + "heavily_compressed_attention": 4, + } + assert cfg.hc_mult == 2 + assert cfg.mlp_layer_types == ["hash_moe", "moe", "moe"] + assert cfg.index_topk == 2 + + +def test_mhc_hooks_preserve_stream_and_collapsed_shapes( + deepseek_v4_case: DeepseekV4Case, +) -> None: + _, cache = deepseek_v4_case.bridge.run_with_cache( + deepseek_v4_case.tokens, + use_cache=False, + ) + + assert cache["blocks.0.hook_in"].shape == (1, 8, 2, 32) + assert cache["blocks.0.attn_hc.hook_post"].shape == (1, 8, 2) + assert cache["blocks.0.attn_hc.hook_comb"].shape == (1, 8, 2, 2) + assert cache["blocks.0.attn_hc.hook_out"].shape == (1, 8, 32) + assert cache["blocks.0.mlp_hc.hook_in"].shape == (1, 8, 2, 32) + assert cache["blocks.0.mlp_hc.hook_out"].shape == (1, 8, 32) + assert cache["blocks.0.hook_out"].shape == (1, 8, 2, 32) + assert cache["hc_head.hook_in"].shape == (1, 8, 2, 32) + assert cache["hc_head.hook_out"].shape == (1, 8, 32) + + +def test_compression_hooks_match_layer_types(deepseek_v4_case: DeepseekV4Case) -> None: + _, cache = deepseek_v4_case.bridge.run_with_cache( + deepseek_v4_case.tokens, + use_cache=False, + ) + + assert not any(key.startswith("blocks.0.attn.compressor") for key in cache) + + for layer in (1, 2): + compressed = cache[f"blocks.{layer}.attn.compressor.hook_out"] + block_bias = cache[f"blocks.{layer}.attn.compressor.hook_block_bias"] + assert compressed.shape[:2] == (1, 1) + assert compressed.shape[-1] == 8 + assert block_bias.shape[:3] == (1, 1, 8) + + indexer = cache["blocks.1.attn.compressor.indexer.hook_out"] + assert indexer.shape == (1, 8, 2) + assert indexer.dtype == torch.long + assert not any(key.startswith("blocks.2.attn.compressor.indexer") for key in cache) + + +def test_hash_and_topk_moe_hooks_fire(deepseek_v4_case: DeepseekV4Case) -> None: + _, cache = deepseek_v4_case.bridge.run_with_cache( + deepseek_v4_case.tokens, + use_cache=False, + ) + + for layer in range(3): + assert cache[f"blocks.{layer}.mlp.gate.hook_out"].shape == (8, 4) + assert cache[f"blocks.{layer}.mlp.experts.hook_out"].shape == (8, 32) + assert cache[f"blocks.{layer}.mlp.shared_experts.hook_out"].shape == ( + 1, + 8, + 32, + ) + + +def test_collapsed_stream_hook_is_patchable(deepseek_v4_case: DeepseekV4Case) -> None: + with torch.no_grad(): + baseline = deepseek_v4_case.bridge(deepseek_v4_case.tokens, use_cache=False) + patched = deepseek_v4_case.bridge.run_with_hooks( + deepseek_v4_case.tokens, + use_cache=False, + fwd_hooks=[ + ( + "blocks.0.attn_hc.hook_out", + lambda activation, hook: torch.zeros_like(activation), + ) + ], + ) + + assert not torch.allclose(baseline, patched) diff --git a/tests/unit/model_bridge/supported_architectures/test_deepseek_v4_adapter.py b/tests/unit/model_bridge/supported_architectures/test_deepseek_v4_adapter.py new file mode 100644 index 000000000..9d3e3fd72 --- /dev/null +++ b/tests/unit/model_bridge/supported_architectures/test_deepseek_v4_adapter.py @@ -0,0 +1,123 @@ +"""Unit tests for the DeepSeek V4 architecture adapter.""" + +import pytest + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge.generalized_components import ( + EmbeddingBridge, + MoEBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) +from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import ( + DeepSeekV4ArchitectureAdapter, + DeepseekV4BlockBridge, + DeepseekV4CompressorBridge, + DeepseekV4HyperConnectionBridge, +) + + +@pytest.fixture +def adapter() -> DeepSeekV4ArchitectureAdapter: + cfg = TransformerBridgeConfig( + d_model=32, + d_head=8, + n_heads=4, + n_layers=3, + n_ctx=32, + d_vocab=64, + d_mlp=16, + n_key_value_heads=1, + architecture="DeepseekV4ForCausalLM", + ) + return DeepSeekV4ArchitectureAdapter(cfg) + + +def test_top_level_mapping(adapter: DeepSeekV4ArchitectureAdapter) -> None: + assert set(adapter.component_mapping) == { + "embed", + "rotary_emb", + "blocks", + "hc_head", + "ln_final", + "unembed", + } + assert isinstance(adapter.component_mapping["embed"], EmbeddingBridge) + assert isinstance(adapter.component_mapping["rotary_emb"], RotaryEmbeddingBridge) + assert isinstance(adapter.component_mapping["unembed"], UnembeddingBridge) + assert adapter.component_mapping["hc_head"].name == "model.hc_head" + + +def test_block_mapping_preserves_mhc_topology(adapter: DeepSeekV4ArchitectureAdapter) -> None: + blocks = adapter.component_mapping["blocks"] + assert isinstance(blocks, DeepseekV4BlockBridge) + assert blocks.name == "model.layers" + assert set(blocks.submodules) == { + "attn_hc", + "ln1", + "attn", + "mlp_hc", + "ln2", + "mlp", + } + assert isinstance(blocks.submodules["attn_hc"], DeepseekV4HyperConnectionBridge) + assert isinstance(blocks.submodules["mlp_hc"], DeepseekV4HyperConnectionBridge) + assert blocks.submodules["mlp_hc"].name == "ffn_hc" + assert blocks.hook_aliases == {} + + +def test_attention_mapping_exposes_compression_surfaces( + adapter: DeepSeekV4ArchitectureAdapter, +) -> None: + attention = adapter.component_mapping["blocks"].submodules["attn"] + assert isinstance(attention, GeneralizedComponent) + assert attention.name == "self_attn" + assert set(attention.submodules) == { + "q_a_proj", + "q_a_norm", + "q_b_proj", + "q_b_norm", + "kv_proj", + "kv_norm", + "compressor", + "o_a_proj", + "o_b_proj", + } + + compressor = attention.submodules["compressor"] + assert isinstance(compressor, DeepseekV4CompressorBridge) + assert compressor.optional + assert compressor.submodules["indexer"].optional + assert set(compressor.submodules["indexer"].submodules) == { + "kv_proj", + "gate_proj", + "kv_norm", + "q_b_proj", + "weights_proj", + "rotary_emb", + } + + +def test_moe_mapping_exposes_both_router_types_and_experts( + adapter: DeepSeekV4ArchitectureAdapter, +) -> None: + mlp = adapter.component_mapping["blocks"].submodules["mlp"] + assert isinstance(mlp, MoEBridge) + assert set(mlp.submodules) == {"gate", "experts", "shared_experts"} + assert mlp.submodules["gate"].name == "gate" + assert mlp.submodules["experts"].name == "experts" + + +def test_config_and_processing_guards(adapter: DeepSeekV4ArchitectureAdapter) -> None: + assert adapter.cfg.normalization_type == "RMS" + assert adapter.cfg.uses_rms_norm + assert adapter.cfg.final_rms + assert not adapter.cfg.rmsnorm_uses_offset + assert adapter.cfg.positional_embedding_type == "rotary" + assert adapter.cfg.gated_mlp + assert adapter.applicable_phases == [2, 4] + assert not adapter.supports_fold_ln + assert not adapter.supports_center_writing_weights diff --git a/transformer_lens/factories/architecture_adapter_factory.py b/transformer_lens/factories/architecture_adapter_factory.py index cef0d22d7..f5283aa36 100644 --- a/transformer_lens/factories/architecture_adapter_factory.py +++ b/transformer_lens/factories/architecture_adapter_factory.py @@ -22,6 +22,7 @@ CohereArchitectureAdapter, DeepSeekV2ArchitectureAdapter, DeepSeekV3ArchitectureAdapter, + DeepSeekV4ArchitectureAdapter, FalconArchitectureAdapter, FalconH1ArchitectureAdapter, Gemma1ArchitectureAdapter, @@ -102,6 +103,7 @@ "CohereForCausalLM": CohereArchitectureAdapter, "DeepseekV2ForCausalLM": DeepSeekV2ArchitectureAdapter, "DeepseekV3ForCausalLM": DeepSeekV3ArchitectureAdapter, + "DeepseekV4ForCausalLM": DeepSeekV4ArchitectureAdapter, "FalconForCausalLM": FalconArchitectureAdapter, "FalconH1ForCausalLM": FalconH1ArchitectureAdapter, "GemmaForCausalLM": Gemma1ArchitectureAdapter, # Default to Gemma1 as it's the original version diff --git a/transformer_lens/model_bridge/sources/_bridge_builder.py b/transformer_lens/model_bridge/sources/_bridge_builder.py index c64930a16..c8f0f7f1e 100644 --- a/transformer_lens/model_bridge/sources/_bridge_builder.py +++ b/transformer_lens/model_bridge/sources/_bridge_builder.py @@ -84,6 +84,20 @@ # Ouro (LoopLM) "total_ut_steps", "early_exit_threshold", + # DeepSeek V4 (mHC + compressed attention) + "compress_rates", + "compress_rope_theta", + "hc_mult", + "hc_sinkhorn_iters", + "hc_eps", + "mlp_layer_types", + "swiglu_limit", + "o_groups", + "o_lora_rank", + "index_n_heads", + "index_head_dim", + "index_topk", + "q_lora_rank", ] diff --git a/transformer_lens/model_bridge/supported_architectures/__init__.py b/transformer_lens/model_bridge/supported_architectures/__init__.py index 6ab2ad971..00c0e9b8c 100644 --- a/transformer_lens/model_bridge/supported_architectures/__init__.py +++ b/transformer_lens/model_bridge/supported_architectures/__init__.py @@ -37,6 +37,9 @@ from transformer_lens.model_bridge.supported_architectures.deepseek_v3 import ( DeepSeekV3ArchitectureAdapter, ) +from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import ( + DeepSeekV4ArchitectureAdapter, +) from transformer_lens.model_bridge.supported_architectures.falcon import ( FalconArchitectureAdapter, ) diff --git a/transformer_lens/model_bridge/supported_architectures/deepseek_v4.py b/transformer_lens/model_bridge/supported_architectures/deepseek_v4.py new file mode 100644 index 000000000..cf634a76d --- /dev/null +++ b/transformer_lens/model_bridge/supported_architectures/deepseek_v4.py @@ -0,0 +1,260 @@ +"""DeepSeek V4 architecture adapter. + +DeepSeek V4 replaces V2/V3's MLA path with a hybrid local/compressed attention +stack and keeps ``hc_mult`` residual streams alive between blocks through +manifold-constrained Hyper-Connections (mHC). The adapter delegates those +architecture-specific calculations to Transformers while exposing the modules +that are useful for interpretability: mHC collapse/mix tensors, compressed KV +states and masks, Lightning Indexer selections, attention projections, and MoE +routing/expert outputs. +""" + +from typing import Any, Dict, Optional + +import torch + +from transformer_lens.hook_points import HookPoint +from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter +from transformer_lens.model_bridge.generalized_components import ( + BlockBridge, + EmbeddingBridge, + GatedMLPBridge, + LinearBridge, + MoEBridge, + RMSNormalizationBridge, + RotaryEmbeddingBridge, + UnembeddingBridge, +) +from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, +) + + +class DeepseekV4HyperConnectionBridge(GeneralizedComponent): + """Bridge an mHC module without discarding its three distinct outputs. + + ``hook_in`` sees the full ``[batch, pos, hc_mult, d_model]`` residual stack. + ``hook_post`` and ``hook_comb`` expose the learned expansion and stream-mix + weights, while ``hook_out`` exposes the collapsed conventional residual that + enters attention or the MLP. + """ + + def __init__( + self, + name: str, + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + ) -> None: + super().__init__(name, config, submodules=submodules or {}) + self.hook_post = HookPoint() + self.hook_comb = HookPoint() + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Run the native mHC module and hook each returned tensor separately.""" + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + + if args and isinstance(args[0], torch.Tensor): + args = (self.hook_in(args[0]),) + args[1:] + elif isinstance(kwargs.get("hidden_streams"), torch.Tensor): + kwargs["hidden_streams"] = self.hook_in(kwargs["hidden_streams"]) + + output = self.original_component(*args, **kwargs) + if not isinstance(output, tuple) or len(output) != 3: + raise RuntimeError( + f"DeepSeek V4 hyper-connection {self.name} returned an unexpected output" + ) + + post, comb, collapsed = output + return self.hook_post(post), self.hook_comb(comb), self.hook_out(collapsed) + + +class DeepseekV4CompressorBridge(GeneralizedComponent): + """Bridge CSA/HCA compression and expose compressed KV plus block bias.""" + + def __init__( + self, + name: str, + config: Optional[Any] = None, + submodules: Optional[Dict[str, GeneralizedComponent]] = None, + optional: bool = False, + ) -> None: + super().__init__( + name, + config, + submodules=submodules or {}, + optional=optional, + ) + self.hook_block_bias = HookPoint() + + def forward(self, *args: Any, **kwargs: Any) -> Any: + """Run the native compressor, preserving and hooking both outputs.""" + if self.original_component is None: + raise RuntimeError( + f"Original component not set for {self.name}. Call set_original_component() first." + ) + + if args and isinstance(args[0], torch.Tensor): + args = (self.hook_in(args[0]),) + args[1:] + elif isinstance(kwargs.get("hidden_states"), torch.Tensor): + kwargs["hidden_states"] = self.hook_in(kwargs["hidden_states"]) + + output = self.original_component(*args, **kwargs) + if not isinstance(output, tuple) or len(output) != 2: + raise RuntimeError(f"DeepSeek V4 compressor {self.name} returned an unexpected output") + + compressed_kv, block_bias = output + compressed_kv = self.hook_out(compressed_kv) + if isinstance(block_bias, torch.Tensor): + block_bias = self.hook_block_bias(block_bias) + return compressed_kv, block_bias + + +class DeepseekV4BlockBridge(BlockBridge): + """Block bridge whose input/output hooks carry the full mHC stream stack. + + Standard residual aliases are intentionally omitted: V4's block boundary is + four-dimensional, and presenting it as a conventional single residual stream + would make otherwise-valid patching code silently target the wrong tensor. + The collapsed attention/MLP inputs are available at ``attn_hc.hook_out`` and + ``mlp_hc.hook_out`` respectively. + """ + + hook_aliases: dict[str, str | list[str]] = {} + maintain_native_attention: bool = True + + +def _compressor_bridge(cfg: Any) -> DeepseekV4CompressorBridge: + """Build the common CSA/HCA compressor mapping, including optional indexer.""" + return DeepseekV4CompressorBridge( + name="compressor", + config=cfg, + optional=True, + submodules={ + "kv_proj": LinearBridge(name="kv_proj"), + "gate_proj": LinearBridge(name="gate_proj"), + "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg), + "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg), + "indexer": GeneralizedComponent( + name="indexer", + optional=True, + submodules={ + "kv_proj": LinearBridge(name="kv_proj"), + "gate_proj": LinearBridge(name="gate_proj"), + "kv_norm": RMSNormalizationBridge(name="kv_norm", config=cfg), + "q_b_proj": LinearBridge(name="q_b_proj"), + "weights_proj": LinearBridge(name="weights_proj"), + "rotary_emb": RotaryEmbeddingBridge(name="rotary_emb", config=cfg), + }, + ), + }, + ) + + +class DeepSeekV4ArchitectureAdapter(ArchitectureAdapter): + """Adapter for ``DeepseekV4ForCausalLM`` (Flash and Pro variants).""" + + # The isolated component harness assumes a three-dimensional residual. V4's + # mHC stack is four-dimensional, so parity is covered by integration tests and + # verify_models' whole-model hook/text phases instead of isolated Phase 1. + applicable_phases: list[int] = [2, 4] + + def __init__(self, cfg: Any) -> None: + super().__init__(cfg) + + self.cfg.normalization_type = "RMS" + self.cfg.uses_rms_norm = True + self.cfg.final_rms = True + self.cfg.rmsnorm_uses_offset = False + self.cfg.positional_embedding_type = "rotary" + self.cfg.gated_mlp = True + self.cfg.attn_implementation = "eager" + + # Folding/centering assumes one additive residual stream. Applying either + # transform to mHC's learned collapse/expand path is not basis preserving. + self.supports_fold_ln = False + self.supports_center_writing_weights = False + self.weight_processing_conversions = {} + + def hyper_connection(name: str) -> DeepseekV4HyperConnectionBridge: + return DeepseekV4HyperConnectionBridge( + name=name, + config=self.cfg, + submodules={ + "input_norm": GeneralizedComponent(name="input_norm"), + }, + ) + + attention = GeneralizedComponent( + name="self_attn", + submodules={ + "q_a_proj": LinearBridge(name="q_a_proj"), + "q_a_norm": RMSNormalizationBridge(name="q_a_norm", config=self.cfg), + "q_b_proj": LinearBridge(name="q_b_proj"), + "q_b_norm": GeneralizedComponent(name="q_b_norm"), + "kv_proj": LinearBridge(name="kv_proj"), + "kv_norm": RMSNormalizationBridge(name="kv_norm", config=self.cfg), + "compressor": _compressor_bridge(self.cfg), + "o_a_proj": GeneralizedComponent(name="o_a_proj"), + "o_b_proj": LinearBridge(name="o_b_proj"), + }, + ) + + mlp = MoEBridge( + name="mlp", + config=self.cfg, + submodules={ + "gate": GeneralizedComponent(name="gate"), + "experts": GeneralizedComponent(name="experts"), + "shared_experts": GatedMLPBridge( + name="shared_experts", + config=self.cfg, + submodules={ + "gate": LinearBridge(name="gate_proj"), + "in": LinearBridge(name="up_proj"), + "out": LinearBridge(name="down_proj"), + }, + ), + }, + ) + + self.component_mapping = { + "embed": EmbeddingBridge(name="model.embed_tokens"), + "rotary_emb": RotaryEmbeddingBridge(name="model.rotary_emb", config=self.cfg), + "blocks": DeepseekV4BlockBridge( + name="model.layers", + config=self.cfg, + submodules={ + "attn_hc": hyper_connection("attn_hc"), + "ln1": RMSNormalizationBridge(name="input_layernorm", config=self.cfg), + "attn": attention, + "mlp_hc": hyper_connection("ffn_hc"), + "ln2": RMSNormalizationBridge(name="post_attention_layernorm", config=self.cfg), + "mlp": mlp, + }, + ), + "hc_head": GeneralizedComponent( + name="model.hc_head", + submodules={ + "input_norm": GeneralizedComponent(name="input_norm"), + }, + ), + "ln_final": RMSNormalizationBridge(name="model.norm", config=self.cfg), + "unembed": UnembeddingBridge(name="lm_head"), + } + + def prepare_loading(self, model_name: str, model_kwargs: dict) -> None: + """Force eager attention so the delegated attention path is deterministic.""" + model_kwargs["attn_implementation"] = "eager" + + def prepare_model(self, hf_model: Any) -> None: + """Force eager attention on a pre-loaded model before installing bridges.""" + if hasattr(hf_model, "config"): + hf_model.config._attn_implementation = "eager" + model = getattr(hf_model, "model", None) + if model is not None and hasattr(model, "layers"): + for layer in model.layers: + if hasattr(layer, "self_attn") and hasattr(layer.self_attn, "config"): + layer.self_attn.config._attn_implementation = "eager" diff --git a/transformer_lens/tools/model_registry/__init__.py b/transformer_lens/tools/model_registry/__init__.py index a304771ef..4fa245c1e 100644 --- a/transformer_lens/tools/model_registry/__init__.py +++ b/transformer_lens/tools/model_registry/__init__.py @@ -58,6 +58,7 @@ "CohereForCausalLM", "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", + "DeepseekV4ForCausalLM", "FalconForCausalLM", "FalconH1ForCausalLM", "GemmaForCausalLM", @@ -137,6 +138,7 @@ "CohereForCausalLM": ["CohereLabs"], "DeepseekV2ForCausalLM": ["deepseek-ai"], "DeepseekV3ForCausalLM": ["deepseek-ai"], + "DeepseekV4ForCausalLM": ["deepseek-ai"], "FalconForCausalLM": ["tiiuae"], "FalconH1ForCausalLM": ["tiiuae"], "Gemma2ForCausalLM": ["google"], diff --git a/transformer_lens/tools/model_registry/generate_report.py b/transformer_lens/tools/model_registry/generate_report.py index 6f514dd6c..c05e9eef3 100644 --- a/transformer_lens/tools/model_registry/generate_report.py +++ b/transformer_lens/tools/model_registry/generate_report.py @@ -41,6 +41,7 @@ "Gemma3nForConditionalGeneration": "Google's Gemma 3n efficient tri-modal model (text-only support)", "Gemma4ForConditionalGeneration": "Google's Gemma 4 multimodal model family (text-only support)", "Gemma4UnifiedForConditionalGeneration": "Google's Gemma 4 unified encoder-free multimodal model (text-only support)", + "DeepseekV4ForCausalLM": "DeepSeek V4 with mHC residual streams and hybrid compressed attention", "GlmMoeDsaForCausalLM": "Z.ai's GLM-5 MoE model with DeepSeek Sparse Attention", "Glm4MoeForCausalLM": "Z.ai's GLM-4.5/4.6/4.7 sparse Mixture-of-Experts causal LM", "Qwen2ForCausalLM": "Alibaba's Qwen2 multilingual model", From a237444ec4e2dd759714b294cdc40941a4dd7685 Mon Sep 17 00:00:00 2001 From: Abhinav Bellapu Date: Sun, 12 Jul 2026 18:47:17 -0700 Subject: [PATCH 2/2] Retrigger CI