diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index e620d07ae..2db3c8822 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/tests/unit/model_bridge/test_boot_native.py @@ -573,6 +573,7 @@ def test_native_adapter_declares_fold_ln_support(): "overrides", [ pytest.param({"normalization_type": "LN"}, id="LN"), + pytest.param({"normalization_type": "LN", "final_rms": True}, id="LN-final-RMS"), pytest.param({"normalization_type": "RMS"}, id="RMS"), pytest.param({"normalization_type": "LNPre"}, id="LNPre"), pytest.param({"normalization_type": "RMSPre"}, id="RMSPre"), @@ -646,6 +647,54 @@ def test_boot_native_weight_processing_preserves_predictions(overrides, recwarn) torch.testing.assert_close(after, before, atol=1e-4, rtol=1e-4) +@pytest.mark.parametrize("nonidentity_scale", [False, True]) +def test_mixed_final_rms_fold_only_preserves_logits_and_unembed(nonidentity_scale): + bridge = TransformerBridge.boot_native( + _cfg(n_layers=2, normalization_type="LN", final_rms=True) + ).eval() + with torch.no_grad(): + if nonidentity_scale: + bridge.ln_final.weight.copy_(torch.linspace(0.5, 1.5, bridge.cfg.d_model)) + weight = bridge.W_U.detach().clone() + scale = bridge.ln_final.weight.detach().clone() + tokens = torch.tensor([[2, 3, 4]]) + before = bridge(tokens) + 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(): + after = bridge(tokens) + torch.testing.assert_close(bridge.W_U, weight * scale[:, None]) + torch.testing.assert_close(bridge.ln_final.weight, torch.ones_like(scale)) + torch.testing.assert_close(after, before, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("fold_ln", [False, True]) +def test_mixed_final_rms_value_bias_folding_preserves_predictions(fold_ln): + bridge = TransformerBridge.boot_native( + _cfg(n_layers=2, normalization_type="LN", final_rms=True) + ).eval() + with torch.no_grad(): + for block in bridge.blocks: + block.attn.v.bias.fill_(0.2) + tokens = torch.tensor([[2, 3, 4]]) + before = bridge(tokens).log_softmax(-1) + bridge.enable_compatibility_mode( + fold_ln=fold_ln, + center_writing_weights=True, + center_unembed=False, + fold_value_biases=True, + refactor_factored_attn_matrices=False, + ) + with torch.no_grad(): + after = bridge(tokens).log_softmax(-1) + torch.testing.assert_close(after, before, atol=1e-5, rtol=1e-5) + + def test_boot_native_fold_ln_zeroes_norm_parameters(): """fold_ln must actually absorb the norm into the downstream weights — an untouched ln1.weight means the fold silently skipped the layer.""" diff --git a/tests/unit/model_bridge/test_mixed_norm_decomposition.py b/tests/unit/model_bridge/test_mixed_norm_decomposition.py new file mode 100644 index 000000000..0c75922c0 --- /dev/null +++ b/tests/unit/model_bridge/test_mixed_norm_decomposition.py @@ -0,0 +1,110 @@ +"""Mixed block-LN/final-RMS predictions and cached decompositions must both reconcile.""" + +import pytest +import torch + +from transformer_lens.config import TransformerBridgeConfig +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis import direct_logit_attribution + + +@pytest.fixture(params=["LN", "LNPre"]) +def mixed_bridge(request): + model = TransformerBridge.boot_native( + TransformerBridgeConfig( + d_model=16, + d_head=8, + n_heads=2, + n_layers=2, + n_ctx=8, + d_vocab=32, + d_mlp=24, + act_fn="gelu", + normalization_type=request.param, + final_rms=True, + seed=3, + ) + ).eval() + return model + + +@pytest.mark.parametrize("processing", ["none", "fold_only", "default"]) +@pytest.mark.parametrize("batchless", [False, True]) +def test_mixed_final_rms_decompositions_reconcile(mixed_bridge, processing, batchless): + model = mixed_bridge + tokens = torch.tensor([[2, 3, 4], [4, 5, 6]]) + with torch.no_grad(): + if processing != "none" and hasattr(model.ln_final, "weight"): + model.ln_final.weight.copy_(torch.linspace(0.5, 1.5, model.cfg.d_model)) + before = model(tokens).log_softmax(-1) + if processing == "none": + model.enable_compatibility_mode(no_processing=True) + elif processing == "fold_only": + model.enable_compatibility_mode( + center_writing_weights=False, center_unembed=False, fold_value_biases=False + ) + else: + model.enable_compatibility_mode() + if batchless: + tokens = tokens[:1] + with torch.no_grad(): + logits, cache = model.run_with_cache(tokens, remove_batch_dim=batchless) + torch.testing.assert_close( + logits.log_softmax(-1), before[: tokens.shape[0]], atol=1e-5, rtol=1e-5 + ) + residual = cache["blocks.1.hook_resid_post"] + normalized = cache.apply_ln_to_stack( + residual.unsqueeze(0), layer=-1, has_batch_dim=not batchless + )[0] + torch.testing.assert_close(normalized, residual / cache["ln_final.hook_scale"]) + direction = model.W_U[:, 7] - model.W_U[:, 8] + expected = logits[..., 7] - logits[..., 8] - (model.b_U[7] - model.b_U[8]) + if batchless: + expected = expected[0] + for expand_neurons in (False, True): + dense = cache.get_full_resid_decomposition(expand_neurons=expand_neurons, apply_ln=True) + projected = cache.get_full_resid_decomposition( + expand_neurons=expand_neurons, apply_ln=True, project_output_onto=direction + ) + torch.testing.assert_close(projected, dense @ direction, atol=1e-5, rtol=1e-5) + torch.testing.assert_close(projected.sum(0), expected, atol=1e-5, rtol=1e-5) + if not batchless: + result = direct_logit_attribution(model, answer_tokens=7, incorrect_tokens=8, cache=cache) + torch.testing.assert_close(result.attribution.sum(0), expected[:, -1], atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("mlp_input", [False, True]) +def test_intermediate_layernorm_still_centers(mixed_bridge, mlp_input): + model = mixed_bridge + model.enable_compatibility_mode(no_processing=True) + with torch.no_grad(): + _, cache = model.run_with_cache(torch.tensor([[2, 3, 4]])) + residual = cache["blocks.1.hook_resid_mid" if mlp_input else "blocks.1.hook_resid_pre"] + expected = (residual - residual.mean(-1, keepdim=True)) / cache[ + f"blocks.1.ln{2 if mlp_input else 1}.hook_scale" + ] + actual = cache.apply_ln_to_stack(residual.unsqueeze(0), layer=1, mlp_input=mlp_input)[0] + torch.testing.assert_close(actual, expected) + direction = torch.linspace(0.1, 1.0, model.cfg.d_model) + dense = cache.stack_neuron_results(layer=1, apply_ln=True, mlp_input=mlp_input) + projected = cache.stack_neuron_results( + layer=1, apply_ln=True, mlp_input=mlp_input, project_output_onto=direction + ) + torch.testing.assert_close(projected, dense @ direction, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("matrix_projection", [False, True]) +def test_final_rms_projected_neurons_match_uncentered_oracle(mixed_bridge, matrix_projection): + model = mixed_bridge + model.enable_compatibility_mode(no_processing=True) + with torch.no_grad(): + _, cache = model.run_with_cache(torch.tensor([[2, 3, 4], [4, 5, 6]])) + direction = torch.linspace(0.1, 1.0, model.cfg.d_model) + if matrix_projection: + direction = torch.stack([direction, direction.square()], dim=-1) + raw = cache.stack_neuron_results(layer=model.cfg.n_layers, pos_slice=-1, apply_ln=False) + expected = (raw / cache["ln_final.hook_scale"][:, -1]) @ direction + actual = cache.stack_neuron_results( + layer=model.cfg.n_layers, pos_slice=-1, apply_ln=True, project_output_onto=direction + ) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5) diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index 15529ffd4..6247dd4c3 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -1179,6 +1179,13 @@ def get_neuron_results( return neuron_acts * projected return neuron_acts[..., None] * projected + def _uses_ln_centering(self, layer: Optional[int]) -> bool: + """Use the target norm's semantics, not just the block normalization family.""" + is_final = layer is None or layer in (-1, self.model.cfg.n_layers) + return self.model.cfg.normalization_type in ["LN", "LNPre"] and not ( + is_final and getattr(self.model.cfg, "final_rms", False) + ) + def _get_cached_ln_scale( self, layer: Optional[int], @@ -1238,7 +1245,7 @@ def _stack_neuron_results_apply_ln_projected( """ scale = self._get_cached_ln_scale(layer, mlp_input=mlp_input, pos_slice=pos_slice) - apply_centering = self.model.cfg.normalization_type in ["LN", "LNPre"] + apply_centering = self._uses_ln_centering(layer) sum_p = project_2d.sum(dim=0) if apply_centering else None # [n_outs] components: list = [] @@ -1508,8 +1515,8 @@ def apply_ln_to_stack( results.append(out.reshape(original_shape)) return torch.stack(results, dim=0) - # Center the stack onlny if the model uses LayerNorm - if self.model.cfg.normalization_type in ["LN", "LNPre"]: + # Final RMSNorm must not inherit the block LayerNorm centering convention. + if self._uses_ln_centering(layer): residual_stack = residual_stack - residual_stack.mean(dim=-1, keepdim=True) # Shape is [batch, position, 1] or [position, 1]; final dim is a dummy for broadcasting. diff --git a/transformer_lens/weight_processing.py b/transformer_lens/weight_processing.py index af3113505..5328a0f04 100644 --- a/transformer_lens/weight_processing.py +++ b/transformer_lens/weight_processing.py @@ -1038,7 +1038,8 @@ def _fold_unembed_layer_norm( ) if alternate_final_w_key != ln_final_w_key and alternate_final_w_key in state_dict: state_dict[alternate_final_w_key] = identity_val - if center_weights: + # RMSNorm outputs are not zero-mean, so centering would change predictions. + if center_weights and not getattr(cfg, "final_rms", False): unembed_weight_centered = ProcessWeights.convert_tensor_to_tl_format( unembed_W_U_key, state_dict, state_dict.get(unembed_W_U_key), cfg, adapter, None ) @@ -1701,19 +1702,18 @@ def process_weights( UserWarning, ) center_writing_weights = False - if center_writing_weights: - if getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"] and ( - not getattr(cfg, "final_rms", False) - ): - state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter) + should_center_writing_weights = ( + center_writing_weights + and getattr(cfg, "normalization_type", "LN") in ["LN", "LNPre"] + and not getattr(cfg, "final_rms", False) + ) + if should_center_writing_weights: + state_dict = ProcessWeights.center_writing_weights(state_dict, cfg, adapter=adapter) if center_unembed: state_dict = ProcessWeights.center_unembed(state_dict, cfg=cfg, adapter=adapter) if fold_value_biases: state_dict = ProcessWeights.fold_value_biases(state_dict, cfg, adapter=adapter) - if center_writing_weights and getattr(cfg, "normalization_type", "LN") in [ - "LN", - "LNPre", - ]: + if should_center_writing_weights: for layer_idx in range(cfg.n_layers): b_O_key = ProcessWeights._get_param_key(f"blocks.{layer_idx}.attn.b_O", adapter) if b_O_key in state_dict: