diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 0a010dd37..2f0069069 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/tests/unit/model_bridge/test_boot_native.py @@ -1,4 +1,5 @@ """Tests for ``TransformerBridge.boot_native`` classmethod.""" + from __future__ import annotations import sys @@ -489,3 +490,40 @@ def test_boot_native_supports_training_step(): ), "No non-zero gradients after backward" optimizer.step() optimizer.zero_grad() + + +def test_boot_native_resolves_initializer_range_sentinel(): + """Regression for #1568 — the resolved initializer range must be used by boot_native.""" + import math + + cfg = _cfg(init_mode="gpt2") + expected = 0.8 / math.sqrt(cfg.d_model) + + assert cfg.initializer_range == pytest.approx(expected) + + bridge = TransformerBridge.boot_native(cfg) + assert bridge.W_E.std().item() == pytest.approx(expected, rel=0.15) + + +def test_boot_native_resolves_non_gpt2_initializer_range_sentinel(): + cfg = _cfg(init_mode="kaiming_normal") + + assert cfg.initializer_range == pytest.approx(1.0) + + +def test_boot_native_kaiming_gain_scales_weights(): + """Regression for #1568 — xavier/kaiming init must use initializer_range + as a multiplicative gain. Without this, the config value is silently + ignored and every kaiming/xavier model gets the same fixed scale + regardless of what the caller asked for.""" + cfg_gain_1 = _cfg(init_mode="kaiming_normal", initializer_range=1.0, seed=0) + cfg_gain_2 = _cfg(init_mode="kaiming_normal", initializer_range=2.0, seed=0) + + bridge_1 = TransformerBridge.boot_native(cfg_gain_1) + bridge_2 = TransformerBridge.boot_native(cfg_gain_2) + + std_1 = bridge_1.W_E.std().item() + std_2 = bridge_2.W_E.std().item() + + # Same seed, only gain differs -> std should scale ~proportionally. + assert std_2 / std_1 == pytest.approx(2.0) diff --git a/transformer_lens/config/transformer_bridge_config.py b/transformer_lens/config/transformer_bridge_config.py index e09d55e75..ac86bb594 100644 --- a/transformer_lens/config/transformer_bridge_config.py +++ b/transformer_lens/config/transformer_bridge_config.py @@ -2,6 +2,7 @@ from typing import Optional +import numpy as np import torch from transformer_lens.utilities.activation_functions import SOFTCAP_DISABLED @@ -219,6 +220,18 @@ def __post_init__(self): ): raise ValueError(f"architecture must be a string, got {type(self.architecture)}") + # Resolve the initializer_range sentinel (-1.0 means "not set by the user"). + # Mirrors HookedTransformerConfig.__post_init__ (hooked_transformer_config.py). + # Guarded with getattr: this method also runs once from the dataclass + # parent's __init__, before self.initializer_range is assigned below. + if getattr(self, "initializer_range", None) is not None: + if self.initializer_range < 0 and self.init_mode == "gpt2": + # Roughly copy the GPT-2 value, but proportional to sqrt(1/d_model) + self.initializer_range = 0.8 / np.sqrt(self.d_model) + if self.initializer_range < 0 and self.init_mode != "gpt2": + # This is the gain parameter for the weight initialisation + self.initializer_range = 1.0 + # Call parent's __post_init__ after our validation if hasattr(super(), "__post_init__"): super().__post_init__() diff --git a/transformer_lens/model_bridge/sources/native/init.py b/transformer_lens/model_bridge/sources/native/init.py index cebfc4171..bc8fdd31e 100644 --- a/transformer_lens/model_bridge/sources/native/init.py +++ b/transformer_lens/model_bridge/sources/native/init.py @@ -72,7 +72,7 @@ def initialize_native_model( generator = None def _staged( - fn: Callable[[torch.Tensor], torch.Tensor] + fn: Callable[[torch.Tensor], torch.Tensor], ) -> Callable[[torch.Tensor], torch.Tensor]: def apply(t: torch.Tensor) -> torch.Tensor: staging = torch.empty(t.shape, dtype=torch.float32) @@ -97,6 +97,14 @@ def apply(t: torch.Tensor) -> torch.Tensor: # std = 0.8/sqrt(d_model), not GPT-2's paper 0.02 — toy-model training # dynamics (e.g. the grokking demo) depend on this scale. std = cfg.initializer_range if cfg.initializer_range > 0 else 0.8 / math.sqrt(cfg.d_model) + + # NOTE: this residual output scaling (1/sqrt(2*n_layers), applied only + # to output projections below) is NOT present in HookedTransformer's + # _init_weights_gpt2 (see transformer_lens/HookedTransformer.py). + # Intentional delta for NativeModel: kept because it follows the + # original GPT-2 paper's residual-scaling convention and improves + # training stability at init for deeper models. Flagged in issue #1568 + # as a maintainer call; kept + documented rather than removed. residual_scale = 1.0 / math.sqrt(2 * cfg.n_layers) weight_init = lambda t: nn.init.normal_( t, mean=0.0, std=std, generator=generator