feat(native): add param-free pre-norm (LNPre/RMSPre) and fold_ln support

tests/unit/model_bridge/test_boot_native.py

@@ -127,6 +127,113 @@ def test_boot_native_distinct_seeds_diverge():
     assert any(diffs), "Two different seeds produced identical params"
 
 
+def test_boot_native_lnpre_param_free():
+    """LNPre builds param-free norm (no learnable weight/bias)."""
+    from transformer_lens.model_bridge.generalized_components.base import (
+        GeneralizedComponent,
+    )
+    from transformer_lens.model_bridge.sources.native.model import NativeLayerNormPre
+
+    bridge = TransformerBridge.boot_native(_cfg(normalization_type="LNPre"))
+    native_model = bridge.original_model
+
+    ln1_wrapped = native_model.layers[0].ln1
+    ln2_wrapped = native_model.layers[0].ln2
+    ln_out_wrapped = native_model.ln_out
+
+    assert isinstance(ln1_wrapped, GeneralizedComponent)
+    assert isinstance(ln2_wrapped, GeneralizedComponent)
+    assert isinstance(ln_out_wrapped, GeneralizedComponent)
+
+    ln1 = ln1_wrapped._original_component
+    ln2 = ln2_wrapped._original_component
+    ln_out = ln_out_wrapped._original_component
+
+    assert isinstance(ln1, NativeLayerNormPre)
+    assert isinstance(ln2, NativeLayerNormPre)
+    assert isinstance(ln_out, NativeLayerNormPre)
+
+    for norm_module in [ln1, ln2, ln_out]:
+        assert not hasattr(norm_module, "weight") or not isinstance(
+            getattr(norm_module, "weight", None), torch.nn.Parameter
+        ), f"LNPre should have no learnable weight"
+        assert not hasattr(norm_module, "bias") or not isinstance(
+            getattr(norm_module, "bias", None), torch.nn.Parameter
+        ), f"LNPre should have no learnable bias"
+
+
+def test_boot_native_rmspre_param_free():
+    """RMSPre builds param-free RMS norm (no learnable weight)."""
+    from transformer_lens.model_bridge.generalized_components.base import (
+        GeneralizedComponent,
+    )
+    from transformer_lens.model_bridge.sources.native.model import NativeRMSNormPre
+
+    bridge = TransformerBridge.boot_native(_cfg(normalization_type="RMSPre"))
+    native_model = bridge.original_model
+
+    ln1_wrapped = native_model.layers[0].ln1
+    ln2_wrapped = native_model.layers[0].ln2
+    ln_out_wrapped = native_model.ln_out
+
+    assert isinstance(ln1_wrapped, GeneralizedComponent)
+    assert isinstance(ln2_wrapped, GeneralizedComponent)
+    assert isinstance(ln_out_wrapped, GeneralizedComponent)
+
+    ln1 = ln1_wrapped._original_component
+    ln2 = ln2_wrapped._original_component
+    ln_out = ln_out_wrapped._original_component
+
+    assert isinstance(ln1, NativeRMSNormPre)
+    assert isinstance(ln2, NativeRMSNormPre)
+    assert isinstance(ln_out, NativeRMSNormPre)
+
+    for norm_module in [ln1, ln2, ln_out]:
+        assert not hasattr(norm_module, "weight") or not isinstance(
+            getattr(norm_module, "weight", None), torch.nn.Parameter
+        ), f"RMSPre should have no learnable weight"
+
+
+def test_boot_native_supports_fold_ln():
+    """Native adapter supports fold_ln and center_writing_weights."""
+    bridge = TransformerBridge.boot_native(_cfg())
+    assert bridge.adapter.supports_fold_ln is True
+    assert bridge.adapter.supports_center_writing_weights is True
+
+
+def test_boot_native_lnpre_forward():
+    """LNPre forward produces correct normalization (param-free LayerNorm)."""
+    from transformer_lens.model_bridge.sources.native.model import NativeLayerNormPre
+
+    norm = NativeLayerNormPre(eps=1e-5)
+    x = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], dtype=torch.float32)
+
+    out = norm(x)
+
+    expected_mean = x.mean(dim=-1, keepdim=True)
+    expected_centered = x - expected_mean
+    expected_scale = (expected_centered.pow(2).mean(dim=-1, keepdim=True) + 1e-5).sqrt()
+    expected = expected_centered / expected_scale
+
+    assert torch.allclose(out, expected, atol=1e-6)
+
+
+def test_boot_native_rmspre_forward():
+    """RMSPre forward produces correct normalization (param-free RMS norm)."""
+    from transformer_lens.model_bridge.sources.native.model import NativeRMSNormPre
+
+    norm = NativeRMSNormPre(eps=1e-5)
+    x = torch.tensor([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]], dtype=torch.float32)
+
+    out = norm(x)
+
+    x_fp32 = x.to(torch.float32)
+    rms_inv = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + 1e-5)
+    expected = (x_fp32 * rms_inv).to(x.dtype)
+
+    assert torch.allclose(out, expected, atol=1e-6)
+
+
 def test_boot_native_forward_and_cache():
     cfg = _cfg()
     bridge = TransformerBridge.boot_native(cfg)
