From 216e2cd3eb55e1bb740685ebd01cc345cdf7ccc7 Mon Sep 17 00:00:00 2001 From: Marco Date: Tue, 18 Aug 2026 20:56:05 -0300 Subject: [PATCH 1/4] Fix native init: resolve initializer_range sentinel, thread gain into xavier/kaiming, document residual scaling delta --- tests/unit/model_bridge/test_boot_native.py | 33 +++++++++++++++++++ .../config/transformer_bridge_config.py | 13 ++++++++ .../model_bridge/sources/native/init.py | 14 +++++++- 3 files changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 0a010dd37..410334b8d 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,35 @@ 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 -1.0 sentinel must be resolved on the + bridge config path the same way HookedTransformerConfig resolves it + (hooked_transformer_config.py), instead of silently falling back to the + unrelated std=0.02 default inside native/init.py.""" + import math + + cfg = _cfg(init_mode="gpt2") + # _cfg() doesn't set initializer_range, so it should still carry the + # -1.0 sentinel until __post_init__ resolves it. + expected = 0.8 / math.sqrt(cfg.d_model) + assert cfg.initializer_range == pytest.approx(expected) + + +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, rel=0.15) 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..befb6727d 100644 --- a/transformer_lens/model_bridge/sources/native/init.py +++ b/transformer_lens/model_bridge/sources/native/init.py @@ -96,7 +96,19 @@ def apply(t: torch.Tensor) -> torch.Tensor: # Default matches the legacy TL scheme: N(0, 0.64/d_model), i.e. # 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) + 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 From 81b940066dfc9cd17b8b2f7ccee7f9a14a86bf03 Mon Sep 17 00:00:00 2001 From: Marco Date: Fri, 21 Aug 2026 16:41:38 -0300 Subject: [PATCH 2/4] Address native init review feedback --- tests/unit/model_bridge/test_boot_native.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 410334b8d..639f8535e 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/tests/unit/model_bridge/test_boot_native.py @@ -493,17 +493,14 @@ def test_boot_native_supports_training_step(): def test_boot_native_resolves_initializer_range_sentinel(): - """Regression for #1568 — the -1.0 sentinel must be resolved on the - bridge config path the same way HookedTransformerConfig resolves it - (hooked_transformer_config.py), instead of silently falling back to the - unrelated std=0.02 default inside native/init.py.""" + """Regression for #1568 — the resolved initializer range must be used by boot_native.""" import math cfg = _cfg(init_mode="gpt2") - # _cfg() doesn't set initializer_range, so it should still carry the - # -1.0 sentinel until __post_init__ resolves it. + bridge = TransformerBridge.boot_native(cfg) + expected = 0.8 / math.sqrt(cfg.d_model) - assert cfg.initializer_range == pytest.approx(expected) + assert bridge.W_E.std().item() == pytest.approx(expected, rel=0.15) def test_boot_native_kaiming_gain_scales_weights(): @@ -521,4 +518,4 @@ def test_boot_native_kaiming_gain_scales_weights(): 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, rel=0.15) + assert std_2 / std_1 == pytest.approx(2.0) From ac4d35258731ef9bbc8a52693aa2bfe9ff2a1f48 Mon Sep 17 00:00:00 2001 From: Marco Date: Fri, 21 Aug 2026 16:52:33 -0300 Subject: [PATCH 3/4] Format native init --- transformer_lens/model_bridge/sources/native/init.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/transformer_lens/model_bridge/sources/native/init.py b/transformer_lens/model_bridge/sources/native/init.py index befb6727d..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) @@ -96,11 +96,7 @@ def apply(t: torch.Tensor) -> torch.Tensor: # Default matches the legacy TL scheme: N(0, 0.64/d_model), i.e. # 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) - ) + 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 From d63b0cec847f598eaee79b45b52af9f2f49eb119 Mon Sep 17 00:00:00 2001 From: Marco Date: Fri, 21 Aug 2026 19:15:06 -0300 Subject: [PATCH 4/4] Add config-level initializer sentinel coverage --- tests/unit/model_bridge/test_boot_native.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tests/unit/model_bridge/test_boot_native.py b/tests/unit/model_bridge/test_boot_native.py index 639f8535e..2f0069069 100644 --- a/tests/unit/model_bridge/test_boot_native.py +++ b/tests/unit/model_bridge/test_boot_native.py @@ -497,12 +497,20 @@ def test_boot_native_resolves_initializer_range_sentinel(): import math cfg = _cfg(init_mode="gpt2") - bridge = TransformerBridge.boot_native(cfg) - 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