Skip to content

Fix native init: resolve initializer_range sentinel, thread gain into xavier/kaiming - #1685

Merged
jlarson4 merged 4 commits into
TransformerLensOrg:devfrom
cjnegao11-cmyk:fix-native-init-range-sentinel
Aug 22, 2026
Merged

Fix native init: resolve initializer_range sentinel, thread gain into xavier/kaiming#1685
jlarson4 merged 4 commits into
TransformerLensOrg:devfrom
cjnegao11-cmyk:fix-native-init-range-sentinel

Conversation

@cjnegao11-cmyk

@cjnegao11-cmyk cjnegao11-cmyk commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Addresses #1568

  • Resolves the initializer_range = -1.0 sentinel in TransformerBridgeConfig.__post_init__, mirroring HookedTransformerConfig
  • Retargets the PR to dev and keeps the existing dev xavier/kaiming gain handling instead of duplicating it
  • Documents the 1/sqrt(2*n_layers) residual output scaling in gpt2-mode as an intentional delta from HookedTransformer._init_weights_gpt2
  • Updates regression coverage so boot_native() is exercised end-to-end and verifies the resulting embedding std
  • Tightens the kaiming gain regression check to require the exact 2x scaling ratio for a fixed seed

Validation:

  • python -m pytest tests/unit/model_bridge/test_boot_native.py -q → 31 passed
  • git diff --check → passed

Note: this issue was previously assigned to @tomatotomata; @jlarson4 asked for a status check before this PR was opened.

@jlarson4

jlarson4 commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator

@cjnegao11-cmyk I just posted 6 fresh issues as well if you want to take a look! Thanks for contributing, it is appreciated.

I will opt for this solution to #1568 & review it tomorrow if we don't here from @tomatotomata by then

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for putting this together @cjnegao11-cmyk and for verifying the gain equivalence rather than just asserting it. A couple small edit requests below:

Comment thread transformer_lens/model_bridge/sources/native/init.py Outdated
Comment thread tests/unit/model_bridge/test_boot_native.py
Comment thread tests/unit/model_bridge/test_boot_native.py Outdated
@cjnegao11-cmyk
cjnegao11-cmyk changed the base branch from main to dev August 21, 2026 18:18
@cjnegao11-cmyk
cjnegao11-cmyk force-pushed the fix-native-init-range-sentinel branch from c8a3423 to 81b9400 Compare August 21, 2026 19:41

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update & rebase @cjnegao11-cmyk! Just one additional comment

optimizer.zero_grad()


def test_boot_native_resolves_initializer_range_sentinel():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These assert weight statistics after boot_native, but dev's init.py already falls back to 0.8/math.sqrt(cfg.d_model) and applies the gain, so both tests pass on dev without your fixes.

Lets keep the end-to-end assertions and add the config-level ones back alongside them — cfg.initializer_range == pytest.approx(0.8 / math.sqrt(cfg.d_model)) for init_mode="gpt2", plus a case for a non-gpt2 mode resolving to 1.0, which is currently untested. Both read -1.0 without this PR.

@jlarson4

Copy link
Copy Markdown
Collaborator

Great work! Merging now

@jlarson4
jlarson4 merged commit 1c55fe0 into TransformerLensOrg:dev Aug 22, 2026
25 checks passed
@cjnegao11-cmyk

Copy link
Copy Markdown
Contributor Author

Thanks for the thorough review and for merging! Happy to keep contributing.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug Report] Native init drifts from HookedTransformer: unresolved initializer_range sentinel and ignored xavier/kaiming gain

2 participants