@@ -323,3 +430,87 @@ def test_boot_native_supports_training_step():
     ), "No non-zero gradients after backward"
     optimizer.step()
     optimizer.zero_grad()
+
+
+def test_boot_native_fold_ln_output_invariant():
+    """fold_ln should not change model output (mathematically equivalent).
+
+    Uses randomized LN weights to actually exercise the folding math, and compares
+    at the logit level for precision.
+    """
+    cfg = _cfg(n_layers=2)
+    bridge = TransformerBridge.boot_native(cfg)
+
+    # Randomize LN weights to actually test folding (identity weights make fold a no-op)
+    torch.manual_seed(42)
+    with torch.no_grad():
+        for name, param in bridge.named_parameters():
+            if "ln" in name.lower() and "weight" in name:
+                param.copy_(torch.randn_like(param) * 0.5 + 1.0)
+            elif "ln" in name.lower() and "bias" in name:
+                param.copy_(torch.randn_like(param) * 0.1)
+
+    inputs = torch.randint(0, cfg.d_vocab, (2, cfg.n_ctx))
+    with torch.no_grad():
+        logits_unfolded = bridge(inputs, return_type="logits")
+
+    bridge.enable_compatibility_mode(
+        fold_ln=True,
+        center_writing_weights=False,
+        center_unembed=False,
+        fold_value_biases=False,
+        refactor_factored_attn_matrices=False,
+    )
+
+    with torch.no_grad():
+        logits_folded = bridge(inputs, return_type="logits")
+
+    assert torch.allclose(
+        logits_folded, logits_unfolded, atol=1e-4, rtol=1e-4
+    ), f"fold_ln should not change logits: max diff={torch.abs(logits_folded - logits_unfolded).max():.6e}"
+
+
+def test_boot_native_lnpre_compatibility_mode():
+    """LNPre models should work with enable_compatibility_mode without crashing.
+
+    LNPre has no weights to fold, so fold_ln should be a no-op. The key is that
+    it doesn't crash (e.g. EinopsError from shape mismatches).
+    """
+    cfg = _cfg(n_layers=2, normalization_type="LNPre")
+    bridge = TransformerBridge.boot_native(cfg)
+
+    inputs = torch.randint(0, cfg.d_vocab, (2, cfg.n_ctx))
+    with torch.no_grad():
+        logits_before = bridge(inputs, return_type="logits")
+
+    # This should not crash — fold_ln on param-free norms is a no-op
+    bridge.enable_compatibility_mode(
+        fold_ln=True,
+        center_writing_weights=False,
+        center_unembed=False,
+        fold_value_biases=False,
+        refactor_factored_attn_matrices=False,
+    )
+
+    with torch.no_grad():
+        logits_after = bridge(inputs, return_type="logits")
+
+    # Since LNPre has no weights, output should be unchanged
+    assert torch.allclose(
+        logits_after, logits_before, atol=1e-6
+    ), f"LNPre fold_ln should be no-op: max diff={torch.abs(logits_after - logits_before).max():.6e}"
+
+
+def test_boot_native_lnpre_has_hooks():
+    """LNPre should expose hook_scale and hook_normalized for ActivationCache compatibility."""
+    cfg = _cfg(normalization_type="LNPre")
+    bridge = TransformerBridge.boot_native(cfg)
+
+    inputs = torch.randint(0, cfg.d_vocab, (2, cfg.n_ctx))
+    _, cache = bridge.run_with_cache(inputs, return_type="logits")
+
+    # Check that the scale and normalized hooks are present
+    assert "blocks.0.ln1.hook_scale" in cache, "LNPre should have hook_scale"
+    assert "blocks.0.ln1.hook_normalized" in cache, "LNPre should have hook_normalized"
+    assert "ln_final.hook_scale" in cache, "ln_final should have hook_scale"
+    assert "ln_final.hook_normalized" in cache, "ln_final should have hook_normalized"

transformer_lens/model_bridge/generalized_components/__init__.py

@@ -81,7 +81,9 @@
     MPTALiBiAttentionBridge,
 )
 from transformer_lens.model_bridge.generalized_components.normalization import (
+    LayerNormPreBridge,
     NormalizationBridge,
+    RMSNormPreBridge,
 )
 from transformer_lens.model_bridge.generalized_components.opaque_block import (
     OpaqueBlockBridge,
@@ -165,8 +167,10 @@
     "ALiBiJointQKVAttentionBridge",
     "RotaryEmbeddingBridge",
     "PosEmbedBridge",
+    "LayerNormPreBridge",
     "NormalizationBridge",
     "RMSNormalizationBridge",
+    "RMSNormPreBridge",
     "JointQKVAttentionBridge",
     "JointQKVPositionEmbeddingsAttentionBridge",
     "JointGateUpMLPBridge",

transformer_lens/model_bridge/generalized_components/normalization.py

@@ -222,3 +222,65 @@ def _hf_autograd_forward_with_hooks(self, x: torch.Tensor) -> torch.Tensor:
             return result
         warnings.warn(NATIVE_PATH_EDIT_FALLBACK_WARNING)
         return self._apply_weight_and_bias(hooked_normalized, input_dtype)
+
+
+class LayerNormPreBridge(NormalizationBridge):
+    """Bridge for param-free LayerNorm (LNPre) — exposes hook_scale/hook_normalized but no weights.
+
+    Used by native models with normalization_type="LNPre" where the normalization has
+    no learnable parameters. The forward pass centers and normalizes without applying
+    any weight or bias.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: Any,
+        submodules: Optional[Dict[str, GeneralizedComponent]] = None,
+        optional: bool = False,
+    ):
+        super().__init__(
+            name,
+            config,
+            submodules=submodules or {},
+            use_native_layernorm_autograd=False,
+            uses_rms_norm=False,
+            optional=optional,
+        )
+
+    def _apply_weight_and_bias(
+        self, hidden_states: torch.Tensor, input_dtype: torch.dtype
+    ) -> torch.Tensor:
+        """No-op for param-free norms — just cast back to input dtype."""
+        return hidden_states.to(input_dtype)
+
+
+class RMSNormPreBridge(NormalizationBridge):
+    """Bridge for param-free RMSNorm (RMSPre) — exposes hook_scale/hook_normalized but no weights.
+
+    Used by native models with normalization_type="RMSPre" where the normalization has
+    no learnable parameters. The forward pass normalizes by RMS without applying
+    any learnable scale.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        config: Any,
+        submodules: Optional[Dict[str, GeneralizedComponent]] = None,
+        optional: bool = False,
+    ):
+        super().__init__(
+            name,
+            config,
+            submodules=submodules or {},
+            use_native_layernorm_autograd=False,
+            uses_rms_norm=True,
+            optional=optional,
+        )
+
+    def _apply_weight_and_bias(
+        self, hidden_states: torch.Tensor, input_dtype: torch.dtype
+    ) -> torch.Tensor:
+        """No-op for param-free norms — just cast back to input dtype."""
+        return hidden_states.to(input_dtype)

transformer_lens/model_bridge/sources/native/init.py

@@ -23,9 +23,11 @@
     NativeAttention,
     NativeBlock,
     NativeGatedMLP,
+    NativeLayerNormPre,
     NativeMLP,
     NativeModel,
     NativeRMSNorm,
+    NativeRMSNormPre,
 )
 
 # Residual-scaled output is gpt2-specific; other modes treat every weight the
@@ -98,6 +100,8 @@ def initialize_native_model(
 def _init_norm(norm: nn.Module) -> None:
     if isinstance(norm, NativeRMSNorm):
         nn.init.ones_(norm.weight)
+    elif isinstance(norm, (NativeRMSNormPre, NativeLayerNormPre)):
+        pass
     elif isinstance(norm, nn.LayerNorm):
         nn.init.ones_(norm.weight)
         nn.init.zeros_(norm.bias)

transformer_lens/model_bridge/sources/native/model.py

@@ -66,11 +66,59 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.weight * normalized
 
 
+class NativeRMSNormPre(nn.Module):
+    """Param-free RMSNorm — normalization only, no learnable scale."""
+
+    def __init__(self, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        x_fp32 = x.to(torch.float32)
+        rms_inv = torch.rsqrt(x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+        return (x_fp32 * rms_inv).to(input_dtype)
+
+
+class NativeLayerNormPre(nn.Module):
+    """Param-free LayerNorm — center + normalize only, no learnable scale/bias.
+
+    Computes in fp32 for numerical stability, matching NativeRMSNormPre and
+    HookedTransformer's LayerNormPre.
+    """
+
+    def __init__(self, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        x_fp32 = x.to(torch.float32)
+        x_fp32 = x_fp32 - x_fp32.mean(dim=-1, keepdim=True)
+        scale = (x_fp32.pow(2).mean(dim=-1, keepdim=True) + self.eps).sqrt()
+        return (x_fp32 / scale).to(input_dtype)
+
+
+def _is_param_free_norm(cfg: TransformerBridgeConfig) -> bool:
+    """Check if the config specifies a param-free normalization type."""
+    return _normalization_type(cfg) in ("RMSPRE", "LNPRE")
+
+
 def _make_norm(cfg: TransformerBridgeConfig, *, force_rms: bool = False) -> nn.Module:
-    if force_rms or _uses_rms_norm(cfg):
+    norm_type = _normalization_type(cfg)
+    is_param_free = _is_param_free_norm(cfg)
+
+    if force_rms or norm_type in ("RMS", "RMSPRE"):
+        if is_param_free or norm_type == "RMSPRE":
+            return NativeRMSNormPre(eps=cfg.eps)
         return NativeRMSNorm(cfg.d_model, eps=cfg.eps)
+
+    if norm_type == "LNPRE":
+        return NativeLayerNormPre(eps=cfg.eps)
+
     if _uses_no_norm(cfg):
         return nn.Identity()
+
     return nn.LayerNorm(cfg.d_model, eps=cfg.eps)