diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 1e11891f6..2fb6946da 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/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" diff --git a/transformer_lens/model_bridge/generalized_components/__init__.py b/transformer_lens/model_bridge/generalized_components/__init__.py index 00225510b..0f34db3dc 100644 --- a/transformer_lens/model_bridge/generalized_components/__init__.py +++ b/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", diff --git a/transformer_lens/model_bridge/generalized_components/normalization.py b/transformer_lens/model_bridge/generalized_components/normalization.py index f500bcb60..6b1ebcc8f 100644 --- a/transformer_lens/model_bridge/generalized_components/normalization.py +++ b/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) diff --git a/transformer_lens/model_bridge/sources/native/init.py b/transformer_lens/model_bridge/sources/native/init.py index 229ac7e27..520ca13d9 100644 --- a/transformer_lens/model_bridge/sources/native/init.py +++ b/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) diff --git a/transformer_lens/model_bridge/sources/native/model.py b/transformer_lens/model_bridge/sources/native/model.py index 5617459e7..5da5a1b6e 100644 --- a/transformer_lens/model_bridge/sources/native/model.py +++ b/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) diff --git a/transformer_lens/model_bridge/supported_architectures/native.py b/transformer_lens/model_bridge/supported_architectures/native.py index 321fdb8cc..aabdb9854 100644 --- a/transformer_lens/model_bridge/supported_architectures/native.py +++ b/transformer_lens/model_bridge/supported_architectures/native.py @@ -12,20 +12,25 @@ BlockBridge, EmbeddingBridge, GatedMLPBridge, + LayerNormPreBridge, LinearBridge, MLPBridge, NormalizationBridge, PosEmbedBridge, RMSNormalizationBridge, + RMSNormPreBridge, UnembeddingBridge, ) -from transformer_lens.model_bridge.generalized_components.base import ( - GeneralizedComponent, -) def _uses_rms(cfg: Any) -> bool: - return (getattr(cfg, "normalization_type", None) or "LN").upper() in ("RMS", "RMSPRE") + norm_type = (getattr(cfg, "normalization_type", None) or "LN").upper() + return norm_type in ("RMS", "RMSPRE") + + +def _uses_param_free_norm(cfg: Any) -> bool: + norm_type = (getattr(cfg, "normalization_type", None) or "LN").upper() + return norm_type in ("RMSPRE", "LNPRE") def _uses_no_norm(cfg: Any) -> bool: @@ -37,9 +42,20 @@ def _is_rotary(cfg: Any) -> bool: def _make_norm_bridge(name: str, cfg: Any, *, force_rms: bool = False): - if force_rms or _uses_rms(cfg): + norm_type = (getattr(cfg, "normalization_type", None) or "LN").upper() + is_param_free = _uses_param_free_norm(cfg) + + if force_rms or norm_type in ("RMS", "RMSPRE"): + if is_param_free or norm_type == "RMSPRE": + return RMSNormPreBridge(name=name, config=cfg) return RMSNormalizationBridge(name=name, config=cfg) + if norm_type == "LNPRE": + return LayerNormPreBridge(name=name, config=cfg) if _uses_no_norm(cfg): + from transformer_lens.model_bridge.generalized_components.base import ( + GeneralizedComponent, + ) + return GeneralizedComponent(name=name, config=cfg) return NormalizationBridge(name=name, config=cfg) @@ -91,13 +107,16 @@ class NativeArchitectureAdapter(ArchitectureAdapter): def __init__(self, cfg: Any) -> None: super().__init__(cfg) - # Native layout already stores Q/K/V split; no rearranges needed. - # Compatibility-mode fold_ln / center_writing_weights aren't wired up, - # so gate the corresponding ProcessWeights paths off — folding without - # the state-dict conversions would mis-place or drop weights. - self.supports_fold_ln = False - self.supports_center_writing_weights = False - self.weight_processing_conversions = {} + # Support fold_ln (LN weight folds into downstream layers) and + # center_writing_weights (residual-stream-writing weights are centered). + self.supports_fold_ln = True + self.supports_center_writing_weights = True + # Native Q/K/V/O stored as nn.Linear [out, in]; fold_ln formulas expect + # TL [head, d_model, d_head] format. Use the shared helper to wire up the + # necessary rearranges so fold_layer_norm sees the right shapes. + self.weight_processing_conversions = { + **self._qkvo_weight_conversions(include_biases=True), + } # Internal attribute names avoid collisions with bridge slot names # ("embed", "blocks", "ln_final", "unembed") — the bridge's __getattr__